diff --git a/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu b/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu index f111d8b5..4ca20569 100644 --- a/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu +++ b/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu @@ -135,7 +135,7 @@ __global__ void gdn_forward_kernel(const scalar_t* __restrict__ q, float beta_value = 0.0F; if (threadIdx.x == 0) { const float gate_raw = load_as_float(gate, gate_index); - gate_value = GateIsExp ? gate_raw : __expf(gate_raw); + { const float gc = fminf(fmaxf(gate_raw, -5.0F), 2.0F); gate_value = GateIsExp ? fminf(gate_raw, 7.389F) : __expf(gc); } beta_value = load_as_float(beta, gate_index); } gate_value = __shfl_sync(0xffffffffU, gate_value, 0); @@ -191,7 +191,7 @@ __global__ void gdn_forward_kernel(const scalar_t* __restrict__ q, for (int c = 0; c < COLS; ++c) { const float new_state = fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); - state_shard[c][r] = new_state; + state_shard[c][r] = fminf(fmaxf(new_state, -65504.0F), 65504.0F); attn_partial[c] += new_state * q_reg[r]; } } @@ -294,7 +294,7 @@ __global__ void gdn_forward_vlk_varlen_kernel( float beta_value = 0.0F; if (threadIdx.x == 0) { const float gate_raw = load_as_float(gate, gate_index); - gate_value = GateIsExp ? gate_raw : __expf(gate_raw); + { const float gc = fminf(fmaxf(gate_raw, -5.0F), 2.0F); gate_value = GateIsExp ? fminf(gate_raw, 7.389F) : __expf(gc); } beta_value = load_as_float(beta, gate_index); } gate_value = __shfl_sync(0xffffffffU, gate_value, 0); @@ -350,7 +350,7 @@ __global__ void gdn_forward_vlk_varlen_kernel( for (int c = 0; c < COLS; ++c) { const float new_state = fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); - state_shard[c][r] = new_state; + state_shard[c][r] = fminf(fmaxf(new_state, -65504.0F), 65504.0F); attn_partial[c] += new_state * q_reg[r]; } } @@ -543,7 +543,7 @@ __global__ void gdn_decode_mixed_qkv_global_state_kernel( for (int c = 0; c < COLS; ++c) { const float new_state = fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); - state_shard[c][r] = new_state; + state_shard[c][r] = fminf(fmaxf(new_state, -65504.0F), 65504.0F); attn_partial[c] += new_state * q_reg[r]; } } @@ -765,7 +765,7 @@ __global__ void gdn_decode_mixed_qkv_ddtree_state_kernel( for (int c = 0; c < COLS; ++c) { const float new_state = fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); - state_shard[c][r] = new_state; + state_shard[c][r] = fminf(fmaxf(new_state, -65504.0F), 65504.0F); attn_partial[c] += new_state * q_reg[r]; } } diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 18b5c7d0..a5acda04 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -581,6 +581,13 @@ class GatedDeltaNet(nn.Module): k_4d = k.unsqueeze(0) # (1, L, Hk, K) v_4d = v_raw.unsqueeze(0) # (1, L, Hv, V) g_3d = gate.unsqueeze(0) # (1, L, Hv) + # Clamp gate to prevent exp() overflow in CUDA kernel. + # gate = -dt * A_log.exp(), typically negative (decay). + # But pathological weights can produce positive values → exp > 1 + # → state grows exponentially over L tokens → inf. + # PyTorch ref clamps g ∈ [-5, 2] before cumsum. + # For recurrent kernel: clamp raw gate so exp(gate) ∈ [exp(-5), exp(2)] + g_3d = g_3d.clamp(-5.0, 2.0) beta_3d = b_seq.unsqueeze(0) # (1, L, Hv) # Initial state from temporal_state @@ -800,6 +807,8 @@ class GatedDeltaNet(nn.Module): k_t.view(BH, self.head_k_dim, 1), delta.view(BH, 1, self.head_v_dim), ) + # Clamp state to prevent gradual drift → NaN over long sequences + temporal_state.clamp_(-65504.0, 65504.0) # Output: core_out = q_t @ updated temporal_state core_out = _ix_bmm(