Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
12 changes: 8 additions & 4 deletions kernels/csrc/cuda/fused/rmsnorm.cu
Original file line number Diff line number Diff line change
Expand Up @@ -204,11 +204,15 @@ __global__ void add_rmsnorm2_q8_kernel(const __nv_bfloat16* __restrict__ x,
__syncthreads();
const float inv_rms = s_warp[0];

// Re-read the bf16-rounded residual sum (exactly as add_rmsnorm2_kernel does), so out_norm
// is bit-identical to the unfused path; then derive Q8_1 from the bf16-rounded out_norm.
// The bf16-rounded residual sum is already in registers (sv, just stored to out_sum):
// round it in-place instead of reloading osum4[t]. bf16(sv) is bit-identical to the
// store+reload rn_pack8/rn_unpack8 round-trip, so out_norm/Q8_1 stay bit-identical to
// the unfused path — but the dependent global load (right after the store, on this
// single-block latency-bound kernel) is removed from the critical path.
const uint4* w4 = reinterpret_cast<const uint4*>(weight);
const uint4* osum4r = reinterpret_cast<const uint4*>(out_sum);
float svb[8], wv[8]; rn_unpack8(__ldg(osum4r + t), svb); rn_unpack8(__ldg(w4 + t), wv);
float svb[8], wv[8]; rn_unpack8(__ldg(w4 + t), wv);
#pragma unroll
for (int j = 0; j < 8; j++) svb[j] = __bfloat162float(__float2bfloat16(sv[j]));
float ov[8], bv[8];
#pragma unroll
for (int j = 0; j < 8; j++) { ov[j] = svb[j] * inv_rms * wv[j]; bv[j] = __bfloat162float(__float2bfloat16(ov[j])); }
Expand Down
6 changes: 5 additions & 1 deletion kernels/csrc/cuda/moe/expert_ffn_q4k.cu
Original file line number Diff line number Diff line change
Expand Up @@ -855,8 +855,12 @@ void launch_moe_expert_ffn_q4k(
if (gu2 < 0) { const char* g2 = getenv("SPARKINFER_GU2"); gu2 = (g2 && g2[0] == '0') ? 0 : 1; }
static int gu_spec = -1;
if (gu_spec < 0) { const char* gs = getenv("SPARKINFER_GU_SPEC"); gu_spec = (gs && gs[0] == '0') ? 0 : 1; }
// gate/up row-packing (2 rows/block, 8 warps). Default OFF: on RTX 5090 (sm_120) the
// 1-row-per-block gate_up_mmvq2_qwen_kernel is measurably faster than the packed variant
// — same 4 warps/row and identical math, but a smaller reduction footprint (sg/su[3][32]
// vs sg/su[2][3][32]) and no group indexing. Set SPARKINFER_GU_PACK2=1 to re-enable packing.
static int gu_pack2 = -1;
if (gu_pack2 < 0) { const char* gp = getenv("SPARKINFER_GU_PACK2"); gu_pack2 = (gp && gp[0] == '0') ? 0 : 1; }
if (gu_pack2 < 0) { const char* gp = getenv("SPARKINFER_GU_PACK2"); gu_pack2 = (gp && gp[0] == '1') ? 1 : 0; }
dim3 gu(num_tokens * top_k, (ffn + WPB - 1) / WPB);
if (mmvq && gu2 && gate_type == 12 && up_type == 12) { // faithful 4-warp mmvq gate/up
const si_block_q8_1* q;
Expand Down
Loading