diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 43ad718eb..02128919a 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -2720,6 +2720,7 @@ def _training_phase(node): # norm stats exist only in TRAINING forward phase "use_qk_l2norm", "checkpoint_every_n_tokens", "use_beta_sigmoid", + "beta_guard", "safe_gate", "gate_lower_bound", "batch_invariant", @@ -2735,7 +2736,7 @@ def _training_phase(node): # norm stats exist only in TRAINING forward phase "gdn2_bwd": dict( node_type=NodeType.GDN2_BWD, inputs=("q", "k", "v", "g", "beta", "w", "cu_seqlens", "dO", "state_checkpoints", "initial_state", "d_final_state", "a_log", "dt_bias"), - attrs=("scale", "use_qk_l2norm", "use_beta_sigmoid", "safe_gate", "gate_lower_bound", "batch_invariant"), + attrs=("scale", "use_qk_l2norm", "use_beta_sigmoid", "beta_guard", "safe_gate", "gate_lower_bound", "batch_invariant"), outputs=("dQ", "dK", "dV", "dG", "dBeta", "dW", "d_initial_state", "d_a_log", "d_dt_bias"), maybe={ "d_initial_state": lambda n: "initial_state" in n.inputs, diff --git a/python/cudnn/linear_attention/frost/common/beta_guard.py b/python/cudnn/linear_attention/frost/common/beta_guard.py new file mode 100644 index 000000000..979ae41cd --- /dev/null +++ b/python/cudnn/linear_attention/frost/common/beta_guard.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Erase-side beta safeguard for GDN-2 (per-token spectral non-expansion). + +Per (token, head) row: with the l2-normalized key k, weights w_d = k_d^2, +n = sum w, a = sum beta*w, nu = sum beta^2*w, and the gate headroom +c^2 = exp(-2*max_d g_d), the erase operator I - k (beta.k)^T stays +non-expansive under the decay budget iff + + n*nu - a^2 <= (c^2 - 1) * (1 - (1 - a)^2 / c^2) + +Violating rows are shrunk toward the key-weighted mean mu = a/n +(a-preserving, so one shot), rounded to the io dtype, re-tested against a +quantization tolerance, and flattened to mu on a re-test failure. + +One device function shared by the prefill, recompute, and bprop row-major +beta blocks so the three kernels produce bitwise-identical beta_eff; the +straight-through backward recomputes it, never differentiates through it. +""" + +import cutlass +import cutlass.cute as cute +import cutlass.experimental.primitives as nvvm + +from cudnn.frost.tile_dsl.swizzle import swizzle_xor_128b + +GUARD_MARGIN = 1.0 / 32 +GUARD_QUANT_TOL_MULT = 4.0 +MACHINE_EPSILON_FP16 = 2.0**-10 +MACHINE_EPSILON_BF16 = 2.0**-7 + + +@cute.jit +def beta_guard( + cfg: cutlass.Constexpr, + raw_beta_regs, + raw_k_regs, + k_inv_norm: cutlass.Float32, + gate_prefix_ptr, + decay_row: cutlass.Int32, + lane_in_row_group: cutlass.Int32, +) -> None: + """Rewrite the 16 per-lane beta registers to beta_eff in place.""" + zero = cutlass.Float32(0.0) + one = cutlass.Float32(1.0) + + # ---- gate headroom --------------------------------------------------------- + prev_row = decay_row - cutlass.Int32(1) + if decay_row == 0: + prev_row = cutlass.Int32(0) + max_ratio = zero + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + for f32_group in cutlass.range_constexpr(2): + f32_dim_base = dim_base + f32_group * 4 + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + row_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) + prev_idx = f32_segment * (cfg.b_t * 32) + prev_row * 32 + swizzle_xor_128b(prev_row, f32_segment_dim, elem_bytes=4) + exp_g_frag = (gate_prefix_ptr + row_idx).load(count=4, alignment=16) + exp_g_prev_frag = (gate_prefix_ptr + prev_idx).load(count=4, alignment=16) + for elem in cutlass.range_constexpr(4): + prev_val = exp_g_prev_frag[elem] + if decay_row == 0: + prev_val = one + ratio = exp_g_frag[elem] * cute.math.rcp(prev_val, approx=True, ftz=True) + max_ratio = cute.math.max(max_ratio, ratio) + max_ratio = cute.math.max(max_ratio, cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, max_ratio, 4, 31, kind=nvvm.Shfl.BFLY))) + max_ratio = cute.math.max(max_ratio, cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, max_ratio, 2, 31, kind=nvvm.Shfl.BFLY))) + max_ratio = cute.math.max(max_ratio, cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, max_ratio, 1, 31, kind=nvvm.Shfl.BFLY))) + + # ---- key-weighted stats ---------------------------------------------------- + n_val = zero + a_val = zero + nu_val = zero + for reg_idx in cutlass.range_constexpr(2 * 8): + k_norm = raw_k_regs[reg_idx] * k_inv_norm + weight = k_norm * k_norm + beta_val = raw_beta_regs[reg_idx] + n_val = n_val + weight + a_val = a_val + beta_val * weight + nu_val = nu_val + (beta_val * beta_val) * weight + n_val = n_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, n_val, 4, 31, kind=nvvm.Shfl.BFLY)) + n_val = n_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, n_val, 2, 31, kind=nvvm.Shfl.BFLY)) + n_val = n_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, n_val, 1, 31, kind=nvvm.Shfl.BFLY)) + a_val = a_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_val, 4, 31, kind=nvvm.Shfl.BFLY)) + a_val = a_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_val, 2, 31, kind=nvvm.Shfl.BFLY)) + a_val = a_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_val, 1, 31, kind=nvvm.Shfl.BFLY)) + nu_val = nu_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_val, 4, 31, kind=nvvm.Shfl.BFLY)) + nu_val = nu_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_val, 2, 31, kind=nvvm.Shfl.BFLY)) + nu_val = nu_val + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_val, 1, 31, kind=nvvm.Shfl.BFLY)) + + # ---- sensor ---------------------------------------------------------------- + inv_c2 = max_ratio * max_ratio + c2 = cute.math.rcp(inv_c2, approx=True, ftz=True) + r2 = cute.math.max(n_val * nu_val - a_val * a_val, zero) + r2_crit = cute.math.max((c2 - one) * (one - ((one - a_val) * (one - a_val)) * inv_c2), zero) + unsafe = cutlass.Boolean(False) + if n_val > cutlass.Float32(1.0e-20): + if r2 > r2_crit: + unsafe = cutlass.Boolean(True) + + # ---- projection ------------------------------------------------------------ + inv_n = cute.math.rcp(cute.math.max(n_val, cutlass.Float32(1.0e-20)), approx=True, ftz=True) + mu = a_val * inv_n + eta = cute.math.sqrt(cutlass.Float32(1.0 - GUARD_MARGIN) * r2_crit * cute.math.rcp(cute.math.max(r2, cutlass.Float32(1.0e-30)), approx=True, ftz=True)) + eta = cute.math.min(cute.math.max(eta, zero), one) + + # ---- quantize + re-test ---------------------------------------------------- + a_q = zero + nu_q = zero + for reg_idx in cutlass.range_constexpr(2 * 8): + candidate = raw_beta_regs[reg_idx] + if unsafe: + candidate = mu + eta * (candidate - mu) + candidate_q = candidate.to(cfg.io_dtype).to(cutlass.Float32) + raw_beta_regs[reg_idx] = candidate_q + k_norm = raw_k_regs[reg_idx] * k_inv_norm + weight = k_norm * k_norm + a_q = a_q + candidate_q * weight + nu_q = nu_q + (candidate_q * candidate_q) * weight + a_q = a_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_q, 4, 31, kind=nvvm.Shfl.BFLY)) + a_q = a_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_q, 2, 31, kind=nvvm.Shfl.BFLY)) + a_q = a_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, a_q, 1, 31, kind=nvvm.Shfl.BFLY)) + nu_q = nu_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_q, 4, 31, kind=nvvm.Shfl.BFLY)) + nu_q = nu_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_q, 2, 31, kind=nvvm.Shfl.BFLY)) + nu_q = nu_q + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, nu_q, 1, 31, kind=nvvm.Shfl.BFLY)) + r2_q = cute.math.max(n_val * nu_q - a_q * a_q, zero) + r2_crit_q = cute.math.max((c2 - one) * (one - ((one - a_q) * (one - a_q)) * inv_c2), zero) + quant_eps = MACHINE_EPSILON_FP16 if cfg.io_dtype == cutlass.Float16 else MACHINE_EPSILON_BF16 + quant_tol = cutlass.Float32(GUARD_QUANT_TOL_MULT * quant_eps) * (n_val * nu_q + a_q * a_q) + fallback = cutlass.Boolean(False) + if unsafe: + if r2_q > r2_crit_q + quant_tol: + fallback = cutlass.Boolean(True) + mu_q = (a_q * inv_n).to(cfg.io_dtype).to(cutlass.Float32) + for reg_idx in cutlass.range_constexpr(2 * 8): + final_val = raw_beta_regs[reg_idx] + if fallback: + final_val = mu_q + raw_beta_regs[reg_idx] = final_val diff --git a/python/cudnn/linear_attention/frost/gdn2_engine.py b/python/cudnn/linear_attention/frost/gdn2_engine.py index e7d734e69..b9fe02d6e 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -57,6 +57,8 @@ def check_support(self, graph) -> None: raise NotImplementedError(f"Gdn2FrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the GDN-2 node (got {checkpoint})") if not facts.gates_at_ho: raise NotImplementedError(f"Gdn2FrostEngine: g/beta/w must carry HO = max(q, v) heads ({facts.h_o})") + if facts.beta_guard and not facts.use_qk_l2norm: + raise NotImplementedError("Gdn2FrostEngine: beta_guard requires use_qk_l2norm (the sensor is defined on the normalized key)") fp32 = cudnn.data_type.FLOAT if facts.io_dtype is not None: for port, got in (("beta", facts.beta_dtype), ("w", facts.w_dtype)): @@ -121,6 +123,7 @@ def __init__(self, node, kernel_module): self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) self.safe_gate = bool(node.params.get("safe_gate", False)) self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) + self.beta_guard = bool(node.params.get("beta_guard", False)) glb = node.params.get("gate_lower_bound") self.gate_lower_bound = float(glb) if glb is not None else kernel_module.DEFAULT_GATE_LOWER_BOUND self.has_final_state = "final_state" in node.outputs @@ -285,6 +288,7 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, use_beta_sigmoid=self.use_beta_sigmoid, + beta_guard=self.beta_guard, work_items=work_items, work_count=work_count, scheduler_counter=scheduler_counter, @@ -324,6 +328,7 @@ def __init__(self, node, bwd_module, recompute_module): self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) self.safe_gate = bool(node.params.get("safe_gate", False)) self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) + self.beta_guard = bool(node.params.get("beta_guard", False)) glb = node.params.get("gate_lower_bound") self.gate_lower_bound = float(glb) if glb is not None else bwd_module.DEFAULT_GATE_LOWER_BOUND self.gate_bwd_blocks = GATE_BWD_BLOCKS @@ -597,6 +602,7 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, use_beta_sigmoid=self.use_beta_sigmoid, + beta_guard=self.beta_guard, work_items=work_items, work_count=work_count, scheduler_counter=scheduler_recompute, @@ -641,6 +647,7 @@ def run(self, views, workspace, stream) -> None: a_log=a_log, dt_bias=dt_bias, use_beta_sigmoid=self.use_beta_sigmoid, + beta_guard=self.beta_guard, work_items=work_items, work_count=work_count, scheduler_counter=scheduler_bwd if self.bwd_dynamic_scheduling else None, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py index 283ad2bee..7897c6159 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py @@ -9,6 +9,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack +from ..common.beta_guard import beta_guard from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr @@ -1842,7 +1843,7 @@ def compute0_warp_group( lane_in_row_group = lane_idx - lane_row_group * 8 decay_row = row_group_start + lane_row_group - g_prefix_ptr = sGate_ptr + gate_prefix_ptr = sGate_ptr channel_dim = cg0_warp * cfg.threads_per_warp + lane_idx # ---- gate prefix scan: cumulative log-gate per key channel --------------- gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) @@ -1972,7 +1973,8 @@ def compute0_warp_group( for dim_offset in cutlass.range_constexpr(8): q_val = raw_q_frag_f32[dim_offset] k_val = raw_k_frag_f32[dim_offset] - raw_q_regs[reg_base + dim_offset] = q_val + if cutlass.const_expr(not cfg.beta_guard): + raw_q_regs[reg_base + dim_offset] = q_val raw_k_regs[reg_base + dim_offset] = k_val beta_val = raw_beta_frag_f32[dim_offset] if cutlass.const_expr(cfg.beta_sigmoid): @@ -1985,9 +1987,11 @@ def compute0_warp_group( qk1_lo, qk1_hi = ffma2(q_val, k_val, q_val, k_val, qk1_lo, qk1_hi) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_q_done[raw_stage].arrive() + if cutlass.const_expr(not cfg.beta_guard): + bars.mb_q_done[raw_stage].arrive() bars.mb_k_done[raw_stage].arrive() - bars.mb_beta_done[raw_stage].arrive() + if cutlass.const_expr(not cfg.beta_guard): + bars.mb_beta_done[raw_stage].arrive() q_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) k_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) @@ -2008,6 +2012,25 @@ def compute0_warp_group( sNorm_raw[(gc % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + cfg.b_t + decay_row] = k_inv_norm q_stage_norm = q_inv_norm * scale + # ---- beta guard ------------------------------------------------------------ + if cutlass.const_expr(cfg.beta_guard): + beta_guard(cfg, raw_beta_regs, raw_k_regs, k_inv_norm, gate_prefix_ptr, decay_row, lane_in_row_group) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + beta_eff_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + beta_eff_pack[pair_idx] = fp32_to_fp16(raw_beta_regs[reg_base + dim0], raw_beta_regs[reg_base + dim0 + 1], dtype=cfg.io_dtype) + beta_eff_vec = cutlass.Vector.from_elements( + (beta_eff_pack[0], beta_eff_pack[1], beta_eff_pack[2], beta_eff_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sBetaP_ptr + raw_f16_idx).store(beta_eff_vec, alignment=16) + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) @@ -2019,9 +2042,9 @@ def compute0_warp_group( f32_segment = f32_dim_base // 32 f32_segment_dim = f32_dim_base - f32_segment * 32 g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) - exp_g_frag = (g_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_frag = (gate_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) - exp_g_last_frag = (g_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) + exp_g_last_frag = (gate_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) f32_reg_base = reg_base + f32_group * 4 exp_g_regs[f32_reg_base] = exp_g_frag[0] exp_g_regs[f32_reg_base + 1] = exp_g_frag[1] @@ -2033,6 +2056,8 @@ def compute0_warp_group( exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] nvvm.fence_proxy("async.shared", space="cta") bars.mb_gate_done[raw_stage].arrive() + if cutlass.const_expr(cfg.beta_guard): + bars.mb_beta_done[raw_stage].arrive() for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 @@ -2080,6 +2105,14 @@ def compute0_warp_group( for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 + if cutlass.const_expr(cfg.beta_guard): + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + raw_q_frag = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_q_frag_f32 = raw_q_frag.to(cutlass.Float32) + for dim_offset in cutlass.range_constexpr(8): + raw_q_regs[reg_base + dim_offset] = raw_q_frag_f32[dim_offset] q_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): @@ -2109,6 +2142,8 @@ def compute0_warp_group( (sK_restore_ptr + op_idx).store(k_restore_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_q_decay_k_restore_ready[decay_stage].arrive() + if cutlass.const_expr(cfg.beta_guard): + bars.mb_q_done[raw_stage].arrive() # ---- state copy: SMEM -> TMEM f16 ---------------------------------------- bars.mb_state_input_done[gc % 2].wait(((gc // 2) + 1) % 2) @@ -2534,6 +2569,7 @@ def compute2_warp_group( lane_idx, tmem_hold, warp_idx, + mBeta, sBeta_raw, sGate_raw, sNorm_raw, @@ -2709,8 +2745,7 @@ def compute2_warp_group( k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] db_regs[t] = k_v * dk_decay beta_v = (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to(cutlass.Float32) - if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) + if cutlass.const_expr(cfg.beta_sigmoid and not cfg.beta_guard): beta_v = sigmoid(beta_v).to(cfg.io_dtype).to(cutlass.Float32) dgate_regs[t] = beta_v * dk_decay dk_n[t] = dk_n[t] + dgate_regs[t] @@ -2730,12 +2765,17 @@ def compute2_warp_group( q_v = q_v * sNorm_raw[norm_base + t] k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] beta_v = (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to(cutlass.Float32) - if cutlass.const_expr(cfg.beta_sigmoid): - half = cutlass.Float32(0.5) + if cutlass.const_expr(cfg.beta_sigmoid and not cfg.beta_guard): beta_v = sigmoid(beta_v).to(cfg.io_dtype).to(cutlass.Float32) dgate_regs[t] = q_v * dq_n[t] + beta_v * db_regs[t] - k_v * (dk_n[t] - dgate_regs[t]) - if cutlass.const_expr(cfg.beta_sigmoid): + if cutlass.const_expr(cfg.beta_sigmoid and not cfg.beta_guard): db_regs[t] = db_regs[t] * (beta_v - beta_v * beta_v) + if cutlass.const_expr(cfg.beta_sigmoid and cfg.beta_guard): + token_idx = chunk_start + cutlass.Int32(t) + sig_val = cutlass.Float32(0.0) + if token_idx < batch_seqlen: + sig_val = sigmoid(mBeta[batch_start + token_idx, head_idx, channel].to(cutlass.Float32)).to(cfg.io_dtype).to(cutlass.Float32) + db_regs[t] = db_regs[t] * (sig_val - sig_val * sig_val) dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + ((dgate_last_acc[0] + dgate_last_acc[1]) + (dgate_last_acc[2] + dgate_last_acc[3])) nvvm.fence_proxy("async.shared", space="cta") bars.mb_beta_done[raw_stage].arrive() @@ -3243,6 +3283,7 @@ def host( state_checkpoints: cute.Tensor, a_log: cute.Tensor | None, dt_bias: cute.Tensor | None, + beta_main: cute.Tensor | None, dgate: cute.Tensor, dbeta: cute.Tensor, dw: cute.Tensor, @@ -3268,6 +3309,7 @@ def host( cu_seqlens, a_log, dt_bias, + beta_main, dgate, dbeta, dw, @@ -3293,6 +3335,7 @@ def kernel( cu_seqlens: cute.Tensor, mA_log: cute.Tensor | None, mDt_bias: cute.Tensor | None, + mBeta: cute.Tensor | None, mDgate: cute.Tensor, mDb: cute.Tensor, mDw_out: cute.Tensor, @@ -3739,6 +3782,7 @@ def kernel( lane_idx, tmem_hold, warp_idx, + mBeta, sBeta_raw, sGate_raw, sNorm_raw, @@ -3789,6 +3833,7 @@ class Gdn2BwdCfg: safe_gate: bool gate_scale_log2: float beta_sigmoid: bool + beta_guard: bool use_initial_state: bool q_ratio: int k_ratio: int @@ -3883,6 +3928,7 @@ def build_cfg( safe_gate: bool, gate_scale_log2: float, beta_sigmoid: bool, + beta_guard: bool = False, use_initial_state: bool, q_ratio: int, k_ratio: int, @@ -3893,6 +3939,8 @@ def build_cfg( ) -> Gdn2BwdCfg: if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + if beta_guard and not l2norm: + raise ValueError("beta_guard requires l2norm (the sensor is defined on the normalized key)") cfg = Gdn2BwdCfg( io_dtype=io_dtype, use_dstate_in=use_dstate_in, @@ -3901,6 +3949,7 @@ def build_cfg( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=beta_sigmoid, + beta_guard=beta_guard, use_initial_state=use_initial_state, q_ratio=q_ratio, k_ratio=k_ratio, @@ -3972,6 +4021,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, + beta_guard: bool, use_initial_state: bool, dynamic_scheduling: bool, order_in_prologue: bool, @@ -4008,6 +4058,7 @@ def chunk_gdn2_bwd_sm100( a_log=None, dt_bias=None, use_beta_sigmoid: bool = False, + beta_guard: bool = False, work_items=None, work_count=None, scheduler_counter=None, @@ -4106,6 +4157,7 @@ def chunk_gdn2_bwd_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, + beta_guard, use_initial_state, dynamic_scheduling, order_in_prologue, @@ -4123,6 +4175,7 @@ def chunk_gdn2_bwd_sm100( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=use_beta_sigmoid, + beta_guard=beta_guard, use_initial_state=use_initial_state, q_ratio=HO // HQ, k_ratio=HO // HK, @@ -4150,6 +4203,9 @@ def chunk_gdn2_bwd_sm100( state_checkpoints_cute = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=len(state_checkpoints.shape) - 1) a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=16) if dt_bias is not None else None + beta_main_cute = None + if beta_guard and use_beta_sigmoid: + beta_main_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=2) dgate_cute = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=len(dgate.shape) - 1) dbeta_cute = from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=len(dbeta.shape) - 1) dw_cute = from_dlpack(dw, assumed_align=16).mark_layout_dynamic(leading_dim=len(dw.shape) - 1) @@ -4159,6 +4215,7 @@ def chunk_gdn2_bwd_sm100( state_checkpoints_cute, a_log_cute, dt_bias_cute, + beta_main_cute, dgate_cute, dbeta_cute, dw_cute, @@ -4255,10 +4312,12 @@ def chunk_gdn2_bwd_sm100( tensormap_workspace, cu_stream, ) + cache["beta_main"] = beta_guard and use_beta_sigmoid cache["compiled"]( state_checkpoints, a_log, dt_bias, + beta if cache["beta_main"] else None, dgate, dbeta, dw, @@ -4336,6 +4395,7 @@ def run_bwd( state_checkpoints, a_log, dt_bias, + beta if cache.get("beta_main") else None, dgate, dbeta, dw, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index b8646ad53..de8d667f5 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -26,6 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack +from ..common.beta_guard import beta_guard from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr @@ -1254,6 +1255,10 @@ def compute0_warp_group( q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + # ---- beta guard ------------------------------------------------------------ + if cutlass.const_expr(cfg.beta_guard): + beta_guard(cfg, raw_beta_regs, raw_k_regs, k_inv_norm, sGate_ptr, decay_row, lane_in_row_group) + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) @@ -2350,6 +2355,7 @@ class Gdn2Cfg: safe_gate: bool gate_scale_log2: float beta_sigmoid: bool + beta_guard: bool q_ratio: int k_ratio: int v_ratio: int @@ -2436,6 +2442,7 @@ def build_cfg( safe_gate: bool, gate_scale_log2: float, beta_sigmoid: bool, + beta_guard: bool = False, q_ratio: int, k_ratio: int, v_ratio: int, @@ -2447,6 +2454,8 @@ def build_cfg( fills the derived TMEM column offsets and SMEM buffer cosizes.""" if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + if beta_guard and not l2norm: + raise ValueError("beta_guard requires l2norm (the sensor is defined on the normalized key)") cfg = Gdn2Cfg( io_dtype=io_dtype, state_dtype=state_dtype, @@ -2457,6 +2466,7 @@ def build_cfg( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=beta_sigmoid, + beta_guard=beta_guard, q_ratio=q_ratio, k_ratio=k_ratio, v_ratio=v_ratio, @@ -2823,6 +2833,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, + beta_guard: bool, dynamic_scheduling: bool, order_gen: bool, ): @@ -2840,6 +2851,7 @@ def compile( safe_gate: bool, gate_scale_log2: float, beta_sigmoid: bool, + beta_guard: bool, q_ratio: int, k_ratio: int, v_ratio: int, @@ -2878,6 +2890,7 @@ def compile( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=beta_sigmoid, + beta_guard=beta_guard, q_ratio=q_ratio, k_ratio=k_ratio, v_ratio=v_ratio, @@ -2932,6 +2945,7 @@ def chunk_gdn2_sm100( a_log=None, dt_bias=None, use_beta_sigmoid: bool = False, + beta_guard: bool = False, work_items=None, work_count=None, scheduler_counter=None, @@ -3037,6 +3051,7 @@ def chunk_gdn2_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, + beta_guard, dynamic_scheduling, order_gen, ) @@ -3083,6 +3098,7 @@ def chunk_gdn2_sm100( safe_gate, gate_scale_log2, use_beta_sigmoid, + beta_guard, q_ratio, k_ratio, v_ratio, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py index 91a007161..4669177b3 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py @@ -26,6 +26,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack +from ..common.beta_guard import beta_guard from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMENTS, ORDER_THREADS, decode_work_item, order_body from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr @@ -1017,6 +1018,10 @@ def compute0_warp_group( norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + # ---- beta guard ------------------------------------------------------------ + if cutlass.const_expr(cfg.beta_guard): + beta_guard(cfg, raw_beta_regs, raw_k_regs, k_inv_norm, sGate_ptr, decay_row, lane_in_row_group) + # ---- decay/restore operands: exp2(+-g) applied per key channel ----------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) @@ -1927,6 +1932,7 @@ class Gdn2RecomputeCfg: safe_gate: bool gate_scale_log2: float beta_sigmoid: bool + beta_guard: bool k_ratio: int v_ratio: int n_heads_out: int @@ -2005,6 +2011,7 @@ def build_cfg( safe_gate: bool, gate_scale_log2: float, beta_sigmoid: bool, + beta_guard: bool = False, k_ratio: int, v_ratio: int, n_heads_out: int, @@ -2015,6 +2022,8 @@ def build_cfg( fills the derived TMEM column offsets and SMEM buffer cosizes.""" if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + if beta_guard and not l2norm: + raise ValueError("beta_guard requires l2norm (the sensor is defined on the normalized key)") cfg = Gdn2RecomputeCfg( io_dtype=io_dtype, state_dtype=state_dtype, @@ -2025,6 +2034,7 @@ def build_cfg( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=beta_sigmoid, + beta_guard=beta_guard, k_ratio=k_ratio, v_ratio=v_ratio, n_heads_out=n_heads_out, @@ -2338,6 +2348,7 @@ def get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, + beta_guard: bool, dynamic_scheduling: bool, order_in_prologue: bool, order_gen: bool, @@ -2357,6 +2368,7 @@ def compile( safe_gate: bool, gate_scale_log2: float, beta_sigmoid: bool, + beta_guard: bool, k_ratio: int, v_ratio: int, n_heads_out: int, @@ -2391,6 +2403,7 @@ def compile( safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, beta_sigmoid=beta_sigmoid, + beta_guard=beta_guard, k_ratio=k_ratio, v_ratio=v_ratio, n_heads_out=n_heads_out, @@ -2438,6 +2451,7 @@ def chunk_gdn2_recompute_sm100( a_log=None, dt_bias=None, use_beta_sigmoid: bool = False, + beta_guard: bool = False, work_items=None, work_count=None, scheduler_counter=None, @@ -2542,6 +2556,7 @@ def chunk_gdn2_recompute_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid, + beta_guard, dynamic_scheduling, order_in_prologue, order_gen, @@ -2588,6 +2603,7 @@ def chunk_gdn2_recompute_sm100( safe_gate, gate_scale_log2, use_beta_sigmoid, + beta_guard, k_ratio, v_ratio, HO, diff --git a/python/cudnn/linear_attention/graph_analyzer.py b/python/cudnn/linear_attention/graph_analyzer.py index 9b8817339..5a0ae873c 100644 --- a/python/cudnn/linear_attention/graph_analyzer.py +++ b/python/cudnn/linear_attention/graph_analyzer.py @@ -107,6 +107,7 @@ class LaGraphFacts: use_qk_l2norm: bool = False safe_gate: bool = False use_beta_sigmoid: bool = False + beta_guard: bool = False gate_lower_bound: Optional[float] = None checkpoint_every_n_tokens: int = 0 batch_invariant: bool = False @@ -236,6 +237,7 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: use_qk_l2norm=bool(params.get("use_qk_l2norm", False)), safe_gate=safe_gate, use_beta_sigmoid=bool(params.get("use_beta_sigmoid", False)), + beta_guard=bool(params.get("beta_guard", False)), gate_lower_bound=float(params["gate_lower_bound"]) if params.get("gate_lower_bound") is not None else None, checkpoint_every_n_tokens=checkpoint, batch_invariant=bool(params.get("batch_invariant", False)), diff --git a/python/cudnn/linear_attention/ops/gdn2.py b/python/cudnn/linear_attention/ops/gdn2.py index 0c37e45c5..f22a4a563 100644 --- a/python/cudnn/linear_attention/ops/gdn2.py +++ b/python/cudnn/linear_attention/ops/gdn2.py @@ -130,6 +130,7 @@ def make_fprop_cache_key( use_qk_l2norm, batch_invariant, use_beta_sigmoid, + beta_guard, safe_gate, gate_lower_bound, has_initial_state, @@ -157,6 +158,7 @@ def make_fprop_cache_key( bool(use_qk_l2norm), bool(batch_invariant), bool(use_beta_sigmoid), + bool(beta_guard), bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, bool(has_initial_state), @@ -188,6 +190,7 @@ def make_bprop_cache_key( use_qk_l2norm, batch_invariant, use_beta_sigmoid, + beta_guard, safe_gate, gate_lower_bound, device, @@ -216,6 +219,7 @@ def make_bprop_cache_key( bool(use_qk_l2norm), bool(batch_invariant), bool(use_beta_sigmoid), + bool(beta_guard), bool(safe_gate), float(gate_lower_bound) if gate_lower_bound is not None else None, device, @@ -249,6 +253,7 @@ def build_fprop_graph( gate_lower_bound, checkpoint, use_beta_sigmoid=False, + beta_guard=False, ): graph = cudnn.pygraph() HO = max(H, HV) @@ -283,6 +288,7 @@ def build_fprop_graph( use_qk_l2norm=use_qk_l2norm, batch_invariant=batch_invariant, use_beta_sigmoid=use_beta_sigmoid or None, + beta_guard=beta_guard or None, safe_gate=safe_gate, gate_lower_bound=gate_lower_bound, checkpoint_every_n_tokens=checkpoint, @@ -325,6 +331,7 @@ def gdn2_fwd( use_qk_l2norm_in_kernel: bool = False, batch_invariant: bool = False, use_beta_sigmoid_in_kernel: bool = False, + beta_guard: bool = False, safe_gate: bool = False, gate_lower_bound: Optional[float] = None, a_log: Optional[torch.Tensor] = None, @@ -401,6 +408,7 @@ def gdn2_fwd( use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, + beta_guard, safe_gate, gate_lower_bound, state0 is not None, @@ -430,6 +438,7 @@ def gdn2_fwd( float(gate_lower_bound) if gate_lower_bound is not None else None, checkpoint, use_beta_sigmoid=bool(use_beta_sigmoid_in_kernel), + beta_guard=bool(beta_guard), ) select_plan(fprop_cache[cache_key][0], plan_name) @@ -479,6 +488,7 @@ def gdn2_fwd_fake( use_qk_l2norm_in_kernel=False, batch_invariant=False, use_beta_sigmoid_in_kernel=False, + beta_guard=False, safe_gate=False, gate_lower_bound=None, a_log=None, @@ -531,6 +541,7 @@ def build_bprop_graph( use_qk_l2norm, batch_invariant, use_beta_sigmoid=False, + beta_guard=False, safe_gate=False, gate_lower_bound=None, ): @@ -576,6 +587,7 @@ def build_bprop_graph( use_qk_l2norm=use_qk_l2norm, batch_invariant=batch_invariant, use_beta_sigmoid=use_beta_sigmoid or None, + beta_guard=beta_guard or None, safe_gate=safe_gate or None, gate_lower_bound=gate_lower_bound, name="gdn2_bwd", @@ -628,6 +640,7 @@ def gdn2_bwd( use_qk_l2norm_in_kernel: bool = False, batch_invariant: bool = False, use_beta_sigmoid_in_kernel: bool = False, + beta_guard: bool = False, safe_gate: bool = False, gate_lower_bound: Optional[float] = None, a_log: Optional[torch.Tensor] = None, @@ -719,6 +732,7 @@ def gdn2_bwd( use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, + beta_guard, safe_gate, gate_lower_bound, device, @@ -744,6 +758,7 @@ def gdn2_bwd( bool(use_qk_l2norm_in_kernel), bool(batch_invariant), use_beta_sigmoid=bool(use_beta_sigmoid_in_kernel), + beta_guard=bool(beta_guard), safe_gate=bool(safe_gate), gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, ) @@ -814,6 +829,7 @@ def gdn2_bwd_fake( use_qk_l2norm_in_kernel=False, batch_invariant=False, use_beta_sigmoid_in_kernel=False, + beta_guard=False, safe_gate=False, gate_lower_bound=None, a_log=None, @@ -858,6 +874,7 @@ def gdn2_setup_context(ctx, inputs, output): use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, + beta_guard, safe_gate, gate_lower_bound, a_log, @@ -878,6 +895,7 @@ def gdn2_setup_context(ctx, inputs, output): ctx.batch_invariant = batch_invariant ctx.plan_name = plan_name ctx.use_beta_sigmoid_in_kernel = bool(use_beta_sigmoid_in_kernel) + ctx.beta_guard = bool(beta_guard) ctx.safe_gate = bool(safe_gate) ctx.gate_lower_bound = gate_lower_bound ctx.set_materialize_grads(False) @@ -914,6 +932,7 @@ def gdn2_backward(ctx, dO, dFinal, dstate_checkpoints): use_qk_l2norm_in_kernel=ctx.use_qk_l2norm_in_kernel, batch_invariant=ctx.batch_invariant, use_beta_sigmoid_in_kernel=ctx.use_beta_sigmoid_in_kernel, + beta_guard=ctx.beta_guard, safe_gate=ctx.safe_gate, gate_lower_bound=ctx.gate_lower_bound, a_log=a_log, @@ -936,6 +955,7 @@ def gdn2_backward(ctx, dO, dFinal, dstate_checkpoints): None, None, None, + None, d_a_log if ctx.safe_gate else None, d_dt_bias if ctx.safe_gate else None, None, @@ -969,6 +989,7 @@ def gated_delta_net_v2( use_qk_l2norm_in_kernel: bool = False, batch_invariant: bool = False, use_beta_sigmoid_in_kernel: bool = False, + beta_guard: bool = False, safe_gate: bool = False, gate_lower_bound: Optional[float] = None, a_log: Optional[torch.Tensor] = None, @@ -1009,6 +1030,11 @@ def gated_delta_net_v2( disables split-K load balancing). use_beta_sigmoid_in_kernel: apply ``sigmoid(beta)`` inside the kernel; the backward returns the raw-logit beta gradient. + beta_guard: apply the erase-side beta safeguard: rows whose per-channel + beta contrast would make the (decayed) erase step expansive are + shrunk toward the key-weighted mean beta before use. The backward + is straight-through (the guard is recomputed, not differentiated). + Requires ``use_qk_l2norm_in_kernel=True``. safe_gate: interpret ``g`` through the safe-gate transform ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))``. Requires ``a_log`` and ``dt_bias``; the backward returns the @@ -1048,6 +1074,7 @@ def gated_delta_net_v2( use_qk_l2norm_in_kernel=bool(use_qk_l2norm_in_kernel), batch_invariant=bool(batch_invariant), use_beta_sigmoid_in_kernel=bool(use_beta_sigmoid_in_kernel), + beta_guard=bool(beta_guard), safe_gate=bool(safe_gate), gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, a_log=a_log, diff --git a/test/python/linear_attention/frost/examples/01_gdn_prefill.py b/test/python/linear_attention/frost/examples/01_gdn_prefill.py index c3637df28..de0934e6e 100644 --- a/test/python/linear_attention/frost/examples/01_gdn_prefill.py +++ b/test/python/linear_attention/frost/examples/01_gdn_prefill.py @@ -26,7 +26,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("gdn_frost")) # pin the FROST entry + g.select_plan(names.index("gdn_frost")) g.check_support() g.build_plans() @@ -50,7 +50,7 @@ def _reference(q, k, v, g, beta, cu, scale): residual = v[t] - a[:, None] * torch.einsum("hd,hdv->hv", k[t], S) S = a[:, None, None] * S + b[:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) - fs[n] = S + fs[n] = S.transpose(-2, -1) return o, fs diff --git a/test/python/linear_attention/frost/examples/02_gdn_backward.py b/test/python/linear_attention/frost/examples/02_gdn_backward.py index 2f4338758..0d66b8a34 100644 --- a/test/python/linear_attention/frost/examples/02_gdn_backward.py +++ b/test/python/linear_attention/frost/examples/02_gdn_backward.py @@ -22,7 +22,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("gdn_frost")) # pin the FROST entry + g.select_plan(names.index("gdn_frost")) g.check_support() g.build_plans() diff --git a/test/python/linear_attention/frost/examples/03_kda_prefill.py b/test/python/linear_attention/frost/examples/03_kda_prefill.py index 393870680..044908046 100644 --- a/test/python/linear_attention/frost/examples/03_kda_prefill.py +++ b/test/python/linear_attention/frost/examples/03_kda_prefill.py @@ -27,7 +27,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("kda_frost")) # pin the FROST entry + g.select_plan(names.index("kda_frost")) g.check_support() g.build_plans() @@ -58,7 +58,7 @@ def _reference(q, k, v, g, beta, cu, scale): residual = v[t] - torch.einsum("hd,hdv->hv", k[t], S) S = S + beta[t][:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) - fs[n] = S + fs[n] = S.transpose(-2, -1) return o, fs diff --git a/test/python/linear_attention/frost/examples/04_kda_backward.py b/test/python/linear_attention/frost/examples/04_kda_backward.py index 4a8b24360..eb3e6f428 100644 --- a/test/python/linear_attention/frost/examples/04_kda_backward.py +++ b/test/python/linear_attention/frost/examples/04_kda_backward.py @@ -23,7 +23,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("kda_frost")) # pin the FROST entry + g.select_plan(names.index("kda_frost")) g.check_support() g.build_plans() diff --git a/test/python/linear_attention/frost/examples/05_gdn2_prefill.py b/test/python/linear_attention/frost/examples/05_gdn2_prefill.py index 19da0445b..3927803db 100644 --- a/test/python/linear_attention/frost/examples/05_gdn2_prefill.py +++ b/test/python/linear_attention/frost/examples/05_gdn2_prefill.py @@ -10,8 +10,18 @@ S_t = S' + k_t^T (w_t . v_t - (beta_t . k_t) S') o_t = q_t S_t -``beta``/``w`` are io-dtype post-sigmoid tensors; ``use_qk_l2norm=False`` -passes q/k through as given, so this example feeds pre-normalized rows. +``beta``/``w`` are io-dtype post-sigmoid tensors. This example runs the +in-kernel q/k L2 norm, the bounded safe gate, and the erase-side beta +safeguard (``beta_guard``). The decay input is a raw pre-activation; with +``safe_gate=True`` the kernel computes the bounded gate + + g = -5 * sigmoid(exp(A_log) * (a + dt_bias)) + +here with ``A_log = 0`` and ``dt_bias = 0``. The guard shrinks tokens whose +per-channel beta contrast would make the decayed erase step expansive toward +the key-weighted mean beta. Even tokens keep real decay headroom (guard +stays quiet); odd tokens carry planted near-zero-decay channels (guard +fires); the fp64 reference applies the same activation and projection. """ from __future__ import annotations @@ -27,7 +37,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("gdn2_frost")) # pin the FROST entry + g.select_plan(names.index("gdn2_frost")) g.check_support() g.build_plans() @@ -44,11 +54,41 @@ def _randu(rows, dim, device): return means + torch.rand(rows, dim, device=device) * 0.5 - 0.25 -def _reference(q, k, v, g, beta, w, cu, scale): - """fp64 per-token recurrence over the packed batch. Returns (o, final_state).""" +def _beta_guard(kn, beta, alpha, io_dtype): + """fp64 mirror of the kernel beta guard: kn l2-normalized, alpha = exp(g).""" + w = kn * kn + n = w.sum(-1) + a = (beta * w).sum(-1) + nu = (beta * beta * w).sum(-1) + r2 = (n * nu - a * a).clamp_min(0.0) + inv_c2 = alpha.amax(-1).pow(2) + c2 = 1.0 / inv_c2 + r2_crit = ((c2 - 1.0) * (1.0 - (1.0 - a).pow(2) * inv_c2)).clamp_min(0.0) + unsafe = (n > 1.0e-20) & (r2 > r2_crit) + mu = a / n.clamp_min(1.0e-20) + eta = ((1.0 - 1.0 / 32) * r2_crit / r2.clamp_min(1.0e-30)).sqrt().clamp(0.0, 1.0) + cand = torch.where(unsafe[..., None], mu[..., None] + eta[..., None] * (beta - mu[..., None]), beta).to(io_dtype).double() + a_q = (cand * w).sum(-1) + nu_q = (cand * cand * w).sum(-1) + r2_q = (n * nu_q - a_q * a_q).clamp_min(0.0) + r2_crit_q = ((c2 - 1.0) * (1.0 - (1.0 - a_q).pow(2) * inv_c2)).clamp_min(0.0) + tol = 4.0 * torch.finfo(io_dtype).eps * (n * nu_q + a_q * a_q) + fallback = unsafe & (r2_q > r2_crit_q + tol) + mu_q = (a_q / n.clamp_min(1.0e-20)).to(io_dtype).double() + return torch.where(fallback[..., None], mu_q[..., None], cand) + + +def _reference(q, k, v, g, beta, w, a_log, dt_bias, cu, scale): + """fp64 per-token recurrence over the packed batch, with the safe-gate + activation, in-kernel L2 norm, and beta guard applied. Returns + (o, final_state).""" total, H, D = q.shape V = v.shape[2] q, k, v, g, beta, w = (x.double() for x in (q, k, v, g, beta, w)) + g = -5.0 * torch.sigmoid(a_log.double().exp()[None, :, None] * (g + dt_bias.double()[None])) + q = torch.nn.functional.normalize(q, dim=-1) + k = torch.nn.functional.normalize(k, dim=-1) + beta = _beta_guard(k, beta, g.exp(), torch.bfloat16) o = torch.zeros(total, H, V, dtype=torch.float64, device=q.device) fs = torch.zeros(cu.numel() - 1, H, D, V, dtype=torch.float64, device=q.device) for n in range(cu.numel() - 1): @@ -59,7 +99,7 @@ def _reference(q, k, v, g, beta, w, cu, scale): v_new = w[t] * v[t] - erase S = S + torch.einsum("hd,hv->hdv", k[t], v_new) o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) - fs[n] = S + fs[n] = S.transpose(-2, -1) return o, fs @@ -69,10 +109,14 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: total, num_seqs = sum(seq_lens), len(seq_lens) scale = 1.0 / math.sqrt(D) - q = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() - k = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() + q = _randu(total * H, D, device).reshape(total, H, D).bfloat16() + k = _randu(total * H, D, device).reshape(total, H, D).bfloat16() v = _randu(total * H, D, device).reshape(total, H, D).bfloat16() - gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + gate = -2.5 + 0.7 * torch.randn(total, H, D, device=device) + gate[1::2, :, :4] = -8.0 + gate = gate.contiguous() + a_log = torch.zeros(H, device=device, dtype=torch.float32) + dt_bias = torch.zeros(H, D, device=device, dtype=torch.float32) beta = (torch.rand(total, H, D, device=device).sigmoid() * 2.0).bfloat16().contiguous() w = torch.rand(total, H, D, device=device).sigmoid().bfloat16().contiguous() cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) @@ -85,6 +129,8 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: beta_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="beta") w_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="w") cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + a_log_t = g.tensor([H], data_type=cudnn.data_type.FLOAT, name="a_log") + dt_bias_t = g.tensor([H, D], data_type=cudnn.data_type.FLOAT, name="dt_bias") O_t, fs_t, _h_t = g.gdn2( q=q_t, k=k_t, @@ -93,9 +139,13 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: beta=beta_t, w=w_t, cu_seqlens=cu_t, + a_log=a_log_t, + dt_bias=dt_bias_t, scale=scale, output_final_state=True, - use_qk_l2norm=False, + use_qk_l2norm=True, + safe_gate=True, + beta_guard=True, name="gdn2", ) O_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) @@ -104,16 +154,16 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: o = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device=device) - pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, O_t: o, fs_t: fs} + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, a_log_t: a_log, dt_bias_t: dt_bias, O_t: o, fs_t: fs} g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) torch.cuda.synchronize() - o_ref, fs_ref = _reference(q, k, v, gate, beta, w, cu, scale) + o_ref, fs_ref = _reference(q, k, v, gate, beta, w, a_log, dt_bias, cu, scale) r_o = _rms_ratio(o, o_ref) assert r_o < 5e-2, f"o rms ratio {r_o:.4g}" r_s = _rms_ratio(fs, fs_ref) assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - print(f"[05] PASS gdn2 prefill seq_lens={list(seq_lens)} H={H} D={D} (fs rms {r_s:.2e})") + print(f"[05] PASS gdn2 prefill (safe gate + beta guard) seq_lens={list(seq_lens)} H={H} D={D} (fs rms {r_s:.2e})") if __name__ == "__main__": diff --git a/test/python/linear_attention/frost/examples/06_gdn2_backward.py b/test/python/linear_attention/frost/examples/06_gdn2_backward.py index 8bfb6a7d4..6cf85967d 100644 --- a/test/python/linear_attention/frost/examples/06_gdn2_backward.py +++ b/test/python/linear_attention/frost/examples/06_gdn2_backward.py @@ -6,8 +6,18 @@ The GDN2_BWD node takes the forward inputs plus ``dO`` and returns ``(dQ, dK, dV, dG, dBeta, dW)`` (per-key-channel ``dG``/``dBeta``, per-value ``dW``; ``dBeta``/``dW`` in io dtype). Without the optional per-chunk ``h`` -input the engine recomputes the forward state pass internally. Gradients are -checked against fp64 autograd through the per-token recurrence. +input the engine recomputes the forward state pass internally. This example +runs the in-kernel q/k L2 norm, the bounded safe gate, and the erase-side +beta safeguard (``beta_guard``). The decay input is a raw pre-activation; +with ``safe_gate=True`` the kernel computes the bounded gate + + g = -5 * sigmoid(exp(A_log) * (a + dt_bias)) + +here with ``A_log = 0`` and ``dt_bias = 0``, ``dG`` is the raw-logit +gradient, and ``d_a_log`` / ``d_dt_bias`` are also produced. The guard is +straight-through, so the reference applies the projection detached (identity +gradient to ``beta``, none to the gate). Gradients are checked against fp64 +autograd through the per-token recurrence. """ from __future__ import annotations @@ -23,7 +33,7 @@ def _build_plans(g) -> None: g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - g.select_plan(names.index("gdn2_frost")) # pin the FROST entry + g.select_plan(names.index("gdn2_frost")) g.check_support() g.build_plans() @@ -33,10 +43,38 @@ def _rms_ratio(out, ref): return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() +def _beta_guard(kn, beta, alpha, io_dtype): + """fp64 mirror of the kernel beta guard: kn l2-normalized, alpha = exp(g).""" + w = kn * kn + n = w.sum(-1) + a = (beta * w).sum(-1) + nu = (beta * beta * w).sum(-1) + r2 = (n * nu - a * a).clamp_min(0.0) + inv_c2 = alpha.amax(-1).pow(2) + c2 = 1.0 / inv_c2 + r2_crit = ((c2 - 1.0) * (1.0 - (1.0 - a).pow(2) * inv_c2)).clamp_min(0.0) + unsafe = (n > 1.0e-20) & (r2 > r2_crit) + mu = a / n.clamp_min(1.0e-20) + eta = ((1.0 - 1.0 / 32) * r2_crit / r2.clamp_min(1.0e-30)).sqrt().clamp(0.0, 1.0) + cand = torch.where(unsafe[..., None], mu[..., None] + eta[..., None] * (beta - mu[..., None]), beta).to(io_dtype).double() + a_q = (cand * w).sum(-1) + nu_q = (cand * cand * w).sum(-1) + r2_q = (n * nu_q - a_q * a_q).clamp_min(0.0) + r2_crit_q = ((c2 - 1.0) * (1.0 - (1.0 - a_q).pow(2) * inv_c2)).clamp_min(0.0) + tol = 4.0 * torch.finfo(io_dtype).eps * (n * nu_q + a_q * a_q) + fallback = unsafe & (r2_q > r2_crit_q + tol) + mu_q = (a_q / n.clamp_min(1.0e-20)).to(io_dtype).double() + return torch.where(fallback[..., None], mu_q[..., None], cand) + + def _reference_o(q, k, v, g, beta, w, cu, scale): - """Differentiable fp64 per-token recurrence; returns o.""" + """Differentiable fp64 per-token recurrence with the in-kernel L2 norm and + the straight-through beta guard; returns o.""" total, H, D = q.shape V = v.shape[2] + q = torch.nn.functional.normalize(q, dim=-1) + k = torch.nn.functional.normalize(k, dim=-1) + beta = beta + (_beta_guard(k.detach(), beta.detach(), g.exp().detach(), torch.bfloat16) - beta.detach()) outs = [] for n in range(cu.numel() - 1): S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) @@ -55,10 +93,14 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: total, num_seqs = sum(seq_lens), len(seq_lens) scale = 1.0 / math.sqrt(D) - q = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() - k = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + q = torch.randn(total, H, D, device=device).bfloat16() + k = torch.randn(total, H, D, device=device).bfloat16() v = torch.randn(total, H, D, device=device).bfloat16() - gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + gate = -2.5 + 0.7 * torch.randn(total, H, D, device=device) + gate[1::2, :, :4] = -8.0 + gate = gate.contiguous() + a_log = torch.zeros(H, device=device, dtype=torch.float32) + dt_bias = torch.zeros(H, D, device=device, dtype=torch.float32) beta = (torch.rand(total, H, D, device=device).sigmoid() * 2.0).bfloat16().contiguous() w = torch.rand(total, H, D, device=device).sigmoid().bfloat16().contiguous() do = torch.randn(total, H, D, device=device).bfloat16() @@ -73,7 +115,9 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: w_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="w") cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - outs = g.gdn2_bwd( + a_log_t = g.tensor([H], data_type=cudnn.data_type.FLOAT, name="a_log") + dt_bias_t = g.tensor([H, D], data_type=cudnn.data_type.FLOAT, name="dt_bias") + dq_t, dk_t, dv_t, dg_t, db_t, dw_t, _dstate0_t, da_log_t, ddt_bias_t = g.gdn2_bwd( q=q_t, k=k_t, v=v_t, @@ -82,11 +126,19 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: w=w_t, cu_seqlens=cu_t, dO=do_t, + a_log=a_log_t, + dt_bias=dt_bias_t, scale=scale, + use_qk_l2norm=True, + safe_gate=True, + beta_guard=True, name="gdn2_bwd", ) - dtypes = (cudnn.data_type.BFLOAT16,) * 3 + (cudnn.data_type.FLOAT, cudnn.data_type.BFLOAT16, cudnn.data_type.BFLOAT16) - grads_t = [out.set_output(True).set_data_type(dt) for out, dt in zip(outs, dtypes)] + io_dt, f32_dt = cudnn.data_type.BFLOAT16, cudnn.data_type.FLOAT + grads_t = [ + out.set_output(True).set_data_type(dt) + for out, dt in ((dq_t, io_dt), (dk_t, io_dt), (dv_t, io_dt), (dg_t, f32_dt), (db_t, io_dt), (dw_t, io_dt), (da_log_t, f32_dt), (ddt_bias_t, f32_dt)) + ] _build_plans(g) dq = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) @@ -95,14 +147,19 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: dg = torch.empty(total, H, D, dtype=torch.float32, device=device) db = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) dw = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) - pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, do_t: do} - pack.update(dict(zip(grads_t, (dq, dk, dv, dg, db, dw)))) + d_a_log = torch.empty(H, dtype=torch.float32, device=device) + d_dt_bias = torch.empty(H, D, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, do_t: do, a_log_t: a_log, dt_bias_t: dt_bias} + pack.update(dict(zip(grads_t, (dq, dk, dv, dg, db, dw, d_a_log, d_dt_bias)))) g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) torch.cuda.synchronize() leaves = [x.double().requires_grad_(True) for x in (q, k, v, gate, beta, w)] - o_ref = _reference_o(*leaves, cu, scale) - grads = torch.autograd.grad((o_ref * do.double()).sum(), leaves) + a_leaf = a_log.double().requires_grad_(True) + dt_leaf = dt_bias.double().requires_grad_(True) + gact = -5.0 * torch.sigmoid(a_leaf.exp()[None, :, None] * (leaves[3] + dt_leaf[None])) + o_ref = _reference_o(leaves[0], leaves[1], leaves[2], gact, leaves[4], leaves[5], cu, scale) + grads = torch.autograd.grad((o_ref * do.double()).sum(), leaves + [a_leaf, dt_leaf]) for name, out, ref in ( ("dQ", dq, grads[0]), ("dK", dk, grads[1]), @@ -110,10 +167,12 @@ def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: ("dG", dg, grads[3]), ("dBeta", db, grads[4]), ("dW", dw, grads[5]), + ("d_a_log", d_a_log, grads[6]), + ("d_dt_bias", d_dt_bias, grads[7]), ): r = _rms_ratio(out, ref) assert r < 5e-2, f"{name} rms ratio {r:.4g}" - print(f"[06] PASS gdn2 backward (recompute) seq_lens={list(seq_lens)} H={H} D={D}") + print(f"[06] PASS gdn2 backward (safe gate + beta guard) seq_lens={list(seq_lens)} H={H} D={D}") if __name__ == "__main__": diff --git a/test/python/linear_attention/reference_gdn2.py b/test/python/linear_attention/reference_gdn2.py index 2335c4bc8..ecbb72f2c 100644 --- a/test/python/linear_attention/reference_gdn2.py +++ b/test/python/linear_attention/reference_gdn2.py @@ -67,6 +67,42 @@ def recurrent_dense(q, k, v, alpha, beta, w, state0): return torch.stack(outs, dim=2), state +def beta_guard_reference( + k: torch.Tensor, + beta: torch.Tensor, + alpha: torch.Tensor, + io_dtype: torch.dtype, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """fp64 mirror of the kernel beta guard (``frost/common/beta_guard.py``). + + ``k`` l2-normalized, ``alpha = exp(g)`` per token; all ``[B, T, HO, K]``. + Returns ``(beta_eff, unsafe, fallback)``; the decision masks let tests + assert the sensor actually fires and compare against other guard + implementations.""" + weight = k * k + n = weight.sum(-1) + a = (beta * weight).sum(-1) + nu = (beta * beta * weight).sum(-1) + r2 = (n * nu - a * a).clamp_min(0.0) + inv_c2 = alpha.amax(-1).pow(2) + c2 = 1.0 / inv_c2 + r2_crit = ((c2 - 1.0) * (1.0 - (1.0 - a).pow(2) * inv_c2)).clamp_min(0.0) + unsafe = (n > 1.0e-20) & (r2 > r2_crit) + mu = a / n.clamp_min(1.0e-20) + eta = ((1.0 - 1.0 / 32) * r2_crit / r2.clamp_min(1.0e-30)).sqrt().clamp(0.0, 1.0) + projected = mu[..., None] + eta[..., None] * (beta - mu[..., None]) + candidate_q = torch.where(unsafe[..., None], projected, beta).to(io_dtype).double() + a_q = (candidate_q * weight).sum(-1) + nu_q = (candidate_q * candidate_q * weight).sum(-1) + r2_q = (n * nu_q - a_q * a_q).clamp_min(0.0) + r2_crit_q = ((c2 - 1.0) * (1.0 - (1.0 - a_q).pow(2) * inv_c2)).clamp_min(0.0) + quant_tol = 4.0 * torch.finfo(io_dtype).eps * (n * nu_q + a_q * a_q) + fallback = unsafe & (r2_q > r2_crit_q + quant_tol) + mu_q = (a_q / n.clamp_min(1.0e-20)).to(io_dtype).double() + beta_eff = torch.where(fallback[..., None], mu_q[..., None], candidate_q) + return beta_eff, unsafe, fallback + + def gdn2_reference( q: torch.Tensor, k: torch.Tensor, @@ -83,6 +119,7 @@ def gdn2_reference( a_log: Optional[torch.Tensor] = None, dt_bias: Optional[torch.Tensor] = None, use_beta_sigmoid: bool = False, + beta_guard: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """GDN-2 reference. @@ -99,6 +136,9 @@ def gdn2_reference( (differentiable; a_log ``[Hg]``, dt_bias ``[Hg, K]``). gate_lower_bound: safe-gate lower bound in log space (default -5.0). use_beta_sigmoid: treat ``beta`` as raw logits; apply ``sigmoid``. + beta_guard: apply the erase-side beta safeguard (straight-through: + the projection is detached, gradients flow as identity to + ``beta`` and not at all to ``g``). Requires l2-normalized ``k``. Returns: ``(o, final_state)`` in fp64: o ``[B, T, HO, V]``, final_state @@ -136,6 +176,10 @@ def gdn2_reference( if w.shape[2] != HO: wf = wf.unsqueeze(3).expand(-1, -1, -1, HO // w.shape[2], -1).reshape(w.shape[0], w.shape[1], HO, -1) + if beta_guard: + beta_eff, _, _ = beta_guard_reference(kf.detach(), betaf.detach(), alphaf.detach(), beta.dtype) + betaf = betaf + (beta_eff - betaf.detach()) + # [B, T, HO, *] -> [B, HO, T, *] qf = qf.permute(0, 2, 1, 3) kf = kf.permute(0, 2, 1, 3) diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py index 0e7984b1e..87002b6e6 100644 --- a/test/python/linear_attention/test_la.py +++ b/test/python/linear_attention/test_la.py @@ -33,7 +33,7 @@ from .conftest import gen_qkv # noqa: E402 from .reference_gdn import gdn_reference, rms_ratio # noqa: E402 -from .reference_gdn2 import gdn2_reference # noqa: E402 +from .reference_gdn2 import beta_guard_reference, gdn2_reference # noqa: E402 from .reference_kda import kda_reference # noqa: E402 pytestmark = [ @@ -254,7 +254,7 @@ def run_fwd(backend, case, *, cu=None, **kw): return pinned_op(backend, case.variant)(*op_args(case, cu=cu), **kw) -def reference(case, *, scale=None, initial_state=None, l2norm=False, cu=None): +def reference(case, *, scale=None, initial_state=None, l2norm=False, cu=None, beta_guard=False): fn = {"gdn": gdn_reference, "kda": kda_reference, "gdn2": gdn2_reference}[case.variant] q, k = case.q, case.k if l2norm: @@ -264,6 +264,8 @@ def reference(case, *, scale=None, initial_state=None, l2norm=False, cu=None): if case.variant == "gdn2": args.append(case.gates["w"]) kwargs = dict(scale=scale, initial_state=initial_state) + if beta_guard: + kwargs["beta_guard"] = True if case.varlen or cu is not None: kwargs["cu_seqlens"] = case.cu if cu is None else cu with torch.no_grad(): @@ -277,13 +279,16 @@ def assert_rms_close(name, out, want, tol): assert r < tol, f"{name} rms ratio {r:.4g} >= {tol}" -def assert_fwd_parity(backend, case, *, scale=None, use_initial_state=False, l2norm=False, seed=SEED + 1): +def assert_fwd_parity(backend, case, *, scale=None, use_initial_state=False, l2norm=False, beta_guard=False, seed=SEED + 1): set_seed(seed) state0 = None if use_initial_state: state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 - o, fs = run_fwd(backend, case, scale=scale, initial_state=state0, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) - o_ref, fs_ref = reference(case, scale=scale, initial_state=state0, l2norm=l2norm) + op_kw = dict(scale=scale, initial_state=state0, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) + if beta_guard: + op_kw["beta_guard"] = True + o, fs = run_fwd(backend, case, **op_kw) + o_ref, fs_ref = reference(case, scale=scale, initial_state=state0, l2norm=l2norm, beta_guard=beta_guard) assert_rms_close("o", o, o_ref, FWD_TOL[case.dtype]) if fs is not None and fs.numel(): assert_rms_close("final_state", fs, fs_ref, STATE_TOL[case.dtype]) @@ -525,7 +530,7 @@ def test_fwd_output_contract(backend, variant): # --------------------------------------------------------------------------- -def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use_dfs=False, l2norm=False, gate_grad_tol=None, seed=SEED + 1): +def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use_dfs=False, l2norm=False, beta_guard=False, gate_grad_tol=None, seed=SEED + 1): variant, tol = case.variant, BWD_TOL[case.dtype] tensors = {"q": case.q, "k": case.k, "v": case.v, "g": case.gates["g"], "beta": case.gates["beta"]} if variant == "gdn2": @@ -544,7 +549,10 @@ def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use if variant == "gdn2": args.append(op_leaves["w"]) args.append(case.cu) - o, fs = pinned_op(backend, variant)(*args, scale=scale, initial_state=state0_op, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) + op_kw = dict(scale=scale, initial_state=state0_op, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) + if beta_guard: + op_kw["beta_guard"] = True + o, fs = pinned_op(backend, variant)(*args, **op_kw) dO = torch.randn_like(o) outputs, grad_outputs = [o], [dO] dFS = None @@ -563,6 +571,8 @@ def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use if variant == "gdn2": ref_args.append(ref_leaves["w"]) ref_kwargs = dict(scale=scale, initial_state=state0_ref) + if beta_guard: + ref_kwargs["beta_guard"] = True if case.varlen: ref_kwargs["cu_seqlens"] = case.cu o_ref, fs_ref = ref_fn(*ref_args, **ref_kwargs) @@ -1052,6 +1062,164 @@ def dbeta(beta, **kw): assert (got - ident).abs().max().item() / scale < 2e-2 +# --------------------------------------------------------------------------- +# GDN-2 beta guard (erase-side safeguard) +# --------------------------------------------------------------------------- + + +def beta_guard_trip_fraction(case): + """Reference sensor trip/fallback fractions on a case with H == HV == HO + (no head expansion) and log-space gates (no safe_gate).""" + kn = F.normalize(case.k.float(), dim=-1).double() + _, unsafe, fallback = beta_guard_reference(kn, case.gates["beta"].double(), case.gates["g"].double().exp(), case.dtype) + return unsafe.double().mean().item(), fallback.double().mean().item() + + +def test_beta_guard_fwd(backend): + """Guard on: parity vs the fp64 guarded reference; the sensor must + actually fire on this data or the parity is vacuous.""" + case = make_case("gdn2", torch.bfloat16, T=256) + trip, _ = beta_guard_trip_fraction(case) + assert trip > 0.01, f"beta guard sensor never fires on this case (trip={trip:.4f})" + assert_fwd_parity(backend, case, l2norm=True, beta_guard=True) + + +def test_beta_guard_fwd_mixed_headroom(backend): + """Tokens with real decay headroom must pass through untouched next to + tripping tokens (exercises the safe path and the per-token gate recovery + at chunk rows 0 and interior rows).""" + case = make_case("gdn2", torch.bfloat16, T=256) + g = case.gates["g"].clone() + g[:, ::2] += math.log(0.5) + case = case.clone(gates=dict(case.gates, g=g)) + trip, _ = beta_guard_trip_fraction(case) + assert 0.01 < trip < 0.99, f"want a mixed safe/unsafe population, got trip={trip:.4f}" + assert_fwd_parity(backend, case, l2norm=True, beta_guard=True) + + +@pytest.mark.parametrize("seq_lens", [[64, 192], [31, 63, 93, 123]], ids=["two", "ragged"]) +def test_beta_guard_fwd_varlen(backend, seq_lens): + assert_fwd_parity(backend, make_case("gdn2", torch.bfloat16, seq_lens=seq_lens), l2norm=True, beta_guard=True) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=DTYPE_IDS.get) +def test_beta_guard_bwd(backend, dtype): + assert_bwd_parity(backend, make_case("gdn2", dtype, T=128), l2norm=True, beta_guard=True) + + +def test_beta_guard_bwd_varlen(backend): + assert_bwd_parity(backend, make_case("gdn2", torch.bfloat16, seq_lens=[31, 63, 93, 123]), l2norm=True, beta_guard=True) + + +def test_beta_guard_bwd_initial_state(backend): + assert_bwd_parity(backend, make_case("gdn2", torch.bfloat16, T=128), l2norm=True, beta_guard=True, use_initial_state=True) + + +def test_beta_guard_recompute_matches_checkpoints(backend): + """Prefill (checkpoint dump) and recompute must apply the same guard: the + gradient gap between the checkpoint-reuse and recompute backward paths + with the guard on must stay at the scale of the guard-off gap.""" + case = make_case("gdn2", torch.bfloat16, T=256) + tensors = (case.q, case.k, case.v, case.gates["g"], case.gates["beta"], case.gates["w"]) + + def path_gap(beta_guard): + grads = {} + dO = None + for ckpt in (16, 0): + leaves = [to_thd(t).detach().clone().requires_grad_(True) for t in tensors] + kw = dict(use_qk_l2norm_in_kernel=True, checkpoint_every_n_tokens=ckpt) + if beta_guard: + kw["beta_guard"] = True + with waive_unsupported(backend, "gdn2"): + out = pinned_op(backend, "gdn2")(*leaves, case.cu, **kw) + o = out[0] + if dO is None: + set_seed(SEED + 23) + dO = torch.randn_like(o) + grads[ckpt] = torch.autograd.grad([o], leaves, [dO]) + return max(rms_ratio(a, b.float()) for a, b in zip(grads[16], grads[0])) + + assert path_gap(True) <= max(4.0 * path_gap(False), 1.0e-3) + + +def test_beta_guard_with_sigmoid_fwd(backend): + """Guard on top of the in-kernel sigmoid: io-dtype logits must match the + post-activation io beta path.""" + case = make_case("gdn2", torch.bfloat16, T=256) + set_seed(SEED + 17) + braw = torch.randn_like(case.gates["beta"].float()).to(case.dtype) + raw_case = case.clone(gates=dict(case.gates, beta=braw)) + eff_case = case.clone(gates=dict(case.gates, beta=torch.sigmoid(braw.float()).to(case.dtype))) + kw = dict(output_final_state=True, use_qk_l2norm_in_kernel=True, beta_guard=True) + o_raw, fs_raw = run_fwd(backend, raw_case, use_beta_sigmoid_in_kernel=True, **kw) + o_eff, fs_eff = run_fwd(backend, eff_case, **kw) + assert_rms_close("o", o_raw, o_eff.double(), 2e-2) + assert rms_ratio(fs_raw, fs_eff) < 2e-2 + + +def test_beta_guard_with_sigmoid_backward(backend): + """Straight-through under the in-kernel sigmoid: dbeta wrt the logits must + equal the post-activation path's dbeta times s*(1-s) at the io-rounded s + (the Jacobian the kernel reads back from the original logits).""" + case = make_case("gdn2", torch.bfloat16, T=256) + set_seed(SEED + 19) + braw = torch.randn_like(case.gates["beta"].float()).to(case.dtype) + s_io = torch.sigmoid(braw.float()).to(case.dtype) + + def dbeta(beta, **kw): + leaf = to_thd(beta).detach().clone().requires_grad_(True) + args = [to_thd(case.q), to_thd(case.k), to_thd(case.v), to_thd(case.gates["g"]), leaf, to_thd(case.gates["w"])] + with waive_unsupported(backend, "gdn2"): + o, _ = pinned_op(backend, "gdn2")(*args, case.cu, use_qk_l2norm_in_kernel=True, beta_guard=True, **kw) + o.sum().backward() + return leaf.grad.double() + + got = dbeta(braw, use_beta_sigmoid_in_kernel=True) + s = to_thd(s_io).double() + ident = dbeta(s_io) * s * (1 - s) + scale = ident.abs().max().item() + assert scale > 1e-3, "dbeta is ~0, the comparison would be vacuous" + assert (got - ident).abs().max().item() / scale < 2e-2 + + +def test_beta_guard_multi_tile(backend): + """B*H well above the SM count with the guard on: several (b, h) tiles per + CTA exercise the moved beta/q stage releases across tile boundaries.""" + case = make_case("gdn2", torch.bfloat16, B=8, T=192, H=64) + assert_fwd_parity(backend, case, l2norm=True, beta_guard=True) + + +def test_beta_guard_bwd_determinism(backend): + """Guard-on backward must stay bitwise repeatable (the FROST determinism + contract) across the relocated mb_beta_done / mb_q_done releases.""" + case = make_case("gdn2", torch.bfloat16, seq_lens=[96, 32, 160, 1]) + tensors = (case.q, case.k, case.v, case.gates["g"], case.gates["beta"], case.gates["w"]) + dO = None + baseline = None + for _ in range(4): + leaves = [to_thd(t).detach().clone().requires_grad_(True) for t in tensors] + with waive_unsupported(backend, "gdn2"): + o, _ = pinned_op(backend, "gdn2")(*leaves, case.cu, use_qk_l2norm_in_kernel=True, beta_guard=True) + if dO is None: + set_seed(SEED + 29) + dO = torch.randn_like(o) + grads = torch.autograd.grad([o], leaves, [dO]) + if baseline is None: + baseline = grads + else: + for name, a, b in zip(("q", "k", "v", "g", "beta", "w"), baseline, grads): + assert torch.equal(a, b), f"d{name} not bitwise repeatable under beta_guard" + + +def test_beta_guard_requires_l2norm(backend): + """The engine must decline beta_guard without the in-kernel l2 norm.""" + case = make_case("gdn2", torch.bfloat16, T=64) + if not backend.engines["gdn2"]: + pytest.skip(f"the {backend.name} backend has no gdn2 engine") + with pytest.raises(cudnn.cudnnGraphNotSupportedError): + pinned_op(backend, "gdn2")(*op_args(case), beta_guard=True) + + @pytest.mark.parametrize("H", (40, 160)) def test_scalar_gate_head_tiling(backend, H): """GDN's scalar gate-parameter reduction tiles heads; the dA_log /