Skip to content

perf(moe): precompute Q4_K scales for Muse concurrent decode #1056

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
shinyuser2345-svg wants to merge 1 commit into from
Closed

perf(moe): precompute Q4_K scales for Muse concurrent decode #1056

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
52 changes: 23 additions & 29 deletions kernels/csrc/cuda/moe/expert_ffn_q4k.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1980,11 +1980,10 @@ __device__ __forceinline__ void si_mma_q4k_scales(const unsigned char* sc12, int
// things here were sized for the 32-row ceiling on EVERY launch, and a packed step narrower than
// 32 rows paid both:
//
// * the shared A tile -- As + Ad + Asum is 10 KB of the CTA's 18.75 KB at MM=32, and at 100 KB
// of shared per SM that 18.75 KB is what holds the kernel to five CTAs per SM. Its own
// __launch_bounds__ asks for eight, and the register allocator already honours that; only
// shared memory was standing in the way. MM=16 is 13.75 KB (seven CTAs), MM=8 is 11.25 KB
// (eight).
// * the shared A tile -- As + Ad + Asum is 10 KB at MM=32. Including the precomputed, padded
// B scale pairs below, total shared storage is 20.25 KB at MM=32, 15.25 KB at MM=16 and
// 12.75 KB at MM=8. Under a 100 KB shared-memory budget these permit four, six and seven
// CTAs respectively (before other resource limits), instead of sizing every width for 32.
// * the M-tile loop, which issued BOTH m16n8k32 tiles unconditionally and then discarded the
// second through `lm < M`. At sixteen rows or fewer the second tile is pure waste: half the
// mma.sync and half the ldmatrix on As.
Expand Down Expand Up @@ -2015,8 +2014,11 @@ void down_q4k_mma_rows_kernel(const unsigned char* __restrict__ down_q,

__shared__ signed char As[MM][256];
__shared__ signed char Bs[SI_MMA_BN][256];
__shared__ unsigned char Ssc[SI_MMA_BN][8], Smn[SI_MMA_BN][8];
__shared__ float2 Wdm[SI_MMA_BN];
// Each loading lane expands one scale group, instead of c==0 serially unpacking all eight.
// A half times a six-bit integer fits exactly in float (at most 17 significant bits), so
// storing these products adds no rounding to the existing fold. The consumer reads columns
// spaced two rows apart; a nine-pair stride avoids their four-way shared-bank aliasing.
__shared__ float2 Bscale[SI_MMA_BN][9];
__shared__ float Ad[MM][8], Asum[MM][8];

constexpr int NT = (MM + 15) / 16; // 16-row mma tiles this width actually needs
Expand Down Expand Up @@ -2058,16 +2060,11 @@ void down_q4k_mma_rows_kernel(const unsigned char* __restrict__ down_q,
*reinterpret_cast<const uint4*>(lo);
*reinterpret_cast<uint4*>(&Bs[r][si_mma_swz(kb + 32, r)]) =
*reinterpret_cast<const uint4*>(hi);
if (c == 0) {
Wdm[r] = __half22float2(b->dm);
#pragma unroll
for (int jj = 0; jj < 4; jj++) {
unsigned char x0, x1, y0, y1;
si_mma_q4k_scales(b->scales, jj, x0, x1, y0, y1);
Ssc[r][2 * jj] = x0; Ssc[r][2 * jj + 1] = x1;
Smn[r][2 * jj] = y0; Smn[r][2 * jj + 1] = y1;
}
}
unsigned char x0, x1, y0, y1;
si_mma_q4k_scales(b->scales, c >> 1, x0, x1, y0, y1);
const float2 dm = __half22float2(b->dm);
Bscale[r][c] = make_float2(dm.x * (float)((c & 1) ? x1 : x0),
dm.y * (float)((c & 1) ? y1 : y0));
}
} else {
for (int u = tid; u < SI_MMA_BN * 16; u += SI_MMA_NW * 32) {
Expand All @@ -2089,15 +2086,12 @@ void down_q4k_mma_rows_kernel(const unsigned char* __restrict__ down_q,
}
const int kb = 64 * j + (sc_ >> 1) * 32 + (sc_ & 1) * 16;
*reinterpret_cast<uint4*>(&Bs[r][si_mma_swz(kb, r)]) = *reinterpret_cast<const uint4*>(out);
if (c == 0) {
Wdm[r] = __half22float2(b->dm);
#pragma unroll
for (int jj = 0; jj < 4; jj++) {
unsigned char x0, x1, y0, y1;
si_mma_q4k_scales(b->scales, jj, x0, x1, y0, y1);
Ssc[r][2 * jj] = x0; Ssc[r][2 * jj + 1] = x1;
Smn[r][2 * jj] = y0; Smn[r][2 * jj + 1] = y1;
}
if (c < 8) {
unsigned char x0, x1, y0, y1;
si_mma_q4k_scales(b->scales, c >> 1, x0, x1, y0, y1);
const float2 dm = __half22float2(b->dm);
Bscale[r][c] = make_float2(dm.x * (float)((c & 1) ? x1 : x0),
dm.y * (float)((c & 1) ? y1 : y0));
}
}
}
Expand All @@ -2116,7 +2110,6 @@ void down_q4k_mma_rows_kernel(const unsigned char* __restrict__ down_q,
__syncthreads();

const int lnA = warp * 8 + tig * 2;
const float2 dmA = Wdm[lnA], dmB = Wdm[lnA + 1];
#pragma unroll 1
for (int g = 0; g < 8; g++) {
const int kk = g * 32;
Expand All @@ -2125,8 +2118,9 @@ void down_q4k_mma_rows_kernel(const unsigned char* __restrict__ down_q,
// accumulator elements is being folded. Fetching them per element cost five shared
// loads per output element per group; hoisting collapses that to a handful per group.
// Ablating this fold-in measured it at 111 us of the kernel's 169.
const float sA = dmA.x * (float)Ssc[lnA][g], mA = dmA.y * (float)Smn[lnA][g];
const float sB = dmB.x * (float)Ssc[lnA + 1][g], mB = dmB.y * (float)Smn[lnA + 1][g];
const float2 scaleA = Bscale[lnA][g], scaleB = Bscale[lnA + 1][g];
const float sA = scaleA.x, mA = scaleA.y;
const float sB = scaleB.x, mB = scaleB.y;
float adv[NT][2], asv[NT][2];
#pragma unroll
for (int ii = 0; ii < NT; ii++)
Expand Down
9 changes: 9 additions & 0 deletions kernels/tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,15 @@ set_target_properties(mmvq_mma_accuracy_gpu_test PROPERTIES CXX_STANDARD 17)
add_executable(down_mma_accuracy_gpu_test down_mma_accuracy_gpu_test.cpp)
target_link_libraries(down_mma_accuracy_gpu_test PRIVATE sparkinfer_kernels CUDA::cudart)
set_target_properties(down_mma_accuracy_gpu_test PROPERTIES CXX_STANDARD 17)
add_executable(down_mma_scales_gpu_test down_mma_scales_gpu_test.cpp)
target_link_libraries(down_mma_scales_gpu_test PRIVATE sparkinfer_kernels CUDA::cudart)
set_target_properties(down_mma_scales_gpu_test PROPERTIES CXX_STANDARD 17)
add_test(NAME down_mma_scales_gpu_test COMMAND down_mma_scales_gpu_test)
add_test(NAME down_mma_scales_legacy_loader_gpu_test COMMAND down_mma_scales_gpu_test)
set_tests_properties(down_mma_scales_gpu_test down_mma_scales_legacy_loader_gpu_test
PROPERTIES SKIP_RETURN_CODE 77)
set_tests_properties(down_mma_scales_gpu_test PROPERTIES ENVIRONMENT "SPARKINFER_MMA_BDEDUP=1")
set_tests_properties(down_mma_scales_legacy_loader_gpu_test PROPERTIES ENVIRONMENT "SPARKINFER_MMA_BDEDUP=0")
# The tensor-core LM head against the dp4a multi-row kernel it replaces. Both arms are exported
# functions, so one process covers both. Not a ctest: it needs ~0.8 GB of device memory.
add_executable(head_mma_accuracy_gpu_test head_mma_accuracy_gpu_test.cpp)
Expand Down
196 changes: 196 additions & 0 deletions kernels/tests/down_mma_scales_gpu_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
// Exercise Q4_K scale/min decoding and the down-MMA scratch lifecycle through the
// production launcher. Logical weights are generated before packing, so the
// oracle never uses the production metadata decoder. Group-constant, exactly
// representable activations make the Q8_1 scale and sum independently known.
#include "sparkinfer/kernels/moe.h"
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <stdexcept>
#include <vector>

namespace {
void check(cudaError_t status) {
if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status));
}
template <class T> struct Buffer {
T* p = nullptr;
explicit Buffer(size_t n) { check(cudaMalloc(&p, n * sizeof(T))); }
~Buffer() { cudaFree(p); }
Buffer(const Buffer&) = delete;
Buffer& operator=(const Buffer&) = delete;
};

struct Inputs {
Buffer<__nv_bfloat16> gate, up, output;
Buffer<float> weights, h, scratch;
Buffer<int> ids;
std::vector<__nv_bfloat16> host_up;
std::vector<float> host_weights;
Inputs(int H, int F) : gate(32u*F), up(32u*F), output(32u*H),
weights(32), h(32u*F), scratch(32u*H), ids(32),
host_up(32u*F), host_weights(32) {
std::vector<__nv_bfloat16> g(32u*F, __float2bfloat16(32.f));
check(cudaMemcpy(gate.p, g.data(), g.size()*sizeof(g[0]), cudaMemcpyHostToDevice));
check(cudaMemset(ids.p, 0, 32*sizeof(int)));
check(cudaMemset(scratch.p, 0, 32u*H*sizeof(float)));
}
};

void run_shape(int H, int F, const std::vector<int>& widths, cudaStream_t* streams) {
const int blocks = F / 256;
const size_t bytes = (size_t)2 * H * blocks * 144;
std::vector<unsigned char> packed(bytes, 0);
std::vector<double> dot_coefficient(2*H, 0), min_coefficient(2*H, 0);
for (int r = 0; r < 2*H; ++r) {
for (int b = 0; b < blocks; ++b) {
auto* p = packed.data() + ((size_t)r*blocks+b)*144;
const float d = std::ldexp(1.f, -10-(r%3));
const float dm = std::ldexp(1.f, -11-(b%2));
const __half hd = __float2half(d), hm = __float2half(dm);
std::memcpy(p, &hd, 2); std::memcpy(p+2, &hm, 2);
unsigned sc[8], mn[8];
for (int g = 0; g < 8; ++g) {
sc[g] = (13*r+7*b+9*g)%63+1;
mn[g] = (17*r+11*b+5*g)%64;
int sum_q = 0;
for (int j = 0; j < 32; ++j) {
const unsigned q = (3*r+5*b+7*g+j)%16;
p[16+(g/2)*32+j] |= (unsigned char)(q << (4*(g&1)));
sum_q += (int)q;
}
// Eight distinct power-of-two activation magnitudes distinguish
// every scale/min group, including a permutation of two groups.
dot_coefficient[r] += (1u<<g)*(double)d*sc[g]*sum_q;
min_coefficient[r] += (1u<<g)*(double)dm*mn[g];
}
// Standard GGML Q4_K six-bit scale packing, independent of the
// kernel's pair-at-a-time uint16 decoder.
for (int g = 0; g < 4; ++g) {
p[4+g] = (unsigned char)(sc[g] | ((sc[g+4] >> 4) << 6));
p[8+g] = (unsigned char)(mn[g] | ((mn[g+4] >> 4) << 6));
p[12+g] = (unsigned char)((sc[g+4]&15) | ((mn[g+4]&15)<<4));
}
}
}
Buffer<unsigned char> down(bytes);
check(cudaMemcpy(down.p, packed.data(), bytes, cudaMemcpyHostToDevice));
Inputs a(H,F), b(H,F);
Inputs* inputs[] = {&a,&b};
std::vector<int> second_expert(32,1);
check(cudaMemcpy(b.ids.p,second_expert.data(),32*sizeof(int),cudaMemcpyHostToDevice));
// Initialization used the default stream; the two nonblocking streams must
// not race its copies, even with pageable host buffers.
check(cudaDeviceSynchronize());
const float values[] = {.25f,-.25f,.5f,-.5f,1.f,-1.f,.125f,-.125f};

auto fill = [&](int lane, int M, bool zero) {
auto& x = *inputs[lane];
for (int row = 0; row < M; ++row) {
const float u = zero ? 0.f : values[row%8]*(lane+1)/128.f;
for (int col = 0; col < F; ++col)
x.host_up[(size_t)row*F+col] = __float2bfloat16(u*(1u<<((col/32)%8)));
x.host_weights[row] = row%3==0 ? 1.f : (row%3==1 ? .5f : -1.f);
}
check(cudaMemcpyAsync(x.up.p, x.host_up.data(), (size_t)M*F*sizeof(__nv_bfloat16),
cudaMemcpyHostToDevice, streams[lane]));
check(cudaMemcpyAsync(x.weights.p, x.host_weights.data(), M*sizeof(float),
cudaMemcpyHostToDevice, streams[lane]));
};
auto launch = [&](int lane, int M) {
auto& x = *inputs[lane];
sparkinfer::kernels::launch_moe_expert_ffn_q4k(
nullptr,nullptr,nullptr,down.p,12,12,12,x.ids.p,x.weights.p,
x.output.p,x.h.p,x.scratch.p,M,1,H,F,nullptr,streams[lane],false,x.gate.p,x.up.p);
check(cudaPeekAtLastError());
};
auto verify = [&](int lane, int M, bool zero) {
auto& x = *inputs[lane];
check(cudaStreamSynchronize(streams[lane]));
std::vector<__nv_bfloat16> out((size_t)M*H);
check(cudaMemcpy(out.data(),x.output.p,out.size()*sizeof(out[0]),cudaMemcpyDeviceToHost));
for (int row = 0; row < M; ++row) {
// SiLU(32) rounds to32 in float; Q8 codes are exactly +/-127.
// The per-group factors were folded into the oracle coefficients.
// Scaling these normal half values by powers of two is exact.
const float hv = zero ? 0.f : 32.f*values[row%8]*(lane+1)/128.f;
const double qd = __half2float(__float2half(std::fabs(hv)/127.f));
const double qs = __half2float(__float2half(hv*32.f));
const int qi = hv>0 ? 127 : (hv<0 ? -127 : 0);
for (int col = 0; col < H; ++col) {
const double want = (qd*qi*dot_coefficient[lane*H+col]-qs*min_coefficient[lane*H+col])
* x.host_weights[row];
const double got = __bfloat162float(out[(size_t)row*H+col]);
// BF16 rounding plus FP32 split-K summation, not a relaxed
// quantization comparison. Zero must remain exactly zero.
const double tolerance = zero ? 0.0 : std::max(.015625, std::fabs(want)*.008);
if (!std::isfinite(got) || std::fabs(got-want)>tolerance) {
std::printf("FAIL H=%d F=%d M=%d stream=%d row=%d col=%d got=%.9g want=%.9g tol=%.9g\n",
H,F,M,lane,row,col,got,want,tolerance);
throw std::runtime_error("down MMA numerical or scratch-lifecycle mismatch");
}
}
}
};
for (int M : widths) {
for (int lane = 0; lane < 2; ++lane) fill(lane,M,false);
for (int lane = 0; lane < 2; ++lane) launch(lane,M);
for (int lane = 0; lane < 2; ++lane) verify(lane,M,false);
cudaGraph_t graphs[2]{};
cudaGraphExec_t execs[2]{};
for (int lane = 0; lane < 2; ++lane) {
check(cudaStreamBeginCapture(streams[lane],cudaStreamCaptureModeGlobal));
launch(lane,M);
check(cudaStreamEndCapture(streams[lane],&graphs[lane]));
check(cudaGraphInstantiate(&execs[lane],graphs[lane],nullptr,nullptr,0));
}
// Reuse the same captured launches across A -> zero -> A. Two streams
// carry distinct inputs; four replays expose missed accumulator resets.
for (bool zero : {false,true,false}) {
for (int lane = 0; lane < 2; ++lane) fill(lane,M,zero);
for (int lane = 0; lane < 2; ++lane) {
for (int repeat = 0; repeat < 4; ++repeat)
check(cudaGraphLaunch(execs[lane],streams[lane]));
}
for (int lane = 0; lane < 2; ++lane) verify(lane,M,zero);
}
for (int lane = 0; lane < 2; ++lane) {
check(cudaGraphExecDestroy(execs[lane]));
check(cudaGraphDestroy(graphs[lane]));
}
std::printf("PASS down MMA scales H=%d F=%d M=%d eager + two-stream graph A/zero/A\n",H,F,M);
}
}
} // namespace

int main() {
int count = 0;
if (cudaGetDeviceCount(&count)!=cudaSuccess || count==0) return 77;
setenv("SPARKINFER_DOWN_MMA","1",1);
setenv("SPARKINFER_DOWN_MMA_MINROWS","8",1);
setenv("SPARKINFER_DOWN_MMVQ","1",1);
setenv("SPARKINFER_DOWN_Q4K","1",1);
setenv("SPARKINFER_DOWN_SPLITK_S_Q4","8",1);
setenv("SPARKINFER_MMA_ASTAGE","1",1);
setenv("SPARKINFER_DOWN_SPLITK_WIDE_ROWS","5",1);
setenv("SPARKINFER_DOWN_SPLITK_WIDE_S","2",1);
cudaStream_t streams[2]{};
try {
check(cudaStreamCreateWithFlags(&streams[0],cudaStreamNonBlocking));
check(cudaStreamCreateWithFlags(&streams[1],cudaStreamNonBlocking));
run_shape(256,256,{8,9,16,17,32},streams); // one K split
run_shape(1024,2304,{8,16,32},streams); // uneven multiple splits
run_shape(6656,19968,{8,9,16,17,32},streams);// real Muse widths
check(cudaStreamDestroy(streams[0]));
check(cudaStreamDestroy(streams[1]));
} catch (const std::exception& e) {
std::fprintf(stderr,"FAIL: %s\n",e.what());
return 1;
}
return 0;
}
Loading