diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 33a14bcfb49..2b63111b126 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -57,6 +57,11 @@ flydsl_pa_mqa_logits_fp4_varqlen, ) from .kernels.qk_norm_rope_quant import flydsl_qk_norm_rope_quant + from .kimi_k3_kda_decode import ( + flydsl_kimi_k3_kda_decode, + flydsl_kimi_k3_kda_decode_with_f_b, + is_flydsl_kimi_k3_kda_decode_supported, + ) from .moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 # from .linear_attention_kernels import flydsl_gdr_decode @@ -68,6 +73,8 @@ "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_hgemm", + "flydsl_kimi_k3_kda_decode", + "flydsl_kimi_k3_kda_decode_with_f_b", "flydsl_moe_stage1", "flydsl_moe_stage2", "flydsl_pa_mqa_logits_fp4", @@ -75,5 +82,6 @@ "flydsl_pa_mqa_logits_fp4_varqlen", "flydsl_preshuffle_gemm_a8", "flydsl_qk_norm_rope_quant", + "is_flydsl_kimi_k3_kda_decode_supported", # "flydsl_gdr_decode", ] diff --git a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py new file mode 100644 index 00000000000..bd03af0342e --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py @@ -0,0 +1,566 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL kernel for the fused Kimi-K3 KDA decode path on gfx950.""" + +import functools +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import gpu as mlir_gpu +from flydsl._mlir.dialects import scf +from flydsl._mlir.dialects import vector as mlir_vector +from flydsl.expr import range_constexpr +from flydsl.expr.typing import T + +from aiter.ops.flydsl.kernels import vector + +from .tensor_shim import GTensor, _to_raw + +_HEADS = 12 +_DIM = 128 +_LOG2E = math.log2(math.e) +_SCALE = _DIM**-0.5 +_BLOCK_THREADS = 256 +_NUM_WARPS = 4 +_WARP_SIZE = 64 +_WARP_THREADS_K = 8 +_VALUES_PER_THREAD_K = 4 +_WARP_TILE_K = _WARP_THREADS_K * _VALUES_PER_THREAD_K +_K_ITERS = _DIM // _WARP_TILE_K +_WARP_THREADS_V = _WARP_SIZE // _WARP_THREADS_K +_V_GROUP_TILE = _NUM_WARPS * _WARP_THREADS_V +_V_ITERS = _DIM // _V_GROUP_TILE +_WAVES_PER_EU = 3 + + +@functools.cache +def create_kimi_k3_kda_decode_kernel(norm_eps: float, lower_bound: float): + """Build the fixed gfx950 BF16 Kimi-K3 decode specialization.""" + + @fx.struct + class SharedStorage: + q: fx.Array[fx.BFloat16, _DIM, 16] + k: fx.Array[fx.BFloat16, _DIM, 16] + v: fx.Array[fx.BFloat16, _DIM, 16] + recurrent_out: fx.Array[fx.BFloat16, _DIM, 16] + norm_partial: fx.Array[fx.Float32, 2, 16] + + @flyc.kernel( + name="kimi_k3_kda_decode_bf16_gfx950", + known_block_size=[_BLOCK_THREADS, 1, 1], + ) + def kernel( + x_mem: fx.Tensor, + weight_mem: fx.Tensor, + conv_state_mem: fx.Tensor, + raw_g_mem: fx.Tensor, + raw_beta_mem: fx.Tensor, + A_log_mem: fx.Tensor, + dt_bias_mem: fx.Tensor, + state_mem: fx.Tensor, + state_indices_mem: fx.Tensor, + output_gate_mem: fx.Tensor, + norm_weight_mem: fx.Tensor, + out_mem: fx.Tensor, + batch_size: fx.Int32, + stride_x_token: fx.Int32, + stride_weight_channel: fx.Int32, + stride_weight_width: fx.Int32, + stride_conv_slot: fx.Int32, + stride_conv_channel: fx.Int32, + stride_conv_width: fx.Int32, + stride_g_token: fx.Int32, + stride_beta_token: fx.Int32, + stride_state_slot: fx.Int32, + stride_gate_token: fx.Int32, + stride_gate_head: fx.Int32, + stride_out_token: fx.Int32, + stride_out_head: fx.Int32, + ): + del batch_size + + x = GTensor(x_mem, dtype=T.bf16, shape=(-1,)) + weight = GTensor(weight_mem, dtype=T.f32, shape=(-1,)) + conv_state = GTensor(conv_state_mem, dtype=T.bf16, shape=(-1,)) + raw_g = GTensor(raw_g_mem, dtype=T.bf16, shape=(-1,)) + raw_beta = GTensor(raw_beta_mem, dtype=T.bf16, shape=(-1,)) + A_log = GTensor(A_log_mem, dtype=T.f32, shape=(-1,)) + dt_bias = GTensor(dt_bias_mem, dtype=T.f32, shape=(-1,)) + state = GTensor(state_mem, dtype=T.f32, shape=(-1,)) + state_indices = GTensor(state_indices_mem, dtype=T.i32, shape=(-1,)) + output_gate = GTensor(output_gate_mem, dtype=T.bf16, shape=(-1,)) + norm_weight = GTensor(norm_weight_mem, dtype=T.bf16, shape=(-1,)) + out = GTensor(out_mem, dtype=T.bf16, shape=(-1,)) + + shared = fx.SharedAllocator().allocate(SharedStorage).peek() + q_lds = shared.q.ptr + k_lds = shared.k.ptr + v_lds = shared.v.ptr + out_lds = shared.recurrent_out.ptr + norm_lds = shared.norm_partial.ptr + + tid = fx.thread_idx.x + block = fx.block_idx.x + batch = block // fx.Int32(_HEADS) + head = block % fx.Int32(_HEADS) + lane = tid % fx.Int32(_WARP_SIZE) + warp = tid // fx.Int32(_WARP_SIZE) + lane_k = lane % fx.Int32(_WARP_THREADS_K) + + state_idx = fx.Int32(state_indices[batch]) + valid = state_idx > fx.Int32(0) + + valid_if = scf.IfOp(_to_raw(valid), results_=[], has_else=True) + with ir.InsertionPoint(valid_if.then_block): + # A workgroup exclusively owns all three convolution channels for + # its (batch, head), so every cache entry is shifted exactly once. + conv_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(conv_if.then_block): + channel_local = tid + q_channel = head * fx.Int32(_DIM) + channel_local + k_channel = ( + fx.Int32(_HEADS * _DIM) + head * fx.Int32(_DIM) + channel_local + ) + v_channel = ( + fx.Int32(2 * _HEADS * _DIM) + head * fx.Int32(_DIM) + channel_local + ) + + def convolve_channel(channel): + cs_base = ( + state_idx * stride_conv_slot + channel * stride_conv_channel + ) + c0 = fx.Float32(conv_state[cs_base]) + c1 = fx.Float32(conv_state[cs_base + stride_conv_width]) + c2 = fx.Float32( + conv_state[cs_base + fx.Int32(2) * stride_conv_width] + ) + current = fx.BFloat16(x[batch * stride_x_token + channel]) + current_f32 = fx.Float32(current) + w_base = channel * stride_weight_channel + acc = c0 * fx.Float32(weight[w_base]) + acc = acc + c1 * fx.Float32(weight[w_base + stride_weight_width]) + acc = acc + c2 * fx.Float32( + weight[w_base + fx.Int32(2) * stride_weight_width] + ) + acc = acc + current_f32 * fx.Float32( + weight[w_base + fx.Int32(3) * stride_weight_width] + ) + silu = acc / ( + fx.Float32(1.0) + fx.math.exp2(-acc * fx.Float32(_LOG2E)) + ) + conv_state.store( + cs_base, + fx.BFloat16(c1), + ) + conv_state.store( + cs_base + stride_conv_width, + fx.BFloat16(c2), + ) + conv_state.store( + cs_base + fx.Int32(2) * stride_conv_width, + current, + ) + return silu.to(fx.BFloat16) + + q_conv = convolve_channel(q_channel) + k_conv = convolve_channel(k_channel) + v_conv = convolve_channel(v_channel) + fx.ptr_store(q_conv, q_lds + tid) + fx.ptr_store(k_conv, k_lds + tid) + fx.ptr_store(v_conv, v_lds + tid) + scf.YieldOp([]) + + fx.gpu.barrier() + + # Four waves split V into 32-row groups. Eight-lane subgroups + # reduce K; each lane issues one aligned f32x4 state transaction. + k_vec_start = lane_k * fx.Int32(_VALUES_PER_THREAD_K) + global_v_start = warp * fx.Int32(_WARP_THREADS_V) + lane // fx.Int32( + _WARP_THREADS_K + ) + vec_f32 = T.vec(_VALUES_PER_THREAD_K, T.f32) + vec_bf16 = T.vec(_VALUES_PER_THREAD_K, T.bf16) + zero_vec = fx.full( + _VALUES_PER_THREAD_K, + 0.0, + fx.Float32, + ) + + q_vecs = [] + k_vecs = [] + decay_vecs = [] + sum_q_partial = fx.Float32(0.0) + sum_k_partial = fx.Float32(0.0) + a = fx.math.exp2(fx.Float32(A_log[head]) * fx.Float32(_LOG2E)) + + for ki in range_constexpr(_K_ITERS): + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + q_bf16 = fx.ptr_load( + q_lds + k_base, + result_type=vec_bf16, + ) + k_bf16 = fx.ptr_load( + k_lds + k_base, + result_type=vec_bf16, + ) + q_f32 = q_bf16.extf(vec_f32) + k_f32 = k_bf16.extf(vec_f32) + q_vecs.append(q_f32) + k_vecs.append(k_f32) + sum_q_vec = q_f32 * q_f32 + sum_k_vec = k_f32 * k_f32 + sum_q_partial = ( + sum_q_partial + + mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_q_vec, + ).dest + ) + sum_k_partial = ( + sum_k_partial + + mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_k_vec, + ).dest + ) + + gate_bf16 = raw_g.vec_load( + (batch * stride_g_token + head * fx.Int32(_DIM) + k_base,), + _VALUES_PER_THREAD_K, + ) + gate_f32 = gate_bf16.extf(vec_f32) + dt = dt_bias.vec_load( + (head * fx.Int32(_DIM) + k_base,), + _VALUES_PER_THREAD_K, + ) + sigmoid_arg = (gate_f32 + dt) * a + gate = fx.Float32(lower_bound) / ( + fx.Float32(1.0) + fx.math.exp2(-sigmoid_arg * fx.Float32(_LOG2E)) + ) + decay_vecs.append(fx.math.exp2(gate * fx.Float32(_LOG2E))) + + width = fx.Int32(_WARP_SIZE) + for offset in (1, 2, 4): + sum_q_partial = ( + sum_q_partial + + mlir_gpu.ShuffleOp( + _to_raw(sum_q_partial), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + sum_k_partial = ( + sum_k_partial + + mlir_gpu.ShuffleOp( + _to_raw(sum_k_partial), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + subgroup_leader = (lane // fx.Int32(_WARP_THREADS_K)) * fx.Int32( + _WARP_THREADS_K + ) + norm_q = mlir_gpu.ShuffleOp( + _to_raw(sum_q_partial), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + norm_k = mlir_gpu.ShuffleOp( + _to_raw(sum_k_partial), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + inv_q = fx.math.rsqrt(fx.Float32(norm_q) + fx.Float32(1e-6)) + inv_k = fx.math.rsqrt(fx.Float32(norm_k) + fx.Float32(1e-6)) + + for ki in range_constexpr(_K_ITERS): + q_vecs[ki] = q_vecs[ki] * fx.Float32(inv_q) * fx.Float32(_SCALE) + k_vecs[ki] = k_vecs[ki] * fx.Float32(inv_k) + + dot_kq_vec = zero_vec + for ki in range_constexpr(_K_ITERS): + dot_kq_vec = mlir_vector.FMAOp( + k_vecs[ki], + q_vecs[ki], + dot_kq_vec, + ).result + dot_kq = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + dot_kq_vec, + ).dest + for offset in (1, 2, 4): + dot_kq = ( + dot_kq + + mlir_gpu.ShuffleOp( + _to_raw(dot_kq), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + beta_value = fx.Float32(raw_beta[batch * stride_beta_token + head]) + beta = fx.Float32(1.0) / ( + fx.Float32(1.0) + fx.math.exp2(-beta_value * fx.Float32(_LOG2E)) + ) + state_head_base = state_idx * stride_state_slot + head * fx.Int32( + _DIM * _DIM + ) + + state_vecs = [] + for vi in range_constexpr(_V_ITERS): + global_v = global_v_start + fx.Int32(vi * _V_GROUP_TILE) + for ki in range_constexpr(_K_ITERS): + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + state_off = state_head_base + global_v * fx.Int32(_DIM) + k_base + state_vecs.append( + state.vec_load( + (state_off,), + _VALUES_PER_THREAD_K, + ) + ) + + for vi in range_constexpr(_V_ITERS): + global_v = global_v_start + fx.Int32(vi * _V_GROUP_TILE) + sum_hk_vec = zero_vec + sum_hq_vec = zero_vec + for ki in range_constexpr(_K_ITERS): + state_pos = vi * _K_ITERS + ki + decayed = state_vecs[state_pos] * decay_vecs[ki] + state_vecs[state_pos] = decayed + sum_hk_vec = mlir_vector.FMAOp( + decayed, + k_vecs[ki], + sum_hk_vec, + ).result + sum_hq_vec = mlir_vector.FMAOp( + decayed, + q_vecs[ki], + sum_hq_vec, + ).result + + sum_hk = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_hk_vec, + ).dest + sum_hq = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_hq_vec, + ).dest + for offset in (1, 2, 4): + sum_hk = ( + sum_hk + + mlir_gpu.ShuffleOp( + _to_raw(sum_hk), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + sum_hq = ( + sum_hq + + mlir_gpu.ShuffleOp( + _to_raw(sum_hq), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + conv_v = fx.Float32(fx.ptr_load(v_lds + global_v)) + v_new = (conv_v - fx.Float32(sum_hk)) * beta + v_new = mlir_gpu.ShuffleOp( + _to_raw(v_new), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + recurrent_value = fx.Float32(sum_hq) + fx.Float32(v_new) * fx.Float32( + dot_kq + ) + v_new_vec = mlir_vector.BroadcastOp( + vec_f32, + _to_raw(v_new), + ).vector + + for ki in range_constexpr(_K_ITERS): + state_pos = vi * _K_ITERS + ki + updated = mlir_vector.FMAOp( + k_vecs[ki], + v_new_vec, + state_vecs[state_pos], + ).result + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + state_off = state_head_base + global_v * fx.Int32(_DIM) + k_base + state.vec_store( + (state_off,), + updated, + _VALUES_PER_THREAD_K, + ) + + if lane_k == fx.Int32(0): + fx.ptr_store( + fx.BFloat16(recurrent_value), + out_lds + global_v, + ) + + fx.gpu.barrier() + + # Preserve the model's BF16 boundary before RMSNorm and gating. + output_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(output_if.then_block): + recurrent_bf16 = fx.ptr_load(out_lds + tid) + recurrent_f32 = fx.Float32(recurrent_bf16) + square = recurrent_f32 * recurrent_f32 + for offset in (32, 16, 8, 4, 2, 1): + square = ( + square + + mlir_gpu.ShuffleOp( + _to_raw(square), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + if lane == fx.Int32(0): + fx.ptr_store(square, norm_lds + warp) + scf.YieldOp([]) + + fx.gpu.barrier() + + output_store_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(output_store_if.then_block): + norm_sum = fx.Float32(fx.ptr_load(norm_lds)) + norm_sum = norm_sum + fx.Float32(fx.ptr_load(norm_lds + fx.Int32(1))) + inv_rms = fx.math.rsqrt( + norm_sum * fx.Float32(1.0 / _DIM) + fx.Float32(norm_eps) + ) + recurrent_f32 = fx.Float32(fx.ptr_load(out_lds + tid)) + norm_w = fx.Float32(norm_weight[tid]) + gate_value = fx.Float32( + output_gate[ + batch * stride_gate_token + head * stride_gate_head + tid + ] + ) + output_sigmoid = fx.Float32(1.0) / ( + fx.Float32(1.0) + fx.math.exp2(-gate_value * fx.Float32(_LOG2E)) + ) + result = recurrent_f32 * inv_rms * norm_w * output_sigmoid + out.store( + batch * stride_out_token + head * stride_out_head + tid, + result.to(fx.BFloat16), + ) + scf.YieldOp([]) + scf.YieldOp([]) + with ir.InsertionPoint(valid_if.else_block): + zero_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(zero_if.then_block): + out.store( + batch * stride_out_token + head * stride_out_head + tid, + fx.BFloat16(0.0), + ) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch( + x_mem: fx.Tensor, + weight_mem: fx.Tensor, + conv_state_mem: fx.Tensor, + raw_g_mem: fx.Tensor, + raw_beta_mem: fx.Tensor, + A_log_mem: fx.Tensor, + dt_bias_mem: fx.Tensor, + state_mem: fx.Tensor, + state_indices_mem: fx.Tensor, + output_gate_mem: fx.Tensor, + norm_weight_mem: fx.Tensor, + out_mem: fx.Tensor, + batch_size: fx.Int32, + stride_x_token: fx.Int32, + stride_weight_channel: fx.Int32, + stride_weight_width: fx.Int32, + stride_conv_slot: fx.Int32, + stride_conv_channel: fx.Int32, + stride_conv_width: fx.Int32, + stride_g_token: fx.Int32, + stride_beta_token: fx.Int32, + stride_state_slot: fx.Int32, + stride_gate_token: fx.Int32, + stride_gate_head: fx.Int32, + stride_out_token: fx.Int32, + stride_out_head: fx.Int32, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + kernel( + x_mem, + weight_mem, + conv_state_mem, + raw_g_mem, + raw_beta_mem, + A_log_mem, + dt_bias_mem, + state_mem, + state_indices_mem, + output_gate_mem, + norm_weight_mem, + out_mem, + batch_size, + stride_x_token, + stride_weight_channel, + stride_weight_width, + stride_conv_slot, + stride_conv_channel, + stride_conv_width, + stride_g_token, + stride_beta_token, + stride_state_slot, + stride_gate_token, + stride_gate_head, + stride_out_token, + stride_out_head, + ).launch( + grid=(batch_size * fx.Int32(_HEADS), 1, 1), + block=(_BLOCK_THREADS, 1, 1), + stream=stream, + ) + + launch.compile_hints = { + "waves_per_eu": _WAVES_PER_EU, + "llvm_options": { + "amdgpu-expert-scheduling-mode": True, + }, + } + return launch + + +__all__ = ["create_kimi_k3_kda_decode_kernel"] diff --git a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py new file mode 100644 index 00000000000..b4a84bea4e5 --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL Kimi-K3 KDA decode with a fused head-local f_b projection.""" + +import functools +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import gpu as mlir_gpu +from flydsl._mlir.dialects import scf +from flydsl._mlir.dialects import vector as mlir_vector +from flydsl.expr import range_constexpr +from flydsl.expr.typing import T + +from aiter.ops.flydsl.kernels import vector + +from .tensor_shim import GTensor, _to_raw + +_HEADS = 12 +_DIM = 128 +_LOG2E = math.log2(math.e) +_SCALE = _DIM**-0.5 +_BLOCK_THREADS = 256 +_NUM_WARPS = 4 +_WARP_SIZE = 64 +_WARP_THREADS_K = 8 +_VALUES_PER_THREAD_K = 4 +_WARP_TILE_K = _WARP_THREADS_K * _VALUES_PER_THREAD_K +_K_ITERS = _DIM // _WARP_TILE_K +_WARP_THREADS_V = _WARP_SIZE // _WARP_THREADS_K +_V_GROUP_TILE = _NUM_WARPS * _WARP_THREADS_V +_V_ITERS = _DIM // _V_GROUP_TILE +_PROJECTION_VECTOR = 4 +_PROJECTION_ITERS = _DIM // _PROJECTION_VECTOR +_WAVES_PER_EU = 2 + + +@functools.cache +def create_kimi_k3_kda_decode_fb_kernel(norm_eps: float, lower_bound: float): + """Build the fixed gfx950 BF16 f_b plus KDA decode specialization.""" + + @fx.struct + class SharedStorage: + q: fx.Array[fx.BFloat16, _DIM, 16] + k: fx.Array[fx.BFloat16, _DIM, 16] + v: fx.Array[fx.BFloat16, _DIM, 16] + gate: fx.Array[fx.BFloat16, _DIM, 16] + recurrent_out: fx.Array[fx.BFloat16, _DIM, 16] + norm_partial: fx.Array[fx.Float32, 2, 16] + + @flyc.kernel( + name="kimi_k3_kda_decode_fb_bf16_gfx950", + known_block_size=[_BLOCK_THREADS, 1, 1], + ) + def kernel( + f_a_mem: fx.Tensor, + f_b_weight_mem: fx.Tensor, + x_mem: fx.Tensor, + weight_mem: fx.Tensor, + conv_state_mem: fx.Tensor, + raw_beta_mem: fx.Tensor, + A_log_mem: fx.Tensor, + dt_bias_mem: fx.Tensor, + state_mem: fx.Tensor, + state_indices_mem: fx.Tensor, + output_gate_mem: fx.Tensor, + norm_weight_mem: fx.Tensor, + out_mem: fx.Tensor, + batch_size: fx.Int32, + stride_f_a_token: fx.Int32, + stride_f_b_head: fx.Int32, + stride_f_b_output: fx.Int32, + stride_x_token: fx.Int32, + stride_weight_channel: fx.Int32, + stride_weight_width: fx.Int32, + stride_conv_slot: fx.Int32, + stride_conv_channel: fx.Int32, + stride_conv_width: fx.Int32, + stride_beta_token: fx.Int32, + stride_state_slot: fx.Int32, + stride_gate_token: fx.Int32, + stride_gate_head: fx.Int32, + stride_out_token: fx.Int32, + stride_out_head: fx.Int32, + ): + del batch_size + + f_a = GTensor(f_a_mem, dtype=T.bf16, shape=(-1,)) + f_b_weight = GTensor(f_b_weight_mem, dtype=T.bf16, shape=(-1,)) + x = GTensor(x_mem, dtype=T.bf16, shape=(-1,)) + weight = GTensor(weight_mem, dtype=T.f32, shape=(-1,)) + conv_state = GTensor(conv_state_mem, dtype=T.bf16, shape=(-1,)) + raw_beta = GTensor(raw_beta_mem, dtype=T.bf16, shape=(-1,)) + A_log = GTensor(A_log_mem, dtype=T.f32, shape=(-1,)) + dt_bias = GTensor(dt_bias_mem, dtype=T.f32, shape=(-1,)) + state = GTensor(state_mem, dtype=T.f32, shape=(-1,)) + state_indices = GTensor(state_indices_mem, dtype=T.i32, shape=(-1,)) + output_gate = GTensor(output_gate_mem, dtype=T.bf16, shape=(-1,)) + norm_weight = GTensor(norm_weight_mem, dtype=T.bf16, shape=(-1,)) + out = GTensor(out_mem, dtype=T.bf16, shape=(-1,)) + + shared = fx.SharedAllocator().allocate(SharedStorage).peek() + q_lds = shared.q.ptr + k_lds = shared.k.ptr + v_lds = shared.v.ptr + gate_lds = shared.gate.ptr + out_lds = shared.recurrent_out.ptr + norm_lds = shared.norm_partial.ptr + + tid = fx.thread_idx.x + block = fx.block_idx.x + batch = block // fx.Int32(_HEADS) + head = block % fx.Int32(_HEADS) + lane = tid % fx.Int32(_WARP_SIZE) + warp = tid // fx.Int32(_WARP_SIZE) + lane_k = lane % fx.Int32(_WARP_THREADS_K) + + state_idx = fx.Int32(state_indices[batch]) + valid = state_idx > fx.Int32(0) + + valid_if = scf.IfOp(_to_raw(valid), results_=[], has_else=True) + with ir.InsertionPoint(valid_if.then_block): + # Threads 0..127 own one output each. Accumulation is FP32 and the + # single BF16 store is the same numerical boundary as F.linear. + projection_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(projection_if.then_block): + vec_f32_projection = T.vec(_PROJECTION_VECTOR, T.f32) + accum = fx.full( + _PROJECTION_VECTOR, + 0.0, + fx.Float32, + ) + f_a_base = batch * stride_f_a_token + f_b_base = head * stride_f_b_head + tid * stride_f_b_output + for projection_iter in range_constexpr(_PROJECTION_ITERS): + projection_offset = fx.Int32(projection_iter * _PROJECTION_VECTOR) + f_a_values = f_a.vec_load( + (f_a_base + projection_offset,), + _PROJECTION_VECTOR, + ).extf(vec_f32_projection) + weight_values = f_b_weight.vec_load( + (f_b_base + projection_offset,), + _PROJECTION_VECTOR, + ).extf(vec_f32_projection) + accum = mlir_vector.FMAOp( + f_a_values, + weight_values, + accum, + ).result + projected = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + accum, + ).dest + fx.ptr_store( + fx.BFloat16(projected), + gate_lds + tid, + ) + scf.YieldOp([]) + + # A workgroup exclusively owns all three convolution channels for + # its (batch, head), so every cache entry is shifted exactly once. + conv_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(conv_if.then_block): + channel_local = tid + q_channel = head * fx.Int32(_DIM) + channel_local + k_channel = ( + fx.Int32(_HEADS * _DIM) + head * fx.Int32(_DIM) + channel_local + ) + v_channel = ( + fx.Int32(2 * _HEADS * _DIM) + head * fx.Int32(_DIM) + channel_local + ) + + def convolve_channel(channel): + cs_base = ( + state_idx * stride_conv_slot + channel * stride_conv_channel + ) + c0 = fx.Float32(conv_state[cs_base]) + c1 = fx.Float32(conv_state[cs_base + stride_conv_width]) + c2 = fx.Float32( + conv_state[cs_base + fx.Int32(2) * stride_conv_width] + ) + current = fx.BFloat16(x[batch * stride_x_token + channel]) + current_f32 = fx.Float32(current) + w_base = channel * stride_weight_channel + acc = c0 * fx.Float32(weight[w_base]) + acc = acc + c1 * fx.Float32(weight[w_base + stride_weight_width]) + acc = acc + c2 * fx.Float32( + weight[w_base + fx.Int32(2) * stride_weight_width] + ) + acc = acc + current_f32 * fx.Float32( + weight[w_base + fx.Int32(3) * stride_weight_width] + ) + silu = acc / ( + fx.Float32(1.0) + fx.math.exp2(-acc * fx.Float32(_LOG2E)) + ) + conv_state.store( + cs_base, + fx.BFloat16(c1), + ) + conv_state.store( + cs_base + stride_conv_width, + fx.BFloat16(c2), + ) + conv_state.store( + cs_base + fx.Int32(2) * stride_conv_width, + current, + ) + return silu.to(fx.BFloat16) + + q_conv = convolve_channel(q_channel) + k_conv = convolve_channel(k_channel) + v_conv = convolve_channel(v_channel) + fx.ptr_store(q_conv, q_lds + tid) + fx.ptr_store(k_conv, k_lds + tid) + fx.ptr_store(v_conv, v_lds + tid) + scf.YieldOp([]) + + # Both projection and convolution LDS values must be visible + # before the recurrent core begins. + fx.gpu.barrier() + + # Four waves split V into 32-row groups. Eight-lane subgroups + # reduce K; each lane issues one aligned f32x4 state transaction. + k_vec_start = lane_k * fx.Int32(_VALUES_PER_THREAD_K) + global_v_start = warp * fx.Int32(_WARP_THREADS_V) + lane // fx.Int32( + _WARP_THREADS_K + ) + vec_f32 = T.vec(_VALUES_PER_THREAD_K, T.f32) + vec_bf16 = T.vec(_VALUES_PER_THREAD_K, T.bf16) + zero_vec = fx.full( + _VALUES_PER_THREAD_K, + 0.0, + fx.Float32, + ) + + q_vecs = [] + k_vecs = [] + decay_vecs = [] + sum_q_partial = fx.Float32(0.0) + sum_k_partial = fx.Float32(0.0) + a = fx.math.exp2(fx.Float32(A_log[head]) * fx.Float32(_LOG2E)) + + for ki in range_constexpr(_K_ITERS): + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + q_bf16 = fx.ptr_load( + q_lds + k_base, + result_type=vec_bf16, + ) + k_bf16 = fx.ptr_load( + k_lds + k_base, + result_type=vec_bf16, + ) + q_f32 = q_bf16.extf(vec_f32) + k_f32 = k_bf16.extf(vec_f32) + q_vecs.append(q_f32) + k_vecs.append(k_f32) + sum_q_vec = q_f32 * q_f32 + sum_k_vec = k_f32 * k_f32 + sum_q_partial = ( + sum_q_partial + + mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_q_vec, + ).dest + ) + sum_k_partial = ( + sum_k_partial + + mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_k_vec, + ).dest + ) + + # The projection is rounded in LDS before the lower-bound gate. + gate_bf16 = fx.ptr_load( + gate_lds + k_base, + result_type=vec_bf16, + ) + gate_f32 = gate_bf16.extf(vec_f32) + dt = dt_bias.vec_load( + (head * fx.Int32(_DIM) + k_base,), + _VALUES_PER_THREAD_K, + ) + sigmoid_arg = (gate_f32 + dt) * a + gate = fx.Float32(lower_bound) / ( + fx.Float32(1.0) + fx.math.exp2(-sigmoid_arg * fx.Float32(_LOG2E)) + ) + decay_vecs.append(fx.math.exp2(gate * fx.Float32(_LOG2E))) + + width = fx.Int32(_WARP_SIZE) + for offset in (1, 2, 4): + sum_q_partial = ( + sum_q_partial + + mlir_gpu.ShuffleOp( + _to_raw(sum_q_partial), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + sum_k_partial = ( + sum_k_partial + + mlir_gpu.ShuffleOp( + _to_raw(sum_k_partial), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + subgroup_leader = (lane // fx.Int32(_WARP_THREADS_K)) * fx.Int32( + _WARP_THREADS_K + ) + norm_q = mlir_gpu.ShuffleOp( + _to_raw(sum_q_partial), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + norm_k = mlir_gpu.ShuffleOp( + _to_raw(sum_k_partial), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + inv_q = fx.math.rsqrt(fx.Float32(norm_q) + fx.Float32(1e-6)) + inv_k = fx.math.rsqrt(fx.Float32(norm_k) + fx.Float32(1e-6)) + + for ki in range_constexpr(_K_ITERS): + q_vecs[ki] = q_vecs[ki] * fx.Float32(inv_q) * fx.Float32(_SCALE) + k_vecs[ki] = k_vecs[ki] * fx.Float32(inv_k) + + dot_kq_vec = zero_vec + for ki in range_constexpr(_K_ITERS): + dot_kq_vec = mlir_vector.FMAOp( + k_vecs[ki], + q_vecs[ki], + dot_kq_vec, + ).result + dot_kq = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + dot_kq_vec, + ).dest + for offset in (1, 2, 4): + dot_kq = ( + dot_kq + + mlir_gpu.ShuffleOp( + _to_raw(dot_kq), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + beta_value = fx.Float32(raw_beta[batch * stride_beta_token + head]) + beta = fx.Float32(1.0) / ( + fx.Float32(1.0) + fx.math.exp2(-beta_value * fx.Float32(_LOG2E)) + ) + state_head_base = state_idx * stride_state_slot + head * fx.Int32( + _DIM * _DIM + ) + + state_vecs = [] + for vi in range_constexpr(_V_ITERS): + global_v = global_v_start + fx.Int32(vi * _V_GROUP_TILE) + for ki in range_constexpr(_K_ITERS): + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + state_off = state_head_base + global_v * fx.Int32(_DIM) + k_base + state_vecs.append( + state.vec_load( + (state_off,), + _VALUES_PER_THREAD_K, + ) + ) + + for vi in range_constexpr(_V_ITERS): + global_v = global_v_start + fx.Int32(vi * _V_GROUP_TILE) + sum_hk_vec = zero_vec + sum_hq_vec = zero_vec + for ki in range_constexpr(_K_ITERS): + state_pos = vi * _K_ITERS + ki + decayed = state_vecs[state_pos] * decay_vecs[ki] + state_vecs[state_pos] = decayed + sum_hk_vec = mlir_vector.FMAOp( + decayed, + k_vecs[ki], + sum_hk_vec, + ).result + sum_hq_vec = mlir_vector.FMAOp( + decayed, + q_vecs[ki], + sum_hq_vec, + ).result + + sum_hk = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_hk_vec, + ).dest + sum_hq = mlir_vector.ReductionOp( + T.f32, + vector.CombiningKind.ADD, + sum_hq_vec, + ).dest + for offset in (1, 2, 4): + sum_hk = ( + sum_hk + + mlir_gpu.ShuffleOp( + _to_raw(sum_hk), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + sum_hq = ( + sum_hq + + mlir_gpu.ShuffleOp( + _to_raw(sum_hq), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + + conv_v = fx.Float32(fx.ptr_load(v_lds + global_v)) + v_new = (conv_v - fx.Float32(sum_hk)) * beta + v_new = mlir_gpu.ShuffleOp( + _to_raw(v_new), + _to_raw(subgroup_leader), + _to_raw(width), + mode="idx", + ).shuffleResult + recurrent_value = fx.Float32(sum_hq) + fx.Float32(v_new) * fx.Float32( + dot_kq + ) + v_new_vec = mlir_vector.BroadcastOp( + vec_f32, + _to_raw(v_new), + ).vector + + for ki in range_constexpr(_K_ITERS): + state_pos = vi * _K_ITERS + ki + updated = mlir_vector.FMAOp( + k_vecs[ki], + v_new_vec, + state_vecs[state_pos], + ).result + k_base = k_vec_start + fx.Int32(ki * _WARP_TILE_K) + state_off = state_head_base + global_v * fx.Int32(_DIM) + k_base + state.vec_store( + (state_off,), + updated, + _VALUES_PER_THREAD_K, + ) + + if lane_k == fx.Int32(0): + fx.ptr_store( + fx.BFloat16(recurrent_value), + out_lds + global_v, + ) + + fx.gpu.barrier() + + # Preserve the model's BF16 boundary before RMSNorm and gating. + output_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(output_if.then_block): + recurrent_bf16 = fx.ptr_load(out_lds + tid) + recurrent_f32 = fx.Float32(recurrent_bf16) + square = recurrent_f32 * recurrent_f32 + for offset in (32, 16, 8, 4, 2, 1): + square = ( + square + + mlir_gpu.ShuffleOp( + _to_raw(square), + _to_raw(fx.Int32(offset)), + _to_raw(width), + mode="xor", + ).shuffleResult + ) + if lane == fx.Int32(0): + fx.ptr_store(square, norm_lds + warp) + scf.YieldOp([]) + + fx.gpu.barrier() + + output_store_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(output_store_if.then_block): + norm_sum = fx.Float32(fx.ptr_load(norm_lds)) + norm_sum = norm_sum + fx.Float32(fx.ptr_load(norm_lds + fx.Int32(1))) + inv_rms = fx.math.rsqrt( + norm_sum * fx.Float32(1.0 / _DIM) + fx.Float32(norm_eps) + ) + recurrent_f32 = fx.Float32(fx.ptr_load(out_lds + tid)) + norm_w = fx.Float32(norm_weight[tid]) + gate_value = fx.Float32( + output_gate[ + batch * stride_gate_token + head * stride_gate_head + tid + ] + ) + output_sigmoid = fx.Float32(1.0) / ( + fx.Float32(1.0) + fx.math.exp2(-gate_value * fx.Float32(_LOG2E)) + ) + result = recurrent_f32 * inv_rms * norm_w * output_sigmoid + out.store( + batch * stride_out_token + head * stride_out_head + tid, + result.to(fx.BFloat16), + ) + scf.YieldOp([]) + scf.YieldOp([]) + with ir.InsertionPoint(valid_if.else_block): + zero_if = scf.IfOp( + _to_raw(tid < fx.Int32(_DIM)), + results_=[], + has_else=False, + ) + with ir.InsertionPoint(zero_if.then_block): + out.store( + batch * stride_out_token + head * stride_out_head + tid, + fx.BFloat16(0.0), + ) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch( + f_a_mem: fx.Tensor, + f_b_weight_mem: fx.Tensor, + x_mem: fx.Tensor, + weight_mem: fx.Tensor, + conv_state_mem: fx.Tensor, + raw_beta_mem: fx.Tensor, + A_log_mem: fx.Tensor, + dt_bias_mem: fx.Tensor, + state_mem: fx.Tensor, + state_indices_mem: fx.Tensor, + output_gate_mem: fx.Tensor, + norm_weight_mem: fx.Tensor, + out_mem: fx.Tensor, + batch_size: fx.Int32, + stride_f_a_token: fx.Int32, + stride_f_b_head: fx.Int32, + stride_f_b_output: fx.Int32, + stride_x_token: fx.Int32, + stride_weight_channel: fx.Int32, + stride_weight_width: fx.Int32, + stride_conv_slot: fx.Int32, + stride_conv_channel: fx.Int32, + stride_conv_width: fx.Int32, + stride_beta_token: fx.Int32, + stride_state_slot: fx.Int32, + stride_gate_token: fx.Int32, + stride_gate_head: fx.Int32, + stride_out_token: fx.Int32, + stride_out_head: fx.Int32, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + kernel( + f_a_mem, + f_b_weight_mem, + x_mem, + weight_mem, + conv_state_mem, + raw_beta_mem, + A_log_mem, + dt_bias_mem, + state_mem, + state_indices_mem, + output_gate_mem, + norm_weight_mem, + out_mem, + batch_size, + stride_f_a_token, + stride_f_b_head, + stride_f_b_output, + stride_x_token, + stride_weight_channel, + stride_weight_width, + stride_conv_slot, + stride_conv_channel, + stride_conv_width, + stride_beta_token, + stride_state_slot, + stride_gate_token, + stride_gate_head, + stride_out_token, + stride_out_head, + ).launch( + grid=(batch_size * fx.Int32(_HEADS), 1, 1), + block=(_BLOCK_THREADS, 1, 1), + stream=stream, + ) + + launch.compile_hints = { + "waves_per_eu": _WAVES_PER_EU, + "llvm_options": { + "amdgpu-expert-scheduling-mode": True, + }, + } + return launch + + +__all__ = ["create_kimi_k3_kda_decode_fb_kernel"] diff --git a/aiter/ops/flydsl/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kimi_k3_kda_decode.py new file mode 100644 index 00000000000..2acdabd799a --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_kda_decode.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""High-level API for the fused Kimi-K3 KDA decode specialization.""" + +from __future__ import annotations + +import functools +from collections.abc import Iterable + +import torch + +from .kernels.kimi_k3_kda_decode import ( + create_kimi_k3_kda_decode_kernel, +) +from .kernels.kimi_k3_kda_decode_fb import ( + create_kimi_k3_kda_decode_fb_kernel, +) +from .kernels.tensor_shim import _run_compiled + +_HEADS = 12 +_DIM = 128 +_CONV_CHANNELS = 3 * _HEADS * _DIM +_CONV_WIDTH = 4 + + +@functools.cache +def _rocm_arch(device: torch.device) -> str | None: + properties = torch.cuda.get_device_properties(device) + arch = getattr(properties, "gcnArchName", None) + return arch.split(":", 1)[0] if arch is not None else None + + +def is_flydsl_kimi_k3_kda_decode_supported( + device: torch.device | str | int | None = None, +) -> bool: + """Return whether ``device`` can run this gfx950-only specialization.""" + if not torch.cuda.is_available(): + return False + try: + resolved = torch.device( + "cuda", + torch.cuda.current_device(), + ) + if device is not None: + resolved = ( + torch.device("cuda", device) + if isinstance(device, int) + else torch.device(device) + ) + if resolved.type != "cuda": + return False + if resolved.index is None: + resolved = torch.device( + "cuda", + torch.cuda.current_device(), + ) + return _rocm_arch(resolved) == "gfx950" + except (AssertionError, RuntimeError, TypeError, ValueError): + return False + + +def _check_tensor( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + inner_strides: tuple[int, ...] = (), +) -> None: + if tensor.shape != shape: + raise ValueError( + f"`{name}` must have shape {list(shape)}, got {list(tensor.shape)}." + ) + if tensor.dtype != dtype: + raise ValueError(f"`{name}` must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"`{name}` must be on {device}, got {tensor.device}.") + if inner_strides and tensor.stride()[-len(inner_strides) :] != inner_strides: + raise ValueError( + f"`{name}` must have inner strides {inner_strides}, got {tensor.stride()}." + ) + + +def _check_same_device( + tensors: Iterable[tuple[str, torch.Tensor]], + device: torch.device, +) -> None: + for name, tensor in tensors: + if not tensor.is_cuda: + raise ValueError(f"`{name}` must be a CUDA tensor.") + if tensor.device != device: + raise ValueError(f"`{name}` must be on {device}, got {tensor.device}.") + + +def _validate_kda_inputs( + *, + api_name: str, + batch_source: str, + device: torch.device, + batch: int, + x: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + state: torch.Tensor, + state_indices: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + out: torch.Tensor | None, +) -> torch.Tensor: + """Validate operands shared by both explicit KDA specializations.""" + if not is_flydsl_kimi_k3_kda_decode_supported(device): + raise RuntimeError(f"`{api_name}` requires a gfx950 GPU.") + if batch <= 0: + raise ValueError(f"`{batch_source}` must have a non-empty batch dimension.") + if conv_bias is not None: + raise ValueError("This specialization requires `conv_bias=None`.") + if lower_bound is None: + raise ValueError("This specialization requires the KDA lower-bound gate.") + + _check_same_device( + ( + ("x", x), + ("conv_weight", conv_weight), + ("conv_state", conv_state), + ("raw_beta", raw_beta), + ("A_log", A_log), + ("dt_bias", dt_bias), + ("state", state), + ("state_indices", state_indices), + ("output_gate", output_gate), + ("norm_weight", norm_weight), + ), + device, + ) + _check_tensor( + "x", + x, + shape=(batch, _CONV_CHANNELS), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "conv_weight", + conv_weight, + shape=(_CONV_CHANNELS, _CONV_WIDTH), + dtype=torch.float32, + device=device, + ) + if conv_state.ndim != 3 or conv_state.shape[1:] != ( + _CONV_CHANNELS, + _CONV_WIDTH - 1, + ): + raise ValueError( + "`conv_state` must have shape [cache, 4608, 3], " + f"got {list(conv_state.shape)}." + ) + if conv_state.dtype != torch.bfloat16: + raise ValueError("`conv_state` must have dtype torch.bfloat16.") + if state.ndim != 4 or state.shape[1:] != ( + _HEADS, + _DIM, + _DIM, + ): + raise ValueError( + f"`state` must have shape [cache, 12, 128, 128], got {list(state.shape)}." + ) + if state.dtype != torch.float32: + raise ValueError("`state` must have dtype torch.float32.") + if state.stride()[-3:] != (_DIM * _DIM, _DIM, 1): + raise ValueError("`state` must be contiguous within each cache slot.") + _check_tensor( + "raw_beta", + raw_beta, + shape=(1, batch, _HEADS), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "A_log", + A_log, + shape=(_HEADS,), + dtype=torch.float32, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "dt_bias", + dt_bias, + shape=(_HEADS * _DIM,), + dtype=torch.float32, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "state_indices", + state_indices, + shape=(batch,), + dtype=torch.int32, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "output_gate", + output_gate, + shape=(batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "norm_weight", + norm_weight, + shape=(_DIM,), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + + if out is None: + return torch.empty( + (1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + ) + _check_same_device((("out", out),), device) + _check_tensor( + "out", + out, + shape=(1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + return out + + +def flydsl_kimi_k3_kda_decode( + x: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + state: torch.Tensor, + state_indices: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Run fused Kimi-K3 KDA decode on MI350-series GPUs. + + This pure-decode specialization fuses the packed width-4 Q/K/V causal + convolution, the FP32 recurrent-state update, and the BF16 + RMSNorm/sigmoid output gate. Slot zero is reserved: non-positive + ``state_indices`` produce zero output without modifying either cache. + + The layout is fixed to Kimi-K3 TP8: 12 local heads and 128-dimensional + key/value state. Call + :func:`is_flydsl_kimi_k3_kda_decode_supported` before dispatching from a + model implementation. + """ + if x.ndim != 2: + raise ValueError(f"`x` must have rank 2, got rank {x.ndim}.") + if not x.is_cuda: + raise ValueError("`x` must be a CUDA tensor.") + device = x.device + batch = x.shape[0] + out = _validate_kda_inputs( + api_name="flydsl_kimi_k3_kda_decode", + batch_source="x", + device=device, + batch=batch, + x=x, + conv_weight=conv_weight, + conv_bias=conv_bias, + conv_state=conv_state, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + state=state, + state_indices=state_indices, + output_gate=output_gate, + norm_weight=norm_weight, + out=out, + ) + _check_same_device((("raw_g", raw_g),), device) + _check_tensor( + "raw_g", + raw_g, + shape=(1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(_DIM, 1), + ) + + executable = create_kimi_k3_kda_decode_kernel( + float(norm_eps), + float(lower_bound), + ) + with torch.cuda.device(device): + stream = torch.cuda.current_stream(device) + _run_compiled( + executable, + x, + conv_weight, + conv_state, + raw_g, + raw_beta, + A_log, + dt_bias, + state, + state_indices, + output_gate, + norm_weight, + out, + batch, + x.stride(0), + conv_weight.stride(0), + conv_weight.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + raw_g.stride(1), + raw_beta.stride(1), + state.stride(0), + output_gate.stride(0), + output_gate.stride(1), + out.stride(1), + out.stride(2), + stream, + ) + return out + + +def flydsl_kimi_k3_kda_decode_with_f_b( + f_a: torch.Tensor, + f_b_weight: torch.Tensor, + x: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + state: torch.Tensor, + state_indices: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Run the explicit gfx950 Kimi-K3 f_b plus KDA decode specialization. + + The kernel consumes ``f_a`` and the head-local ``f_b_weight`` directly, + accumulates the projection in FP32, and rounds once to BF16 before the KDA + lower-bound decay gate. It does not materialize the projected raw-g tensor + in global memory. + """ + if f_a.ndim != 2: + raise ValueError(f"`f_a` must have rank 2, got rank {f_a.ndim}.") + if not f_a.is_cuda: + raise ValueError("`f_a` must be a CUDA tensor.") + device = f_a.device + batch = f_a.shape[0] + _check_same_device((("f_b_weight", f_b_weight),), device) + _check_tensor( + "f_a", + f_a, + shape=(batch, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(1,), + ) + _check_tensor( + "f_b_weight", + f_b_weight, + shape=(_HEADS, _DIM, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(_DIM, 1), + ) + out = _validate_kda_inputs( + api_name="flydsl_kimi_k3_kda_decode_with_f_b", + batch_source="f_a", + device=device, + batch=batch, + x=x, + conv_weight=conv_weight, + conv_bias=conv_bias, + conv_state=conv_state, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + state=state, + state_indices=state_indices, + output_gate=output_gate, + norm_weight=norm_weight, + out=out, + ) + + executable = create_kimi_k3_kda_decode_fb_kernel( + float(norm_eps), + float(lower_bound), + ) + with torch.cuda.device(device): + stream = torch.cuda.current_stream(device) + _run_compiled( + executable, + f_a, + f_b_weight, + x, + conv_weight, + conv_state, + raw_beta, + A_log, + dt_bias, + state, + state_indices, + output_gate, + norm_weight, + out, + batch, + f_a.stride(0), + f_b_weight.stride(0), + f_b_weight.stride(1), + x.stride(0), + conv_weight.stride(0), + conv_weight.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + raw_beta.stride(1), + state.stride(0), + output_gate.stride(0), + output_gate.stride(1), + out.stride(1), + out.stride(2), + stream, + ) + return out + + +__all__ = [ + "flydsl_kimi_k3_kda_decode", + "flydsl_kimi_k3_kda_decode_with_f_b", + "is_flydsl_kimi_k3_kda_decode_supported", +] diff --git a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py new file mode 100644 index 00000000000..87e25cd78b2 --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness tests for the fused FlyDSL Kimi-K3 KDA decode path.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F + +pytest.importorskip("flydsl") +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.kimi_k3_kda_decode import ( + flydsl_kimi_k3_kda_decode, + flydsl_kimi_k3_kda_decode_with_f_b, + is_flydsl_kimi_k3_kda_decode_supported, +) +from aiter.ops.flydsl.utils import is_flydsl_available + + +def _gfx950_flydsl_available() -> bool: + if not torch.cuda.is_available() or not is_flydsl_available(): + return False + try: + return get_gfx() == "gfx950" + except (AssertionError, KeyError, RuntimeError): + return False + + +pytestmark = pytest.mark.skipif( + not _gfx950_flydsl_available(), + reason="gfx950 FlyDSL required", +) + +_DEVICE = torch.device("cuda") +_HEADS = 12 +_DIM = 128 +_CHANNELS = 3 * _HEADS * _DIM +_CONV_WIDTH = 4 +_LOWER_BOUND = -5.0 +_NORM_EPS = 1e-5 + + +@dataclass +class Inputs: + x: torch.Tensor + conv_weight: torch.Tensor + conv_state: torch.Tensor + raw_g: torch.Tensor + raw_beta: torch.Tensor + A_log: torch.Tensor + dt_bias: torch.Tensor + state: torch.Tensor + state_indices: torch.Tensor + output_gate: torch.Tensor + norm_weight: torch.Tensor + + +def _make_inputs(batch: int, seed: int = 20260728) -> Inputs: + generator = torch.Generator(device=_DEVICE).manual_seed(seed + batch) + slots = batch + 2 + + x_storage = torch.randn( + (batch, _CHANNELS + 17), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + x = x_storage[:, :_CHANNELS] + conv_weight = 0.1 * torch.randn( + (_CHANNELS, _CONV_WIDTH), + dtype=torch.float32, + device=_DEVICE, + generator=generator, + ) + + # Kimi's hybrid cache can pad the slot stride. Exercise that layout + # explicitly while keeping each slot's inner dimensions contiguous. + conv_storage = torch.randn( + (slots, _CHANNELS * (_CONV_WIDTH - 1) + 19), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + conv_state = conv_storage[:, : _CHANNELS * (_CONV_WIDTH - 1)].view( + slots, _CHANNELS, _CONV_WIDTH - 1 + ) + state_storage = 0.01 * torch.randn( + (slots, _HEADS * _DIM * _DIM + 23), + dtype=torch.float32, + device=_DEVICE, + generator=generator, + ) + state = state_storage[:, : _HEADS * _DIM * _DIM].view(slots, _HEADS, _DIM, _DIM) + + raw_beta_storage = torch.randn( + (1, batch, _HEADS + 1), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + output_gate_storage = torch.randn( + (batch, _HEADS * _DIM + 7), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + return Inputs( + x=x, + conv_weight=conv_weight, + conv_state=conv_state, + raw_g=torch.randn( + (1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ), + raw_beta=raw_beta_storage[:, :, :_HEADS], + A_log=0.5 + * torch.randn( + (_HEADS,), + dtype=torch.float32, + device=_DEVICE, + generator=generator, + ), + dt_bias=0.1 + * torch.randn( + (_HEADS * _DIM,), + dtype=torch.float32, + device=_DEVICE, + generator=generator, + ), + state=state, + state_indices=torch.arange( + 1, + batch + 1, + dtype=torch.int32, + device=_DEVICE, + ), + output_gate=output_gate_storage[:, : _HEADS * _DIM].view(batch, _HEADS, _DIM), + norm_weight=torch.randn( + (_DIM,), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ), + ) + + +def _copy_inputs(inputs: Inputs) -> Inputs: + def clone_preserving_strides(tensor: torch.Tensor) -> torch.Tensor: + clone = torch.empty_strided( + tensor.shape, + tensor.stride(), + dtype=tensor.dtype, + device=tensor.device, + ) + clone.copy_(tensor) + return clone + + return Inputs( + **{ + name: clone_preserving_strides(getattr(inputs, name)) + for name in Inputs.__dataclass_fields__ + } + ) + + +def _reference(inputs: Inputs) -> torch.Tensor: + batch = inputs.x.shape[0] + output = torch.zeros( + (1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=_DEVICE, + ) + dt_bias = inputs.dt_bias.view(_HEADS, _DIM) + + for batch_idx in range(batch): + state_idx = int(inputs.state_indices[batch_idx]) + if state_idx <= 0: + continue + + history = inputs.conv_state[state_idx] + conv_values = torch.cat( + ( + history.float(), + inputs.x[batch_idx, :, None].float(), + ), + dim=-1, + ) + packed_qkv = F.silu((conv_values * inputs.conv_weight).sum(dim=-1)).to( + torch.bfloat16 + ) + history[:, 0].copy_(history[:, 1]) + history[:, 1].copy_(history[:, 2]) + history[:, 2].copy_(inputs.x[batch_idx]) + + q, k, v = packed_qkv.view( + 3, + _HEADS, + _DIM, + ).unbind(0) + for head_idx in range(_HEADS): + q_head = q[head_idx].float() + k_head = k[head_idx].float() + q_head = q_head * torch.rsqrt(q_head.square().sum() + 1e-6) + q_head = q_head * (_DIM**-0.5) + k_head = k_head * torch.rsqrt(k_head.square().sum() + 1e-6) + + a = inputs.A_log[head_idx].exp() + decay = ( + _LOWER_BOUND + * torch.sigmoid( + ( + inputs.raw_g[ + 0, + batch_idx, + head_idx, + ].float() + + dt_bias[head_idx] + ) + * a + ) + ).exp() + decayed_state = inputs.state[state_idx, head_idx] * decay[None, :] + state_dot_k = decayed_state @ k_head + state_dot_q = decayed_state @ q_head + k_dot_q = torch.dot(k_head, q_head) + beta = torch.sigmoid( + inputs.raw_beta[ + 0, + batch_idx, + head_idx, + ].float() + ) + v_new = (v[head_idx].float() - state_dot_k) * beta + inputs.state[state_idx, head_idx].copy_( + decayed_state + v_new[:, None] * k_head[None, :] + ) + + # The model materializes recurrent output in BF16 before the + # normalization/gating operation. + recurrent = (state_dot_q + v_new * k_dot_q).to(torch.bfloat16) + recurrent_f32 = recurrent.float() + inv_rms = torch.rsqrt(recurrent_f32.square().mean() + _NORM_EPS) + output[0, batch_idx, head_idx] = ( + recurrent_f32 + * inv_rms + * inputs.norm_weight.float() + * torch.sigmoid( + inputs.output_gate[ + batch_idx, + head_idx, + ].float() + ) + ).to(torch.bfloat16) + return output + + +def _relative_rmse( + reference: torch.Tensor, + actual: torch.Tensor, +) -> float: + delta = actual.float() - reference.float() + return float( + delta.square().mean().sqrt() / (reference.float().square().mean().sqrt() + 1e-8) + ) + + +def _run(inputs: Inputs) -> torch.Tensor: + return flydsl_kimi_k3_kda_decode( + x=inputs.x, + conv_weight=inputs.conv_weight, + conv_bias=None, + conv_state=inputs.conv_state, + raw_g=inputs.raw_g, + raw_beta=inputs.raw_beta, + A_log=inputs.A_log, + dt_bias=inputs.dt_bias, + lower_bound=_LOWER_BOUND, + state=inputs.state, + state_indices=inputs.state_indices, + output_gate=inputs.output_gate, + norm_weight=inputs.norm_weight, + norm_eps=_NORM_EPS, + ) + + +def _make_fb_inputs( + batch: int, + seed: int = 20260728, +) -> tuple[torch.Tensor, torch.Tensor, Inputs]: + generator = torch.Generator(device=_DEVICE).manual_seed(seed + 10_000 + batch) + f_a_storage = torch.randn( + (batch, _DIM + 5), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + f_a = f_a_storage[:, :_DIM] + f_b_weight = ( + 0.05 + * torch.randn( + (_HEADS, _DIM, _DIM), + dtype=torch.bfloat16, + device=_DEVICE, + generator=generator, + ) + ).to(torch.bfloat16) + inputs = _make_inputs(batch, seed) + projected = F.linear( + f_a.float(), + f_b_weight.view(_HEADS * _DIM, _DIM).float(), + ).to(torch.bfloat16) + inputs.raw_g.copy_(projected.view(1, batch, _HEADS, _DIM)) + return f_a, f_b_weight, inputs + + +def _run_with_f_b( + f_a: torch.Tensor, + f_b_weight: torch.Tensor, + inputs: Inputs, +) -> torch.Tensor: + return flydsl_kimi_k3_kda_decode_with_f_b( + f_a=f_a, + f_b_weight=f_b_weight, + x=inputs.x, + conv_weight=inputs.conv_weight, + conv_bias=None, + conv_state=inputs.conv_state, + raw_beta=inputs.raw_beta, + A_log=inputs.A_log, + dt_bias=inputs.dt_bias, + lower_bound=_LOWER_BOUND, + state=inputs.state, + state_indices=inputs.state_indices, + output_gate=inputs.output_gate, + norm_weight=inputs.norm_weight, + norm_eps=_NORM_EPS, + ) + + +def test_public_api_and_support_predicate() -> None: + import aiter.ops.flydsl as flydsl_ops + + assert flydsl_ops.flydsl_kimi_k3_kda_decode is flydsl_kimi_k3_kda_decode + assert ( + flydsl_ops.is_flydsl_kimi_k3_kda_decode_supported + is is_flydsl_kimi_k3_kda_decode_supported + ) + assert is_flydsl_kimi_k3_kda_decode_supported(0) + assert not is_flydsl_kimi_k3_kda_decode_supported("cpu") + + +def test_f_b_public_api() -> None: + import aiter.ops.flydsl as flydsl_ops + + assert ( + flydsl_ops.flydsl_kimi_k3_kda_decode_with_f_b + is flydsl_kimi_k3_kda_decode_with_f_b + ) + + +@pytest.mark.parametrize("batch", [1, 8, 16]) +def test_kimi_k3_kda_decode_matches_reference(batch: int) -> None: + seed = _make_inputs(batch) + reference_inputs = _copy_inputs(seed) + actual_inputs = _copy_inputs(seed) + + reference = _reference(reference_inputs) + actual = _run(actual_inputs) + torch.cuda.synchronize() + + assert is_flydsl_kimi_k3_kda_decode_supported(_DEVICE) + assert not torch.isnan(actual).any() + assert _relative_rmse(reference, actual) < 1e-3 + assert ( + _relative_rmse( + reference_inputs.state, + actual_inputs.state, + ) + < 1e-3 + ) + assert torch.equal( + reference_inputs.conv_state, + actual_inputs.conv_state, + ) + + +def test_non_positive_slots_do_not_modify_caches() -> None: + inputs = _make_inputs(batch=2) + inputs.state_indices.copy_(torch.tensor([0, -1], dtype=torch.int32, device=_DEVICE)) + conv_before = inputs.conv_state.clone() + state_before = inputs.state.clone() + + actual = _run(inputs) + torch.cuda.synchronize() + + assert torch.count_nonzero(actual) == 0 + assert torch.equal(inputs.conv_state, conv_before) + assert torch.equal(inputs.state, state_before) + + +@pytest.mark.parametrize("batch", [1, 8, 16]) +def test_kimi_k3_kda_decode_with_f_b_matches_reference(batch: int) -> None: + f_a, f_b_weight, seed = _make_fb_inputs(batch) + reference_inputs = _copy_inputs(seed) + actual_inputs = _copy_inputs(seed) + + reference = _reference(reference_inputs) + actual = _run_with_f_b(f_a, f_b_weight, actual_inputs) + torch.cuda.synchronize() + + assert not torch.isnan(actual).any() + assert _relative_rmse(reference, actual) < 1e-3 + assert _relative_rmse(reference_inputs.state, actual_inputs.state) < 1e-3 + assert torch.equal(reference_inputs.conv_state, actual_inputs.conv_state) + + +def test_f_b_non_positive_slots_do_not_modify_caches() -> None: + f_a, f_b_weight, inputs = _make_fb_inputs(batch=2) + inputs.state_indices.copy_(torch.tensor([0, -1], dtype=torch.int32, device=_DEVICE)) + conv_before = inputs.conv_state.clone() + state_before = inputs.state.clone() + + actual = _run_with_f_b(f_a, f_b_weight, inputs) + torch.cuda.synchronize() + + assert torch.count_nonzero(actual) == 0 + assert torch.equal(inputs.conv_state, conv_before) + assert torch.equal(inputs.state, state_before) + + +def test_f_b_api_rejects_invalid_projection_inputs() -> None: + f_a, f_b_weight, inputs = _make_fb_inputs(batch=1) + + with pytest.raises(ValueError, match="`f_a` must have rank 2"): + _run_with_f_b(f_a.unsqueeze(0), f_b_weight, inputs) + with pytest.raises(ValueError, match="`f_a` must have dtype"): + _run_with_f_b(f_a.float(), f_b_weight, inputs) + with pytest.raises(ValueError, match="`f_b_weight` must have shape"): + _run_with_f_b(f_a, f_b_weight[:, :, :-1], inputs) + with pytest.raises(ValueError, match="`f_b_weight` must have inner strides"): + _run_with_f_b(f_a, f_b_weight.transpose(1, 2), inputs) + + +def test_decode_api_rejects_invalid_input_rank() -> None: + inputs = _make_inputs(batch=1) + inputs.x = inputs.x.unsqueeze(0) + + with pytest.raises(ValueError, match="`x` must have rank 2"): + _run(inputs)