Skip to content

route-resident: support 256-expert top-8 models (Qwen3.6-35B-A3B) #2321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
atomicmilkshake wants to merge 1 commit into from
Closed

route-resident: support 256-expert top-8 models (Qwen3.6-35B-A3B) #2321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions include/strata/kernels/native_router.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,16 @@ void native_router_top10_multi(const float* logits, int32_t* ids, float* weights
/// 256 / 8 (Qwen3.6-35B-A3B). logits [n_tok, n_expert], ids / weights [n_tok, k]. CUDA / HIP.
bool native_router_topk_supported(int n_expert, int k);
void native_router_topk(const float* logits, int32_t* ids, float* weights, int n_expert, int k, int n_tok, void* stream);
/// STRATA_ROUTE_RESIDENT: swap non-resident tail-rank experts for resident ones within `margin` logits (see the .cu).
/// STRATA_ROUTE_RESIDENT (opt-in, changes output): swap non-resident tail-rank experts for resident ones within
/// `margin` logits. Supports only (n_expert,k) = (512,10) or (256,8). Reads logits [n_tok,n_expert] and residency
/// [n_expert] (negative = non-resident); updates IDs/weights [n_tok,k]. IDs must be distinct and in [0,n_expert).
/// Ranks are zero-based and require 0 <= rank_lo <= rank_hi < k; margin must be finite.
/// No swap leaves IDs/weights bit-identical. Requires aligned, disjoint buffers and a nonnull ordered stream.
/// `stats` is nullptr or 4 device uint64 counters: tail entries seen, swaps, non-resident before, non-resident after.
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok, float margin,
int rank_lo, int rank_hi, unsigned long long* stats, void* stream);
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok,
int n_expert, int k, float margin, int rank_lo, int rank_hi,
unsigned long long* stats, void* stream);
/// Backward-compatible overload for 512/10 models.
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok,
float margin, int rank_lo, int rank_hi, unsigned long long* stats, void* stream);
}
23 changes: 18 additions & 5 deletions src/core/verify.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,12 +88,14 @@ using Clock = std::chrono::steady_clock;
double ms_since(Clock::time_point t) { return std::chrono::duration<double, std::milli>(Clock::now() - t).count(); }
const bool g_dbg = std::getenv("STRATA_VERIFY_DEBUG") != nullptr;
inline bool g_lfuse() { static const bool on = [] { const char* v = std::getenv("STRATA_LFUSE"); return v != nullptr && v[0] == '1'; }(); return on; }
// STRATA_ROUTE_RESIDENT=<margin logits> (+ STRATA_ROUTE_RESIDENT_RANKS=lo-hi, default 6-9): residency-biased routing.
// STRATA_ROUTE_RESIDENT=<margin logits> (+ STRATA_ROUTE_RESIDENT_RANKS=lo-hi): residency-biased routing.
// Zero-based tail ranks default to 6-9 for 512/10, 5-7 for 256/8. Off by default; swaps change the output.
struct RouteResidentCfg;
RouteResidentCfg& route_resident_cfg();
struct RouteResidentCfg {
float margin = 0.0f;
int lo = 6, hi = 9;
bool ranks_override = false;
// #1578: one counter block per GPU. A layer split inits a verifier per stage; a block allocated on whichever device
// was current at the first init is unreachable from the other stage's kernels without P2P (illegal memory access).
// The Verifier asks for its device's block in init (outside graph capture) and keeps the pointer: a launch never
Expand Down Expand Up @@ -144,7 +146,10 @@ RouteResidentCfg& route_resident_cfg() {
static RouteResidentCfg c;
static const bool parsed = [] {
if (const char* v = std::getenv("STRATA_ROUTE_RESIDENT")) c.margin = (float) std::atof(v);
if (const char* v = std::getenv("STRATA_ROUTE_RESIDENT_RANKS")) std::sscanf(v, "%d-%d", &c.lo, &c.hi);
if (const char* v = std::getenv("STRATA_ROUTE_RESIDENT_RANKS")) {
c.ranks_override = true;
if (std::sscanf(v, "%d-%d", &c.lo, &c.hi) != 2) c.lo = c.hi = -1; // rejected before launching
}
return true;
}();
(void) parsed;
Expand Down Expand Up @@ -563,8 +568,14 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
if (slot.load() == this || slot.compare_exchange_strong(none, this)) break;
}
release_gpu_fn().store(&release_live_verifiers);
RouteResidentCfg& rr = route_resident_cfg();
if (!rr.ranks_override) {
const bool rr256 = g.n_expert == 256 && ss.k == 8;
rr.lo = rr256 ? 5 : 6;
rr.hi = rr256 ? 7 : 9;
}
cudaGetDevice(&device_); // a layer split's stage on another GPU: its streams, graphs and buffers live there
rr_stats_ = route_resident_cfg().stats_on(device_); // STRATA_ROUTE_RESIDENT: this GPU's counters, outside graph capture
rr_stats_ = rr.stats_on(device_); // STRATA_ROUTE_RESIDENT: this GPU's counters, outside graph capture
strata::kernels::fused_gr_check(); // once per card: which bitwise-equal hyper-connection read runs there
wt_ = &wt;
g_ = &g;
Expand Down Expand Up @@ -1501,11 +1512,13 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
static const float route_prior = [] { const char* v = std::getenv("STRATA_ROUTE_PRIOR"); return v ? (float) std::atof(v) : 0.0f; }();
if (route_prior > 0.0f && NE == 512 && K == 10 && hits_.d_res != nullptr)
route_prior_top10(logits_ + tb * NE, hits_.d_res + l * g.n_expert, route_prior, ids_ + tb * K, w_ + tb * K, n, cs);
if (route_resident_cfg().margin > 0.0f && NE == 512 && K == 10 && hits_.d_res != nullptr) {
if (route_resident_cfg().margin > 0.0f && ((NE == 512 && K == 10) || (NE == 256 && K == 8)) &&
hits_.d_res != nullptr) {
// STRATA_ROUTE_RESIDENT: after the router, before the plan/doorbell read ids_/w_ (opt-in, changes the output)
try {
native_route_resident(logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, hits_.d_res + l * g.n_expert, n,
route_resident_cfg().margin, route_resident_cfg().lo, route_resident_cfg().hi,
(int) NE, (int) K, route_resident_cfg().margin,
route_resident_cfg().lo, route_resident_cfg().hi,
rr_stats_ != nullptr ? rr_stats_ + (n <= 8 ? 4 : 0) : nullptr, cs);
} catch (const std::exception& e) { err = "verify route-resident: " + std::string(e.what()); return false; }
}
Expand Down
112 changes: 73 additions & 39 deletions src/kernels/cuda/native_router.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,41 +250,44 @@ void native_router_topk(const float* logits, int32_t* ids, float* weights, int n
if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
}

// STRATA_ROUTE_RESIDENT (EXPERIMENTAL, opt-in, changes the output): residency-biased routing for 512 x 10.
// After the top-10: for the ranks lo..hi a non-resident expert is replaced by the best resident expert not already
// STRATA_ROUTE_RESIDENT (EXPERIMENTAL, opt-in, changes the output): residency-biased routing.
// After the top-k: for the ranks lo..hi a non-resident expert is replaced by the best resident expert not already
// picked when its logit is within `margin` of the replaced one; the weights are then the softmax of the selected
// logits (the router's renormalisation). Tokens without a swap keep the router's exact bits.
// stats (device, 4 x uint64): [0] tail entries seen, [1] swaps, [2] non-resident entries before, [3] after.
//
// One warp per token. This used to be one THREAD per token scanning all 512 experts (4 ranks x 512 x 12 scalar
// ops on up to 6 threads): 45 us per layer, 2.2 ms of a 19.4 ms window, which is exactly what the option saved in
// PCIe experts (~2.4 ms) - the swap cost as much as the misses it removed. The ranks are still visited in order
// (a swap changes the set already picked), but the scan is coalesced (lane j reads f = j, j + 32, ...) and reduced
// within the warp. The bits are unchanged: a strictly greater logit wins, the smallest expert index breaks ties,
// and the softmax is the same expression.
__global__ void route_resident_k(const float* __restrict__ logits, int32_t* __restrict__ ids, float* __restrict__ weights,
const int32_t* __restrict__ res, int n_tok, float margin, int lo, int hi,
unsigned long long* __restrict__ stats) {
const int t = blockIdx.x;
// One warp per token (PR #1737, 95dd807). The serial scan cost as much as the PCIe misses it removed. Ranks are
// still visited in order, because a swap changes the set already picked. Lane j reads experts j, j + 32, ... and
// the warp reduces with __shfl_down_sync: a strictly greater logit wins, and the smallest expert index breaks a
// tie. Both instantiated widths (512 and 256) are multiples of the warp, so every lane takes the same number of
// steps and the shuffle stays convergent. The 512/10 entry point keeps its symbol; the 256/8 instance is the
// other template. Lane 0 writes ids, weights and counters. The qwen4exp top-10 router above is not this kernel.
template <int NE, int K>
__device__ __forceinline__ void route_resident_impl(const float* __restrict__ logits, int32_t* __restrict__ ids,
float* __restrict__ weights, const int32_t* __restrict__ res,
int n_tok, float margin, int lo, int hi,
unsigned long long* __restrict__ stats) {
static_assert(NE % 32 == 0, "route-resident warp scan requires NE to be a multiple of the warp");
const int t = (int) blockIdx.x;
if (t >= n_tok) return;
const int lane = (int) threadIdx.x;
const float* l = logits + (size_t) t * 512;
int32_t* id = ids + (size_t) t * 10;
int32_t my[10];
for (int r = 0; r < 10; ++r) my[r] = id[r];
const float* l = logits + (size_t) t * NE;
int32_t* id = ids + (size_t) t * K;
int32_t my[K];
for (int r = 0; r < K; ++r) my[r] = id[r];
int swaps = 0, tail = 0, before = 0;
if (lane == 0)
for (int r = 0; r < 10; ++r) before += res[my[r]] < 0;
for (int r = lo; r <= hi && r < 10; ++r) {
for (int r = 0; r < K; ++r) before += res[my[r]] < 0;
for (int r = lo; r <= hi && r < K; ++r) {
const int e = my[r];
if (res[e] >= 0) continue; // already resident: uniform across the warp
if (lane == 0) ++tail;
float bv = -INFINITY;
int bi = -1;
for (int f = lane; f < 512; f += 32) {
for (int f = lane; f < NE; f += 32) {
if (res[f] < 0) continue;
bool used = false;
for (int q = 0; q < 10; ++q) used |= my[q] == f;
for (int q = 0; q < K; ++q) used |= my[q] == f;
if (used) continue;
if (l[f] > bv) { bv = l[f]; bi = f; }
}
Expand All @@ -295,35 +298,66 @@ __global__ void route_resident_k(const float* __restrict__ logits, int32_t* __re
if (ov > bv || (ov == bv && oi >= 0 && (bi < 0 || oi < bi))) { bv = ov; bi = oi; }
}
if (lane == 0 && bi >= 0 && l[e] - bv <= margin) { my[r] = bi; ++swaps; }
my[r] = __shfl_sync(0xffffffffu, my[r], 0); // all lanes keep the same picked set
}
int after = 0;
if (lane == 0)
for (int r = 0; r < 10; ++r) after += res[my[r]] < 0;
const int sw = __shfl_sync(0xffffffffu, swaps, 0);
if (sw) {
float m = -INFINITY;
for (int r = 0; r < 10; ++r) m = fmaxf(m, l[my[r]]);
float ex[10], sum = 0.0f;
for (int r = 0; r < 10; ++r) { ex[r] = expf(l[my[r]] - m); sum += ex[r]; }
for (int r = 0; r < 10; ++r) weights[(size_t) t * 10 + r] = ex[r] / sum;
my[r] = __shfl_sync(0xffffffffu, my[r], 0); // every lane keeps the picked set for the next rank
}
if (lane == 0) {
for (int r = 0; r < 10; ++r) id[r] = my[r];
int after = 0;
for (int r = 0; r < K; ++r) after += res[my[r]] < 0;
if (swaps) {
float m = -INFINITY;
for (int r = 0; r < K; ++r) m = fmaxf(m, l[my[r]]);
float ex[K], sum = 0.0f;
for (int r = 0; r < K; ++r) { ex[r] = expf(l[my[r]] - m); sum += ex[r]; }
for (int r = 0; r < K; ++r) weights[(size_t) t * K + r] = ex[r] / sum;
}
for (int r = 0; r < K; ++r) id[r] = my[r];
if (stats) {
atomicAdd(stats + 0, (unsigned long long) tail);
atomicAdd(stats + 1, (unsigned long long) sw);
atomicAdd(stats + 1, (unsigned long long) swaps);
atomicAdd(stats + 2, (unsigned long long) before);
atomicAdd(stats + 3, (unsigned long long) after);
}
}
}
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok, float margin,
int rank_lo, int rank_hi, unsigned long long* stats, void* stream) {
if (!stream || n_tok < 1 || !res_layer) throw std::invalid_argument("native_route_resident: bad arguments");
route_resident_k<<<(unsigned) n_tok, 32, 0, static_cast<cudaStream_t>(stream)>>>(
logits, ids, weights, res_layer, n_tok, margin, rank_lo, rank_hi, stats);
// 512/10 keeps the original entry-point symbol. 256/8 is the other instantiation of the same warp reduction.
__global__ void route_resident_k(const float* __restrict__ logits, int32_t* __restrict__ ids, float* __restrict__ weights,
const int32_t* __restrict__ res, int n_tok, float margin, int lo, int hi,
unsigned long long* __restrict__ stats) {
route_resident_impl<512, 10>(logits, ids, weights, res, n_tok, margin, lo, hi, stats);
}
template <int NE, int K>
__global__ void route_resident_geometry_k(const float* __restrict__ logits, int32_t* __restrict__ ids,
float* __restrict__ weights, const int32_t* __restrict__ res,
int n_tok, float margin, int lo, int hi, unsigned long long* __restrict__ stats) {
route_resident_impl<NE, K>(logits, ids, weights, res, n_tok, margin, lo, hi, stats);
}
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok,
int n_expert, int k, float margin, int rank_lo, int rank_hi,
unsigned long long* stats, void* stream) {
if (!stream || n_tok < 1 || !((n_expert == 512 && k == 10) || (n_expert == 256 && k == 8)) ||
!(margin >= -FLT_MAX && margin <= FLT_MAX) || rank_lo < 0 || rank_hi < rank_lo || rank_hi >= k)
throw std::invalid_argument("native_route_resident: requires 512/10 or 256/8, finite margin and ranks in [0,k)");
const size_t logit_bytes = (size_t) n_tok * n_expert * sizeof(float);
const size_t out_bytes = (size_t) n_tok * k * sizeof(int32_t);
const size_t res_bytes = (size_t) n_expert * sizeof(int32_t);
if (!valid(logits, logit_bytes) || !valid(ids, out_bytes) || !valid(weights, out_bytes) || !valid(res_layer, res_bytes) ||
overlap(logits, logit_bytes, ids, out_bytes) || overlap(logits, logit_bytes, weights, out_bytes) ||
overlap(ids, out_bytes, weights, out_bytes) || overlap(res_layer, res_bytes, ids, out_bytes) ||
overlap(res_layer, res_bytes, weights, out_bytes))
throw std::invalid_argument("native_route_resident: requires aligned spans and disjoint outputs");
// One warp per token. blockIdx.x is the token; all 32 lanes participate in the reduction.
const unsigned blocks = (unsigned) n_tok;
if (n_expert == 512)
route_resident_k<<<blocks, 32, 0, static_cast<cudaStream_t>(stream)>>>(
logits, ids, weights, res_layer, n_tok, margin, rank_lo, rank_hi, stats);
else
route_resident_geometry_k<256, 8><<<blocks, 32, 0, static_cast<cudaStream_t>(stream)>>>(
logits, ids, weights, res_layer, n_tok, margin, rank_lo, rank_hi, stats);
const auto error = cudaGetLastError();
if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
}
void native_route_resident(const float* logits, int32_t* ids, float* weights, const int32_t* res_layer, int n_tok, float margin,
int rank_lo, int rank_hi, unsigned long long* stats, void* stream) {
native_route_resident(logits, ids, weights, res_layer, n_tok, 512, 10, margin, rank_lo, rank_hi, stats, stream);
}
}
Loading