diff --git a/dsv4/kernels/cuda/compressor_reduce.cu b/dsv4/kernels/cuda/compressor_reduce.cu index 4abc06b7..a6e07caf 100644 --- a/dsv4/kernels/cuda/compressor_reduce.cu +++ b/dsv4/kernels/cuda/compressor_reduce.cu @@ -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; diff --git a/dsv4/kernels/gemm/fused_swiglu.py b/dsv4/kernels/gemm/fused_swiglu.py index aace00ce..1c4b1d07 100644 --- a/dsv4/kernels/gemm/fused_swiglu.py +++ b/dsv4/kernels/gemm/fused_swiglu.py @@ -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) diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py index faa51fc8..6d7bf149 100644 --- a/dsv4/layers/moe.py +++ b/dsv4/layers/moe.py @@ -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)