biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert compressor pos bias fix + SwiGLU clamp ordering from commit 3320abf

Browse Source 
These changes were made based on CORRECTNESS_FIX_ATTEMPTS.md but the
PyTorch reference (which works correctly) doesn't have them. Reverting
to match the reference behavior:
- Compressor: position bias added to BOTH gate logits AND KV content
- SwiGLU: clamp after silu(gate), not before
This commit is contained in:
biondizzle
2026-06-03 13:57:16 +00:00
parent f05ee6cd69
commit 89510601f5
3 changed files with 21 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
dsv4/kernels/cuda/compressor_reduce.cu
View File

@@ -124,13 +124,15 @@ __global__ void csa_compress_reduce_kernel(
float g = gate_proj[token_idx * kv_dim + gate_offset + c];
float kv_val = kv_proj[token_idx * kv_dim + kv_offset + c];
// Position bias added ONLY to gate (softmax logit), NOT to KV content.
// Paper eq. 11-12: compressed = softmax(Z + B) * C — bias B is on the
// compression weights/logits, not on the KV content C.
// Position bias: same (m, 2*hd) bias added to every block
// Added to BOTH gate (softmax logit) and kv (content) per reference
if (position_bias != nullptr) {
int pos_bias_row = (block_i > 0 && t < m) ? t : (block_i > 0 ? (t - m) : t);
if (pos_bias_row >= 0 && pos_bias_row < m) {
g += position_bias[pos_bias_row * kv_dim + gate_offset + c];
float pb = position_bias[pos_bias_row * kv_dim + gate_offset + c];
g += pb;
// kv_offset matches gate_offset for CSA: both are 0 (a-stream) or hd (b-stream)
kv_val += position_bias[pos_bias_row * kv_dim + kv_offset + c];
}
}
float e = expf(g - local_max[ci]);
@@ -190,11 +192,12 @@ __global__ void hca_compress_reduce_kernel(
if (token_idx >= T) break;
float g = gate_proj[token_idx * hd + c];
float kv_val = kv_proj[token_idx * hd + c];
// Position bias added ONLY to gate (softmax logit), NOT to KV content.
// Paper eq. 9-10: compressed = softmax(Z + B) * C — bias B is on the
// compression weights/logits, not on the KV content C.
// Position bias: same (m, hd) bias added to every block
// Added to BOTH gate (softmax logit) and kv (content) per reference
if (position_bias != nullptr && t < m) {
g += position_bias[t * hd + c];
float pb = position_bias[t * hd + c];
g += pb;
kv_val += pb;
}
float e = expf(g - local_max);
local_denom += e;

7
dsv4/kernels/gemm/fused_swiglu.py
View File

@@ -2196,11 +2196,12 @@ class FusedSwiGLUScaledGroupedGemmKernel:
neg_acc = acc_vec * cutlass.Float32(-1.0)
exp_neg = cute.exp(neg_acc)
sigmoid = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + exp_neg)
# Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
silu_result = acc_vec * sigmoid
# Paper §4.2.3: gate component capped at swiglu_limit
# CuTe DSL clamp: min(x, limit) = cute.where(x > limit, limit, x)
if cutlass.const_expr(self.swiglu_limit > 0.0):
limit = cutlass.Float32(self.swiglu_limit)
acc_vec = cute.where(acc_vec > limit, limit, acc_vec)
silu_result = acc_vec * sigmoid
silu_result = cute.where(silu_result > limit, limit, silu_result)
silu_result = silu_result.to(self.c_dtype)
silu_gate_buf.store(silu_result)
# Keep acc_vec in BF16 (same type as the up branch)

14
dsv4/layers/moe.py
View File

@@ -512,11 +512,10 @@ class Nvfp4MoE:
l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
gate = l1_deil[:, :self.intermediate_size]
up = l1_deil[:, self.intermediate_size:]
# Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
if self._swiglu_limit is not None:
gate = gate.clamp(max=self._swiglu_limit)
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
gate_silu = torch.nn.functional.silu(gate)
if self._swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=self._swiglu_limit)
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
activated = gate_silu * up
_, _, l2_gs = quantize_to_nvfp4(activated)
@@ -652,11 +651,10 @@ class Nvfp4MoE:
l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
gate = l1_deil[:, :self.intermediate_size]
up = l1_deil[:, self.intermediate_size:]
# Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
if self._swiglu_limit is not None:
gate = gate.clamp(max=self._swiglu_limit)
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
gate_silu = torch.nn.functional.silu(gate)
if self._swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=self._swiglu_limit)
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
activated = gate_silu * up
# Compute runtime gsa for L2 from activated output (non-fused path)