mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-16 21:01:14 +08:00
yyj
41613064c7
* feat: support SwiGLU limit for block FP8 experts * perf(fp8): add native layerwise batch transport * feat(fp8): support layerwise transport for TP1/2/4/8 * release: prepare GLM-5.3-flash support * docs(glm): refine GLM-5.3-flash documentation * docs(glm): simplify release announcement
744 lines
31 KiB
C++
744 lines
31 KiB
C++
#include "fp8_layerwise_transport.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <condition_variable>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <limits>
|
|
#include <mutex>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
#include <thread>
|
|
#include <type_traits>
|
|
#include <utility>
|
|
|
|
#if defined(KTRANSFORMERS_USE_CUDA) || defined(USE_CUDA)
|
|
#include <cuda_runtime_api.h>
|
|
#define KT_FP8_LAYERWISE_HAS_CUDA 1
|
|
#else
|
|
#define KT_FP8_LAYERWISE_HAS_CUDA 0
|
|
#endif
|
|
|
|
namespace kt::layerwise {
|
|
namespace {
|
|
|
|
using Clock = std::chrono::steady_clock;
|
|
|
|
constexpr std::uint64_t kControlMagic = 0x4b544650384c5731ULL; // "KTFP8LW1"
|
|
constexpr std::uint64_t kControlVersion = 2;
|
|
constexpr std::uint64_t kPoisonClaimed = std::numeric_limits<std::uint64_t>::max();
|
|
|
|
enum ErrorCode : int {
|
|
kErrorNone = 0,
|
|
kErrorProtocol = 1,
|
|
kErrorTimeout = 2,
|
|
kErrorCudaInitialization = 3,
|
|
kErrorCudaCopy = 4,
|
|
kErrorWriter = 5,
|
|
kErrorClosedWhileActive = 6,
|
|
};
|
|
|
|
struct alignas(64) AtomicCell {
|
|
std::uint64_t value;
|
|
std::uint8_t padding[56];
|
|
};
|
|
static_assert(sizeof(AtomicCell) == 64);
|
|
|
|
struct alignas(64) ControlHeader {
|
|
std::uint64_t magic;
|
|
std::uint64_t version;
|
|
std::uint64_t struct_bytes;
|
|
std::uint64_t tp_size;
|
|
std::uint8_t padding[32];
|
|
};
|
|
static_assert(sizeof(ControlHeader) == 64);
|
|
|
|
struct alignas(64) ReadySlot {
|
|
AtomicCell sequence;
|
|
AtomicCell epoch;
|
|
AtomicCell expert_id;
|
|
};
|
|
|
|
// This type is intentionally POD. It lives in mmap'ed process-shared memory;
|
|
// all synchronization goes through the __atomic helpers below.
|
|
struct alignas(64) FP8LayerwiseControl {
|
|
ControlHeader header;
|
|
AtomicCell poison_state; // 0 healthy, UINT64_MAX claimed/writing, 1 poisoned
|
|
AtomicCell error_code;
|
|
AtomicCell error_rank;
|
|
alignas(64) char error_message[256];
|
|
|
|
AtomicCell active_epoch;
|
|
AtomicCell layer_id;
|
|
AtomicCell expert_count;
|
|
AtomicCell next_sequence;
|
|
AtomicCell producer_done_epoch;
|
|
|
|
AtomicCell join_epoch[kFP8LayerwiseMaxTPSize];
|
|
AtomicCell consumer_done_epoch[kFP8LayerwiseMaxTPSize];
|
|
AtomicCell wait_returned_epoch[kFP8LayerwiseMaxTPSize];
|
|
|
|
ReadySlot ready[kFP8LayerwiseHostSlots];
|
|
AtomicCell ack_sequence[kFP8LayerwiseMaxTPSize][kFP8LayerwiseHostSlots];
|
|
|
|
AtomicCell layer_start_ns;
|
|
AtomicCell writer_ns;
|
|
AtomicCell slot_wait_ns;
|
|
AtomicCell producer_total_ns;
|
|
AtomicCell h2d_ns[kFP8LayerwiseMaxTPSize];
|
|
AtomicCell bytes[kFP8LayerwiseMaxTPSize];
|
|
};
|
|
|
|
static_assert(std::is_trivial_v<FP8LayerwiseControl>);
|
|
static_assert(std::is_standard_layout_v<FP8LayerwiseControl>);
|
|
static_assert(sizeof(FP8LayerwiseControl) <= kFP8LayerwiseControlBytes,
|
|
"FP8 layerwise control must fit in the shared control region");
|
|
|
|
std::uint64_t load_acquire(const AtomicCell& cell) {
|
|
return __atomic_load_n(&cell.value, __ATOMIC_ACQUIRE);
|
|
}
|
|
|
|
std::uint64_t load_relaxed(const AtomicCell& cell) {
|
|
return __atomic_load_n(&cell.value, __ATOMIC_RELAXED);
|
|
}
|
|
|
|
void store_release(AtomicCell& cell, std::uint64_t value) {
|
|
__atomic_store_n(&cell.value, value, __ATOMIC_RELEASE);
|
|
}
|
|
|
|
void store_relaxed(AtomicCell& cell, std::uint64_t value) {
|
|
__atomic_store_n(&cell.value, value, __ATOMIC_RELAXED);
|
|
}
|
|
|
|
std::uint64_t fetch_add_relaxed(AtomicCell& cell, std::uint64_t value) {
|
|
return __atomic_fetch_add(&cell.value, value, __ATOMIC_RELAXED);
|
|
}
|
|
|
|
bool compare_exchange(AtomicCell& cell, std::uint64_t& expected, std::uint64_t desired) {
|
|
return __atomic_compare_exchange_n(&cell.value, &expected, desired, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQUIRE);
|
|
}
|
|
|
|
std::uint64_t now_ns() {
|
|
return static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now().time_since_epoch()).count());
|
|
}
|
|
|
|
double ns_to_ms(std::uint64_t ns) { return static_cast<double>(ns) / 1.0e6; }
|
|
|
|
std::uint64_t signed_to_bits(std::int64_t value) { return static_cast<std::uint64_t>(value); }
|
|
|
|
std::int64_t bits_to_signed(std::uint64_t value) { return static_cast<std::int64_t>(value); }
|
|
|
|
void validate_control_region(std::uintptr_t control_ptr, std::size_t control_size) {
|
|
if (control_ptr == 0) throw std::invalid_argument("FP8 layerwise control pointer is null");
|
|
if ((control_ptr & 63U) != 0) throw std::invalid_argument("FP8 layerwise control pointer must be 64-byte aligned");
|
|
if (control_size < kFP8LayerwiseControlBytes) {
|
|
throw std::invalid_argument("FP8 layerwise control region must be at least 8192 bytes");
|
|
}
|
|
}
|
|
|
|
FP8LayerwiseControl* checked_control(std::uintptr_t control_ptr, std::size_t control_size, int tp_size) {
|
|
validate_control_region(control_ptr, control_size);
|
|
auto* control = reinterpret_cast<FP8LayerwiseControl*>(control_ptr);
|
|
const auto magic = __atomic_load_n(&control->header.magic, __ATOMIC_ACQUIRE);
|
|
if (magic != kControlMagic || control->header.version != kControlVersion ||
|
|
control->header.struct_bytes != sizeof(FP8LayerwiseControl)) {
|
|
throw std::runtime_error("FP8 layerwise control is uninitialized or has an incompatible ABI");
|
|
}
|
|
if (control->header.tp_size != static_cast<std::uint64_t>(tp_size)) {
|
|
throw std::invalid_argument("FP8 layerwise control TP size does not match transport TP size");
|
|
}
|
|
return control;
|
|
}
|
|
|
|
std::string poison_message(FP8LayerwiseControl* control) {
|
|
// The process which wins the poison claim fills diagnostics before changing
|
|
// kPoisonClaimed to 1. Wait briefly so readers normally avoid a partial
|
|
// string, but never let a process dying mid-publication hang every rank.
|
|
const auto claim_deadline = Clock::now() + std::chrono::milliseconds(10);
|
|
while (load_acquire(control->poison_state) == kPoisonClaimed && Clock::now() < claim_deadline) {
|
|
std::this_thread::yield();
|
|
}
|
|
if (load_acquire(control->poison_state) == kPoisonClaimed) {
|
|
return "FP8 layerwise transport poisoned, but poison publication did not complete (claim owner may have died)";
|
|
}
|
|
std::ostringstream out;
|
|
out << "FP8 layerwise transport poisoned (code=" << load_relaxed(control->error_code)
|
|
<< ", rank=" << bits_to_signed(load_relaxed(control->error_rank)) << ")";
|
|
if (control->error_message[0] != '\0') out << ": " << control->error_message;
|
|
return out.str();
|
|
}
|
|
|
|
void publish_poison(FP8LayerwiseControl* control, int code, int rank, const std::string& message) noexcept {
|
|
std::uint64_t expected = 0;
|
|
if (!compare_exchange(control->poison_state, expected, kPoisonClaimed)) return;
|
|
store_relaxed(control->error_code, static_cast<std::uint64_t>(code));
|
|
store_relaxed(control->error_rank, signed_to_bits(rank));
|
|
std::snprintf(control->error_message, sizeof(control->error_message), "%s", message.c_str());
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control->poison_state, 1);
|
|
}
|
|
|
|
void throw_if_poisoned(FP8LayerwiseControl* control) {
|
|
if (load_acquire(control->poison_state) != 0) throw std::runtime_error(poison_message(control));
|
|
}
|
|
|
|
std::uint64_t fold_progress(std::uint64_t seed, std::uint64_t value) {
|
|
// A cheap mixer used only to notice progress and refresh a timeout deadline.
|
|
return seed ^ (value + 0x9e3779b97f4a7c15ULL + (seed << 6U) + (seed >> 2U));
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void initialize_fp8_layerwise_control(std::uintptr_t control_ptr, std::size_t control_size, int tp_size) {
|
|
validate_control_region(control_ptr, control_size);
|
|
if (tp_size <= 0 || tp_size > kFP8LayerwiseMaxTPSize) {
|
|
throw std::invalid_argument("FP8 layerwise transport supports TP sizes 1 through 8");
|
|
}
|
|
auto* control = reinterpret_cast<FP8LayerwiseControl*>(control_ptr);
|
|
std::memset(control, 0, kFP8LayerwiseControlBytes);
|
|
control->header.version = kControlVersion;
|
|
control->header.struct_bytes = sizeof(FP8LayerwiseControl);
|
|
control->header.tp_size = static_cast<std::uint64_t>(tp_size);
|
|
store_relaxed(control->error_rank, signed_to_bits(-1));
|
|
__atomic_store_n(&control->header.magic, kControlMagic, __ATOMIC_RELEASE);
|
|
}
|
|
|
|
class FP8LayerwiseTransport::Impl {
|
|
public:
|
|
Impl(std::uintptr_t control_ptr, std::size_t control_size, int rank, int tp_size, int cuda_device,
|
|
const std::vector<std::uintptr_t>& local_host_ptrs, const std::vector<std::uintptr_t>& local_gpu_ptrs,
|
|
const std::vector<std::uintptr_t>& all_rank_host_ptrs, const std::vector<std::size_t>& expert_nbytes,
|
|
int num_experts, std::uint64_t timeout_ms)
|
|
: control_(checked_control(control_ptr, control_size, tp_size)),
|
|
rank_(rank),
|
|
tp_size_(tp_size),
|
|
cuda_device_(cuda_device),
|
|
num_experts_(num_experts),
|
|
timeout_(timeout_ms) {
|
|
if (rank < 0 || rank >= tp_size) throw std::invalid_argument("FP8 layerwise rank is outside TP range");
|
|
if (num_experts <= 0) throw std::invalid_argument("FP8 layerwise num_experts must be positive");
|
|
if (timeout_ms == 0) throw std::invalid_argument("FP8 layerwise timeout_ms must be positive");
|
|
if (local_host_ptrs.size() != kFP8LayerwiseHostSlots * kFP8LayerwiseBufferKinds) {
|
|
throw std::invalid_argument("FP8 layerwise local_host_ptrs must contain 8 pointers in [slot][kind] order");
|
|
}
|
|
if (local_gpu_ptrs.size() != kFP8LayerwiseBufferKinds) {
|
|
throw std::invalid_argument("FP8 layerwise local_gpu_ptrs must contain 4 pointers");
|
|
}
|
|
if (expert_nbytes.size() != kFP8LayerwiseBufferKinds) {
|
|
throw std::invalid_argument("FP8 layerwise expert_nbytes must contain 4 sizes");
|
|
}
|
|
const std::size_t expected_all_rank = kFP8LayerwiseHostSlots * tp_size * kFP8LayerwiseBufferKinds;
|
|
if ((rank == 0 && all_rank_host_ptrs.size() != expected_all_rank) || (rank != 0 && !all_rank_host_ptrs.empty())) {
|
|
throw std::invalid_argument(
|
|
"FP8 layerwise all_rank_host_ptrs must be [slot][rank][kind] on rank zero and empty on other ranks");
|
|
}
|
|
if (std::any_of(local_host_ptrs.begin(), local_host_ptrs.end(), [](auto ptr) { return ptr == 0; }) ||
|
|
std::any_of(local_gpu_ptrs.begin(), local_gpu_ptrs.end(), [](auto ptr) { return ptr == 0; }) ||
|
|
std::any_of(expert_nbytes.begin(), expert_nbytes.end(), [](auto size) { return size == 0; }) ||
|
|
(rank == 0 &&
|
|
std::any_of(all_rank_host_ptrs.begin(), all_rank_host_ptrs.end(), [](auto ptr) { return ptr == 0; }))) {
|
|
throw std::invalid_argument("FP8 layerwise buffer pointers and expert byte sizes must be non-zero");
|
|
}
|
|
|
|
std::copy(local_host_ptrs.begin(), local_host_ptrs.end(), local_host_ptrs_.begin());
|
|
std::copy(local_gpu_ptrs.begin(), local_gpu_ptrs.end(), local_gpu_ptrs_.begin());
|
|
std::copy(expert_nbytes.begin(), expert_nbytes.end(), expert_nbytes_.begin());
|
|
all_rank_host_ptrs_ = all_rank_host_ptrs;
|
|
for (int slot = 0; slot < kFP8LayerwiseHostSlots; ++slot) {
|
|
last_ready_sequence_[slot] = load_acquire(control_->ready[slot].sequence);
|
|
}
|
|
|
|
#if KT_FP8_LAYERWISE_HAS_CUDA
|
|
worker_ = std::thread(&Impl::consumer_main, this);
|
|
std::unique_lock lock(init_mutex_);
|
|
init_cv_.wait(lock, [this] { return init_complete_; });
|
|
if (!init_error_.empty()) {
|
|
lock.unlock();
|
|
stopping_.store(true, std::memory_order_release);
|
|
if (worker_.joinable()) worker_.join();
|
|
throw std::runtime_error(init_error_);
|
|
}
|
|
#else
|
|
throw std::runtime_error("FP8 layerwise transport requires a CUDA-enabled kt-kernel build");
|
|
#endif
|
|
}
|
|
|
|
~Impl() { close(); }
|
|
|
|
void join(std::uint64_t epoch, std::int64_t layer_id, int expert_count) {
|
|
ensure_open();
|
|
throw_if_poisoned(control_);
|
|
validate_call(epoch, layer_id, expert_count);
|
|
|
|
if (rank_ == 0) {
|
|
const std::uint64_t previous_epoch = load_acquire(control_->active_epoch);
|
|
if (epoch <= previous_epoch) {
|
|
poison_and_throw(kErrorProtocol, "rank zero attempted to reuse or decrease a layerwise epoch");
|
|
}
|
|
if (previous_epoch != 0) {
|
|
wait_until(
|
|
"all ranks returning from the previous layer",
|
|
[&] {
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
if (load_acquire(control_->wait_returned_epoch[rank]) < previous_epoch) return false;
|
|
}
|
|
return true;
|
|
},
|
|
[&] {
|
|
std::uint64_t progress = 0;
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
progress = fold_progress(progress, load_relaxed(control_->wait_returned_epoch[rank]));
|
|
}
|
|
return progress;
|
|
});
|
|
}
|
|
|
|
store_relaxed(control_->layer_id, signed_to_bits(layer_id));
|
|
store_relaxed(control_->expert_count, static_cast<std::uint64_t>(expert_count));
|
|
store_relaxed(control_->layer_start_ns, now_ns());
|
|
store_relaxed(control_->writer_ns, 0);
|
|
store_relaxed(control_->slot_wait_ns, 0);
|
|
store_relaxed(control_->producer_total_ns, 0);
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
store_relaxed(control_->h2d_ns[rank], 0);
|
|
store_relaxed(control_->bytes[rank], 0);
|
|
}
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control_->active_epoch, epoch);
|
|
} else {
|
|
wait_until(
|
|
"rank zero publishing the requested epoch",
|
|
[&] {
|
|
const auto active = load_acquire(control_->active_epoch);
|
|
if (active > epoch) poison_and_throw(kErrorProtocol, "rank joined an epoch older than the active epoch");
|
|
return active == epoch;
|
|
},
|
|
[&] { return load_relaxed(control_->active_epoch); });
|
|
if (bits_to_signed(load_relaxed(control_->layer_id)) != layer_id ||
|
|
load_relaxed(control_->expert_count) != static_cast<std::uint64_t>(expert_count)) {
|
|
poison_and_throw(kErrorProtocol, "ranks supplied inconsistent layer_id or expert_count");
|
|
}
|
|
}
|
|
|
|
store_release(control_->join_epoch[rank_], epoch);
|
|
}
|
|
|
|
FP8LayerwiseStats wait(std::uint64_t epoch) {
|
|
ensure_open();
|
|
throw_if_poisoned(control_);
|
|
if (epoch == 0 || load_acquire(control_->join_epoch[rank_]) != epoch) {
|
|
poison_and_throw(kErrorProtocol, "wait called before this rank joined the epoch");
|
|
}
|
|
|
|
wait_until(
|
|
"local H2D consumer and producer completion",
|
|
[&] {
|
|
return load_acquire(control_->consumer_done_epoch[rank_]) >= epoch &&
|
|
load_acquire(control_->producer_done_epoch) >= epoch;
|
|
},
|
|
[&] {
|
|
return fold_progress(load_relaxed(control_->consumer_done_epoch[rank_]),
|
|
load_relaxed(control_->producer_done_epoch));
|
|
});
|
|
|
|
FP8LayerwiseStats stats;
|
|
stats.epoch = epoch;
|
|
stats.layer_id = bits_to_signed(load_relaxed(control_->layer_id));
|
|
stats.expert_count = static_cast<int>(load_relaxed(control_->expert_count));
|
|
stats.rank = rank_;
|
|
stats.writer_ms = ns_to_ms(load_relaxed(control_->writer_ns));
|
|
stats.slot_wait_ms = ns_to_ms(load_relaxed(control_->slot_wait_ns));
|
|
stats.h2d_ms = ns_to_ms(load_relaxed(control_->h2d_ns[rank_]));
|
|
stats.total_ms = ns_to_ms(load_relaxed(control_->producer_total_ns));
|
|
stats.bytes = load_relaxed(control_->bytes[rank_]);
|
|
stats.poisoned = load_acquire(control_->poison_state) != 0;
|
|
stats.error_code = static_cast<int>(load_relaxed(control_->error_code));
|
|
stats.error_rank = static_cast<int>(bits_to_signed(load_relaxed(control_->error_rank)));
|
|
stats.error_message = control_->error_message;
|
|
store_release(control_->wait_returned_epoch[rank_], epoch);
|
|
return stats;
|
|
}
|
|
|
|
void run_producer(std::uint64_t epoch, std::int64_t layer_id, int expert_count, const Writer& writer) {
|
|
ensure_open();
|
|
throw_if_poisoned(control_);
|
|
if (rank_ != 0) throw std::invalid_argument("run_layerwise_fp8_batch is rank-zero only");
|
|
if (!writer) throw std::invalid_argument("FP8 layerwise writer callback is empty");
|
|
validate_call(epoch, layer_id, expert_count);
|
|
if (load_acquire(control_->active_epoch) != epoch || bits_to_signed(load_relaxed(control_->layer_id)) != layer_id ||
|
|
load_relaxed(control_->expert_count) != static_cast<std::uint64_t>(expert_count) ||
|
|
load_acquire(control_->join_epoch[0]) != epoch) {
|
|
poison_and_throw(kErrorProtocol, "producer arguments do not match the joined layer");
|
|
}
|
|
|
|
const auto producer_begin = Clock::now();
|
|
try {
|
|
wait_until(
|
|
"all ranks joining the layer",
|
|
[&] {
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
if (load_acquire(control_->join_epoch[rank]) != epoch) return false;
|
|
}
|
|
return true;
|
|
},
|
|
[&] {
|
|
std::uint64_t progress = 0;
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
progress = fold_progress(progress, load_relaxed(control_->join_epoch[rank]));
|
|
}
|
|
return progress;
|
|
});
|
|
|
|
std::array<std::uint64_t, kFP8LayerwiseHostSlots> published_sequence{};
|
|
std::array<std::vector<std::uintptr_t>, kFP8LayerwiseBufferKinds> ptrs;
|
|
for (auto& values : ptrs) values.resize(tp_size_);
|
|
std::uint64_t writer_ns = 0;
|
|
std::uint64_t slot_wait_ns = 0;
|
|
for (int expert = 0; expert < expert_count; ++expert) {
|
|
const int slot = expert % kFP8LayerwiseHostSlots;
|
|
if (published_sequence[slot] != 0) {
|
|
const auto wait_begin = Clock::now();
|
|
wait_for_slot_ack(slot, published_sequence[slot]);
|
|
slot_wait_ns += static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - wait_begin).count());
|
|
}
|
|
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
for (int kind = 0; kind < kFP8LayerwiseBufferKinds; ++kind) {
|
|
ptrs[kind][rank] =
|
|
all_rank_host_ptrs_[((slot * tp_size_ + rank) * kFP8LayerwiseBufferKinds) + kind];
|
|
}
|
|
}
|
|
|
|
const auto writer_begin = Clock::now();
|
|
writer(expert, ptrs[0], ptrs[1], ptrs[2], ptrs[3]);
|
|
writer_ns += static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - writer_begin).count());
|
|
|
|
const std::uint64_t sequence = fetch_add_relaxed(control_->next_sequence, 1) + 1;
|
|
store_relaxed(control_->ready[slot].epoch, epoch);
|
|
store_relaxed(control_->ready[slot].expert_id, static_cast<std::uint64_t>(expert));
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control_->ready[slot].sequence, sequence);
|
|
published_sequence[slot] = sequence;
|
|
}
|
|
|
|
for (int slot = 0; slot < kFP8LayerwiseHostSlots; ++slot) {
|
|
if (published_sequence[slot] == 0) continue;
|
|
const auto wait_begin = Clock::now();
|
|
wait_for_slot_ack(slot, published_sequence[slot]);
|
|
slot_wait_ns += static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - wait_begin).count());
|
|
}
|
|
|
|
const auto total_ns = static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - producer_begin).count());
|
|
store_relaxed(control_->writer_ns, writer_ns);
|
|
store_relaxed(control_->slot_wait_ns, slot_wait_ns);
|
|
store_relaxed(control_->producer_total_ns, total_ns);
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control_->producer_done_epoch, epoch);
|
|
} catch (const std::exception& error) {
|
|
if (load_acquire(control_->poison_state) == 0) publish_poison(control_, kErrorWriter, rank_, error.what());
|
|
throw;
|
|
} catch (...) {
|
|
publish_poison(control_, kErrorWriter, rank_, "unknown exception in FP8 layerwise writer");
|
|
throw;
|
|
}
|
|
}
|
|
|
|
void close() noexcept {
|
|
bool expected = false;
|
|
if (!closed_.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) return;
|
|
const auto active = load_acquire(control_->active_epoch);
|
|
if (active != 0 && load_acquire(control_->consumer_done_epoch[rank_]) < active &&
|
|
load_acquire(control_->poison_state) == 0) {
|
|
publish_poison(control_, kErrorClosedWhileActive, rank_, "transport closed while a layer was active");
|
|
}
|
|
stopping_.store(true, std::memory_order_release);
|
|
if (worker_.joinable()) worker_.join();
|
|
}
|
|
|
|
int rank() const noexcept { return rank_; }
|
|
int tp_size() const noexcept { return tp_size_; }
|
|
int num_experts() const noexcept { return num_experts_; }
|
|
bool closed() const noexcept { return closed_.load(std::memory_order_acquire); }
|
|
|
|
private:
|
|
void validate_call(std::uint64_t epoch, std::int64_t layer_id, int expert_count) const {
|
|
if (epoch == 0) throw std::invalid_argument("FP8 layerwise epoch zero is reserved");
|
|
if (layer_id < 0) throw std::invalid_argument("FP8 layerwise layer_id must be non-negative");
|
|
if (expert_count <= 0 || expert_count > num_experts_) {
|
|
throw std::invalid_argument("FP8 layerwise expert_count is outside the configured expert range");
|
|
}
|
|
}
|
|
|
|
void ensure_open() const {
|
|
if (closed_.load(std::memory_order_acquire)) throw std::runtime_error("FP8 layerwise transport is closed");
|
|
}
|
|
|
|
[[noreturn]] void poison_and_throw(int code, const std::string& message) {
|
|
publish_poison(control_, code, rank_, message);
|
|
throw std::runtime_error(message);
|
|
}
|
|
|
|
template <typename Predicate, typename Progress>
|
|
void wait_until(const char* description, Predicate&& predicate, Progress&& progress) {
|
|
auto last_progress_time = Clock::now();
|
|
std::uint64_t last_progress = progress();
|
|
std::uint32_t spins = 0;
|
|
while (!predicate()) {
|
|
throw_if_poisoned(control_);
|
|
if (stopping_.load(std::memory_order_acquire)) throw std::runtime_error("FP8 layerwise transport is stopping");
|
|
const auto current_progress = progress();
|
|
if (current_progress != last_progress) {
|
|
last_progress = current_progress;
|
|
last_progress_time = Clock::now();
|
|
} else if (Clock::now() - last_progress_time >= timeout_) {
|
|
std::ostringstream message;
|
|
message << "timed out after " << timeout_.count() << " ms without progress while waiting for " << description;
|
|
poison_and_throw(kErrorTimeout, message.str());
|
|
}
|
|
if (++spins < 2048) {
|
|
std::this_thread::yield();
|
|
} else {
|
|
spins = 0;
|
|
std::this_thread::sleep_for(std::chrono::microseconds(50));
|
|
}
|
|
}
|
|
throw_if_poisoned(control_);
|
|
}
|
|
|
|
void wait_for_slot_ack(int slot, std::uint64_t sequence) {
|
|
wait_until(
|
|
"all ranks acknowledging a reused host slot",
|
|
[&] {
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
if (load_acquire(control_->ack_sequence[rank][slot]) < sequence) return false;
|
|
}
|
|
return true;
|
|
},
|
|
[&] {
|
|
std::uint64_t progress = 0;
|
|
for (int rank = 0; rank < tp_size_; ++rank) {
|
|
progress = fold_progress(progress, load_relaxed(control_->ack_sequence[rank][slot]));
|
|
}
|
|
return progress;
|
|
});
|
|
}
|
|
|
|
#if KT_FP8_LAYERWISE_HAS_CUDA
|
|
static void check_cuda(cudaError_t status, const char* operation) {
|
|
if (status == cudaSuccess) return;
|
|
std::ostringstream message;
|
|
message << operation << " failed: " << cudaGetErrorString(status);
|
|
throw std::runtime_error(message.str());
|
|
}
|
|
|
|
void initialize_cuda_worker() {
|
|
check_cuda(cudaSetDevice(cuda_device_), "cudaSetDevice");
|
|
check_cuda(cudaStreamCreateWithFlags(©_stream_, cudaStreamNonBlocking), "cudaStreamCreateWithFlags");
|
|
for (int slot = 0; slot < kFP8LayerwiseHostSlots; ++slot) {
|
|
check_cuda(cudaEventCreateWithFlags(&slot_events_[slot], cudaEventDisableTiming), "cudaEventCreateWithFlags");
|
|
}
|
|
}
|
|
|
|
void destroy_cuda_worker() noexcept {
|
|
for (auto& event : slot_events_) {
|
|
if (event != nullptr) {
|
|
cudaEventDestroy(event);
|
|
event = nullptr;
|
|
}
|
|
}
|
|
if (copy_stream_ != nullptr) {
|
|
cudaStreamDestroy(copy_stream_);
|
|
copy_stream_ = nullptr;
|
|
}
|
|
}
|
|
|
|
void copy_expert(int expert, int slot) {
|
|
const auto begin = Clock::now();
|
|
for (int kind = 0; kind < kFP8LayerwiseBufferKinds; ++kind) {
|
|
if (expert_nbytes_[kind] > std::numeric_limits<std::uintptr_t>::max() / static_cast<std::size_t>(num_experts_)) {
|
|
throw std::overflow_error("FP8 layerwise GPU expert offset overflows uintptr_t");
|
|
}
|
|
const auto destination = local_gpu_ptrs_[kind] + static_cast<std::uintptr_t>(expert) * expert_nbytes_[kind];
|
|
const auto source = local_host_ptrs_[slot * kFP8LayerwiseBufferKinds + kind];
|
|
check_cuda(cudaMemcpyAsync(reinterpret_cast<void*>(destination), reinterpret_cast<const void*>(source),
|
|
expert_nbytes_[kind], cudaMemcpyHostToDevice, copy_stream_),
|
|
"cudaMemcpyAsync(H2D)");
|
|
}
|
|
check_cuda(cudaEventRecord(slot_events_[slot], copy_stream_), "cudaEventRecord");
|
|
wait_until(
|
|
"the local H2D completion event",
|
|
[&] {
|
|
const auto status = cudaEventQuery(slot_events_[slot]);
|
|
if (status == cudaSuccess) return true;
|
|
if (status == cudaErrorNotReady) return false;
|
|
check_cuda(status, "cudaEventQuery");
|
|
return false;
|
|
},
|
|
[] { return std::uint64_t{0}; });
|
|
local_h2d_ns_ += static_cast<std::uint64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - begin).count());
|
|
for (auto size : expert_nbytes_) local_bytes_ += size;
|
|
}
|
|
#endif
|
|
|
|
void consumer_main() noexcept {
|
|
#if KT_FP8_LAYERWISE_HAS_CUDA
|
|
try {
|
|
initialize_cuda_worker();
|
|
{
|
|
std::lock_guard lock(init_mutex_);
|
|
init_complete_ = true;
|
|
}
|
|
init_cv_.notify_one();
|
|
|
|
std::uint64_t completed_epoch = load_acquire(control_->consumer_done_epoch[rank_]);
|
|
while (!stopping_.load(std::memory_order_acquire)) {
|
|
const auto epoch = load_acquire(control_->active_epoch);
|
|
if (epoch == 0 || epoch <= completed_epoch || load_acquire(control_->join_epoch[rank_]) != epoch) {
|
|
throw_if_poisoned(control_);
|
|
std::this_thread::sleep_for(std::chrono::microseconds(50));
|
|
continue;
|
|
}
|
|
|
|
const int expert_count = static_cast<int>(load_relaxed(control_->expert_count));
|
|
if (expert_count <= 0 || expert_count > num_experts_) {
|
|
poison_and_throw(kErrorProtocol, "consumer observed an invalid expert_count");
|
|
}
|
|
local_h2d_ns_ = 0;
|
|
local_bytes_ = 0;
|
|
|
|
for (int expert = 0; expert < expert_count; ++expert) {
|
|
const int slot = expert % kFP8LayerwiseHostSlots;
|
|
std::uint64_t observed_sequence = 0;
|
|
wait_until(
|
|
"the next expert generation",
|
|
[&] {
|
|
const auto active = load_acquire(control_->active_epoch);
|
|
if (active > epoch) poison_and_throw(kErrorProtocol, "active epoch advanced before consumer completed");
|
|
const auto sequence = load_acquire(control_->ready[slot].sequence);
|
|
if (sequence == last_ready_sequence_[slot]) return false;
|
|
if (load_relaxed(control_->ready[slot].epoch) != epoch ||
|
|
load_relaxed(control_->ready[slot].expert_id) != static_cast<std::uint64_t>(expert)) {
|
|
return false;
|
|
}
|
|
observed_sequence = sequence;
|
|
return true;
|
|
},
|
|
[&] { return load_relaxed(control_->ready[slot].sequence); });
|
|
|
|
copy_expert(expert, slot);
|
|
last_ready_sequence_[slot] = observed_sequence;
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control_->ack_sequence[rank_][slot], observed_sequence);
|
|
}
|
|
|
|
store_relaxed(control_->h2d_ns[rank_], local_h2d_ns_);
|
|
store_relaxed(control_->bytes[rank_], local_bytes_);
|
|
__atomic_thread_fence(__ATOMIC_RELEASE);
|
|
store_release(control_->consumer_done_epoch[rank_], epoch);
|
|
completed_epoch = epoch;
|
|
}
|
|
} catch (const std::exception& error) {
|
|
bool initialization_failure = false;
|
|
{
|
|
std::lock_guard lock(init_mutex_);
|
|
if (!init_complete_) {
|
|
initialization_failure = true;
|
|
init_error_ = error.what();
|
|
init_complete_ = true;
|
|
}
|
|
}
|
|
if (initialization_failure) init_cv_.notify_one();
|
|
if (!stopping_.load(std::memory_order_acquire)) {
|
|
publish_poison(control_, initialization_failure ? kErrorCudaInitialization : kErrorCudaCopy, rank_,
|
|
error.what());
|
|
}
|
|
} catch (...) {
|
|
bool initialization_failure = false;
|
|
{
|
|
std::lock_guard lock(init_mutex_);
|
|
if (!init_complete_) {
|
|
initialization_failure = true;
|
|
init_error_ = "unknown CUDA worker initialization error";
|
|
init_complete_ = true;
|
|
}
|
|
}
|
|
if (initialization_failure) init_cv_.notify_one();
|
|
if (!stopping_.load(std::memory_order_acquire)) {
|
|
publish_poison(control_, initialization_failure ? kErrorCudaInitialization : kErrorCudaCopy, rank_,
|
|
initialization_failure ? "unknown CUDA worker initialization error"
|
|
: "unknown exception in FP8 layerwise H2D consumer");
|
|
}
|
|
}
|
|
destroy_cuda_worker();
|
|
#endif
|
|
}
|
|
|
|
FP8LayerwiseControl* control_;
|
|
int rank_;
|
|
int tp_size_;
|
|
int cuda_device_;
|
|
int num_experts_;
|
|
std::chrono::milliseconds timeout_;
|
|
std::array<std::uintptr_t, kFP8LayerwiseHostSlots * kFP8LayerwiseBufferKinds> local_host_ptrs_{};
|
|
std::array<std::uintptr_t, kFP8LayerwiseBufferKinds> local_gpu_ptrs_{};
|
|
std::vector<std::uintptr_t> all_rank_host_ptrs_;
|
|
std::array<std::size_t, kFP8LayerwiseBufferKinds> expert_nbytes_{};
|
|
std::array<std::uint64_t, kFP8LayerwiseHostSlots> last_ready_sequence_{};
|
|
|
|
std::atomic<bool> stopping_{false};
|
|
std::atomic<bool> closed_{false};
|
|
std::thread worker_;
|
|
std::mutex init_mutex_;
|
|
std::condition_variable init_cv_;
|
|
bool init_complete_ = false;
|
|
std::string init_error_;
|
|
std::uint64_t local_h2d_ns_ = 0;
|
|
std::uint64_t local_bytes_ = 0;
|
|
|
|
#if KT_FP8_LAYERWISE_HAS_CUDA
|
|
cudaStream_t copy_stream_ = nullptr;
|
|
std::array<cudaEvent_t, kFP8LayerwiseHostSlots> slot_events_{};
|
|
#endif
|
|
};
|
|
|
|
FP8LayerwiseTransport::FP8LayerwiseTransport(
|
|
std::uintptr_t control_ptr, std::size_t control_size, int rank, int tp_size, int cuda_device,
|
|
const std::vector<std::uintptr_t>& local_host_ptrs, const std::vector<std::uintptr_t>& local_gpu_ptrs,
|
|
const std::vector<std::uintptr_t>& all_rank_host_ptrs, const std::vector<std::size_t>& expert_nbytes,
|
|
int num_experts, std::uint64_t timeout_ms)
|
|
: impl_(std::make_unique<Impl>(control_ptr, control_size, rank, tp_size, cuda_device, local_host_ptrs,
|
|
local_gpu_ptrs, all_rank_host_ptrs, expert_nbytes, num_experts, timeout_ms)) {}
|
|
|
|
FP8LayerwiseTransport::~FP8LayerwiseTransport() = default;
|
|
|
|
void FP8LayerwiseTransport::join(std::uint64_t epoch, std::int64_t layer_id, int expert_count) {
|
|
impl_->join(epoch, layer_id, expert_count);
|
|
}
|
|
|
|
FP8LayerwiseStats FP8LayerwiseTransport::wait(std::uint64_t epoch) { return impl_->wait(epoch); }
|
|
|
|
void FP8LayerwiseTransport::close() { impl_->close(); }
|
|
|
|
void FP8LayerwiseTransport::run_producer(std::uint64_t epoch, std::int64_t layer_id, int expert_count,
|
|
const Writer& writer) {
|
|
impl_->run_producer(epoch, layer_id, expert_count, writer);
|
|
}
|
|
|
|
int FP8LayerwiseTransport::rank() const noexcept { return impl_->rank(); }
|
|
|
|
int FP8LayerwiseTransport::tp_size() const noexcept { return impl_->tp_size(); }
|
|
|
|
int FP8LayerwiseTransport::num_experts() const noexcept { return impl_->num_experts(); }
|
|
|
|
bool FP8LayerwiseTransport::closed() const noexcept { return impl_->closed(); }
|
|
|
|
} // namespace kt::layerwise
|