diff --git a/python/sglang/kernels/jit/csrc/distributed/communicator.cuh b/python/sglang/kernels/jit/csrc/distributed/communicator.cuh deleted file mode 100644 index 2d2108cd6..000000000 --- a/python/sglang/kernels/jit/csrc/distributed/communicator.cuh +++ /dev/null @@ -1,117 +0,0 @@ -#include -#include -#include - -#include - -#include -#include - -#include -#include -#include -#include -#include - -namespace sglang { - -namespace host::distributed { - -inline CommunicatorObj::CommunicatorObj( - const uint32_t rank, - const uint32_t world_size, - std::vector push_workspaces, - std::vector pull_workspaces, - std::vector pull_semaphores, - TensorView push_counter, - const std::optional 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(); - for (uint32_t i = 0; i < world_size; ++i) { - TensorMatcher({2 * world_size, push_bytes}).with_dtype().with_device(device).verify(push_workspaces[i]); - TensorMatcher({pull_bytes}) // - .with_dtype() - .with_device(device) - .verify(pull_workspaces[i]); - TensorMatcher({num_pull_blocks, static_cast(sizeof(Semaphore))}) - .with_dtype() - .with_device(device) - .verify(pull_semaphores[i]); - this->push_workspaces[i] = static_cast(push_workspaces[i].data_ptr()); - this->pull_workspaces[i] = static_cast(pull_workspaces[i].data_ptr()); - this->pull_semaphores[i] = static_cast(pull_semaphores[i].data_ptr()); - } - TensorMatcher({num_push_blocks, static_cast(sizeof(Counter))}) - .with_dtype() - .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(static_cast(pull_mc_workspace_ptr.value())); - } else { - this->pull_mc_workspace = nullptr; - } - - // push config - this->push_counter = static_cast(push_counter.data_ptr()); - this->push_bytes = push_bytes.unwrap(); - this->num_push_blocks = static_cast(num_push_blocks.unwrap()); - - // pull config - this->pull_bytes = pull_bytes.unwrap(); - this->num_pull_blocks = static_cast(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 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() - .def( - refl::init< - uint32_t, - uint32_t, - std::vector, - std::vector, - std::vector, - TensorView, - std::optional>(), - "__init__") - .def_ro("world_size", &Class::world_size) - .def_ro("rank", &Class::rank) - .def("_config", &Class::config); -} - -} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh index c3b3ac179..627218dfd 100644 --- a/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh @@ -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 #include #include +#include #include #include #include @@ -22,475 +30,235 @@ #include #include +#include -#include -#include #include #include #include -#include 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 -struct fp_trait {}; - -template <> -struct fp_trait { - 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 { - 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 { - using type = uint32_t; - [[maybe_unused]] - static constexpr uint32_t pos_zero = 0x00000000u; - [[maybe_unused]] - static constexpr uint32_t neg_zero = 0x80000000u; -}; - -template -SGL_DEVICE void clear_pos_zero(DType& val) { - using Trait = fp_trait; - const auto ptr = reinterpret_cast(&val); - if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero; -} - -template -SGL_DEVICE bool is_pos_zero(const DType& val) { - using Trait = fp_trait; - const auto ptr = reinterpret_cast(&val); - return *ptr == Trait::pos_zero; -} - -template -SGL_DEVICE DType get_pos_zero() { - using Trait = fp_trait; - const auto value = Trait::pos_zero; - return *reinterpret_cast(&value); -} - -template -SGL_DEVICE auto reduce(device::AlignedVector (&vec)[M]) -> device::AlignedVector { - 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(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 out_vec; -#pragma unroll - for (size_t j = 0; j < N; ++j) { - out_vec[j] = device::cast(acc_vec[j]); - } - return out_vec; -} - -template -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(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(&val); -} - -template -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(&x); - addr = static_cast(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 -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(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(&val); -} - -template -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(&x); - addr = static_cast(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 -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(mc_addr) + vec_offset * 16; - if constexpr (std::is_same_v>) { - 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(&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>) { - 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>); // 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(&val); - } -#else - assert(false && "multimem load is only supported on Hopper or later architecture"); -#endif -} - -template -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(&x); - mc_addr = static_cast(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 -struct AllReduceParams { +struct AllReducePushParams { const void* __restrict__ input; void* __restrict__ output; - uint32_t num_elements; + uint32_t num_vecs; + uint32_t rank; + PushWorkSpace ws; +}; + +template +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 ws; }; -template -struct AllReducePushImpl { - private: - using T2 = packed_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; - static_assert(kWorldSize <= kMaxWorldSize); - - static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) { - device::PDLWaitPrimary(); - return (params.push_counter[blockIdx.x].get() % 2) != 0; - } - - static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) { - device::PDLTriggerSecondary(); - __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(); -#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 +struct LoadStoreImpl { public: - static SGL_DEVICE void forward_1shot(const AllReduceParams& 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& params, uint32_t vec_offset = 0) { + if constexpr (kUseGraph) { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + m_data[i] = reinterpret_cast(params.graph_params[i]) + vec_offset; + } + } else { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + m_data[i] = reinterpret_cast(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 -struct AllReducePullImpl { private: - using T2 = packed_t; - static constexpr uint32_t kVecSize = 16 / sizeof(T2); - static constexpr uint32_t kElemsPerVec = 16 / sizeof(T); - using vec_t = device::AlignedVector; - static_assert(kWorldSize <= kMaxWorldSize); - - template - static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams& params) { - uint32_t current_counter_val = 0; - if (const auto tx = threadIdx.x; tx < kWorldSize) { - device::PDLWaitPrimary(); - const auto bx = blockIdx.x; - const auto semaphore = ¶ms.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 - static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) { - device::PDLTriggerSecondary(); - __syncthreads(); - if (const auto tx = threadIdx.x; tx < kWorldSize) { - const auto bx = blockIdx.x; - const auto semaphore = ¶ms.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 - 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& 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(params); - reduce_impl(total_num_vecs, data, params.output, params.pull_mc_workspace); - sync_exit_pull(params, counter); - } - - static SGL_DEVICE void forward_2shot(const AllReduceParams& 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(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(params.pull_workspaces[i]) + local_vec_bias; - } - } - const auto counter = sync_enter_pull(params); - const auto mc_addr = reinterpret_cast(params.pull_mc_workspace) + local_vec_bias; - reduce_impl(local_num_vecs, data, params.output, mc_addr); - sync_exit_pull(params, counter); - } + V* m_data[kWorldSize]; }; -template -__global__ __launch_bounds__(1024, 1) // - void all_reduce_kernel(const __grid_constant__ AllReduceParams 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 +struct MultiCastImpl { + public: + static_assert(kUseGraph == false); + static constexpr uint32_t size() { + return kWorldSize; } + + SGL_DEVICE MultiCastImpl(const AllReducePullParams& params, uint32_t vec_offset = 0) + : m_multicast_ptr(reinterpret_cast(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 +ALL_REDUCE_KERNEL void all_reduce_1shot_push_kernel(const __grid_constant__ AllReducePushParams params) { + using namespace device; + constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2); + using vec_t = AlignedVector, kVecSize>; + using Lamport = distributed::LamportTrait; + 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(); + const auto epoch = distributed::PushEpoch{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(); + __syncthreads(); + epoch.flip(); +} + +template +ALL_REDUCE_KERNEL void all_reduce_1shot_pull_kernel(const __grid_constant__ AllReducePullParams params) { + using namespace device; + constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2); + using vec_t = AlignedVector, kVecSize>; + const auto num_vecs = params.num_vecs; + const auto impl = Impl{params}; + PDLWaitPrimary(); + const auto barrier = distributed::Barrier{ + 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(); + __syncthreads(); + barrier.arrive_relaxed(/*n=*/1); +} + +template +ALL_REDUCE_KERNEL void all_reduce_2shot_pull_kernel(const __grid_constant__ AllReducePullParams params) { + using namespace device; + constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2); + using vec_t = AlignedVector, 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(); + const auto barrier = distributed::Barrier{ + 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(); + __syncthreads(); + barrier.arrive_rel_acq(/*n=*/1); } template @@ -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(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 - static constexpr auto kernel_pull = - all_reduce_kernel, kWorldSize, kShot>; - template - static constexpr auto kernel_push = all_reduce_kernel, kWorldSize, kShot>; + using PushParams = AllReducePushParams; + using PullParams = AllReducePullParams; + using vec_t = device::AlignedVector, 16 / (sizeof(T) * 2)>; public: - static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant pull_arg) { + static Tensor + run(const CommunicatorRef comm_ref, + const Tensor in, + const std::string algo, + const tvm::ffi::Optional 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(in_.dtype()), "Input dtype mismatch"); - RuntimeCheck(in_.device().device_type == kDLCUDA, "Only CUDA device is supported"); - RuntimeCheck(std::bit_cast(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(in.dtype()) && in.device().device_type == kDLCUDA); + const auto num_elems_int64 = in.numel(); const auto num_elems = static_cast(num_elems_int64); - RuntimeCheck(static_cast(num_elems) == num_elems_int64, "Number of items exceeds 4G limit"); - const bool use_graph = std::holds_alternative(pull_arg); - const auto graph_ptr = use_graph ? std::get(pull_arg).data_ptr() : nullptr; - const bool inplace = use_graph && algo == "2shot_pull"; - Tensor out = inplace ? in_ : ffi::empty_like(in_); - AllReduceParams params{ - .input = in_.data_ptr(), - .output = out.data_ptr(), - .num_elements = num_elems, - .rank = data.rank, - .graph_params = static_cast(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(num_elems_int64 * sizeof(T)); + CHECK_HOST(static_cast(num_elems) == num_elems_int64) << num_elems_int64; + CHECK_HOST(reinterpret_cast(in.data_ptr()) % 16 == 0 && nbytes % 16 == 0); + const auto num_vecs = static_cast(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(nbytes), + }; + using Impl = LoadStoreImpl; + const auto kernel = all_reduce_1shot_push_kernel; // 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(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(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(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; + using LS_GRAPH = LoadStoreImpl; + using MC = MultiCastImpl; 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 + : all_reduce_1shot_pull_kernel; // 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; + // 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 + : all_reduce_2shot_pull_kernel; 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 -tvm::ffi::Tensor custom_all_reduce( - CommunicatorRef comm, tvm::ffi::Tensor input, std::string algo, std::variant pull_arg) { - return AllReduceKernel::run(comm, input, algo, pull_arg); -} - } // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/distributed/registry.cuh b/python/sglang/kernels/jit/csrc/distributed/registry.cuh new file mode 100644 index 000000000..2dd77e16d --- /dev/null +++ b/python/sglang/kernels/jit/csrc/distributed/registry.cuh @@ -0,0 +1,146 @@ +#include + +#include + +#include +#include + +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 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(); + for (uint32_t i = 0; i < world_size; ++i) { + TensorMatcher({2 * world_size, N}) // + .with_dtype() + .with_device(device_sym) + .verify(workspaces[i]); + } + TensorMatcher({M, static_cast(sizeof(Counter))}) // + .with_dtype() + .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(M.unwrap()); + this->slot_bytes = N.unwrap(); + this->counter = static_cast(counter.data_ptr()); + for (uint32_t i = 0; i < world_size; ++i) { + this->workspaces[i] = static_cast(workspaces[i].data_ptr()); + } + this->mc_workspace = reinterpret_cast(mc_workspace_ptr); +} + +PullPlaneObj::PullPlaneObj( + uint32_t rank, + uint32_t world_size, + std::vector workspaces, // world_size * [num_bytes] + std::vector 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(); + for (uint32_t i = 0; i < world_size; ++i) { + TensorMatcher({N}) // + .with_dtype() + .with_device(device_sym) + .verify(workspaces[i]); + TensorMatcher({M, static_cast(sizeof(Semaphore))}) // + .with_dtype() + .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(M.unwrap()); + this->num_bytes = N.unwrap(); + for (uint32_t i = 0; i < world_size; ++i) { + this->semaphores[i] = static_cast(semaphores[i].data_ptr()); + this->workspaces[i] = static_cast(workspaces[i].data_ptr()); + } + this->mc_semaphore = reinterpret_cast(mc_semaphore_ptr); + this->mc_workspace = reinterpret_cast(mc_workspace_ptr); +} + +CommunicatorObj::CommunicatorObj(Optional push, Optional 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; + + refl::ObjectDef() + .def(refl::init(), "__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() + .def(refl::init(), "__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() + .def(refl::init, dist::Optional>(), "__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 diff --git a/python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh b/python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh index 8873a9c43..842a3a5f7 100644 --- a/python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh +++ b/python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh @@ -36,23 +36,6 @@ struct ParallelQKNormParams { uint32_t num_clean_up_count = 0; }; -template -SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) { - static_assert(alignof(T) == 8 && sizeof(T) == 8); - addr = device::pointer::offset(addr, offset); - uint2 val; - asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); - x = *reinterpret_cast(&val); -} - -template -SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) { - static_assert(alignof(T) == 8 && sizeof(T) == 8); - const uint2 val = *reinterpret_cast(&x); - addr = device::pointer::offset(addr, offset); - asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr)); -} - [[maybe_unused]] SGL_DEVICE float sync_float(float x) { return __shfl_sync(0xffffffffu, x, 0); @@ -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(comm.push_bytes); + params.epoch_bytes = static_cast(push.slot_bytes); params.num_clean_up_count = num_clean; const auto needed_buffer_bytes = static_cast(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(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned"); - RuntimeCheck(needed_buffer_bytes <= 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() { diff --git a/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh b/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh index ab4f9deb1..c47110ea1 100644 --- a/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh @@ -15,6 +15,7 @@ limitations under the License. #pragma once +#include #include #include @@ -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); diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh index 912f06f31..b25347e83 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh @@ -9,9 +9,12 @@ #include #include +#include +#include + #include +#include -#include "../../distributed/custom_all_reduce.cuh" #include #include #include @@ -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 -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 -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= 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(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(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(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 __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 = ¶ms.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(); - 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(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(smem_raw)); __syncthreads(); - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.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(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 __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 = ¶ms.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(smem_raw)); __syncthreads(); - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.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(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(T_.unwrap()); const auto H = static_cast(H_.unwrap()); const auto NB = static_cast(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(num_sm) * kOccupancy, num_tokens, max_blocks, - static_cast(data.num_pull_blocks)}); + static_cast(ref.get()->get_pull_blocks())}); const auto local_elems = num_tokens * H; const auto params = AttnResTMAParams{ .prefix_sum = static_cast(prefix_out.data_ptr()), @@ -847,14 +817,14 @@ struct AttnResFusedTmaKernel { .out = static_cast(out.data_ptr()), .prefix_dst = nullptr, .input_mc = reinterpret_cast(static_cast(input_mc_ptr)) + - data.rank * local_elems * sizeof(bf16_t), + pull.rank * local_elems * sizeof(bf16_t), .residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr, .prefix_out = static_cast(prefix_out.data_ptr()), - .sem_local = data.pull_semaphores[data.rank], - .sem_mc = reinterpret_cast(static_cast(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(eps), .num_tokens = static_cast(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(GT_.unwrap()); const auto H = static_cast(H_.unwrap()); const auto NB = static_cast(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(num_sm) * kOccupancy, num_tokens, max_blocks, - static_cast(data.num_pull_blocks)}); + static_cast(ref.get()->get_pull_blocks())}); const auto params = AttnResTMAParams{ .prefix_sum = static_cast(prefix_sum.data_ptr()), .bank = static_cast(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(static_cast(sem_mc_ptr)), + .sem_local = pull.semaphores[pull.rank], + .sem_mc = pull.mc_semaphore, .output_mc = reinterpret_cast(static_cast(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(eps), .num_tokens = static_cast(num_tokens), diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh index 26ebd56e0..e0a8a9577 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh @@ -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 + #include +#include #include #include #include @@ -38,15 +41,23 @@ #include // TODO: remove dependency on the custom_all_reduce, move out common utilities -#include "../../distributed/custom_all_reduce.cuh" -#include "ptx_sys.cuh" +#include +#include + +#include 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; 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::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(&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(), get_pos_zero()}); + 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(&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::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(&vec); - if (bits.x == 0) bits.x = fp_trait::neg_zero; - if (bits.y == 0) bits.y = fp_trait::neg_zero; - if (bits.z == 0) bits.z = fp_trait::neg_zero; - if (bits.w == 0) bits.w = fp_trait::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(), get_pos_zero()}); + 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(&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(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(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 -SGL_DEVICE uint32_t pull_barrier_enter(const PullParams& params) { - uint32_t current = 0; - if (threadIdx.x == 0) { - const auto semaphore = ¶ms.sem_local[blockIdx.x]; - const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size); - current = reserved + params.world_size; - device::PDLWaitPrimary(); - 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. -template -SGL_DEVICE void pull_barrier_exit(const PullParams& params, uint32_t current) { - device::PDLTriggerSecondary(); - __syncthreads(); - if (threadIdx.x == 0) { - const auto semaphore = ¶ms.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(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(); + 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(vid, num_vecs, step, mc_ptr, res_ptr); - pull_barrier_exit(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(); + __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(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(); + 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(wvec[j]); vec[u][j] = cast(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(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(); + __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 struct AllReduceFusionKernel { @@ -624,8 +601,9 @@ struct AllReduceFusionKernel { static constexpr auto res_push_kernel = all_reduce_push_res_kernel; static FusionParams - make_params(const host::distributed::CommunicatorObj& data, TensorView input, std::optional residual) { + make_params(const host::distributed::CommunicatorObj& comm, TensorView input, std::optional residual) { using namespace host; + const auto& push = comm.get_push_obj(); SymbolicSize N = {"num_elements"}; SymbolicDevice device; device.set_options(); @@ -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(input.data_ptr()); params.residual = residual.has_value() ? static_cast(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(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(); TensorMatcher({kNormDim}).with_dtype().with_device(device).verify(weight); @@ -679,18 +658,18 @@ struct AllReduceFusionKernel { params.norm_weight = static_cast(weight.data_ptr()); params.norm_eps = static_cast(eps); params.num_norm_rows = static_cast(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 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(); @@ -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(static_cast(input_mc_ptr)); params.residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr; - params.sem_local = data.pull_semaphores[data.rank]; - params.sem_mc = reinterpret_cast(static_cast(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(num_elems / 8); return params; } @@ -765,33 +744,26 @@ struct AllReduceFusionKernel { } public: - static void push_res(CommunicatorRef ref, TensorView input, std::optional 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 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(static_cast(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 : res_push_kernel; - 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(static_cast(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(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, 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().with_device(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(static_cast(ws_mc_base)); + CHECK_HOST(nbytes <= push.slot_bytes) << "output size " << nbytes << " exceeds push slot size " << push.slot_bytes; params.fin_gemm2 = static_cast(gemm2_out.data_ptr()); params.fin_idx = static_cast(permuted_idx.data_ptr()); params.fin_weights = static_cast(expert_weights.data_ptr()); constexpr uint32_t kMaxClusters = 96; const auto num_row_clusters = std::max(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, 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 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(num_blocks, data.num_pull_blocks); + num_blocks = std::min(num_blocks, ref.get()->get_pull_blocks()); if (residual.has_value()) { launch_pull_res(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(num_blocks, data.num_pull_blocks); + num_blocks = std::min(num_blocks, ref.get()->get_pull_blocks()); using namespace host; SymbolicDevice device; device.set_options(); diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh index de7636302..e50b84dfb 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh @@ -11,7 +11,6 @@ #include #include -#include "ptx_sys.cuh" #include #include #include @@ -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(params.ws_mc + phase * params.half_bytes); const auto dst = reinterpret_cast(base + elem); - multimem_store_relaxed(dst, bits); + ptx::multimem_store_relaxed(dst, bits); } PDLTriggerSecondary(); } @@ -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(&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{}); static void - run(CommunicatorRef ref, - TensorView x, - TensorView weight, - TensorView b, - std::optional c, - TensorView out, - intptr_t ws_mc_base) { + run(CommunicatorRef ref, TensorView x, TensorView weight, TensorView b, std::optional 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().with_device(device).verify(out); const auto num_tokens = static_cast(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(ws_mc_base), - .counter = data.push_counter, - .half_bytes = static_cast(data.push_bytes * data.world_size), - .rank = data.rank, + .ws_mc = push.mc_workspace, + .counter = push.counter, + .half_bytes = static_cast(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(data.push_bytes * data.world_size), + .ws_local = push.workspaces[push.rank], + .counter = push.counter, + .num_counters = push.num_blocks, + .half_bytes = static_cast(push.slot_bytes * push.world_size), .b = static_cast(b.data_ptr()), .c = c.has_value() ? static_cast(c.value().data_ptr()) : nullptr, .out = static_cast(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; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh index 6bada89e3..ea3498c8f 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh @@ -17,7 +17,6 @@ #include #include -#include "ptx_sys.cuh" #include #include #include @@ -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(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(prm.uc_base[r] + foff) + ring, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(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(prm.uc_base[r] + doff2) + off, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(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(a_smem + uint32_t(k2) * 32); - const uint64_t db = ptx::mma_smem_desc_k_major(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(a_smem + uint32_t(k2) * 32); + const uint64_t db = device::ptx::mma_smem_desc_k_major(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(__int_as_float(r1), __int_as_float(r0)); - v.y = ptx::cvt_pack_f32x2_to(__int_as_float(r3), __int_as_float(r2)); - v.z = ptx::cvt_pack_f32x2_to(__int_as_float(r5), __int_as_float(r4)); - v.w = ptx::cvt_pack_f32x2_to(__int_as_float(r7), __int_as_float(r6)); + v.x = device::ptx::cvt_pack_f32x2_to(__int_as_float(r1), __int_as_float(r0)); + v.y = device::ptx::cvt_pack_f32x2_to(__int_as_float(r3), __int_as_float(r2)); + v.z = device::ptx::cvt_pack_f32x2_to(__int_as_float(r5), __int_as_float(r4)); + v.w = device::ptx::cvt_pack_f32x2_to(__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(); 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(prm.uc_base[r] + foff) + ring, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(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(prm.uc_base[r] + foff + 256) + ring, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + foff + 256) + ring, 1); } } - while (load_acquire_sys(reinterpret_cast(prm.uc_base[prm.my_rank] + foff + 256) + ring) < + while (device::ptx::load_acquire_sys(reinterpret_cast(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(prm.uc_base[r] + doff2) + off, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(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(prm.uc_base[r] + doff2) + off, 1); + device::ptx::red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + doff2) + off, 1); } } } diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh deleted file mode 100644 index 59eafa897..000000000 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh +++ /dev/null @@ -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 - -#include - -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 diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh index 16b0de31c..908456647 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh @@ -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 + #include #include -#include "../../distributed/custom_all_reduce.cuh" +#include +#include + +#include 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 -SGL_DEVICE void make_nonzero(Vec& vec) { - constexpr uint32_t kNegZeroPair = 0x8000u; - auto& bits = *reinterpret_cast(&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(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; template -SGL_DEVICE bool has_empty_marker(const Vec& vec) { - const auto bits = *reinterpret_cast(&vec); - return bits.x == 0 || bits.y == 0 || bits.z == 0 || bits.w == 0; -} - -template -SGL_DEVICE Vec zero_vec() { - Vec zero; - zero.fill(bf16x2_t{get_pos_zero(), get_pos_zero()}); - return zero; +SGL_DEVICE Vec empty_vec() { + Vec vec; + Lamport::fill_pos_zero(vec.data()); + return vec; } template @@ -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(); @@ -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(); + const auto zero = empty_vec(); 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(); const auto poll_base = params.push_workspaces[params.rank] + phase_offset; - const auto zero = zero_vec(); + const auto zero = empty_vec(); 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 __global__ void all_gather_direct_kernel(const __grid_constant__ Params params) { using vec_t = device::AlignedVector; // 16 B - uint32_t exit_base = 0; - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.sem_local[blockIdx.x]; - const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize); - exit_base = reserved + kWorldSize; - device::PDLWaitPrimary(); - 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(); + 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(); __syncthreads(); - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.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; // 16 B using SumOp = device::ReductionTrait; - uint32_t exit_base = 0; - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.sem_local[blockIdx.x]; - const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize); - exit_base = reserved + kWorldSize; - device::PDLWaitPrimary(); - 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(); + 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(); __syncthreads(); - if (threadIdx.x == 0) { - auto* semaphore = ¶ms.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 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 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().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(input.data_ptr()), .output = static_cast(output.data_ptr()), .residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr, .push_workspaces = {}, - .push_ws_mc = reinterpret_cast(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(local_elems.unwrap() * sizeof(bf16_t) / 16), .residual_is_local = static_cast(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 : sp_collective::reduce_scatter_res_kernel; 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(input.data_ptr()); params.output = static_cast(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(input.data_ptr()); params.output = static_cast(output.data_ptr()); params.output_mc = reinterpret_cast(static_cast(output_mc_ptr)); - params.sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)); host::LaunchKernel(num_blocks, block_size, input.device()) .enable_pdl(kUsePDL)(sp_collective::all_gather_direct_kernel, params); } @@ -421,17 +385,14 @@ struct SPCollectiveKernel { std::optional 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(static_cast(input_mc_ptr)); - params.sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)); const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_pull_kernel : sp_collective::reduce_scatter_pull_kernel; host::LaunchKernel(num_blocks, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params); diff --git a/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh b/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh index 19bc87098..411e0c3c2 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh @@ -1,16 +1,21 @@ #pragma once +#include #include #include +#include + #include #include +#include +#include #include #include -#include +#include #include -#include +#include #include 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 +struct PushWorkSpace { + std::array 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 +struct PullWorkSpace { + std::array semaphores; + std::array 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 +struct FloatTrait {}; + +template <> +struct FloatTrait { + using type = uint16_t; + static constexpr uint16_t kNegZero = 0x8000u; +}; + +template <> +struct FloatTrait { + using type = uint16_t; + static constexpr uint16_t kNegZero = 0x8000u; +}; + +template <> +struct FloatTrait { + using type = uint32_t; + static constexpr uint32_t kNegZero = 0x80000000u; +}; + +template +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>; + static constexpr Packed kNegZero = FloatTrait::kNegZero; + static constexpr uint32_t kNumPacked = N / (kAtom / sizeof(T)); + + SGL_DEVICE static void clear_pos_zero(void* val) { + const auto ptr = static_cast(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(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(val); +#pragma unroll + for (uint32_t i = 0; i < kNumPacked; ++i) { + ptr[i] = 0; + } + } +}; + +template +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 + 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(n); + } + + SGL_DEVICE void arrive_rel_acq(uint32_t n) const { + return this->arrive(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 +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& 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 + SGL_DEVICE void arrive(uint32_t n) const { + return arrive_at(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(n); + } + + SGL_DEVICE void arrive_rel_acq(uint32_t n) const { + return this->arrive(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 + 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 +using Optional = tvm::ffi::Optional; + 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 push_workspaces, - std::vector pull_workspaces, - std::vector pull_semaphores, - TensorView push_counter, - std::optional pull_mc_workspace_ptr); - - void config(std::map 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 pull_workspaces; // symmetric memory - std::array push_workspaces; // symmetric memory - std::array 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 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 workspaces; // symmetric memory + uint8_t* mc_workspace; // multicast VA of the local workspace (may be null) + + template + PushWorkSpace 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::max()) + << 2 * N * slot_bytes << " bytes of push plane exceeds the 32-bit offset range"; + PushWorkSpace ws{{}, counter, offset_mc(mc_workspace, offset), static_cast(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 workspaces, // world_size * [num_bytes] + std::vector 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 semaphores; // symmetric memory + std::array 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 + PullWorkSpace 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 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 push, Optional 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 get_push() const { + return m_push; + } + Optional 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 num_blocks) { + m_pull_blocks = num_blocks; + } + void set_pull_multicast_blocks(std::optional 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 m_pull_blocks; + std::optional m_pull_multicast_blocks; + Optional m_push; + Optional m_pull; +}; + +SGLANG_REGISTER_FFI_REFERENCE_CLASS(CommunicatorRef, CommunicatorObj); + } // namespace host::distributed } // namespace sglang diff --git a/python/sglang/kernels/jit/include/sgl_kernel/distributed/ptx.cuh b/python/sglang/kernels/jit/include/sgl_kernel/distributed/ptx.cuh new file mode 100644 index 000000000..0b2983da9 --- /dev/null +++ b/python/sglang/kernels/jit/include/sgl_kernel/distributed/ptx.cuh @@ -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 + +#include +#include +#include + +#include +#include + +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 +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(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(&val); +} + +template +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(&x); + addr = static_cast(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 +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(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(&val); +} + +template +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(&x); + addr = static_cast(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 +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(mc_addr) + vec_offset * 16; + if constexpr (std::is_same_v>) { + 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(&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>) { + 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>); // 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(&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 +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(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(&val); +} + +template +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(&x); + addr = static_cast(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 +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(&x); + mc_addr = static_cast(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 diff --git a/python/sglang/kernels/jit/include/sgl_kernel/ffi.h b/python/sglang/kernels/jit/include/sgl_kernel/ffi.h index 8d619afbb..97c17805d 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/ffi.h +++ b/python/sglang/kernels/jit/include/sgl_kernel/ffi.h @@ -103,6 +103,23 @@ inline Tensor from_blob_like( return from_blob(data, t.shape(), t.dtype(), t.device(), std::forward(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(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 diff --git a/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh b/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh index a2796a454..2bbcfd9c2 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh @@ -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 @@ -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 -static SGL_DEVICE uint32_t to_shared(T* ptr) { +SGL_DEVICE uint32_t to_shared(T* ptr) { return static_cast(__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 diff --git a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh index 5c7ca2941..28b9245e3 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh @@ -15,7 +15,6 @@ #pragma once -#include #include #include @@ -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(required_bytes)}; - return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); -} - /** * \brief Kernel launcher with automatic stream resolution and PDL support. * diff --git a/python/sglang/kernels/jit/include/sgl_kernel/vec.cuh b/python/sglang/kernels/jit/include/sgl_kernel/vec.cuh index e9882339f..4c224d07f 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/vec.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/vec.cuh @@ -7,6 +7,7 @@ /// 256 bits (32 bytes), matching CUDA's widest vector load. #pragma once +#include #include #include @@ -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 +SGL_DEVICE auto reduce_vec(device::AlignedVector (&vec)[M]) -> device::AlignedVector { + static_assert(DTypeTrait::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(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 out_vec; +#pragma unroll + for (size_t j = 0; j < N; ++j) { + out_vec[j] = cast(acc[j]); + } + return out_vec; + } else { + using SumOp = ReductionTrait; + device::AlignedVector 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 diff --git a/python/sglang/kernels/ops/communication/all_reduce.py b/python/sglang/kernels/ops/communication/all_reduce.py index 65c5d1474..d7f39e33b 100644 --- a/python/sglang/kernels/ops/communication/all_reduce.py +++ b/python/sglang/kernels/ops/communication/all_reduce.py @@ -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 diff --git a/python/sglang/kernels/ops/kimi_k3/all_reduce.py b/python/sglang/kernels/ops/kimi_k3/all_reduce.py index 7e7b7a52d..f50f95cff 100644 --- a/python/sglang/kernels/ops/kimi_k3/all_reduce.py +++ b/python/sglang/kernels/ops/kimi_k3/all_reduce.py @@ -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, diff --git a/python/sglang/kernels/ops/kimi_k3/attn_res.py b/python/sglang/kernels/ops/kimi_k3/attn_res.py index 33b936975..d8c96a18d 100644 --- a/python/sglang/kernels/ops/kimi_k3/attn_res.py +++ b/python/sglang/kernels/ops/kimi_k3/attn_res.py @@ -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, ) diff --git a/python/sglang/kernels/ops/kimi_k3/gemm_ag.py b/python/sglang/kernels/ops/kimi_k3/gemm_ag.py index 280cdad3a..e108082c7 100644 --- a/python/sglang/kernels/ops/kimi_k3/gemm_ag.py +++ b/python/sglang/kernels/ops/kimi_k3/gemm_ag.py @@ -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 diff --git a/python/sglang/kernels/ops/kimi_k3/sp_collective.py b/python/sglang/kernels/ops/kimi_k3/sp_collective.py index d3a1bcd8d..abb93b474 100644 --- a/python/sglang/kernels/ops/kimi_k3/sp_collective.py +++ b/python/sglang/kernels/ops/kimi_k3/sp_collective.py @@ -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, ) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 3011c9958..5b3ea5a6d 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -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] diff --git a/python/sglang/srt/layers/k3_ar_fusion.py b/python/sglang/srt/layers/k3_ar_fusion.py index 86556c27f..e2f72c008 100644 --- a/python/sglang/srt/layers/k3_ar_fusion.py +++ b/python/sglang/srt/layers/k3_ar_fusion.py @@ -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, ) diff --git a/python/sglang/srt/layers/k3_sp_collective.py b/python/sglang/srt/layers/k3_sp_collective.py index 6a48d43ad..b63c75b8c 100644 --- a/python/sglang/srt/layers/k3_sp_collective.py +++ b/python/sglang/srt/layers/k3_sp_collective.py @@ -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": diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 0b655a872..431ff22b6 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -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, diff --git a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py index 585610b53..35c239a30 100644 --- a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py +++ b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py @@ -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) diff --git a/test/registered/kernels/ops/kimi_k3/test_collectives.py b/test/registered/kernels/ops/kimi_k3/test_collectives.py index faab5e8bc..10a99b106 100644 --- a/test/registered/kernels/ops/kimi_k3/test_collectives.py +++ b/test/registered/kernels/ops/kimi_k3/test_collectives.py @@ -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)