chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
178
third_party/sglang/sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip
vendored
Normal file
178
third_party/sglang/sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip
vendored
Normal file
@@ -0,0 +1,178 @@
|
||||
// Deterministic All-Reduce for ROCm/HIP
|
||||
//
|
||||
// This is a wrapper that forces the use of the existing 1-stage all-reduce kernel
|
||||
// (cross_device_reduce_1stage) which is inherently deterministic due to fixed
|
||||
// accumulation ordering (no atomics, no race conditions).
|
||||
//
|
||||
// How the 1-stage kernel works:
|
||||
// - Each GPU reads ALL data from ALL other GPUs via direct memory access
|
||||
// - Each GPU reduces the data locally in a fixed order
|
||||
// - Result: every GPU has the complete reduced output
|
||||
//
|
||||
// This is NOT a reduce-scatter + all-gather (RS+AG) approach.
|
||||
// The 2-stage kernel (cross_device_reduce_2stage) implements RS+AG but may have
|
||||
// non-deterministic behavior, so we explicitly avoid it here.
|
||||
|
||||
#include <ATen/hip/Exceptions.h>
|
||||
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
|
||||
#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "custom_all_reduce_hip.cuh"
|
||||
|
||||
using fptr_t = int64_t;
|
||||
static_assert(sizeof(void*) == sizeof(fptr_t));
|
||||
|
||||
// Helper function for weak contiguity check
|
||||
bool _is_weak_contiguous_det(torch::Tensor& t) {
|
||||
return t.is_contiguous() ||
|
||||
(t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
|
||||
}
|
||||
|
||||
// Deterministic all-reduce for registered buffers (ROCm)
|
||||
// Uses the 1-stage kernel which is deterministic (fixed ordering)
|
||||
void deterministic_all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out) {
|
||||
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
|
||||
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
|
||||
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
|
||||
TORCH_CHECK_EQ(inp.numel(), out.numel());
|
||||
TORCH_CHECK(_is_weak_contiguous_det(out));
|
||||
TORCH_CHECK(_is_weak_contiguous_det(inp));
|
||||
|
||||
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
|
||||
|
||||
// For ROCm, manually call the 1-stage kernel to ensure deterministic ordering
|
||||
// Get rank data pointer
|
||||
sglang::RankData* ptrs;
|
||||
hipStreamCaptureStatus status;
|
||||
AT_CUDA_CHECK(hipStreamIsCapturing(stream, &status));
|
||||
if (status == hipStreamCaptureStatusActive) {
|
||||
ptrs = fa->d_rank_data_base_ + fa->graph_unreg_buffers_.size();
|
||||
fa->graph_unreg_buffers_.push_back(inp.data_ptr());
|
||||
} else {
|
||||
auto it = fa->buffers_.find(inp.data_ptr());
|
||||
if (it == fa->buffers_.end()) {
|
||||
throw std::runtime_error("buffer not registered!");
|
||||
}
|
||||
ptrs = it->second;
|
||||
}
|
||||
|
||||
int size = out.numel();
|
||||
int threads = 512;
|
||||
|
||||
switch (out.scalar_type()) {
|
||||
case at::ScalarType::Float: {
|
||||
using T = float;
|
||||
using P = typename sglang::packed_t<T>::P;
|
||||
auto d = P::size;
|
||||
if (size % d != 0) {
|
||||
throw std::runtime_error("size must be multiple of " + std::to_string(d));
|
||||
}
|
||||
size /= d;
|
||||
int blocks = std::min(16, (size + threads - 1) / threads);
|
||||
// Always use 1-stage kernel for determinism
|
||||
switch (fa->world_size_) {
|
||||
case 2:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 4:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 6:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 8:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("world_size must be in (2,4,6,8)");
|
||||
}
|
||||
break;
|
||||
}
|
||||
case at::ScalarType::Half: {
|
||||
using T = half;
|
||||
using P = typename sglang::packed_t<T>::P;
|
||||
auto d = P::size;
|
||||
if (size % d != 0) {
|
||||
throw std::runtime_error("size must be multiple of " + std::to_string(d));
|
||||
}
|
||||
size /= d;
|
||||
int blocks = std::min(16, (size + threads - 1) / threads);
|
||||
switch (fa->world_size_) {
|
||||
case 2:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 4:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 6:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 8:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("world_size must be in (2,4,6,8)");
|
||||
}
|
||||
break;
|
||||
}
|
||||
#if (__HIP_ARCH__ >= 800 || !defined(__HIP_ARCH__))
|
||||
case at::ScalarType::BFloat16: {
|
||||
using T = nv_bfloat16;
|
||||
using P = typename sglang::packed_t<T>::P;
|
||||
auto d = P::size;
|
||||
if (size % d != 0) {
|
||||
throw std::runtime_error("size must be multiple of " + std::to_string(d));
|
||||
}
|
||||
size /= d;
|
||||
int blocks = std::min(16, (size + threads - 1) / threads);
|
||||
switch (fa->world_size_) {
|
||||
case 2:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 4:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 6:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
case 8:
|
||||
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
|
||||
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("world_size must be in (2,4,6,8)");
|
||||
}
|
||||
break;
|
||||
}
|
||||
#endif
|
||||
default:
|
||||
throw std::runtime_error("deterministic allreduce only supports float32, float16 and bfloat16");
|
||||
}
|
||||
}
|
||||
|
||||
// Deterministic all-reduce for unregistered buffers (ROCm)
|
||||
void deterministic_all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out) {
|
||||
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
|
||||
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
|
||||
|
||||
auto input_size = inp.numel() * inp.element_size();
|
||||
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
|
||||
TORCH_CHECK_EQ(inp.numel(), out.numel());
|
||||
TORCH_CHECK(input_size <= reg_buffer.numel() * reg_buffer.element_size(),
|
||||
"registered buffer is too small to contain the input");
|
||||
AT_CUDA_CHECK(hipMemcpyAsync(reg_buffer.data_ptr(), inp.data_ptr(),
|
||||
input_size, hipMemcpyDeviceToDevice, stream));
|
||||
deterministic_all_reduce_reg(_fa, reg_buffer, out);
|
||||
}
|
||||
Reference in New Issue
Block a user