From aa09a9d81198c0838d878bbbc314ca72bf11ce1a Mon Sep 17 00:00:00 2001 From: Rifky Bujana Bisri Date: Tue, 16 Jun 2026 16:38:38 +0700 Subject: [PATCH 1/3] feat(gdn): Hopper fused-recurrent decode + SGLang verify kernels (TileLang) Gemm-free Gated Delta Rule fused-recurrent kernels for NVIDIA Hopper (SM90), complementing the chunked-prefill kernels with the single-step / draft-verify regime used by speculative decoding. Drop-in faster alternative to the FLA Triton recurrent/verify kernels. - Decode (recurrent_gated_delta_rule / fused_recurrent_gdr_fwd): state kept [block_DV, DK] fp32 in registers; the two GEMVs via T.reduce_sum over K and the rank-1 via a T.Parallel outer product (no tensor-core M-padding). GQA, ragged seqlens, optional initial/final state. block_DV=128 on the final-state path coalesces the K-major transposed store (~2x at all batch sizes). - Verify (recurrent_gated_delta_rule_verify): paged V-major bf16 state-pool gather/scatter (state_indices, -1=skip), per-token intermediate states, no-commit, varlen cu_seqlens, host- and in-kernel fused gating (g=-exp(A_log)*softplus(a+dt_bias), beta=sigmoid(b), qk-l2norm), CUDA-graph safe. A gating+l2norm dedup pre-pass (regime-gated) gives +8-24% at large batch (1.13-1.18x vs FLA at draft length T=12). Requires SM90+, CUDA 12.8+, PyTorch 2.8+, tilelang==0.1.8; dispatch hard-gates on cc=="9.0". Co-Authored-By: Claude Opus 4.8 (1M context) --- flash_qla/__init__.py | 10 + flash_qla/ops/gated_delta_rule/__init__.py | 10 +- .../fused_recurrent/__init__.py | 163 ++++++ .../fused_recurrent/hopper/__init__.py | 5 + .../hopper/fused_recurrent_fwd.py | 357 ++++++++++++ .../hopper/fused_recurrent_verify.py | 538 ++++++++++++++++++ 6 files changed, 1082 insertions(+), 1 deletion(-) create mode 100644 flash_qla/ops/gated_delta_rule/fused_recurrent/__init__.py create mode 100644 flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/__init__.py create mode 100644 flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_fwd.py create mode 100644 flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_verify.py diff --git a/flash_qla/__init__.py b/flash_qla/__init__.py index 640f684..cf9061b 100644 --- a/flash_qla/__init__.py +++ b/flash_qla/__init__.py @@ -8,9 +8,19 @@ chunk_gated_delta_rule_bwd, chunk_gated_delta_rule, ) +from flash_qla.ops.gated_delta_rule.fused_recurrent import ( + fused_recurrent_gdr_fwd, + recurrent_gated_delta_rule, + fused_recurrent_gdr_verify_fwd, + recurrent_gated_delta_rule_verify, +) __all__ = [ "chunk_gated_delta_rule_fwd", "chunk_gated_delta_rule_bwd", "chunk_gated_delta_rule", + "fused_recurrent_gdr_fwd", + "recurrent_gated_delta_rule", + "fused_recurrent_gdr_verify_fwd", + "recurrent_gated_delta_rule_verify", ] diff --git a/flash_qla/ops/gated_delta_rule/__init__.py b/flash_qla/ops/gated_delta_rule/__init__.py index ea07eeb..91a2aa5 100644 --- a/flash_qla/ops/gated_delta_rule/__init__.py +++ b/flash_qla/ops/gated_delta_rule/__init__.py @@ -2,6 +2,14 @@ # Licensed under The MIT License [see LICENSE for details] from .chunk import chunk_gated_delta_rule +from .fused_recurrent import ( + recurrent_gated_delta_rule, + recurrent_gated_delta_rule_verify, +) -__all__ = ["chunk_gated_delta_rule"] +__all__ = [ + "chunk_gated_delta_rule", + "recurrent_gated_delta_rule", + "recurrent_gated_delta_rule_verify", +] diff --git a/flash_qla/ops/gated_delta_rule/fused_recurrent/__init__.py b/flash_qla/ops/gated_delta_rule/fused_recurrent/__init__.py new file mode 100644 index 0000000..0dd5a98 --- /dev/null +++ b/flash_qla/ops/gated_delta_rule/fused_recurrent/__init__.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] +import torch +import torch.nn.functional as F +import tilelang + +from flash_qla.utils import l2norm + +if tilelang.contrib.nvcc.get_target_compute_version() == "9.0": + from .hopper import fused_recurrent_gdr_fwd # noqa: F401 + from .hopper.fused_recurrent_verify import ( # noqa: F401 + fused_recurrent_gdr_verify_fwd, + fused_recurrent_gdr_verify_gated_fwd, + fused_recurrent_gdr_verify_prepass, + get_prepass_scratch, + should_use_prepass, + ) +else: + raise ValueError("FlashQLA now support sm90 only.") + +__all__ = [ + "fused_recurrent_gdr_fwd", + "recurrent_gated_delta_rule", + "fused_recurrent_gdr_verify_fwd", + "fused_recurrent_gdr_verify_gated_fwd", + "recurrent_gated_delta_rule_verify", +] + + +def recurrent_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=None, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + seqlens=None, + head_first=False, + head_batch=None, +): + assert q.dtype == k.dtype == v.dtype and q.dtype != torch.float32 + assert not head_first, "head_first=True is not supported." + assert v.shape[2] % k.shape[2] == 0 and q.shape[-1] == v.shape[-1] == 128 + if scale is None: + scale = k.shape[-1] ** -0.5 + if use_qk_l2norm_in_kernel: + q = l2norm(q) + k = l2norm(k) + o, final_state = fused_recurrent_gdr_fwd( + q, + k, + v, + g, + beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + seqlens=seqlens, + head_batch=head_batch, + ) + return o.to(q.dtype), final_state + + +def gdn_sigmoid_gate(A_log, a, dt_bias, b, allow_neg_eigval=False): + """Host-side GDN gating (the sigmoid_gating family): g = -exp(A_log)*softplus(a+dt_bias), + beta = sigmoid(b) (x2 if allow_neg_eigval). A_log,dt_bias:[H]; a,b:[...,H]. Returns fp32.""" + g = -torch.exp(A_log.float())[(None,) * (a.dim() - 1)] * F.softplus( + a.float() + dt_bias.float()[(None,) * (a.dim() - 1)] + ) + beta = torch.sigmoid(b.float()) + if allow_neg_eigval: + beta = beta * 2 + return g, beta + + +def recurrent_gated_delta_rule_verify( + A_log, + a, + dt_bias, + q, + k, + v, + b, + ssm_states, + cache_indices, + query_start_loc, + intermediate_states_buffer, + intermediate_state_indices, + cache_steps=None, + o=None, + scale=None, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + allow_neg_eigval=False, + fuse_gating=False, + prepass=None, # H1 dedup gating+l2norm pre-pass: None=auto (regime-gated), True/False=force + retrieve_parent_token=None, # accepted-and-IGNORED (DFlash width-1; tree path not built) +): + """High-level SGLang DFlash verify entry. q,k:[1,T,Hk,128] v:[1,T,Hv,128]; a,b:[1,T,Hv]; + A_log,dt_bias:[Hv]; ssm_states pool V-major [num_slots,Hv,128,128]. + + CUDA-graph note: for capture, use ``fuse_gating=True`` (computes g/beta + qk-l2norm INSIDE + the kernel from raw a,b,A_log,dt_bias -- no PyTorch gating/l2norm, no allocation when ``o`` + is provided -> fully capture-safe). The default ``fuse_gating=False`` path computes g/beta + + qk-l2norm in PyTorch (l2norm is ``@torch.compile``'d) and allocates them; run it OUTSIDE + capture or prefer ``fuse_gating=True`` inside it. + + H1 (``fuse_gating=True``): in the large-batch / multi-draft-token regime the gating + qk-l2norm + are split into a tiny dedup PRE-PASS kernel (computed ONCE per (token,K-head) instead of once + per (token,V-head,V-tile) in the hot loop) feeding the host-gated main kernel -- measured + ~12-15% faster at T=12 large batch. Regime-gated by ``should_use_prepass`` (the single + in-kernel-gated kernel A is kept for latency-bound single-request, where a 2nd launch is a net + loss). The prepass scratch is a persistent, never-evicting cache -> capture-safe after warmup. + ``prepass`` forces the choice (None=auto, True=prepass+main, False=single gated kernel A).""" + assert q.dtype == k.dtype == v.dtype and q.dtype != torch.float32 + assert q.shape[-1] == v.shape[-1] == 128 and v.shape[2] % k.shape[2] == 0 + if cache_steps is not None: # static shape check (capture-safe; no value read) + assert intermediate_states_buffer.shape[1] >= cache_steps, ( + f"intermediate_states_buffer cache-steps dim {intermediate_states_buffer.shape[1]} " + f"< cache_steps {cache_steps}" + ) + scale = scale if scale is not None else q.shape[-1] ** -0.5 + if o is None: + o = torch.empty(1, q.shape[1], v.shape[2], v.shape[-1], device=q.device, dtype=q.dtype) + + if fuse_gating: + N, total_tokens, Hk = cache_indices.shape[0], q.shape[1], k.shape[2] + H = v.shape[2] + use_prepass = should_use_prepass(N, H, total_tokens) if prepass is None else prepass + if use_prepass: + # H1: dedup gating + qk-l2norm into a run-once pre-pass, then the host-gated main + # kernel (no in-loop transcendentals). Persistent scratch -> capture-safe after warmup. + q_n, k_n, g_pp, beta_pp = get_prepass_scratch(total_tokens, Hk, H, q.device, q.dtype) + fused_recurrent_gdr_verify_prepass( + q, k, a, b, A_log, dt_bias, q_n, k_n, g_pp, beta_pp, + allow_neg_eigval=allow_neg_eigval, + ) + fused_recurrent_gdr_verify_fwd( + q_n, k_n, v, g_pp, beta_pp, ssm_states, cache_indices, query_start_loc, + intermediate_states_buffer, intermediate_state_indices, o, + scale=scale, disable_state_update=disable_state_update, + ) + return o + fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, ssm_states, cache_indices, query_start_loc, + intermediate_states_buffer, intermediate_state_indices, o, + scale=scale, disable_state_update=disable_state_update, allow_neg_eigval=allow_neg_eigval, + ) + return o + + if use_qk_l2norm_in_kernel: + q = l2norm(q) + k = l2norm(k) + g, beta = gdn_sigmoid_gate(A_log, a, dt_bias, b, allow_neg_eigval) + fused_recurrent_gdr_verify_fwd( + q, k, v, g, beta, ssm_states, cache_indices, query_start_loc, + intermediate_states_buffer, intermediate_state_indices, o, + scale=scale, disable_state_update=disable_state_update, + ) + return o diff --git a/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/__init__.py b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/__init__.py new file mode 100644 index 0000000..a8fd7ca --- /dev/null +++ b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] +from .fused_recurrent_fwd import fused_recurrent_gdr_fwd + +__all__ = ["fused_recurrent_gdr_fwd"] diff --git a/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_fwd.py b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_fwd.py new file mode 100644 index 0000000..be60039 --- /dev/null +++ b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_fwd.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] +import torch +import tilelang +import tilelang.language as T + +MULTI_PROCESSOR_COUNT = torch.cuda.get_device_properties().multi_processor_count +TARGET_NUM_CTAS = int(MULTI_PROCESSOR_COUNT * 0.7) + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def tilelang_fused_recurrent_gdr_fwd( + H, + Hg, + DK, + DV, + scale, + accum_dtype, + qkva_dtype, + g_dtype, + b_dtype, + h0_dtype, + ht_dtype, + o_dtype, + seqlen_dtype, + use_initial_state, + store_final_state, + has_seqlens, + block_DV=128, + threads=128, +): + batch_size = T.dynamic("batch_size") + num_tokens = T.dynamic("num_tokens") # = q_len (D); D fixed per call + q_shape = (batch_size, num_tokens, Hg, DK) + v_shape = (batch_size, num_tokens, H, DV) + g_shape = (batch_size, num_tokens, H) + s_shape = (batch_size, H, DK, DV) + n_vt = (DV + block_DV - 1) // block_DV + + @T.prim_func + def kernel( + q: T.Tensor(q_shape, qkva_dtype), + k: T.Tensor(q_shape, qkva_dtype), + v: T.Tensor(v_shape, qkva_dtype), + g: T.Tensor(g_shape, g_dtype), + b: T.Tensor(g_shape, b_dtype), + h0: T.Tensor(s_shape, h0_dtype), + seqlens: T.Tensor([batch_size], seqlen_dtype), + o: T.Tensor(v_shape, o_dtype), + ht: T.Tensor(s_shape, ht_dtype), + ): + # Gemm-free, memory-bound decode recurrence. One CTA owns (sequence bb, V-head bh, + # V-column-tile bv). State S is kept [block_DV, DK] (V-major rows) in an fp32 fragment; + # the two GEMVs are reductions over the last (DK) dim and the rank-1 is a T.Parallel + # outer product. No tensor-core gemm -> no M-padding, no warp-partition constraint. + with T.Kernel(n_vt * batch_size * H, threads=threads) as (bbhv,): + bbh = bbhv // n_vt + bv = bbhv % n_vt + bb = bbh // H + bh = bbh % H + bhg = bh // (H // Hg) + v0 = bv * block_DV + + L = T.alloc_var("int32") + L = seqlens[bb] if has_seqlens else num_tokens + + S = T.alloc_fragment((block_DV, DK), accum_dtype) # state [V-tile, K], fp32 + prod = T.alloc_fragment((block_DV, DK), accum_dtype) + q_s = T.alloc_shared((1, DK), qkva_dtype) # 2-D staging (1-D slice copies fail layout) + k_s = T.alloc_shared((1, DK), qkva_dtype) + v_s = T.alloc_shared((1, block_DV), qkva_dtype) + o_sh = T.alloc_shared((1, block_DV), o_dtype) + kS = T.alloc_fragment((block_DV,), accum_dtype) + oo = T.alloc_fragment((block_DV,), accum_dtype) + vnew = T.alloc_fragment((block_DV,), accum_dtype) + decay = T.alloc_fragment((1,), accum_dtype) + bt = T.alloc_fragment((1,), accum_dtype) + + if use_initial_state: + # h0 is K-major [B,H,DK,DV]; load transposed into S[v, dk] = h0[bb,bh,dk,v0+v] + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] = h0[bb, bh, j_k, v0 + j_v] + else: + T.clear(S) + + for t in T.serial(L): + T.copy(q[bb, t : t + 1, bhg, 0:DK], q_s) + T.copy(k[bb, t : t + 1, bhg, 0:DK], k_s) + T.copy(v[bb, t : t + 1, bh, v0 : v0 + block_DV], v_s) + decay[0] = T.exp2(g[bb, t, bh] * 1.442695) # raw g, exp2 + bt[0] = b[bb, t, bh] + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] *= decay[0] + # kS[v] = sum_dk k[dk] * S[v, dk] + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = k_s[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, kS, dim=1) + # v_new[v] = beta * (v[v] - kS[v]) + for j_v in T.Parallel(block_DV): + vnew[j_v] = bt[0] * (v_s[0, j_v] - kS[j_v]) + # rank-1: S[v, dk] += k[dk] * v_new[v] + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] += k_s[0, j_k] * vnew[j_v] + # o[v] = scale * sum_dk q[dk] * S[v, dk] (post-update read) + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = q_s[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, oo, dim=1) + for j_v in T.Parallel(block_DV): + o_sh[0, j_v] = oo[j_v] * scale + T.copy(o_sh, o[bb, t : t + 1, bh, v0 : v0 + block_DV]) + + if store_final_state: + # ht is K-major [B,H,DK,DV]; store transposed ht[bb,bh,dk,v0+v] = S[v, dk] + for j_v, j_k in T.Parallel(block_DV, DK): + ht[bb, bh, j_k, v0 + j_v] = S[j_v, j_k] + + return kernel + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def tilelang_fused_recurrent_gdr_fwd_hb( + H, + Hg, + DK, + DV, + scale, + accum_dtype, + qkva_dtype, + g_dtype, + b_dtype, + h0_dtype, + ht_dtype, + o_dtype, + seqlen_dtype, + use_initial_state, + store_final_state, + has_seqlens, + block_DV=64, + threads=256, +): + """Head-batched GQA specialization of the gemm-free decode kernel. One CTA owns + (sequence bb, K/Q head-group hg, V-col-tile bv) and processes ALL grp = H//Hg V-heads + h = hg*grp + i that share K/Q head hg, loading q/k ONCE (dedup). State is row-stacked + S[grp*block_DV, DK]: row gv -> head-band i = gv//block_DV, v-row jv = gv%block_DV. + Layout-inference rule learned on H100: every [M] fragment (S/prod/decay_f/.../oo) MUST be + accessed over the FULL Parallel(M,...) range -- a partial Parallel(block_DV) write at offset + i*block_DV makes the fragment's affine map non-invertible (TVM InverseAffineIterMap check + fails). So ALL per-head divergence is routed through GLOBAL reads/writes indexed by the + derived head hg*grp + gv//block_DV and channel v0 + gv%block_DV (global needs no fragment + inversion); only q/k stage through shared (loaded once for the group). + threads = grp*128 keeps per-thread register footprint identical to the per-head kernel.""" + grp = H // Hg + M = grp * block_DV + batch_size = T.dynamic("batch_size") + num_tokens = T.dynamic("num_tokens") + q_shape = (batch_size, num_tokens, Hg, DK) + v_shape = (batch_size, num_tokens, H, DV) + g_shape = (batch_size, num_tokens, H) + s_shape = (batch_size, H, DK, DV) + n_vt = (DV + block_DV - 1) // block_DV + + @T.prim_func + def kernel( + q: T.Tensor(q_shape, qkva_dtype), + k: T.Tensor(q_shape, qkva_dtype), + v: T.Tensor(v_shape, qkva_dtype), + g: T.Tensor(g_shape, g_dtype), + b: T.Tensor(g_shape, b_dtype), + h0: T.Tensor(s_shape, h0_dtype), + seqlens: T.Tensor([batch_size], seqlen_dtype), + o: T.Tensor(v_shape, o_dtype), + ht: T.Tensor(s_shape, ht_dtype), + ): + with T.Kernel(n_vt * batch_size * Hg, threads=threads) as (bbhv,): + bbh = bbhv // n_vt # flattened (bb, hg) + bv = bbhv % n_vt + bb = bbh // Hg + hg = bbh % Hg + v0 = bv * block_DV + + L = T.alloc_var("int32") + L = seqlens[bb] if has_seqlens else num_tokens + + S = T.alloc_fragment((M, DK), accum_dtype) # row-stacked state [grp*V-tile, K] + prod = T.alloc_fragment((M, DK), accum_dtype) + q_s = T.alloc_shared((1, DK), qkva_dtype) # shared across the grp heads + k_s = T.alloc_shared((1, DK), qkva_dtype) + decay_f = T.alloc_fragment((M,), accum_dtype) # row-aligned bands (full-M access) + b_f = T.alloc_fragment((M,), accum_dtype) + kS = T.alloc_fragment((M,), accum_dtype) + oo = T.alloc_fragment((M,), accum_dtype) + vnew = T.alloc_fragment((M,), accum_dtype) + + if use_initial_state: + # full-M gather; K-major h0[B,H,DK,DV] -> S[gv, K] with derived head/channel + for gv, j_k in T.Parallel(M, DK): + S[gv, j_k] = h0[bb, hg * grp + gv // block_DV, j_k, v0 + gv % block_DV] + else: + T.clear(S) + + for t in T.serial(L): + T.copy(q[bb, t : t + 1, hg, 0:DK], q_s) # loaded ONCE for the whole group + T.copy(k[bb, t : t + 1, hg, 0:DK], k_s) + # per-band decay/beta into FULL-M fragments (global g/beta at derived head). A + # shared [grp] band read by gv//block_DV does NOT lower in TileLang (the Parallel + # layout inferencer rejects the shared derived-index read), so we materialize the + # full M; the redundant exp2 is cheap vs the FMA floor on this memory-bound kernel. + for gv in T.Parallel(M): + decay_f[gv] = T.exp2(g[bb, t, hg * grp + gv // block_DV] * 1.442695) + b_f[gv] = b[bb, t, hg * grp + gv // block_DV] + for gv, j_k in T.Parallel(M, DK): + S[gv, j_k] *= decay_f[gv] + for gv, j_k in T.Parallel(M, DK): + prod[gv, j_k] = k_s[0, j_k] * S[gv, j_k] + T.reduce_sum(prod, kS, dim=1) + # v read straight from global (derived head/channel); no shared staging + for gv in T.Parallel(M): + vnew[gv] = b_f[gv] * ( + v[bb, t, hg * grp + gv // block_DV, v0 + gv % block_DV] - kS[gv] + ) + for gv, j_k in T.Parallel(M, DK): + S[gv, j_k] += k_s[0, j_k] * vnew[gv] + for gv, j_k in T.Parallel(M, DK): + prod[gv, j_k] = q_s[0, j_k] * S[gv, j_k] + T.reduce_sum(prod, oo, dim=1) + for gv in T.Parallel(M): # o written straight to global (derived head/channel) + o[bb, t, hg * grp + gv // block_DV, v0 + gv % block_DV] = oo[gv] * scale + + if store_final_state: + for gv, j_k in T.Parallel(M, DK): + ht[bb, hg * grp + gv // block_DV, j_k, v0 + gv % block_DV] = S[gv, j_k] + + return kernel + + +def _fused_recurrent_gdr_fwd_hb( + q, k, v, g, beta, scale, initial_state, output_final_state, seqlens, grp +): + """Head-batched dispatch helper. Grid collapses to n_vt*B*Hg; block_DV keys off the + POST-collapse supply (B*Hg); threads = grp*128 (cap 512).""" + B, Tq, Hg, K = k.shape + _, _, H, V = v.shape + + # block_DV from the head-batched grid (B*Hg CTAs, not B*H): the collapse already cost + # CTAs, so the V-split must work harder to refill. threads scale with grp to hold the + # per-thread footprint constant (== the per-head kernel). + grid_base = B * Hg + block_DV = 64 if grid_base * 2 >= TARGET_NUM_CTAS else 32 + threads = min(512, grp * 128) + + use_initial_state = initial_state is not None + if initial_state is None: + initial_state = torch.empty((B, H, K, V), dtype=torch.float32, device=k.device) + final_state = torch.empty((B, H, K, V), dtype=torch.float32, device=k.device) + o = torch.empty_like(v) + + has_seqlens = seqlens is not None + if seqlens is None: + seqlens = torch.empty((B,), dtype=torch.int32, device=k.device) + + kern = tilelang_fused_recurrent_gdr_fwd_hb( + H, Hg, K, V, scale, + accum_dtype="float32", + qkva_dtype=q.dtype, + g_dtype=g.dtype, + b_dtype=beta.dtype, + h0_dtype=initial_state.dtype, + ht_dtype=final_state.dtype, + o_dtype=o.dtype, + seqlen_dtype=seqlens.dtype, + use_initial_state=use_initial_state, + store_final_state=output_final_state, + has_seqlens=has_seqlens, + block_DV=block_DV, + threads=threads, + ) + kern(q, k, v, g, beta, initial_state, seqlens, o, final_state) + return o, (final_state if output_final_state else None) + + +def fused_recurrent_gdr_fwd( + q, + k, + v, + g, + beta, + scale=None, + initial_state=None, + output_final_state=False, + seqlens=None, + head_batch=None, +): + B, Tq, Hg, K = k.shape + _, _, H, V = v.shape + assert K == V == 128 and H % Hg == 0 + scale = scale or K ** -0.5 + + # Head-batched GQA: one CTA processes all grp = H//Hg V-heads sharing a K/Q head, loading + # q/k once. Auto OFF here (the host-gated decode path saves only the q/k load ~ sub-1%, + # below noise); exposed as a forceable flag for benchmarking/tests. Restricted to grp in + # {2,4} (threads <= 512) -- grp>4 risks register spills / 1024-thread occupancy loss. + grp = H // Hg + if head_batch is None: + head_batch = False + if head_batch: + assert grp in (2, 4), f"head_batch supports grp=H//Hg in {{2,4}}, got grp={grp}" + return _fused_recurrent_gdr_fwd_hb( + q, k, v, g, beta, scale, initial_state, output_final_state, seqlens, grp + ) + + # Tile selection. When writing the final state (`output_final_state`, the decode default), the + # K-major transposed store `ht[bb,bh,jk,v0+jv] = S[jv,jk]` is the dominant cost: at block_DV<128 + # it is catastrophically uncoalesced (~0.4 TB/s), but at block_DV=128 (full-V tile, n_vt=1) the + # transpose coalesces -> measured ~2x faster at EVERY batch size (1.35x @ B=1 .. 3.0x @ B=8-16 + # .. 2.0x @ B=256; benchmark/probe_h2_blockdv_crossover.py, final_state bit-identical). So force + # block_DV=128 whenever we store the final state. (This inverts the verify kernel's V-major + # choice of 64 -- that store needs no transpose, so there occupancy wins; here the write does.) + # No final-state write: keep the occupancy ladder (64/32) -- block_DV is then perf-neutral. + grid_base = B * H + if output_final_state: + block_DV = 128 + else: + block_DV = 64 if grid_base * 2 >= TARGET_NUM_CTAS else 32 + + use_initial_state = initial_state is not None + if initial_state is None: + initial_state = torch.empty((B, H, K, V), dtype=torch.float32, device=k.device) + final_state = torch.empty((B, H, K, V), dtype=torch.float32, device=k.device) + o = torch.empty_like(v) + + has_seqlens = seqlens is not None + if seqlens is None: + seqlens = torch.empty((B,), dtype=torch.int32, device=k.device) + seqlen_dtype = seqlens.dtype + + kern = tilelang_fused_recurrent_gdr_fwd( + H, + Hg, + K, + V, + scale, + accum_dtype="float32", + qkva_dtype=q.dtype, + g_dtype=g.dtype, + b_dtype=beta.dtype, + h0_dtype=initial_state.dtype, + ht_dtype=final_state.dtype, + o_dtype=o.dtype, + seqlen_dtype=seqlen_dtype, + use_initial_state=use_initial_state, + store_final_state=output_final_state, + has_seqlens=has_seqlens, + block_DV=block_DV, + threads=128, + ) + kern(q, k, v, g, beta, initial_state, seqlens, o, final_state) + return o, (final_state if output_final_state else None) diff --git a/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_verify.py b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_verify.py new file mode 100644 index 0000000..1bcacd8 --- /dev/null +++ b/flash_qla/ops/gated_delta_rule/fused_recurrent/hopper/fused_recurrent_verify.py @@ -0,0 +1,538 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] +"""SGLang verify kernel: gemm-free GDN recurrence + paged V-major (bf16) state pool, +per-token intermediate states, no-commit, varlen cu_seqlens. Host-side gating (g/beta +pre-activated, q/k pre-l2normed by the wrapper). CUDA-graph safe (no host sync / no alloc +in the captured entry; all buffers caller-provided).""" +import os + +import torch +import tilelang +import tilelang.language as T + +MULTI_PROCESSOR_COUNT = torch.cuda.get_device_properties().multi_processor_count +TARGET_NUM_CTAS = int(MULTI_PROCESSOR_COUNT * 0.7) + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def tilelang_fused_recurrent_gdr_verify( + H, + Hg, + DK, + DV, + scale, + accum_dtype, + qkva_dtype, + g_dtype, + b_dtype, + pool_dtype, + o_dtype, + seqlen_dtype, + idx_dtype, + store_intermediate, + disable_state_update, + block_DV=128, + threads=128, +): + total_tokens = T.dynamic("total_tokens") + N = T.dynamic("N") # number of requests + num_slots = T.dynamic("num_slots") + num_cache_slots = T.dynamic("num_cache_slots") + cache_steps = T.dynamic("cache_steps") + q_shape = (1, total_tokens, Hg, DK) + v_shape = (1, total_tokens, H, DV) + g_shape = (1, total_tokens, H) + pool_shape = (num_slots, H, DV, DK) # V-major [., H, V, K] + ibuf_shape = (num_cache_slots, cache_steps, H, DV, DK) # V-major + n_vt = (DV + block_DV - 1) // block_DV + + @T.prim_func + def kernel( + q: T.Tensor(q_shape, qkva_dtype), + k: T.Tensor(q_shape, qkva_dtype), + v: T.Tensor(v_shape, qkva_dtype), + g: T.Tensor(g_shape, g_dtype), + b: T.Tensor(g_shape, b_dtype), + pool: T.Tensor(pool_shape, pool_dtype), + state_indices: T.Tensor([N], idx_dtype), + cu_seqlens: T.Tensor([N + 1], seqlen_dtype), + intermediate_state_indices: T.Tensor([N], idx_dtype), + o: T.Tensor(v_shape, o_dtype), + ibuf: T.Tensor(ibuf_shape, pool_dtype), + ): + with T.Kernel(n_vt * N * H, threads=threads) as (bbhv,): + bbh = bbhv // n_vt + bv = bbhv % n_vt + bb = bbh // H # request index + bh = bbh % H + bhg = bh // (H // Hg) + v0 = bv * block_DV + + slot = T.alloc_var("int32") + cslot = T.alloc_var("int32") + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + slot = state_indices[bb] + cslot = intermediate_state_indices[bb] + seq_start = cu_seqlens[bb] + seq_end = cu_seqlens[bb + 1] + + S = T.alloc_fragment((block_DV, DK), accum_dtype) # state [V-tile, K] fp32 + prod = T.alloc_fragment((block_DV, DK), accum_dtype) + q_s = T.alloc_shared((1, DK), qkva_dtype) + k_s = T.alloc_shared((1, DK), qkva_dtype) + v_s = T.alloc_shared((1, block_DV), qkva_dtype) + o_sh = T.alloc_shared((1, block_DV), o_dtype) + kS = T.alloc_fragment((block_DV,), accum_dtype) + oo = T.alloc_fragment((block_DV,), accum_dtype) + vnew = T.alloc_fragment((block_DV,), accum_dtype) + decay = T.alloc_fragment((1,), accum_dtype) + bt = T.alloc_fragment((1,), accum_dtype) + + # gather V-major pool[slot, bh, v0:v0+block_DV, :] directly into S[v, dk] + T.clear(S) + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] = pool[slot, bh, v0 + j_v, j_k] + + for t in T.serial(seq_end - seq_start): + tt = seq_start + t # absolute token position in the flattened layout + T.copy(q[0, tt : tt + 1, bhg, 0:DK], q_s) + T.copy(k[0, tt : tt + 1, bhg, 0:DK], k_s) + T.copy(v[0, tt : tt + 1, bh, v0 : v0 + block_DV], v_s) + decay[0] = T.exp2(g[0, tt, bh] * 1.442695) + bt[0] = b[0, tt, bh] + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] *= decay[0] + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = k_s[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, kS, dim=1) + for j_v in T.Parallel(block_DV): + vnew[j_v] = bt[0] * (v_s[0, j_v] - kS[j_v]) + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] += k_s[0, j_k] * vnew[j_v] + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = q_s[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, oo, dim=1) + for j_v in T.Parallel(block_DV): + o_sh[0, j_v] = oo[j_v] * scale + T.copy(o_sh, o[0, tt : tt + 1, bh, v0 : v0 + block_DV]) + # per-token intermediate (V-major), gated by the POOL slot mask + if store_intermediate: + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + ibuf[cslot, t, bh, v0 + j_v, j_k] = S[j_v, j_k] + + # commit final state to the pool unless no-commit (verify) + if not disable_state_update: + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + pool[slot, bh, v0 + j_v, j_k] = S[j_v, j_k] + + return kernel + + +def fused_recurrent_gdr_verify_fwd( + q, + k, + v, + g, + beta, + pool, + state_indices, + cu_seqlens, + intermediate_states_buffer, + intermediate_state_indices, + o, + scale=None, + disable_state_update=True, +): + """Graph-safe low-level verify entry. ALL buffers caller-preallocated (o, pool, ibuf); + no host sync, no allocation. g/beta pre-activated and q/k pre-l2normed host-side.""" + _, total_tokens, Hg, K = k.shape + _, _, H, V = v.shape + N = state_indices.shape[0] + assert K == V == 128 and H % Hg == 0 + scale = scale or K ** -0.5 + store_intermediate = intermediate_states_buffer is not None + + # block_DV=64 (2 V-tiles) @ threads=128 is the bandwidth sweet spot (autotuned, H100); + # 32 (4 V-tiles) for the low-CTA tail. block_DV=128 is occupancy-starved -> never used. + grid_base = N * H + block_DV = 64 if grid_base * 2 >= TARGET_NUM_CTAS else 32 + + kern = tilelang_fused_recurrent_gdr_verify( + H, + Hg, + K, + V, + scale, + accum_dtype="float32", + qkva_dtype=q.dtype, + g_dtype=g.dtype, + b_dtype=beta.dtype, + pool_dtype=pool.dtype, + o_dtype=o.dtype, + seqlen_dtype=cu_seqlens.dtype, + idx_dtype=state_indices.dtype, + store_intermediate=store_intermediate, + disable_state_update=disable_state_update, + block_DV=block_DV, + threads=128, + ) + kern(q, k, v, g, beta, pool, state_indices, cu_seqlens, + intermediate_state_indices, o, intermediate_states_buffer) + return o + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def tilelang_fused_recurrent_gdr_verify_gated( + H, Hg, DK, DV, scale, accum_dtype, qkva_dtype, ab_dtype, gate_dtype, pool_dtype, + o_dtype, seqlen_dtype, idx_dtype, store_intermediate, disable_state_update, + l2norm_eps=1e-6, softplus_thr=20.0, allow_neg_eigval=False, block_DV=128, threads=128, +): + """In-kernel fused-gating verify kernel (req #5): takes raw a,b,A_log,dt_bias and + computes g=-exp(A_log)*softplus(a+dt_bias), beta=sigmoid(b), and qk-l2norm in-kernel.""" + total_tokens = T.dynamic("total_tokens") + N = T.dynamic("N") + num_slots = T.dynamic("num_slots") + num_cache_slots = T.dynamic("num_cache_slots") + cache_steps = T.dynamic("cache_steps") + qk_shape = (1, total_tokens, Hg, DK) + v_shape = (1, total_tokens, H, DV) + ab_shape = (1, total_tokens, H) + pool_shape = (num_slots, H, DV, DK) + ibuf_shape = (num_cache_slots, cache_steps, H, DV, DK) + n_vt = (DV + block_DV - 1) // block_DV + beta_mul = 2.0 if allow_neg_eigval else 1.0 + + @T.prim_func + def kernel( + q: T.Tensor(qk_shape, qkva_dtype), + k: T.Tensor(qk_shape, qkva_dtype), + v: T.Tensor(v_shape, qkva_dtype), + a: T.Tensor(ab_shape, ab_dtype), + b: T.Tensor(ab_shape, ab_dtype), + A_log: T.Tensor([H], gate_dtype), + dt_bias: T.Tensor([H], gate_dtype), + pool: T.Tensor(pool_shape, pool_dtype), + state_indices: T.Tensor([N], idx_dtype), + cu_seqlens: T.Tensor([N + 1], seqlen_dtype), + intermediate_state_indices: T.Tensor([N], idx_dtype), + o: T.Tensor(v_shape, o_dtype), + ibuf: T.Tensor(ibuf_shape, pool_dtype), + ): + with T.Kernel(n_vt * N * H, threads=threads) as (bbhv,): + bbh = bbhv // n_vt + bv = bbhv % n_vt + bb = bbh // H + bh = bbh % H + bhg = bh // (H // Hg) + v0 = bv * block_DV + + slot = T.alloc_var("int32") + cslot = T.alloc_var("int32") + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + slot = state_indices[bb] + cslot = intermediate_state_indices[bb] + seq_start = cu_seqlens[bb] + seq_end = cu_seqlens[bb + 1] + a_log_h = T.alloc_var(accum_dtype) + dt_b_h = T.alloc_var(accum_dtype) + a_log_h = A_log[bh] + dt_b_h = dt_bias[bh] + + S = T.alloc_fragment((block_DV, DK), accum_dtype) + prod = T.alloc_fragment((block_DV, DK), accum_dtype) + q_s = T.alloc_shared((1, DK), qkva_dtype) + k_s = T.alloc_shared((1, DK), qkva_dtype) + q_n = T.alloc_shared((1, DK), qkva_dtype) + k_n = T.alloc_shared((1, DK), qkva_dtype) + v_s = T.alloc_shared((1, block_DV), qkva_dtype) + o_sh = T.alloc_shared((1, block_DV), o_dtype) + qsq = T.alloc_fragment((1, DK), accum_dtype) + ksq = T.alloc_fragment((1, DK), accum_dtype) + ssq = T.alloc_fragment((1,), accum_dtype) + kS = T.alloc_fragment((block_DV,), accum_dtype) + oo = T.alloc_fragment((block_DV,), accum_dtype) + vnew = T.alloc_fragment((block_DV,), accum_dtype) + decay = T.alloc_fragment((1,), accum_dtype) + bt = T.alloc_fragment((1,), accum_dtype) + + T.clear(S) + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] = pool[slot, bh, v0 + j_v, j_k] + + for t in T.serial(seq_end - seq_start): + tt = seq_start + t + T.copy(q[0, tt : tt + 1, bhg, 0:DK], q_s) + T.copy(k[0, tt : tt + 1, bhg, 0:DK], k_s) + T.copy(v[0, tt : tt + 1, bh, v0 : v0 + block_DV], v_s) + # in-kernel qk l2norm: x_n = x / sqrt(sum(x^2) + eps) + for _i, j_k in T.Parallel(1, DK): + qsq[0, j_k] = q_s[0, j_k] * q_s[0, j_k] + T.reduce_sum(qsq, ssq, dim=1) + for _i, j_k in T.Parallel(1, DK): + q_n[0, j_k] = q_s[0, j_k] * T.rsqrt(ssq[0] + l2norm_eps) + for _i, j_k in T.Parallel(1, DK): + ksq[0, j_k] = k_s[0, j_k] * k_s[0, j_k] + T.reduce_sum(ksq, ssq, dim=1) + for _i, j_k in T.Parallel(1, DK): + k_n[0, j_k] = k_s[0, j_k] * T.rsqrt(ssq[0] + l2norm_eps) + # in-kernel gating: g = -exp(A_log)*softplus(a+dt_bias); beta = sigmoid(b) + x = a[0, tt, bh] + dt_b_h + sp = T.if_then_else(x > softplus_thr, x, T.log(1.0 + T.exp(x))) + decay[0] = T.exp2((-T.exp(a_log_h) * sp) * 1.442695) + bt[0] = beta_mul * T.sigmoid(b[0, tt, bh]) + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] *= decay[0] + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = k_n[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, kS, dim=1) + for j_v in T.Parallel(block_DV): + vnew[j_v] = bt[0] * (v_s[0, j_v] - kS[j_v]) + for j_v, j_k in T.Parallel(block_DV, DK): + S[j_v, j_k] += k_n[0, j_k] * vnew[j_v] + for j_v, j_k in T.Parallel(block_DV, DK): + prod[j_v, j_k] = q_n[0, j_k] * S[j_v, j_k] + T.reduce_sum(prod, oo, dim=1) + for j_v in T.Parallel(block_DV): + o_sh[0, j_v] = oo[j_v] * scale + T.copy(o_sh, o[0, tt : tt + 1, bh, v0 : v0 + block_DV]) + if store_intermediate: + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + ibuf[cslot, t, bh, v0 + j_v, j_k] = S[j_v, j_k] + + if not disable_state_update: + with T.If(slot >= 0): + with T.Then(): + for j_v, j_k in T.Parallel(block_DV, DK): + pool[slot, bh, v0 + j_v, j_k] = S[j_v, j_k] + + return kernel + + +def fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, pool, state_indices, cu_seqlens, + intermediate_states_buffer, intermediate_state_indices, o, + scale=None, disable_state_update=True, allow_neg_eigval=False, +): + """In-kernel fused-gating verify entry: raw a,b,A_log,dt_bias; computes g/beta + qk-l2norm + inside the kernel. Graph-safe (all buffers caller-provided).""" + _, total_tokens, Hg, K = k.shape + _, _, H, V = v.shape + N = state_indices.shape[0] + assert K == V == 128 and H % Hg == 0 + scale = scale or K ** -0.5 + store_intermediate = intermediate_states_buffer is not None + + grid_base = N * H # bandwidth sweet spot (autotuned, H100): block_DV=64 @ threads=128 + block_DV = 64 if grid_base * 2 >= TARGET_NUM_CTAS else 32 + + kern = tilelang_fused_recurrent_gdr_verify_gated( + H, Hg, K, V, scale, + accum_dtype="float32", qkva_dtype=q.dtype, ab_dtype=a.dtype, gate_dtype=A_log.dtype, + pool_dtype=pool.dtype, o_dtype=o.dtype, seqlen_dtype=cu_seqlens.dtype, + idx_dtype=state_indices.dtype, store_intermediate=store_intermediate, + disable_state_update=disable_state_update, allow_neg_eigval=allow_neg_eigval, + block_DV=block_DV, threads=max(128, block_DV * 2), + ) + kern(q, k, v, a, b, A_log, dt_bias, pool, state_indices, cu_seqlens, + intermediate_state_indices, o, intermediate_states_buffer) + return o + + +# ---------------------------------------------------------------------------------------------- +# H1: gating + qk-l2norm DEDUP PRE-PASS. +# The in-kernel-gated kernel above recomputes g/beta + qk-l2norm INSIDE the per-token hot loop, +# once per (token, V-head, V-tile) CTA -> l2norm redundant grp*n_vt times per (token,K-head), +# gating redundant n_vt times per (token,V-head). This pre-pass computes each ONCE (grid = +# total_tokens*Hk, one CTA per (token, K-head)) and feeds the HOST-gated main kernel +# (fused_recurrent_gdr_verify_fwd) whose hot loop then has NO transcendentals. Measured ceiling +# (gated-vs-host-gated, H100): 14-16% at T=12 large batch. Regime-gated (see should_use_prepass): +# at single-request the second-launch tax exceeds the small ceiling, so variant A is kept there. +# ---------------------------------------------------------------------------------------------- + +# Regime gate tunables. CALIBRATED FOR THE CUDA-GRAPH (production) PATH (H100, +# benchmark/bench_prepass.py bench_graph, _time_graph with 50+ warmup). Prepass wins when (a) drafts +# are long enough that the per-token gating/l2norm recompute it dedups is a meaningful fraction +# (T_avg >= 4; at T=1 it is paid once -> nothing to dedup across tokens, measured neutral/loss even +# under graphs), AND (b) the main-kernel work N*H*(1+T_avg) clears a SMALL floor. +# +# KEY: under CUDA-graph REPLAY the prepass's 2nd launch shrinks to a graph node, so the eager +# second-launch tax (~13-15us) VANISHES and the dedup win dominates down to SINGLE-REQUEST. Measured +# (graph, Hk=16 Hv=32): N=1,T=12 work=416 -> 1.18x WIN (eager was 0.60x LOSS); N=2,T=12 -> 1.08x; +# N=4,T=12 -> 1.18x; every T=12 point N>=1 wins 1.08-1.24x. The only graph losses are work<=320 +# (N=1,T=4=160 0.96x; N=2,T=4=320 0.93x). So the gate window for the verify (always T=12) is +# work in (320, 416]; MIN_WORK=384 fires the prepass for the ENTIRE T=12 verify path incl. N=1, +# capturing the 1.08-1.24x the old eager-conservative 3000 was leaving on the table at small batch. +# TRADEOFF: this assumes the CUDA-graph deployment (qwen36 captures all batch sizes). Under EAGER +# (warmup / --disable-cuda-graph), small-N now pays the 2nd-launch tax (N=1,T=12 eager 0.60x) -- but +# steady-state production verify is always captured, and warmup is untimed. Metric scales with H, so +# small-head configs self-raise the batch needed (gate stays safe for untested H<32). +PREPASS_MIN_T = 4.0 +# graph-calibrated default 384 (was 3000, eager-conservative); floor below N=1/T=12 work=416. +# Env-overridable for controlled A/B (set FLASHQLA_PREPASS_MIN_WORK=3000 to reproduce the old gate). +PREPASS_MIN_WORK = int(os.environ.get("FLASHQLA_PREPASS_MIN_WORK", "384")) + + +def should_use_prepass(N, H, total_tokens): + """Static (capture-safe; shape-only) decision: run the dedup prepass + host-gated main kernel + (True) vs the single in-kernel-gated kernel A (False). Both produce identical output; this + only picks the faster path per regime. CALIBRATED FOR THE CUDA-GRAPH (production) path, where + the prepass's 2nd launch is a free graph node and it wins down to single-request T>=4 (the eager + second-launch tax that favored variant A at small batch does NOT exist under graph replay).""" + if N <= 0: + return False + t_avg = total_tokens / N + work = H * (N + total_tokens) # == N*H*(1 + t_avg), proportional to the main-kernel runtime + return (t_avg >= PREPASS_MIN_T) and (work >= PREPASS_MIN_WORK) + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def tilelang_gdr_verify_prepass( + Hk, Hv, DK, accum_dtype, qk_dtype, ab_dtype, gate_dtype, g_out_dtype, b_out_dtype, + l2norm_eps=1e-6, softplus_thr=20.0, allow_neg_eigval=False, threads=128, +): + """Dedup pre-pass: ONE CTA per token tt computes q_n/k_n = l2norm(q/k) for all Hk K-heads at + once via a parallel [Hk,DK] reduce (the proven reduce_sum(...,dim=1) idiom), and + g = -exp(A_log[h])*softplus(a+dt_bias) (RAW log-decay; the main kernel applies exp2) + + beta = beta_mul*sigmoid(b) for all Hv V-heads. No recurrence, no cu_seqlens (token-local): + the main kernel only consumes tokens within cu_seqlens, so normalizing trailing padding is + harmless. + + q/k are read into a fragment and q_n/k_n written DIRECTLY from a full-range T.Parallel (the + head-batch global-write idiom) -- no per-row [1,DK] T.copy, so no serial Hk loop (that was a + measured ~6x slowdown) and no Hopper copy-layout trap. The gate write is the full contiguous + [1,Hv] row (Hv>=4 -> >=128-bit fp32 extent; a per-K-head [1,grp] write fails layout inference + for grp<4). Grid = total_tokens balances occupancy (~total_tokens CTAs).""" + beta_mul = 2.0 if allow_neg_eigval else 1.0 + total_tokens = T.dynamic("total_tokens") + qk_shape = (1, total_tokens, Hk, DK) + ab_shape = (1, total_tokens, Hv) + + @T.prim_func + def kernel( + q: T.Tensor(qk_shape, qk_dtype), + k: T.Tensor(qk_shape, qk_dtype), + a: T.Tensor(ab_shape, ab_dtype), + b: T.Tensor(ab_shape, ab_dtype), + A_log: T.Tensor([Hv], gate_dtype), + dt_bias: T.Tensor([Hv], gate_dtype), + q_n: T.Tensor(qk_shape, qk_dtype), + k_n: T.Tensor(qk_shape, qk_dtype), + g_out: T.Tensor(ab_shape, g_out_dtype), + beta_out: T.Tensor(ab_shape, b_out_dtype), + ): + with T.Kernel(total_tokens, threads=threads) as (tt,): + xf = T.alloc_fragment((Hk, DK), accum_dtype) # raw q/k rows (reused q then k) + sq = T.alloc_fragment((Hk, DK), accum_dtype) # squared, reduce input + ssq = T.alloc_fragment((Hk,), accum_dtype) # per-K-head sum-of-squares + g_f = T.alloc_fragment((1, Hv), accum_dtype) + b_f = T.alloc_fragment((1, Hv), accum_dtype) + g_sh = T.alloc_shared((1, Hv), g_out_dtype) + b_sh = T.alloc_shared((1, Hv), b_out_dtype) + + # q l2norm: load [Hk,DK] -> square -> reduce over DK -> normalize + write (all parallel) + for i, j in T.Parallel(Hk, DK): + xf[i, j] = q[0, tt, i, j] + for i, j in T.Parallel(Hk, DK): + sq[i, j] = xf[i, j] * xf[i, j] + T.reduce_sum(sq, ssq, dim=1) + for i, j in T.Parallel(Hk, DK): + q_n[0, tt, i, j] = xf[i, j] * T.rsqrt(ssq[i] + l2norm_eps) + + # k l2norm (reuse xf/sq/ssq) + for i, j in T.Parallel(Hk, DK): + xf[i, j] = k[0, tt, i, j] + for i, j in T.Parallel(Hk, DK): + sq[i, j] = xf[i, j] * xf[i, j] + T.reduce_sum(sq, ssq, dim=1) + for i, j in T.Parallel(Hk, DK): + k_n[0, tt, i, j] = xf[i, j] * T.rsqrt(ssq[i] + l2norm_eps) + + # gating for all Hv V-heads (fragment idiom, inlined exprs, direct global reads), + # staged through [1,Hv] shared and written as one contiguous row (>=128-bit extent) + for _i, h in T.Parallel(1, Hv): + g_f[0, h] = -T.exp(A_log[h]) * T.if_then_else( + a[0, tt, h] + dt_bias[h] > softplus_thr, + a[0, tt, h] + dt_bias[h], + T.log(1.0 + T.exp(a[0, tt, h] + dt_bias[h])), + ) + b_f[0, h] = beta_mul * T.sigmoid(b[0, tt, h]) + for _i, h in T.Parallel(1, Hv): + g_sh[0, h] = g_f[0, h] + b_sh[0, h] = b_f[0, h] + T.copy(g_sh, g_out[0, tt : tt + 1, 0:Hv]) + T.copy(b_sh, beta_out[0, tt : tt + 1, 0:Hv]) + + return kernel + + +def fused_recurrent_gdr_verify_prepass( + q, k, a, b, A_log, dt_bias, + q_n=None, k_n=None, g_out=None, beta_out=None, allow_neg_eigval=False, +): + """Dedup pre-pass dispatch. Computes q_n=l2norm(q), k_n=l2norm(k), g=raw log-decay, + beta=beta_mul*sigmoid(b) ONCE, feeding the host-gated verify main kernel. Buffers are + caller-provided for CUDA-graph capture safety (like o/pool/ibuf); allocate-if-None is an + EAGER convenience for tests/bench only -- it must NOT run inside a captured region (no alloc + in capture). g/beta are fp32 (matching gdn_sigmoid_gate); q_n/k_n match q/k dtype.""" + _, total_tokens, Hk, K = q.shape + Hv = a.shape[2] + assert Hv % Hk == 0, f"num_v_heads {Hv} must be divisible by num_k_heads {Hk}" + assert total_tokens > 0 + if q_n is None: + q_n = torch.empty_like(q) + if k_n is None: + k_n = torch.empty_like(k) + if g_out is None: + g_out = torch.empty(1, total_tokens, Hv, device=q.device, dtype=torch.float32) + if beta_out is None: + beta_out = torch.empty(1, total_tokens, Hv, device=q.device, dtype=torch.float32) + + kern = tilelang_gdr_verify_prepass( + Hk, Hv, K, + accum_dtype="float32", qk_dtype=q.dtype, ab_dtype=a.dtype, gate_dtype=A_log.dtype, + g_out_dtype=g_out.dtype, b_out_dtype=beta_out.dtype, + allow_neg_eigval=allow_neg_eigval, threads=128, + ) + kern(q, k, a, b, A_log, dt_bias, q_n, k_n, g_out, beta_out) + return q_n, k_n, g_out, beta_out + + +# Module-level UNBOUNDED, never-evicting scratch cache for the prepass outputs. Unbounded is the +# capture-safe choice: each captured graph bakes the device pointers of its own (total_tokens,...) +# entry; an evicting LRU (e.g. utils.tensor_cache) could free a still-referenced storage and make +# a later replay read freed memory. One live entry per captured (N,T); entries are never freed. +_PREPASS_SCRATCH = {} + + +def get_prepass_scratch(total_tokens, Hk, Hv, device, qk_dtype): + """Persistent prepass scratch (q_n,k_n bf16 [1,T,Hk,128]; g,beta fp32 [1,T,Hv]). Allocated + lazily on the first (warmup) call and reused thereafter so addresses are stable across + CUDA-graph capture/replay. Raises (rather than allocating + aborting capture) on a cold miss + during capture -- warmup must exercise every (N,T) that will be captured.""" + key = (total_tokens, Hk, Hv, device, qk_dtype) + buf = _PREPASS_SCRATCH.get(key) + if buf is None: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"prepass scratch not warmed for shape {key}; run an eager warmup pass before " + "CUDA-graph capture (the scratch must be allocated outside the captured region)." + ) + q_n = torch.empty(1, total_tokens, Hk, 128, device=device, dtype=qk_dtype) + k_n = torch.empty(1, total_tokens, Hk, 128, device=device, dtype=qk_dtype) + g_out = torch.empty(1, total_tokens, Hv, device=device, dtype=torch.float32) + beta_out = torch.empty(1, total_tokens, Hv, device=device, dtype=torch.float32) + buf = (q_n, k_n, g_out, beta_out) + _PREPASS_SCRATCH[key] = buf + return buf From 6a3dc8a211a6b4441c3d644d64bd87a7a304f2da Mon Sep 17 00:00:00 2001 From: Rifky Bujana Bisri Date: Tue, 16 Jun 2026 16:39:06 +0700 Subject: [PATCH 2/3] test(gdn): decode/verify/prepass/head-batch correctness + reference 67 tests against a reference recurrence and the FLA kernel: GQA, ragged batches, host/in-kernel gating, the dedup pre-pass path, CUDA-graph capture/replay, and negative controls (step order, GQA mapping, V-major transpose). The accuracy check runs the kernel 1000x asserting <=2% drift (a deliberate race/nondeterminism detector). ref_gdr.py adds decode_recur and verify_ref. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/ref_gdr.py | 111 ++++++++++++++++ tests/test_decode_gdr.py | 137 ++++++++++++++++++++ tests/test_head_batch_gdr.py | 135 +++++++++++++++++++ tests/test_prepass_gdr.py | 153 ++++++++++++++++++++++ tests/test_verify_gdr.py | 243 +++++++++++++++++++++++++++++++++++ 5 files changed, 779 insertions(+) create mode 100644 tests/test_decode_gdr.py create mode 100644 tests/test_head_batch_gdr.py create mode 100644 tests/test_prepass_gdr.py create mode 100644 tests/test_verify_gdr.py diff --git a/tests/ref_gdr.py b/tests/ref_gdr.py index fde5335..be80742 100644 --- a/tests/ref_gdr.py +++ b/tests/ref_gdr.py @@ -698,3 +698,114 @@ def chunk_gated_delta_rule_bwd( dk = torch.sum(dk.reshape(B, T, Hg, -1, K), dim=3) dg = torch_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens) return dq, dk, dv, db, dg, dh0 + + +def decode_recur( + q, + k, + v, + g, + beta, # q,k:[B,T,Hk,128] v:[B,T,Hv,128] g,beta:[B,T,Hv] + scale=None, + initial_state=None, # initial_state: [B,Hv,128,128] fp32 or None + seqlens=None, # [B] int32 accepted lengths (default: all T) + read_pre_update=False, # NEGATIVE CONTROL: read o before the rank-1 (wrong) + gqa_mod=False, # NEGATIVE CONTROL: hg = h % Hk (wrong GQA mapping) + band_perm=False, # NEGATIVE CONTROL: cyclically swap V-heads WITHIN each GQA group (wrong) +): + """Ground-truth GDN decode recurrence (spec 2): per (b,h), per token tband + mis-mapping bug. (Use a per-head DISTINCT gate so the permuted result is far from correct.)""" + B, T, Hk, K = k.shape + _, _, Hv, V = v.shape + assert K == V == 128 and Hv % Hk == 0 + scale = scale if scale is not None else K ** -0.5 + grp = Hv // Hk + dev = k.device + S = ( + initial_state.clone().float() + if initial_state is not None + else torch.zeros(B, Hv, K, V, device=dev, dtype=torch.float32) + ) + o = torch.zeros(B, T, Hv, V, device=dev, dtype=torch.float32) + if seqlens is None: + seqlens = torch.full((B,), T, device=dev, dtype=torch.int32) + for b in range(B): + L = int(seqlens[b]) + for t in range(L): + for h in range(Hv): + hg = (h % Hk) if gqa_mod else (h // grp) + qt = q[b, t, hg].float() + kt = k[b, t, hg].float() + vt = v[b, t, h].float() + decay = torch.exp(g[b, t, h].float()) + Sh = S[b, h] * decay # [K,V] + kS = kt @ Sh # [V] + v_new = beta[b, t, h].float() * (vt - kS) # [V] + if read_pre_update: + o[b, t, h] = scale * (qt @ Sh) # WRONG: pre rank-1 + Sh = Sh + torch.outer(kt, v_new) # [K,V] + S[b, h] = Sh + if not read_pre_update: + o[b, t, h] = scale * (qt @ Sh) # [V] (post-update, correct) + if band_perm and grp > 1: # cyclic within-group head swap -> wrong (head-batch band control) + perm = torch.arange(Hv, device=dev) + for hgi in range(Hk): + base = hgi * grp + perm[base : base + grp] = base + (torch.arange(grp, device=dev) + 1) % grp + o = o[:, :, perm, :].contiguous() + S = S[:, perm, :, :].contiguous() + return o, S # o:[B,T,Hv,V] (only [:, :L_b] valid per b), final_state S:[B,Hv,K,V] + + +def verify_ref( + q, + k, + v, + g, + beta, # q,k:[1,T,Hk,128] v:[1,T,Hv,128] g,beta:[1,T,Hv] + pool, # [num_slots, Hv, V=128, K=128] V-major (state_v_first) + state_indices, # [N] int32; <0 => skip (gather + commit) + cu_seqlens, # [N+1] int32; request bb owns flattened tokens [cu_seqlens[bb], cu_seqlens[bb+1]) + intermediate_states_buffer, # [num_cache_slots, cache_steps, Hv, V, K] V-major + intermediate_state_indices, # [N] int32 dense (never -1); ibuf write gated by pool slot + scale=None, + disable_state_update=True, # no-commit +): + """Reference for the SGLang verify kernel: paged V-major bf16 pool, per-token + intermediate states, varlen cu_seqlens, no-commit. Mirrors decode_recur per request.""" + K = V = 128 + Hk, Hv = k.shape[2], v.shape[2] + grp = Hv // Hk + scale = scale if scale is not None else K ** -0.5 + N = len(cu_seqlens) - 1 + dev = k.device + o = torch.zeros(1, q.shape[1], Hv, V, device=dev, dtype=torch.float32) + pool_out = pool.clone() + ibuf_out = intermediate_states_buffer.clone() + for bb in range(N): + slot = int(state_indices[bb]) + cslot = int(intermediate_state_indices[bb]) + s0, s1 = int(cu_seqlens[bb]), int(cu_seqlens[bb + 1]) + for h in range(Hv): + hg = h // grp + # gather V-major pool[slot,h] [V,K] -> decode state S [K,V] + S = (pool[slot, h].float().t().clone() if slot >= 0 + else torch.zeros(K, V, device=dev, dtype=torch.float32)) + for ti, t in enumerate(range(s0, s1)): + decay = torch.exp(g[0, t, h].float()) + S = S * decay + kS = k[0, t, hg].float() @ S + v_new = beta[0, t, h].float() * (v[0, t, h].float() - kS) + S = S + torch.outer(k[0, t, hg].float(), v_new) + o[0, t, h] = scale * (q[0, t, hg].float() @ S) + if slot >= 0: # per-token intermediate, V-major [V,K]; gated by POOL slot mask + ibuf_out[cslot, ti, h] = S.t().to(ibuf_out.dtype) + if slot >= 0 and not disable_state_update: + pool_out[slot, h] = S.t().to(pool_out.dtype) + return o, pool_out, ibuf_out diff --git a/tests/test_decode_gdr.py b/tests/test_decode_gdr.py new file mode 100644 index 0000000..1652c15 --- /dev/null +++ b/tests/test_decode_gdr.py @@ -0,0 +1,137 @@ +# tests/test_decode_gdr.py +import pytest +import torch + +from ref_gdr import decode_recur +from ref_gdr import chunk_gated_delta_rule_fwd as chunk_fwd_ref +from flash_qla.utils import l2norm + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _mk(B, T, Hk, Hv, seed=0, dtype=torch.float32): + torch.manual_seed(seed) + q = l2norm(torch.randn(B, T, Hk, 128, device="cuda", dtype=dtype)) + k = l2norm(torch.randn(B, T, Hk, 128, device="cuda", dtype=dtype)) + v = torch.randn(B, T, Hv, 128, device="cuda", dtype=dtype) + g = torch.nn.functional.logsigmoid(torch.randn(B, T, Hv, device="cuda")) / 16 + beta = torch.randn(B, T, Hv, device="cuda").sigmoid() + return q, k, v, g, beta + + +def _ref_bf16_inputs(B, T, Hk, Hv, seed=0): + return _mk(B, T, Hk, Hv, seed=seed, dtype=torch.bfloat16) + + +@CUDA +def test_decode_recur_matches_chunk_at_cs64(): + # A length-L single sequence: decode_recur must match the chunk reference on the L-prefix. + B, T, H = 1, 50, 8 + q, k, v, g, beta = _mk(B, T, H, H, dtype=torch.float32) + o_dec, s_dec = decode_recur(q, k, v, g, beta) + g_ref, o_ref, A_ref, h_ref, s_ref = chunk_fwd_ref( + q=q.double(), k=k.double(), v=v.double(), g=g.double(), beta=beta.double(), + scale=128 ** -0.5, initial_state=None, cu_seqlens=None) + assert (o_dec - o_ref.float()).abs().max() / o_ref.abs().max() < 1e-3 + assert (s_dec - s_ref.float()).abs().max() / s_ref.abs().max() < 1e-3 + + +@CUDA +@pytest.mark.parametrize("D", [1, 8]) +@pytest.mark.parametrize("Hk,Hv", [(8, 8), (2, 8), (1, 8)]) +@pytest.mark.parametrize("use_h0", [False, True]) +def test_kernel_matches_reference(D, Hk, Hv, use_h0): + from flash_qla import recurrent_gated_delta_rule + B = 1 + q, k, v, g, beta = _ref_bf16_inputs(B, D, Hk, Hv) + h0 = (torch.randn(B, Hv, 128, 128, device="cuda", dtype=torch.float32) if use_h0 else None) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, initial_state=h0) + o_qla, s_qla = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, initial_state=h0, output_final_state=True) + assert (o_qla.float() - o_ref).abs().max() / o_ref.abs().max().clamp_min(1e-6) <= 0.02 + assert (s_qla - s_ref).abs().max() / s_ref.abs().max().clamp_min(1e-6) <= 0.02 + + +@CUDA +def test_kernel_g0_swa_heads(): + from flash_qla import recurrent_gated_delta_rule + B, D, H = 1, 8, 8 + q, k, v, g, beta = _ref_bf16_inputs(B, D, H, H) + g[:, :, :H // 2] = 0.0 # half the heads have no decay + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_qla, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5) + assert (o_qla.float() - o_ref).abs().max() / o_ref.abs().max() <= 0.02 + + +@CUDA +def test_kernel_ragged_seqlens(): + from flash_qla import recurrent_gated_delta_rule + B, D, H = 3, 8, 8 + q, k, v, g, beta = _ref_bf16_inputs(B, D, H, H) + seqlens = torch.tensor([1, 5, 8], device="cuda", dtype=torch.int32) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, seqlens=seqlens) + o_qla, s_qla = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, seqlens=seqlens, output_final_state=True) + for b in range(B): + L = int(seqlens[b]) + assert (o_qla[b, :L].float() - o_ref[b, :L]).abs().max() / o_ref[b, :L].abs().max() <= 0.02 + assert (s_qla[b] - s_ref[b]).abs().max() / s_ref[b].abs().max() <= 0.02 + + +@CUDA +def test_kernel_low_occupancy_vsplit(): + from flash_qla import recurrent_gated_delta_rule + # B*H small => wrapper picks block_DV in {64,32}; result must still match. + B, D, H = 1, 4, 4 + q, k, v, g, beta = _ref_bf16_inputs(B, D, H, H) + o_ref, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_qla, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5) + assert (o_qla.float() - o_ref).abs().max() / o_ref.abs().max() <= 0.02 + + +@CUDA +@pytest.mark.parametrize("B,H", [(8, 8), (16, 8)]) # B*H=64 -> block_DV=64; 128 -> block_DV=128 +def test_kernel_high_occupancy(B, H): + from flash_qla import recurrent_gated_delta_rule + D = 4 + q, k, v, g, beta = _ref_bf16_inputs(B, D, H, H) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_qla, s_qla = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, output_final_state=True) + assert (o_qla.float() - o_ref).abs().max() / o_ref.abs().max() <= 0.02 + assert (s_qla - s_ref).abs().max() / s_ref.abs().max() <= 0.02 + + +@CUDA +def test_negctrl_step_order(): + # discriminating: kernel matches the post-update reference and DIFFERS from a pre-update one + from flash_qla import recurrent_gated_delta_rule + B, D, H = 1, 6, 8 + q, k, v, g, beta = _ref_bf16_inputs(B, D, H, H) + o_post, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_pre, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, read_pre_update=True) + o_qla, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5) + assert (o_qla.float() - o_post).abs().max() / o_post.abs().max() <= 0.02 + assert (o_qla.float() - o_pre).abs().max() / o_pre.abs().max() > 0.2 # must NOT match wrong order + + +@CUDA +def test_negctrl_gqa_mapping(): + # discriminating: kernel uses hg=h//grp and must DIFFER from the hg=h%Hk mapping + from flash_qla import recurrent_gated_delta_rule + B, D, Hk, Hv = 1, 6, 2, 8 + q, k, v, g, beta = _ref_bf16_inputs(B, D, Hk, Hv) + o_div, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_mod, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, gqa_mod=True) + o_qla, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5) + assert (o_qla.float() - o_div).abs().max() / o_div.abs().max() <= 0.02 + assert (o_qla.float() - o_mod).abs().max() / o_mod.abs().max() > 0.2 # must NOT match wrong GQA + + +def test_signature_contract(): + import inspect + from flash_qla import recurrent_gated_delta_rule + sig = inspect.signature(recurrent_gated_delta_rule) + for p in ["q", "k", "v", "g", "beta", "scale", "initial_state", + "output_final_state", "use_qk_l2norm_in_kernel", "seqlens"]: + assert p in sig.parameters diff --git a/tests/test_head_batch_gdr.py b/tests/test_head_batch_gdr.py new file mode 100644 index 0000000..ac88073 --- /dev/null +++ b/tests/test_head_batch_gdr.py @@ -0,0 +1,135 @@ +# tests/test_head_batch_gdr.py +# Head-batched GQA specialization of the gemm-free decode kernel: one CTA processes all +# grp = Hv//Hk V-heads sharing a K/Q head. Validates correctness + a within-group band-swap +# negative control + direct equality with the per-head path. Forces head_batch=True so the +# new code is actually exercised (auto stays OFF for the host-gated decode path). +import pytest +import torch + +from ref_gdr import decode_recur +from flash_qla.utils import l2norm + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _mk(B, T, Hk, Hv, seed=0, distinct_gate=True): + """bf16 inputs. distinct_gate gives each V-head its own decay magnitude so a within-group + head-band swap produces a far-off result (mandatory for the band-swap negative control).""" + torch.manual_seed(seed) + q = l2norm(torch.randn(B, T, Hk, 128, device="cuda", dtype=torch.bfloat16)) + k = l2norm(torch.randn(B, T, Hk, 128, device="cuda", dtype=torch.bfloat16)) + v = torch.randn(B, T, Hv, 128, device="cuda", dtype=torch.bfloat16) + g = torch.nn.functional.logsigmoid(torch.randn(B, T, Hv, device="cuda")) / 16 + if distinct_gate: + g = g * (1 + torch.arange(Hv, device="cuda").float()[None, None, :]) + beta = torch.randn(B, T, Hv, device="cuda").sigmoid() + return q, k, v, g, beta + + +def _rel(a, b): + return (a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6) + + +@CUDA +@pytest.mark.parametrize("Hk,Hv", [(4, 8), (2, 8)]) # grp = 2, 4 +def test_head_batch_compiles_and_runs(Hk, Hv): + # smoke: the factory must JIT-build for both grp values (isolates a lowering failure). + from flash_qla import recurrent_gated_delta_rule + q, k, v, g, beta = _mk(1, 8, Hk, Hv) + o, s = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, output_final_state=True, head_batch=True) + assert o.shape == (1, 8, Hv, 128) and s.shape == (1, Hv, 128, 128) + assert torch.isfinite(o.float()).all() + + +@CUDA +@pytest.mark.parametrize("D", [1, 8, 12]) +@pytest.mark.parametrize("Hk,Hv", [(4, 8), (2, 8)]) # grp = 2, 4 +@pytest.mark.parametrize("use_h0", [False, True]) +def test_head_batch_matches_reference(D, Hk, Hv, use_h0): + from flash_qla import recurrent_gated_delta_rule + B = 1 + q, k, v, g, beta = _mk(B, D, Hk, Hv) + h0 = (torch.randn(B, Hv, 128, 128, device="cuda", dtype=torch.float32) if use_h0 else None) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, initial_state=h0) + o_hb, s_hb = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, initial_state=h0, + output_final_state=True, head_batch=True) + assert _rel(o_hb, o_ref) <= 0.02 + assert _rel(s_hb, s_ref) <= 0.02 + + +@CUDA +@pytest.mark.parametrize("Hk,Hv", [(4, 8), (2, 8)]) +def test_head_batch_equals_per_head(Hk, Hv): + # the two paths must agree on identical inputs (cheap regression catch; no fp32 ref). + from flash_qla import recurrent_gated_delta_rule + q, k, v, g, beta = _mk(1, 8, Hk, Hv) + o_hb, s_hb = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, output_final_state=True, head_batch=True) + o_ph, s_ph = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, output_final_state=True, head_batch=False) + assert _rel(o_hb, o_ph) <= 0.02 + assert _rel(s_hb, s_ph) <= 0.02 + + +@CUDA +def test_head_batch_low_occupancy_block_dv32(): + # B*Hg small -> head-batch wrapper picks block_DV=32 (the low-CTA tail); must still match. + from flash_qla import recurrent_gated_delta_rule + q, k, v, g, beta = _mk(1, 8, 2, 8) # Hg=2 -> grid_base=2 -> block_DV=32 + o_ref, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_hb, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5, head_batch=True) + assert _rel(o_hb, o_ref) <= 0.02 + + +@CUDA +def test_head_batch_high_occupancy_block_dv64(): + # B*Hg large -> head-batch wrapper picks block_DV=64; must still match. + from flash_qla import recurrent_gated_delta_rule + B, Hk, Hv = 16, 4, 8 # Hg=4 -> grid_base=64 -> block_DV=64, grp=2 + q, k, v, g, beta = _mk(B, 4, Hk, Hv) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_hb, s_hb = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, output_final_state=True, head_batch=True) + assert _rel(o_hb, o_ref) <= 0.02 + assert _rel(s_hb, s_ref) <= 0.02 + + +@CUDA +def test_head_batch_ragged_seqlens(): + from flash_qla import recurrent_gated_delta_rule + B, Hk, Hv = 3, 2, 8 + q, k, v, g, beta = _mk(B, 8, Hk, Hv) + seqlens = torch.tensor([1, 5, 8], device="cuda", dtype=torch.int32) + o_ref, s_ref = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, seqlens=seqlens) + o_hb, s_hb = recurrent_gated_delta_rule( + q, k, v, g, beta, scale=128 ** -0.5, seqlens=seqlens, + output_final_state=True, head_batch=True) + for b in range(B): + L = int(seqlens[b]) + assert _rel(o_hb[b, :L], o_ref[b, :L]) <= 0.02 + assert _rel(s_hb[b], s_ref[b]) <= 0.02 + + +@CUDA +def test_negctrl_head_band_swap(): + # discriminating: head-batch must match the correct band layout and DIFFER from a + # within-group V-head swap. gqa_mod can't express this (it only re-routes the K/Q source). + from flash_qla import recurrent_gated_delta_rule + B, D, Hk, Hv = 1, 8, 2, 8 # grp=4 + q, k, v, g, beta = _mk(B, D, Hk, Hv) # distinct per-head gate (mandatory) + o_ok, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5) + o_swap, _ = decode_recur(q, k, v, g, beta, scale=128 ** -0.5, band_perm=True) + o_hb, _ = recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5, head_batch=True) + assert _rel(o_hb, o_ok) <= 0.02 + assert _rel(o_hb, o_swap) > 0.2 # must NOT match a swapped-band layout + + +@CUDA +def test_head_batch_rejects_unsupported_grp(): + # grp not in {2,4} when forced must raise (thread/register cap), not silently mis-run. + from flash_qla import recurrent_gated_delta_rule + q, k, v, g, beta = _mk(1, 4, 1, 8) # grp=8 + with pytest.raises(AssertionError): + recurrent_gated_delta_rule(q, k, v, g, beta, scale=128 ** -0.5, head_batch=True) diff --git a/tests/test_prepass_gdr.py b/tests/test_prepass_gdr.py new file mode 100644 index 0000000..fa1824a --- /dev/null +++ b/tests/test_prepass_gdr.py @@ -0,0 +1,153 @@ +# tests/test_prepass_gdr.py +# H1: the gating + qk-l2norm PRE-PASS kernel must produce g/beta/q_n/k_n bit-equivalent (bf16 +# noise) to the reference gdn_sigmoid_gate + l2norm that the host-gated main verify kernel +# consumes. This is the correctness gate for the dedup pre-pass (replaces the in-hot-loop +# recompute of variant A). +import pytest +import torch + +from flash_qla.utils import l2norm +from flash_qla.ops.gated_delta_rule.fused_recurrent import gdn_sigmoid_gate + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _rel(a, b): + return ((a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6)).item() + + +@CUDA +@pytest.mark.parametrize("Hk,Hv", [(8, 8), (2, 8), (16, 32), (4, 16)]) +@pytest.mark.parametrize("neg", [False, True]) +def test_prepass_matches_reference(Hk, Hv, neg): + from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_prepass, + ) + N, D = 4, 5 + total = N * D + torch.manual_seed(7) + q = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(Hv, device="cuda").abs().log() + dt_bias = torch.randn(Hv, device="cuda") + + q_n, k_n, g, beta = fused_recurrent_gdr_verify_prepass( + q, k, a, b, A_log, dt_bias, allow_neg_eigval=neg) + g_ref, beta_ref = gdn_sigmoid_gate(A_log, a, dt_bias, b, allow_neg_eigval=neg) + + assert _rel(q_n, l2norm(q)) <= 0.02, "q l2norm mismatch" + assert _rel(k_n, l2norm(k)) <= 0.02, "k l2norm mismatch" + assert _rel(g, g_ref) <= 0.02, "g (raw log-decay) mismatch" + assert _rel(beta, beta_ref) <= 0.02, "beta mismatch" + + +@CUDA +def test_prepass_negctrl_raw_g_not_exp2(): + # discriminating: g must be the RAW log-decay (g<=0), NOT pre-exp2'd (which would be in (0,1]). + # The main kernel applies decay=exp2(g*1.442695); a pre-exp2'd g would be silently wrong. + from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_prepass, + ) + N, D, Hk, Hv = 3, 4, 4, 8 + total = N * D + torch.manual_seed(11) + q = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(Hv, device="cuda").abs().log() + dt_bias = torch.randn(Hv, device="cuda") + _, _, g, _ = fused_recurrent_gdr_verify_prepass(q, k, a, b, A_log, dt_bias) + assert (g <= 1e-4).all(), "g must be raw log-decay (<=0), not pre-exp2'd" + + +@CUDA +@pytest.mark.parametrize("Hk,Hv", [(16, 32), (2, 8)]) +def test_prepass_plus_hostgated_matches_in_kernel_gated(Hk, Hv): + # end-to-end: prepass + host-gated main kernel must match the in-kernel-gated kernel on the + # SAME raw inputs (both compute the identical recurrence; only WHERE gating/l2norm happen + # differs). Within in-kernel-gating tolerance (0.03). + from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_prepass, + fused_recurrent_gdr_verify_fwd, + fused_recurrent_gdr_verify_gated_fwd, + ) + N, D = 4, 6 + total = N * D + torch.manual_seed(13) + q = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(Hv, device="cuda").abs().log() + dt_bias = torch.randn(Hv, device="cuda") + pool = torch.randn(N, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf_a = torch.zeros(N + 1, D, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + ibuf_b = torch.zeros(N + 1, D, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + o_gated = torch.empty(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + o_pp = torch.empty(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + + # in-kernel-gated (the current path) + fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, pool, si, cu, ibuf_a, ci, o_gated, disable_state_update=True) + # prepass + host-gated (the H1 path) + q_n, k_n, g, beta = fused_recurrent_gdr_verify_prepass(q, k, a, b, A_log, dt_bias) + fused_recurrent_gdr_verify_fwd( + q_n, k_n, v, g, beta, pool, si, cu, ibuf_b, ci, o_pp, disable_state_update=True) + + assert _rel(o_pp, o_gated) <= 0.03, "prepass+host-gated o must match in-kernel-gated o" + assert _rel(ibuf_b, ibuf_a) <= 0.03, "prepass+host-gated ibuf must match in-kernel-gated ibuf" + + +@CUDA +def test_prepass_path_cuda_graph(): + # the auto-prepass path (large-batch regime) must be CUDA-graph capturable: the persistent + # scratch is warmed by the eager warmup runs, so capture sees no allocation, and replay + # reuses the same buffers -> must match eager bit-for-bit. + from flash_qla import recurrent_gated_delta_rule_verify + N, D, Hk, Hv = 64, 4, 8, 8 # N*Hv*n_vt comfortably saturates SMs -> should_use_prepass True + total = N * D + torch.manual_seed(5) + q = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(Hv, device="cuda").abs().log() + dt_bias = torch.randn(Hv, device="cuda") + pool = torch.randn(N, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf = torch.zeros(N + 1, D, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + o = torch.empty(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + + def run(): + recurrent_gated_delta_rule_verify( + A_log, a, dt_bias, q, k, v, b, pool, si, cu, ibuf, ci, + o=o, fuse_gating=True, prepass=True, disable_state_update=True) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): # warmup: JIT-compile prepass+main and allocate persistent scratch + run() + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): # would raise if the prepass allocated/synced inside capture + run() + graph.replay() + torch.cuda.synchronize() + o_graph = o.clone() + + o_eager = torch.empty_like(o) + recurrent_gated_delta_rule_verify( + A_log, a, dt_bias, q, k, v, b, pool, si, cu, ibuf, ci, + o=o_eager, fuse_gating=True, prepass=True, disable_state_update=True) + assert torch.equal(o_graph, o_eager), "prepass-path graph replay must match eager" diff --git a/tests/test_verify_gdr.py b/tests/test_verify_gdr.py new file mode 100644 index 0000000..a72f41b --- /dev/null +++ b/tests/test_verify_gdr.py @@ -0,0 +1,243 @@ +# tests/test_verify_gdr.py +import pytest +import torch + +from ref_gdr import verify_ref +from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_fwd, +) +from flash_qla.utils import l2norm + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _mk(N, D, Hk, Hv, num_slots, seed=0, ragged=False, distinct_gate=False): + torch.manual_seed(seed) + Ls = (torch.randint(1, D + 1, (N,)) if ragged else torch.full((N,), D)).to(torch.int32) + cu = torch.zeros(N + 1, dtype=torch.int32) + cu[1:] = torch.cumsum(Ls, 0) + total = int(cu[-1]) + q = l2norm(torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16)) + k = l2norm(torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16)) + v = torch.randn(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + if distinct_gate: # per-head distinct decay (catches a V-major transpose bug) + g = (torch.nn.functional.logsigmoid(torch.randn(1, total, Hv, device="cuda")) / 16 + * (1 + torch.arange(Hv, device="cuda").float()[None, None, :])) + else: + g = torch.nn.functional.logsigmoid(torch.randn(1, total, Hv, device="cuda")) / 16 + beta = torch.randn(1, total, Hv, device="cuda").sigmoid() + pool = torch.randn(num_slots, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) # V-major + ibuf = torch.zeros(N + 1, D, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + return q, k, v, g, beta, pool, si, cu.cuda(), ibuf, ci + + +def _rel(a, b): + return ((a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6)).item() + + +@CUDA +@pytest.mark.parametrize("ragged", [False, True]) +@pytest.mark.parametrize("Hk,Hv", [(8, 8), (2, 8)]) +def test_verify_nocommit(ragged, Hk, Hv): + N, D = 4, 4 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, Hk, Hv, num_slots=N, ragged=ragged) + o = torch.empty(1, q.shape[1], Hv, 128, device="cuda", dtype=torch.bfloat16) + pool0 = pool.clone() + o_ref, pool_ref, ibuf_ref = verify_ref(q, k, v, g, beta, pool, si, cu, ibuf, ci, disable_state_update=True) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + assert _rel(o, o_ref) <= 0.02 + assert _rel(ibuf, ibuf_ref) <= 0.02 + assert torch.equal(pool, pool0), "no-commit must not touch the pool" + + +@CUDA +def test_verify_commit(): + N, D, H = 4, 4, 8 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, H, H, num_slots=N) + o = torch.empty(1, q.shape[1], H, 128, device="cuda", dtype=torch.bfloat16) + pool_k = pool.clone() + o_ref, pool_ref, ibuf_ref = verify_ref(q, k, v, g, beta, pool, si, cu, ibuf, ci, disable_state_update=False) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool_k, si, cu, ibuf, ci, o, disable_state_update=False) + assert _rel(o, o_ref) <= 0.02 + assert _rel(pool_k, pool_ref) <= 0.02, "committed final state mismatch" + + +@CUDA +def test_verify_distinct_gate_vmajor(): + # per-head-distinct gates: a wrong V-major store would diverge from the reference + N, D, H = 4, 4, 8 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, H, H, num_slots=N, distinct_gate=True) + o = torch.empty(1, q.shape[1], H, 128, device="cuda", dtype=torch.bfloat16) + o_ref, pool_ref, ibuf_ref = verify_ref(q, k, v, g, beta, pool, si, cu, ibuf, ci, disable_state_update=True) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + assert _rel(o, o_ref) <= 0.02 + assert _rel(ibuf, ibuf_ref) <= 0.02 + + +@CUDA +def test_verify_wrapper_gating(): + # high-level wrapper: host-side sigmoid gating + l2norm must match a hand-built reference + from flash_qla import recurrent_gated_delta_rule_verify + from flash_qla.ops.gated_delta_rule.fused_recurrent import gdn_sigmoid_gate + + N, D, H = 4, 4, 8 + total = N * D + torch.manual_seed(1) + q = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, H, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, H, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(H, device="cuda").abs().log() + dt_bias = torch.randn(H, device="cuda") + pool = torch.randn(N, H, 128, 128, device="cuda", dtype=torch.bfloat16) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf = torch.zeros(N + 1, D, H, 128, 128, device="cuda", dtype=torch.bfloat16) + + g_ref, beta_ref = gdn_sigmoid_gate(A_log, a, dt_bias, b) + o_ref, _, ibuf_ref = verify_ref( + l2norm(q), l2norm(k), v, g_ref, beta_ref, pool, si, cu, ibuf, ci, disable_state_update=True) + o = recurrent_gated_delta_rule_verify( + A_log, a, dt_bias, q, k, v, b, pool, si, cu, ibuf, ci, disable_state_update=True) + assert _rel(o, o_ref) <= 0.02 + assert _rel(ibuf, ibuf_ref) <= 0.02 + + +@CUDA +@pytest.mark.parametrize("Hk,Hv", [(8, 8), (2, 8)]) +def test_verify_in_kernel_gating(Hk, Hv): + # in-kernel g/beta/l2norm must match the host-gating reference on the same raw inputs + from flash_qla.ops.gated_delta_rule.fused_recurrent import gdn_sigmoid_gate + from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_gated_fwd, + ) + + N, D = 4, 4 + total = N * D + torch.manual_seed(2) + q = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, Hk, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, Hv, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(Hv, device="cuda").abs().log() + dt_bias = torch.randn(Hv, device="cuda") + pool = torch.randn(N, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf = torch.zeros(N + 1, D, Hv, 128, 128, device="cuda", dtype=torch.bfloat16) + o = torch.empty(1, total, Hv, 128, device="cuda", dtype=torch.bfloat16) + + g_ref, beta_ref = gdn_sigmoid_gate(A_log, a, dt_bias, b) + o_ref, _, ibuf_ref = verify_ref( + l2norm(q), l2norm(k), v, g_ref, beta_ref, pool, si, cu, ibuf, ci, disable_state_update=True) + fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, pool, si, cu, ibuf, ci, o, disable_state_update=True) + assert _rel(o, o_ref) <= 0.03 + assert _rel(ibuf, ibuf_ref) <= 0.03 + + +@CUDA +def test_verify_cuda_graph(): + # the low-level entry must be CUDA-graph capturable (no host sync / no alloc) and replay correctly + N, D, H = 4, 4, 8 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, H, H, num_slots=N) + o = torch.empty(1, q.shape[1], H, 128, device="cuda", dtype=torch.bfloat16) + + def run(): + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + graph.replay() + torch.cuda.synchronize() + + o_eager = torch.empty_like(o) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o_eager, disable_state_update=True) + assert torch.equal(o, o_eager), "graph replay must match eager" + + +@CUDA +def test_verify_negctrl_vmajor_transpose(): + # discriminating: the V-major ibuf must DIFFER from a transposed-major reference + N, D, H = 4, 4, 8 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, H, H, num_slots=N, distinct_gate=True) + o = torch.empty(1, q.shape[1], H, 128, device="cuda", dtype=torch.bfloat16) + o_ref, _, ibuf_ref = verify_ref(q, k, v, g, beta, pool, si, cu, ibuf, ci, disable_state_update=True) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + assert _rel(ibuf, ibuf_ref) <= 0.02 + ibuf_wrong = ibuf_ref.transpose(-1, -2).contiguous() # K/V swapped (the silent-bug case) + assert _rel(ibuf, ibuf_wrong) > 0.2 # the test discriminates a wrong major-order + + +@CUDA +def test_verify_gated_cuda_graph(): + # the fully in-kernel gated path (no PyTorch gating/l2norm) is the capture-safe SGLang entry + from flash_qla.ops.gated_delta_rule.fused_recurrent.hopper.fused_recurrent_verify import ( + fused_recurrent_gdr_verify_gated_fwd, + ) + N, D, H = 4, 4, 8 + total = N * D + torch.manual_seed(3) + q = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(1, total, H, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, total, H, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(H, device="cuda").abs().log() + dt_bias = torch.randn(H, device="cuda") + pool = torch.randn(N, H, 128, 128, device="cuda", dtype=torch.bfloat16) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf = torch.zeros(N + 1, D, H, 128, 128, device="cuda", dtype=torch.bfloat16) + o = torch.empty(1, total, H, 128, device="cuda", dtype=torch.bfloat16) + + def run(): + fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, pool, si, cu, ibuf, ci, o, disable_state_update=True) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + graph.replay() + torch.cuda.synchronize() + o_eager = torch.empty_like(o) + fused_recurrent_gdr_verify_gated_fwd( + q, k, v, a, b, A_log, dt_bias, pool, si, cu, ibuf, ci, o_eager, disable_state_update=True) + assert torch.equal(o, o_eager), "in-kernel gated graph replay must match eager" + + +@CUDA +def test_verify_skip_slot(): + # a -1 pool slot must skip gather/commit/intermediate writes for that request + N, D, H = 3, 4, 8 + q, k, v, g, beta, pool, si, cu, ibuf, ci = _mk(N, D, H, H, num_slots=N) + si[1] = -1 # request 1 has no pool slot + o = torch.empty(1, q.shape[1], H, 128, device="cuda", dtype=torch.bfloat16) + ibuf0 = ibuf.clone() + o_ref, pool_ref, ibuf_ref = verify_ref(q, k, v, g, beta, pool, si, cu, ibuf, ci, disable_state_update=True) + fused_recurrent_gdr_verify_fwd(q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + # request 1's tokens: o still produced (from zero state), ibuf row left untouched + s0, s1 = int(cu[1]), int(cu[2]) + assert _rel(o[:, s0:s1], o_ref[:, s0:s1]) <= 0.02 + assert torch.equal(ibuf[int(ci[1])], ibuf0[int(ci[1])]), "skipped slot must not write ibuf" From 929da66cc5fb2bf81eb59ad008447cdf6c5637c0 Mon Sep 17 00:00:00 2001 From: Rifky Bujana Bisri Date: Wed, 1 Jul 2026 00:18:08 +0700 Subject: [PATCH 3/3] bench: standalone verify-kernel bandwidth benchmark FLA-free; measures wall time + achieved HBM bandwidth across SGLang-relevant server-batched and single-request/TP regimes. Co-Authored-By: Claude Opus 4.8 --- benchmark/bench_recurrent_gdr.py | 62 ++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 benchmark/bench_recurrent_gdr.py diff --git a/benchmark/bench_recurrent_gdr.py b/benchmark/bench_recurrent_gdr.py new file mode 100644 index 0000000..fe5fc9d --- /dev/null +++ b/benchmark/bench_recurrent_gdr.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] +"""Benchmark the FlashQLA GDN verify kernel (memory-bound): wall time + achieved HBM +bandwidth vs the theoretical state-I/O floor, across SGLang-relevant regimes.""" +import torch + +from flash_qla import fused_recurrent_gdr_verify_fwd +from flash_qla.utils import l2norm + +PEAK_TBS = 3.35 # H100 HBM3 ~3.35 TB/s + + +def _time(fn, iters=50, warmup=10): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters # ms + + +def bench(N, H, Hg, D, dtype=torch.bfloat16, store_intermediate=True): + total = N * D + q = l2norm(torch.randn(1, total, Hg, 128, device="cuda", dtype=dtype)) + k = l2norm(torch.randn(1, total, Hg, 128, device="cuda", dtype=dtype)) + v = torch.randn(1, total, H, 128, device="cuda", dtype=dtype) + g = torch.nn.functional.logsigmoid(torch.randn(1, total, H, device="cuda")) / 16 + beta = torch.randn(1, total, H, device="cuda").sigmoid() + pool = torch.randn(N, H, 128, 128, device="cuda", dtype=dtype) + cu = torch.arange(0, total + 1, D, dtype=torch.int32, device="cuda") + si = torch.arange(N, dtype=torch.int32, device="cuda") + ci = torch.arange(N, dtype=torch.int32, device="cuda") + ibuf = (torch.zeros(N + 1, D, H, 128, 128, device="cuda", dtype=dtype) + if store_intermediate else None) + o = torch.empty(1, total, H, 128, device="cuda", dtype=dtype) + + fn = lambda: fused_recurrent_gdr_verify_fwd( + q, k, v, g, beta, pool, si, cu, ibuf, ci, o, disable_state_update=True) + ms = _time(fn) + + es = 2 # bf16 state element size + # dominant state I/O: 1 gather + D intermediate writes per (request, head) + state_bytes = N * H * (1 + D) * 128 * 128 * es + io_bytes = (q.numel() + k.numel() + v.numel() + o.numel()) * es # tiny qkvo + gbs = (state_bytes + io_bytes) / (ms * 1e-3) / 1e9 + print(f" N={N:<4d} H={H} Hg={Hg} D={D:<2d} {ms*1e3:8.1f} us " + f"{gbs:7.1f} GB/s ({100*gbs/1000/PEAK_TBS:4.1f}% peak) state={state_bytes/1e6:6.1f} MB") + + +if __name__ == "__main__": + print(f"device: {torch.cuda.get_device_name()} | verify kernel (bf16 pool, per-token intermediates)") + print("server / batched decode (H=32, Hg=16):") + for N in (8, 64, 256, 512): + for D in (1, 4, 12): + bench(N, 32, 16, D) + print("single-request / TP (N=1, varying H = TP1..TP8):") + for H in (64, 32, 16, 8): + bench(1, H, max(1, H // 4), 12)