diff --git a/aiter/ops/triton/_gluon_kernels/gfx1250/attention/kda_decode.py b/aiter/ops/triton/_gluon_kernels/gfx1250/attention/kda_decode.py index f2b748e37de..71333bbb9d4 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx1250/attention/kda_decode.py +++ b/aiter/ops/triton/_gluon_kernels/gfx1250/attention/kda_decode.py @@ -2,6 +2,8 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import math + from triton.experimental import gluon from triton.experimental.gluon import language as gl from triton.language.extra.hip import libdevice @@ -12,6 +14,7 @@ softplus, ) from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir _TDM = gl.amd.gfx1250.tdm @@ -211,27 +214,64 @@ def _conv_qkv_load( ).to(gl.float32), hist_off, mask, + x_off, ), ) return loads +@gluon.jit +def _conv_qkv_roll(q_ptr, k_ptr, v_ptr, loads): + """Advance each q/k/v conv window one token in registers; load the next input.""" + rolled = () + for qkv in gl.static_range(3): + x, _h0, h1, h2, w0, w1, w2, w3, hist_off, mask, x_off = loads[qkv] + if qkv == 2: + x_ptr = v_ptr + else: + x_ptr = q_ptr if qkv == 0 else k_ptr + rolled = rolled + ( + ( + gl.amd.gfx1250.buffer_load(x_ptr, x_off, mask=mask, other=0), + h1, + h2, + x.to(h2.dtype), + w0, + w1, + w2, + w3, + hist_off, + mask, + x_off, + ), + ) + return rolled + + @gluon.jit def _conv_qkv_finish( - raw, loads, hist_ptr, pos_stride, K_LAYOUT: gl.constexpr, V_LAYOUT: gl.constexpr + raw, + loads, + hist_ptr, + pos_stride, + pos, + ROLL: gl.constexpr, + K_LAYOUT: gl.constexpr, + V_LAYOUT: gl.constexpr, ): - # causal conv + SiLU for q, k, v, and roll each history window by one token + """Causal conv + SiLU of q/k/v; the input goes to tap pos, ROLL shifts the taps.""" conv_out = () for qkv in gl.static_range(3): - x, h0, h1, h2, w0, w1, w2, w3, hist_off, mask = loads[qkv] + x, h0, h1, h2, w0, w1, w2, w3, hist_off, mask, _x_off = loads[qkv] acc = h0.to(gl.float32) * w0 + h1.to(gl.float32) * w1 acc += h2.to(gl.float32) * w2 + x.to(gl.float32) * w3 - gl.amd.gfx1250.buffer_store(h1, hist_ptr, hist_off, mask=mask) - gl.amd.gfx1250.buffer_store(h2, hist_ptr, hist_off + pos_stride, mask=mask) + if ROLL: + gl.amd.gfx1250.buffer_store(h1, hist_ptr, hist_off, mask=mask) + gl.amd.gfx1250.buffer_store(h2, hist_ptr, hist_off + pos_stride, mask=mask) gl.amd.gfx1250.buffer_store( x.to(hist_ptr.dtype.element_ty), hist_ptr, - hist_off + 2 * pos_stride, + hist_off + pos * pos_stride, mask=mask, ) conv_out = conv_out + (acc * sigmoid(acc),) @@ -406,6 +446,10 @@ def fused_recurrent_kda_packed_decode_kernel( "USE_CONV/USE_RMS_GATE run on the register-prefetch path", ) gl.static_assert((not USE_CONV) or USE_INITIAL_STATE, "USE_CONV needs a state slot") + gl.static_assert( + (not (USE_CONV and IS_SPEC_DECODING)) or IS_CONTINUOUS_BATCHING, + "spec decoding's conv window lives in paged state slots", + ) CR: gl.constexpr = (32 * NUM_WARPS) // K if 32 * NUM_WARPS > K else 1 gl.static_assert( (not USE_CONV) @@ -451,11 +495,9 @@ def fused_recurrent_kda_packed_decode_kernel( ) if pad_slot <= 0: zero = gl.full([BV], 0.0, o_ptr.dtype.element_ty, V_LAYOUT) - gl.amd.gfx1250.buffer_store( - zero, - o_ptr + bos.to(gl.int64) * stride_o_token + i_hv * V + i_v * BV, - off_v, - ) + o_z = o_ptr + bos.to(gl.int64) * stride_o_token + i_hv * V + i_v * BV + for t in range(n_tok): + gl.amd.gfx1250.buffer_store(zero, o_z + t * stride_o_token, off_v) return col = 0 if STATE_V_FIRST else i_v * BV @@ -598,13 +640,20 @@ def fused_recurrent_kda_packed_decode_kernel( [CR, NUM_WARPS // CR], [1, 0], ) - hb = conv_state_ptr + slot.to(gl.int64) * stride_cs_slot - # conv loads first, then the state stream, the conv math and its history stores run under the state load + if IS_SPEC_DECODING: + cslot = gl.load(state_indices_ptr + i_n * stride_indices_seq).to( + gl.int32 + ) + hb = conv_state_ptr + cslot.to(gl.int64) * stride_cs_slot + hb_in = hb + seed * stride_cs_pos + else: + hb = conv_state_ptr + slot.to(gl.int64) * stride_cs_slot + hb_in = hb conv_loads = _conv_qkv_load( q_p, k_p, v_p, - hb, + hb_in, conv_weight_ptr, i_h, i_hv, @@ -650,7 +699,7 @@ def fused_recurrent_kda_packed_decode_kernel( _TDM.async_load(desc_in, [row_in, col], smem_in) if USE_CONV: nxt = _conv_qkv_finish( - nxt, conv_loads, hb, stride_cs_pos, K_LAYOUT, V_LAYOUT + nxt, conv_loads, hb, stride_cs_pos, W - 2, True, K_LAYOUT, V_LAYOUT ) _TDM.async_wait(0) S = smem_in.load(STATE_LAYOUT) @@ -660,7 +709,7 @@ def fused_recurrent_kda_packed_decode_kernel( ).to(gl.float32) if USE_CONV: nxt = _conv_qkv_finish( - nxt, conv_loads, hb, stride_cs_pos, K_LAYOUT, V_LAYOUT + nxt, conv_loads, hb, stride_cs_pos, W - 2, True, K_LAYOUT, V_LAYOUT ) if CACHE_STATE_UPDATES and IS_SPEC_DECODING: for j in range(1, seed + 1): @@ -761,27 +810,47 @@ def fused_recurrent_kda_packed_decode_kernel( q_p, k_p, v_p, g_p, b_p, off_k, off_v, IS_BETA_HEADWISE ) if USE_CONV: - conv_loads = _conv_qkv_load( - q_p, - k_p, - v_p, - hb, - conv_weight_ptr, - i_h, - i_hv, - stride_cs_dim, - stride_cs_pos, - H * K, - W, - CR, - K, - V, - FLATK, - FLATV, - ) - nxt = _conv_qkv_finish( - nxt, conv_loads, hb, stride_cs_pos, K_LAYOUT, V_LAYOUT - ) + if IS_SPEC_DECODING: # token t+1's input lands at W-1+t + conv_loads = _conv_qkv_roll(q_p, k_p, v_p, conv_loads) + nxt = _conv_qkv_finish( + nxt, + conv_loads, + hb, + stride_cs_pos, + W - 1 + t, + False, + K_LAYOUT, + V_LAYOUT, + ) + else: + conv_loads = _conv_qkv_load( + q_p, + k_p, + v_p, + hb, + conv_weight_ptr, + i_h, + i_hv, + stride_cs_dim, + stride_cs_pos, + H * K, + W, + CR, + K, + V, + FLATK, + FLATV, + ) + nxt = _conv_qkv_finish( + nxt, + conv_loads, + hb, + stride_cs_pos, + W - 2, + True, + K_LAYOUT, + V_LAYOUT, + ) if USE_RMS_GATE: og_p += stride_og_token og_nxt = gl.amd.gfx1250.buffer_load(og_p, off_v) @@ -917,3 +986,43 @@ def fused_recurrent_kda_packed_decode_kernel( state_out_ptr + row_out.to(gl.int64) * ROWLEN + col, off_s, ) + + +def get_kda_config( + avg_T: int, + num_seqs: int, + HV: int, + K: int, + V: int, + overrides: dict | None = None, + fused: bool = False, +) -> dict: + cfg_dir = resolve_config_dir("attention", "KDA_DECODE", backend="gluon") + tuned = load_config_json(f"{cfg_dir}/DEFAULT.json") + num_seq_heads = num_seqs * HV + aligned = K % 32 == 0 and V % 32 == 0 + if fused: + bucket = "fused" + elif avg_T > 1: + if aligned and num_seq_heads >= 3072 and V % 128 == 0: + bucket = "t_gt1_seq_heads_geq_3072" + elif aligned and num_seq_heads >= 384: + bucket = "t_gt1_seq_heads_geq_384" + elif K % 16 == 0: + bucket = "t_gt1_default" + else: + bucket = "default" + elif aligned and 256 <= num_seq_heads <= 512: + bucket = "t1_seq_heads_256_to_512" + else: + bucket = "default" + config = dict(tuned[bucket]) + overrides = overrides or {} + config.update(overrides) + if "SK" not in overrides: + config["SK"] = min(config["SK"], math.gcd(32, K)) + if "num_warps" not in overrides: + config["num_warps"] = max( + 1, min(config["num_warps"], config["BV"] * config["SK"] // 32) + ) + return config diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/kda_decode.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/kda_decode.py index 5626e78bc61..512dad988b6 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/kda_decode.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/kda_decode.py @@ -2,6 +2,8 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import math + from triton.experimental import gluon from triton.experimental.gluon import language as gl @@ -11,6 +13,20 @@ softplus, ) from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir + + +@gluon.jit +def _conv_load(hb, wb, ho, wo, m, s_pos, LP: gl.constexpr): + """The three history taps and four conv weights of this head's q/k/v rows.""" + h0 = gl.amd.cdna4.buffer_load(hb, ho, mask=m, cache=".cg") + h1 = gl.amd.cdna4.buffer_load(hb, ho + s_pos, mask=m, cache=".cg") + h2 = gl.amd.cdna4.buffer_load(hb, ho + 2 * s_pos, mask=m, cache=".cg") + w0 = gl.amd.cdna4.buffer_load(wb, wo, mask=m).to(gl.float32) + w1 = gl.amd.cdna4.buffer_load(wb, wo + LP, mask=m).to(gl.float32) + w2 = gl.amd.cdna4.buffer_load(wb, wo + 2 * LP, mask=m).to(gl.float32) + w3 = gl.amd.cdna4.buffer_load(wb, wo + 3 * LP, mask=m).to(gl.float32) + return h0, h1, h2, w0, w1, w2, w3 @gluon.jit @@ -36,13 +52,7 @@ def _token( # g has its own dtype: load it apart and splice it into row 2 x = gl.where(m, x.to(gl.float32), gl.load(gp, mask=~m).to(gl.float32)) if USE_CONV: - h0 = gl.amd.cdna4.buffer_load(hb, ho, mask=m, cache=".cg") - h1 = gl.amd.cdna4.buffer_load(hb, ho + s_pos, mask=m, cache=".cg") - h2 = gl.amd.cdna4.buffer_load(hb, ho + 2 * s_pos, mask=m, cache=".cg") - w0 = gl.amd.cdna4.buffer_load(wb, wo, mask=m).to(gl.float32) - w1 = gl.amd.cdna4.buffer_load(wb, wo + LP, mask=m).to(gl.float32) - w2 = gl.amd.cdna4.buffer_load(wb, wo + 2 * LP, mask=m).to(gl.float32) - w3 = gl.amd.cdna4.buffer_load(wb, wo + 3 * LP, mask=m).to(gl.float32) + h0, h1, h2, w0, w1, w2, w3 = _conv_load(hb, wb, ho, wo, m, s_pos, LP) acc = h0.to(gl.float32) * w0 + h1.to(gl.float32) * w1 acc += h2.to(gl.float32) * w2 + x.to(gl.float32) * w3 gl.amd.cdna4.buffer_store(h1, hb, ho, mask=m, cache=".cs") @@ -52,6 +62,24 @@ def _token( ) else: acc = x.to(gl.float32) + return _finish( + x, acc, m, a_exp, dt, gate_c, USE_CONV, USE_GATE_IN_KERNEL, USE_LOWER_BOUND + ) + + +@gluon.jit +def _finish( + x, + acc, + m, + a_exp, + dt, + gate_c, + USE_CONV: gl.constexpr, + USE_GATE_IN_KERNEL: gl.constexpr, + USE_LOWER_BOUND: gl.constexpr, +): + """Gate chain on row 2 and SiLU on the convolved q/k/v rows.""" x = x.to(gl.float32) if USE_GATE_IN_KERNEL and USE_LOWER_BOUND: sg = sigmoid(gl.where(m, acc, a_exp * (x + dt))) @@ -69,7 +97,15 @@ def _token( _fused_recurrent_kda_packed_decode_repr = make_kernel_repr( "fused_recurrent_kda_packed_decode_kernel", - ["BV", "SK", "NUM_WARPS", "PAD_SLOT_GUARD", "USE_CONV", "USE_RMS_GATE"], + [ + "BV", + "SK", + "NUM_WARPS", + "PAD_SLOT_GUARD", + "USE_CONV", + "USE_RMS_GATE", + "RESIDENT_STATE", + ], ) @@ -142,6 +178,7 @@ def fused_recurrent_kda_packed_decode_kernel( USE_CONV: gl.constexpr = False, USE_RMS_GATE: gl.constexpr = False, TDM_STORE_BUFS: gl.constexpr = 2, + RESIDENT_STATE: gl.constexpr = False, # multi-token: state held in registers ): gl.static_assert(STATE_V_FIRST, "gfx950 kda decode keeps the state [V, K]") gl.static_assert(BV == V and K == V, "gfx950 kda decode needs BV == V == K") @@ -162,14 +199,15 @@ def fused_recurrent_kda_packed_decode_kernel( gl.static_assert((not USE_CONV) or W == 4, "USE_CONV needs W == 4") gl.static_assert( not ( - IS_SPEC_DECODING - or CACHE_STATE_UPDATES - or USE_TDM_STORE - or USE_TDM_LOAD - or USE_TDM_FUSED_LOAD + CACHE_STATE_UPDATES or USE_TDM_STORE or USE_TDM_LOAD or USE_TDM_FUSED_LOAD ), "not supported on gfx950", ) + gl.static_assert( + (not IS_SPEC_DECODING) + or (RESIDENT_STATE and IS_CONTINUOUS_BATCHING and INPLACE_FINAL_STATE), + "spec decoding runs the paged, in-place, register-resident path", + ) gl.static_assert( (not PAD_SLOT_GUARD) or IS_CONTINUOUS_BATCHING, "PAD_SLOT_GUARD is a paged-mode contract", @@ -206,6 +244,11 @@ def fused_recurrent_kda_packed_decode_kernel( else: bos = i_n * T n_tok = T + # Kept nested: RESIDENT_STATE is constexpr, so only resident variants get the + # runtime early return (and the control-flow rewrite it brings). + if RESIDENT_STATE: # noqa: SIM102 + if n_tok == 0: + return tok0 = bos.to(gl.int64) o_p = o_ptr + tok0 * stride_o_token + i_hv * V @@ -259,17 +302,19 @@ def fused_recurrent_kda_packed_decode_kernel( og_p = out_gate_ptr + tok0 * stride_og_token + i_hv * V nw = gl.amd.cdna4.buffer_load(norm_weight_ptr, off_o).to(gl.float32) if IS_CONTINUOUS_BATCHING: - slot = gl.load(state_indices_ptr + i_n * stride_indices_seq).to(gl.int32) + if IS_SPEC_DECODING: + seed = gl.maximum(gl.load(num_accepted_ptr + i_n).to(gl.int32) - 1, 0) + else: + seed = 0 + slot = gl.load(state_indices_ptr + i_n * stride_indices_seq + seed).to(gl.int32) else: + seed = 0 slot = i_n - if PAD_SLOT_GUARD: # noqa: SIM102 (constexpr guard, runtime test) - if slot <= 0: - gl.store( - o_p + off_o, - gl.full([V], 0.0, o_ptr.dtype.element_ty, OUT), - mask=off_o < n_tok * V, - ) - return + if PAD_SLOT_GUARD and slot <= 0: + zero = gl.full([V], 0.0, o_ptr.dtype.element_ty, OUT) + for t in range(n_tok): + gl.store(o_p + t * stride_o_token + off_o, zero) + return if USE_INITIAL_STATE: s_src = state_ptr + (slot * stride_state_slot_rows + i_hv * V).to(gl.int64) * K @@ -278,16 +323,57 @@ def fused_recurrent_kda_packed_decode_kernel( else: s_src = state_out_ptr if USE_CONV: - hb = conv_state_ptr + slot.to(gl.int64) * stride_cs_slot + if IS_SPEC_DECODING: + cs_slot = gl.load(state_indices_ptr + i_n * stride_indices_seq).to(gl.int32) + else: + cs_slot = slot + hb = conv_state_ptr + cs_slot.to(gl.int64) * stride_cs_slot else: hb = conv_state_ptr qkav_smem = gl.allocate_shared_memory(gl.float32, [4 * K], SMEM1) obuf = gl.allocate_shared_memory(gl.float32, [V], SMEM1) if USE_RMS_GATE: sbuf = gl.allocate_shared_memory(gl.float32, [RC], SMEM1) + if RESIDENT_STATE: + st = () + for i in gl.static_range(NC): + if USE_INITIAL_STATE: + st = st + ( + gl.amd.cdna4.buffer_load(s_src + i * K, off_s, cache=".cg").to( + gl.float32 + ), + ) + else: + st = st + (gl.zeros([RC, K], gl.float32, CHUNK),) + if USE_CONV: + h0, h1, h2, w0, w1, w2, w3 = _conv_load( + hb, conv_weight_ptr, ho + seed * stride_cs_pos, wo, m, stride_cs_pos, LP + ) + if IS_SPEC_DECODING: + keep1 = h1 + keep2 = h2 + x_nxt = gl.load(xp) + if gp is not None: + g_nxt = gl.load(gp, mask=~m) + if not IS_BETA_HEADWISE: + b_nxt = gl.load(b_p) + if USE_RMS_GATE: + og_nxt = gl.amd.cdna4.buffer_load(og_p, off_o) + if INPLACE_FINAL_STATE and IS_CONTINUOUS_BATCHING: + slot_nxt = gl.load(state_indices_ptr + i_n * stride_indices_seq) for t in range(n_tok): - if INPLACE_FINAL_STATE and IS_CONTINUOUS_BATCHING and n_tok > 1: + if RESIDENT_STATE: + more = (t + 1 < n_tok).to(gl.int32) + if RESIDENT_STATE and INPLACE_FINAL_STATE and IS_CONTINUOUS_BATCHING: + nxt = slot_nxt.to(gl.int32) + slot_nxt = gl.load(state_indices_ptr + i_n * stride_indices_seq + t + more) + if PAD_SLOT_GUARD and not IS_SPEC_DECODING: + nxt = gl.where(nxt > 0, nxt, slot) + slot = nxt + row_out = slot * stride_state_out_slot_rows + i_hv * V + s_dst = state_out_ptr + row_out.to(gl.int64) * K + elif INPLACE_FINAL_STATE and IS_CONTINUOUS_BATCHING and n_tok > 1: nxt = gl.load(state_indices_ptr + i_n * stride_indices_seq + t).to(gl.int32) if PAD_SLOT_GUARD: nxt = gl.where(nxt > 0, nxt, slot) @@ -303,35 +389,73 @@ def fused_recurrent_kda_packed_decode_kernel( if not USE_INITIAL_STATE: s_mask = (off_s >= 0) & (t > 0) s_other = 0.0 - if not IS_BETA_HEADWISE: - b = gl.load(b_p).to(gl.float32) - if USE_RMS_GATE: - og = gl.amd.cdna4.buffer_load(og_p, off_o) - bufs = () - for i in gl.static_range(NUM_BUFFERS): - bufs = bufs + ( - gl.amd.cdna4.buffer_load( - s_src + i * K, off_s, mask=s_mask, other=s_other, cache=".cg" - ), + if RESIDENT_STATE: + x = x_nxt + x_nxt = gl.load(xp + xs * more) + if gp is not None: + # g has its own dtype: loaded apart and spliced into row 2 + x = gl.where(m, x.to(gl.float32), g_nxt.to(gl.float32)) + g_nxt = gl.load(gp + stride_g_token * more, mask=~m) + if not IS_BETA_HEADWISE: + b = b_nxt.to(gl.float32) + b_nxt = gl.load(b_p + stride_beta_token * more) + if USE_RMS_GATE: + og = og_nxt + og_nxt = gl.amd.cdna4.buffer_load(og_p + stride_og_token * more, off_o) + else: + if not IS_BETA_HEADWISE: + b = gl.load(b_p).to(gl.float32) + if USE_RMS_GATE: + og = gl.amd.cdna4.buffer_load(og_p, off_o) + if RESIDENT_STATE: + if USE_CONV: + acc = h0.to(gl.float32) * w0 + h1.to(gl.float32) * w1 + acc += h2.to(gl.float32) * w2 + x.to(gl.float32) * w3 + h0 = h1 + h1 = h2 + h2 = x.to(hb.dtype.element_ty) + if IS_SPEC_DECODING: # this token's input lands at W-2+t + gl.amd.cdna4.buffer_store( + h2, hb, ho + (W - 2 + t) * stride_cs_pos, mask=m, cache=".cs" + ) + else: + acc = x.to(gl.float32) + y = _finish( + x, + acc, + m, + a_exp, + dt, + gate_c, + USE_CONV, + USE_GATE_IN_KERNEL, + USE_LOWER_BOUND, + ) + else: + bufs = () + for i in gl.static_range(NUM_BUFFERS): + bufs = bufs + ( + gl.amd.cdna4.buffer_load( + s_src + i * K, off_s, mask=s_mask, other=s_other, cache=".cg" + ), + ) + y = _token( + xp, + gp, + hb, + conv_weight_ptr, + ho, + wo, + m, + stride_cs_pos, + a_exp, + dt, + gate_c, + LP, + USE_CONV, + USE_GATE_IN_KERNEL, + USE_LOWER_BOUND, ) - - y = _token( - xp, - gp, - hb, - conv_weight_ptr, - ho, - wo, - m, - stride_cs_pos, - a_exp, - dt, - gate_c, - LP, - USE_CONV, - USE_GATE_IN_KERNEL, - USE_LOWER_BOUND, - ) if not IS_BETA_HEADWISE and APPLY_BETA_SIGMOID: b = sigmoid(b) if ALLOW_NEG_EIGVAL: @@ -351,21 +475,35 @@ def fused_recurrent_kda_packed_decode_kernel( qv = qv * scale qk = qk * scale + if RESIDENT_STATE and not IS_CONTINUOUS_BATCHING: + s_keep = (off_s >= 0) & (t == n_tok - 1) + elif IS_SPEC_DECODING and PAD_SLOT_GUARD: + s_keep = (off_s >= 0) & (slot > 0) + else: + s_keep = None ss = gl.zeros([RC], gl.float32, R_LAYOUT) + if RESIDENT_STATE: + upd = () + vvs = () for i in gl.static_range(NC): - S = bufs[0] - if i + NUM_BUFFERS < NC: - nxt = gl.amd.cdna4.buffer_load( - s_src + (i + NUM_BUFFERS) * K, - off_s, - mask=s_mask, - other=s_other, - cache=".cg", - ) - bufs = bufs[1:] + (nxt,) + vvs = vvs + (qkav_smem.slice(3 * K, K).gather(rows + i, 0),) + for i in gl.static_range(NC): + if RESIDENT_STATE: + S = st[i] else: - bufs = bufs[1:] - vv = qkav_smem.slice(3 * K, K).gather(rows + i, 0) + S = bufs[0] + if i + NUM_BUFFERS < NC: + nxt = gl.amd.cdna4.buffer_load( + s_src + (i + NUM_BUFFERS) * K, + off_s, + mask=s_mask, + other=s_other, + cache=".cg", + ) + bufs = bufs[1:] + (nxt,) + else: + bufs = bufs[1:] + vv = vvs[i] if IS_BETA_HEADWISE: b = gl.load(b_p + rows + i).to(gl.float32) if APPLY_BETA_SIGMOID: @@ -380,10 +518,18 @@ def fused_recurrent_kda_packed_decode_kernel( S = S + err[:, None] * kv[None, :] # write: S += err (x) k gl.amd.cdna4.buffer_store( - S.to(s_dst.dtype.element_ty), s_dst + i * K, off_s, cache=".wt" + S.to(s_dst.dtype.element_ty), + s_dst + i * K, + off_s, + mask=s_keep, + cache=".wt", ) + if RESIDENT_STATE: + upd = upd + (S,) obuf.slice(i * RC, RC).store(o) ss += o * o + if RESIDENT_STATE: + st = upd if USE_RMS_GATE: sbuf.store(ss) o = obuf.gather(o_idx, 0) @@ -398,3 +544,57 @@ def fused_recurrent_kda_packed_decode_kernel( gp += stride_g_token b_p += stride_beta_token o_p += stride_o_token + + if RESIDENT_STATE and USE_CONV: + if IS_SPEC_DECODING: + gl.amd.cdna4.buffer_store(keep1, hb, ho, mask=m, cache=".cs") + gl.amd.cdna4.buffer_store( + keep2, hb, ho + stride_cs_pos, mask=m, cache=".cs" + ) + else: + gl.amd.cdna4.buffer_store(h0, hb, ho, mask=m, cache=".cs") + gl.amd.cdna4.buffer_store(h1, hb, ho + stride_cs_pos, mask=m, cache=".cs") + gl.amd.cdna4.buffer_store( + h2, hb, ho + 2 * stride_cs_pos, mask=m, cache=".cs" + ) + + +def get_kda_config( + avg_T: int, + num_seqs: int, + HV: int, + K: int, + V: int, + overrides: dict | None = None, + fused: bool = False, +) -> dict: + cfg_dir = resolve_config_dir("attention", "KDA_DECODE", backend="gluon") + tuned = load_config_json(f"{cfg_dir}/DEFAULT.json") + num_seq_heads = num_seqs * HV + aligned = K % 32 == 0 and V % 32 == 0 + if fused: + multi = avg_T > 1 and num_seq_heads < 3072 + bucket = "fused_t_gt1" if multi else "fused" + elif avg_T > 1: + if aligned and num_seq_heads >= 3072 and V % 128 == 0: + bucket = "t_gt1_seq_heads_geq_3072" + elif aligned and num_seq_heads >= 384: + bucket = "t_gt1_seq_heads_geq_384" + elif K % 16 == 0: + bucket = "t_gt1_default" + else: + bucket = "default" + elif aligned and 256 <= num_seq_heads <= 512: + bucket = "t1_seq_heads_256_to_512" + else: + bucket = "default" + config = dict(tuned[bucket]) + overrides = overrides or {} + config.update(overrides) + if "SK" not in overrides: + config["SK"] = min(config["SK"], math.gcd(32, K)) + if "num_warps" not in overrides: + config["num_warps"] = max( + 1, min(config["num_warps"], config["BV"] * config["SK"] // 32) + ) + return config diff --git a/aiter/ops/triton/attention/kda.py b/aiter/ops/triton/attention/kda.py index 99577be061d..33ad7ad312a 100644 --- a/aiter/ops/triton/attention/kda.py +++ b/aiter/ops/triton/attention/kda.py @@ -2,13 +2,11 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import logging -import math import torch import triton from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -17,56 +15,19 @@ _ARCH = arch_info.get_arch() fused_recurrent_kda_packed_decode_kernel = None +get_kda_config = None if _ARCH == "gfx1250": from aiter.ops.triton._gluon_kernels.gfx1250.attention.kda_decode import ( fused_recurrent_kda_packed_decode_kernel, + get_kda_config, ) elif _ARCH == "gfx950": from aiter.ops.triton._gluon_kernels.gfx950.attention.kda_decode import ( fused_recurrent_kda_packed_decode_kernel, + get_kda_config, ) -def get_kda_config( - avg_T: int, - num_seqs: int, - HV: int, - K: int, - V: int, - overrides: dict | None = None, - fused: bool = False, -) -> dict: - cfg_dir = resolve_config_dir("attention", "KDA_DECODE", backend="gluon") - tuned = load_config_json(f"{cfg_dir}/DEFAULT.json") - num_seq_heads = num_seqs * HV - aligned = K % 32 == 0 and V % 32 == 0 - if fused: - bucket = "fused" - elif avg_T > 1: - if aligned and num_seq_heads >= 3072 and V % 128 == 0: - bucket = "t_gt1_seq_heads_geq_3072" - elif aligned and num_seq_heads >= 384: - bucket = "t_gt1_seq_heads_geq_384" - elif K % 16 == 0: - bucket = "t_gt1_default" - else: - bucket = "default" - elif aligned and 256 <= num_seq_heads <= 512: - bucket = "t1_seq_heads_256_to_512" - else: - bucket = "default" - config = dict(tuned[bucket]) - overrides = overrides or {} - config.update(overrides) - if "SK" not in overrides: - config["SK"] = min(config["SK"], math.gcd(32, K)) - if "num_warps" not in overrides: - config["num_warps"] = max( - 1, min(config["num_warps"], config["BV"] * config["SK"] // 32) - ) - return config - - def fused_recurrent_kda( q: torch.Tensor, k: torch.Tensor, @@ -167,7 +128,9 @@ def fused_recurrent_kda( W = conv_weight.shape[1] if use_conv else 4 if use_conv: assert conv_weight.shape == (3, W, H * K) and conv_weight.is_contiguous() - assert conv_state.shape[1:] == (3 * H * K, W - 1) and HV * V == H * K + assert conv_state.shape[1] == 3 * H * K and HV * V == H * K + if num_accepted_tokens is None: + assert conv_state.shape[2] >= W - 1 if use_rms_gate: assert out_gate.shape == v.shape and out_gate.stride()[2:] == (V, 1) assert B == 1 or out_gate.stride(0) == T * out_gate.stride(1) @@ -240,6 +203,9 @@ def fused_recurrent_kda( use_tdm_fused_load = config.get("use_tdm_fused_load", False) tdm_store_bufs = config.get("tdm_store_bufs", 2) sched_strategy = config.get("sched_strategy") + resident_state = config.get( + "resident_state", avg_T > 1 or num_accepted_tokens is not None + ) assert V % BV == 0, f"BV={BV} must divide V={V}" assert 32 % SK == 0 and K % SK == 0, f"SK={SK} must divide 32 and K={K}" assert (BV * SK) % ( @@ -266,6 +232,13 @@ def fused_recurrent_kda( stride_indices_seq = 1 assert num_accepted_tokens is None, "spec decoding requires ssm_state_indices" + if use_conv and num_accepted_tokens is not None: + # spec window: taps read from num_accepted-1 (< T); token t lands at W-2+t + max_tok = ssm_state_indices.shape[1] if ssm_state_indices.ndim > 1 else 1 + assert ( + conv_state.shape[2] >= W - 2 + max_tok + ), f"spec conv_state needs {W - 2 + max_tok} taps, got {conv_state.shape[2]}" + if initial_state is not None: assert initial_state.dtype == torch.float32 assert ( @@ -407,6 +380,7 @@ def _rows(t): W=W, USE_CONV=use_conv, USE_RMS_GATE=use_rms_gate, + **({"RESIDENT_STATE": bool(resident_state)} if _ARCH == "gfx950" else {}), num_warps=num_warps, **( {"llvm_fn_attrs": f"amdgpu-sched-strategy={sched_strategy}"} diff --git a/aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json index a0e96b8b649..83767c67303 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json @@ -5,6 +5,12 @@ "num_warps": 4, "num_buffers": 8 }, + "fused_t_gt1": { + "BV": 128, + "SK": 16, + "num_warps": 8, + "num_buffers": 2 + }, "t_gt1_seq_heads_geq_3072": { "BV": 128, "SK": 16, @@ -14,8 +20,8 @@ "t_gt1_seq_heads_geq_384": { "BV": 128, "SK": 16, - "num_warps": 4, - "num_buffers": 4 + "num_warps": 8, + "num_buffers": 2 }, "t1_seq_heads_256_to_512": { "BV": 128, @@ -26,8 +32,8 @@ "t_gt1_default": { "BV": 128, "SK": 16, - "num_warps": 4, - "num_buffers": 4 + "num_warps": 8, + "num_buffers": 2 }, "default": { "BV": 128, diff --git a/op_tests/op_benchmarks/triton/bench_kda.py b/op_tests/op_benchmarks/triton/bench_kda.py index 64b2741717b..1d677f9a809 100644 --- a/op_tests/op_benchmarks/triton/bench_kda.py +++ b/op_tests/op_benchmarks/triton/bench_kda.py @@ -74,9 +74,11 @@ def make_inputs( extra = {} if fused: + # spec decoding keeps vLLM's W-2+T tap window, read from num_accepted-1 + taps = W - 1 if accepted is None else W - 2 + T extra = { "conv_state": torch.randn( - state.shape[0], 3 * H * D, W - 1, dtype=dtype, device=device + state.shape[0], 3 * H * D, taps, dtype=dtype, device=device ), "conv_weight": torch.randn(3, W, H * D, dtype=torch.float32, device=device) * 0.1, @@ -328,8 +330,8 @@ def parse_args(): "--gluon_configs", type=str, nargs="+", - default=["32,4,32"], - help="BV,NUM_WARPS,SK triples, one benchmark line each", + default=None, + help="BV,NUM_WARPS,SK triples, one benchmark line each (default per arch)", ) parser.add_argument( "--num_buffers", @@ -448,6 +450,10 @@ def parse_args(): def run_bench(args): if arch_info.get_arch() not in ("gfx950", "gfx1250"): sys.exit(f"KDA gluon decode needs gfx950/gfx1250, got {arch_info.get_arch()}") + if args.gluon_configs is None: + # The gfx950 kernel holds the whole [V, K] state: BV == V. + gfx950 = arch_info.get_arch() == "gfx950" + args.gluon_configs = [f"{args.head_dim},4,16"] if gfx950 else ["32,4,32"] if "fla" in args.backends and not HAS_FLA: print("fla not importable -- dropping the upstream backend.") print(" PYTHONPATH=/path/to/flash-linear-attention to enable it") diff --git a/op_tests/triton_tests/attention/test_kda.py b/op_tests/triton_tests/attention/test_kda.py index ea1316dc2e0..acd27fea646 100644 --- a/op_tests/triton_tests/attention/test_kda.py +++ b/op_tests/triton_tests/attention/test_kda.py @@ -22,14 +22,13 @@ arch not in ("gfx950", "gfx1250"), reason=f"KDA gluon decode needs gfx950/gfx1250, got {arch}", ) -# gfx950 runs the vLLM decode path only: V-first state, head_dim 128, no spec -# decoding, cache_state_updates or tiling knobs + gfx1250_only = pytest.mark.skipif(arch != "gfx1250", reason="gfx1250-only feature") -def skip_on_gfx950(D=128, spec=False): - if arch == "gfx950" and (D != 128 or spec): - pytest.skip("gfx950 KDA decode needs head_dim 128 and no spec decoding") +def skip_on_gfx950(D=128): + if arch == "gfx950" and D != 128: + pytest.skip("gfx950 KDA decode needs head_dim 128") DEVICE = "cuda" @@ -444,7 +443,6 @@ def test_fused_recurrent_varlen(lens): @pytest.mark.parametrize("spec", [False, True]) def test_fused_recurrent_vllm_decode(B, T, spec): """Continuous batching with paged state, poisoning untouched slots.""" - skip_on_gfx950(spec=spec) H, D = 8, 128 pool_kv, indices, untouched = make_pool(B, T, H, D) @@ -1325,8 +1323,7 @@ def test_pad_slot_guard(): assert_close(f"o[{n}]", ref, o[:, b:e]) assert_close(f"ht[{n}]", ref_ht[0], pool_out[int(indices[n, T - 1])]) - assert torch.all(o[:, T] == 0), "pad seq token 0 must be zeroed" - assert torch.all(o[:, T + 1 : 2 * T] == 123.0), "pad seq must not be computed" + assert torch.all(o[:, T : 2 * T] == 0), "every pad seq token must be zeroed" assert torch.all(pool_out[0] == POISON), "the pad slot was written" @@ -1501,6 +1498,88 @@ def test_fused_conv_rms_gate(N, T, H): assert_close("ht", S2, S1) +@pytest.mark.parametrize("N, k", [(3, 3), (2, 7), (5, 1)]) +@pytest.mark.parametrize("padded", [False, True]) +def test_fused_conv_rms_gate_spec(N, k, padded): + """vLLM's spec conv window: W-1+k taps in slot column 0, read from + num_accepted-1; the step leaves taps 1..W-2 followed by its inputs. A padded + sequence (slot 0) zeroes all of its tokens and writes no state.""" + D, W, eps, lb, H = 128, 4, 1e-6, -5.0, 8 + T = k + 1 + lp, TT, NS = H * D, N * T, N * T + 4 + torch.manual_seed(N * 11 + k) + mixed = torch.randn(TT, 3 * lp, dtype=torch.bfloat16, device=DEVICE) + cw = torch.randn(3 * lp, W, dtype=torch.bfloat16, device=DEVICE) * 0.1 + cs = torch.randn(NS, W - 1 + k, 3 * lp, dtype=torch.bfloat16, device=DEVICE) + cs = cs.transpose(-1, -2) * 0.1 # vLLM's SD storage, seen as (dim, taps) + g = torch.randn(1, TT, H, D, dtype=torch.bfloat16, device=DEVICE) * 0.5 + beta = torch.randn(1, TT, H, dtype=torch.bfloat16, device=DEVICE) + og = torch.randn(1, TT, H, D, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.1 + dt_bias = torch.randn(lp, dtype=torch.float32, device=DEVICE) * 0.1 + S = torch.randn(NS, H, D, D, dtype=torch.float32, device=DEVICE) * 0.01 + nw = torch.rand(D, dtype=torch.float32, device=DEVICE) + 0.5 + idx = (torch.randperm(NS - 1, device=DEVICE)[:TT] + 1).int().view(N, T) + if padded: + idx[0] = 0 + acc = torch.randint(1, T + 1, (N,), dtype=torch.int32, device=DEVICE) + kw = { + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "lower_bound": lb, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "use_beta_sigmoid_in_kernel": True, + "pad_slot_guard": True, + "cu_seqlens": torch.arange(0, TT + 1, T, dtype=torch.int64, device=DEVICE), + "ssm_state_indices": idx, + "num_accepted_tokens": acc, + } + view = lambda off: mixed.as_strided((1, TT, H, D), (TT * 3 * lp, 3 * lp, D, 1), off) + S1, cs1 = S.clone(), cs.clone() + o1, _ = fused_recurrent_kda( + q=view(0), + k=view(lp), + v=view(2 * lp), + initial_state=S1, + conv_state=cs1, + conv_weight=cw.view(3, lp, W).permute(0, 2, 1).contiguous().float(), + out_gate=og, + norm_weight=nw, + norm_eps=eps, + out=torch.full_like(og, float("nan")), + **kw, + ) + # reference: torch conv under vLLM's window rules, then the recurrence alone + x, w, S2, cs2 = mixed.float(), cw.float(), S.clone(), cs.clone() + y = torch.zeros(TT, 3 * lp, device=DEVICE) + for n in range(N): + c, off, rows = int(idx[n, 0]), int(acc[n]) - 1, slice(n * T, (n + 1) * T) + if c == 0: + continue + h = cs2[c].float() + taps = h[:, off : off + W - 1] + for t in range(T): + y[n * T + t] = (taps * w[:, : W - 1]).sum(-1) + x[n * T + t] * w[:, W - 1] + taps = torch.cat([taps[:, 1:], x[n * T + t, :, None]], -1) + window = torch.cat([h[:, off + 1 : off + W - 1], x[rows].T], -1) + cs2[c][:, : W - 2 + T] = window.to(cs.dtype) + y = y * torch.sigmoid(y) + q2, k2, v2 = ( + y[None, :, i * lp : (i + 1) * lp].reshape(1, TT, H, D) for i in range(3) + ) + o2, _ = fused_recurrent_kda(q=q2, k=k2, v=v2, initial_state=S2, **kw) + ob = o2.to(torch.bfloat16).float() + ref = ob * torch.rsqrt(ob.square().mean(-1, keepdim=True) + eps) * nw + ref = ref * torch.sigmoid(og.float()) + assert torch.equal(cs1, cs2), "conv window" + assert not padded or o1[0, :T].eq(0).all(), "padded sequence output" + assert_close("o", ref, o1) + assert_close("ht", S2, S1) + + def test_fused_recurrent_kda_packed_decode(): N, H, D, W, NS, pad = 16, 24, 128, 4, 40, 16 lp = H * D