diff --git a/kernels/csrc/cuda/fused/rmsnorm.cu b/kernels/csrc/cuda/fused/rmsnorm.cu index 38f7159..5d045e2 100644 --- a/kernels/csrc/cuda/fused/rmsnorm.cu +++ b/kernels/csrc/cuda/fused/rmsnorm.cu @@ -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(weight); - const uint4* osum4r = reinterpret_cast(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])); } diff --git a/kernels/csrc/cuda/moe/expert_ffn_q4k.cu b/kernels/csrc/cuda/moe/expert_ffn_q4k.cu index 8ace63d..e9906ec 100644 --- a/kernels/csrc/cuda/moe/expert_ffn_q4k.cu +++ b/kernels/csrc/cuda/moe/expert_ffn_q4k.cu @@ -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;