From fb2d7253322f5d4d795f7ae90c04a57afeb00915 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 03:32:06 +0000 Subject: [PATCH 1/6] perf(flydsl): fuse Kimi-K3 KDA decode on gfx950 Signed-off-by: Yanyuan Qin Assisted-by: OpenAI Codex --- aiter/ops/flydsl/__init__.py | 6 + .../ops/flydsl/kernels/kimi_k3_kda_decode.py | 564 ++++++++++++++++++ aiter/ops/flydsl/kimi_k3_kda_decode.py | 305 ++++++++++ .../flydsl_tests/test_kimi_k3_kda_decode.py | 340 +++++++++++ 4 files changed, 1215 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py create mode 100644 aiter/ops/flydsl/kimi_k3_kda_decode.py create mode 100644 op_tests/flydsl_tests/test_kimi_k3_kda_decode.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 33a14bcfb49..d148e18c41c 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -48,6 +48,10 @@ from .kernels.fp8_mqa_logits import ( flydsl_fp8_mqa_logits, ) + from .kimi_k3_kda_decode import ( + flydsl_kimi_k3_kda_decode, + is_flydsl_kimi_k3_kda_decode_supported, + ) from .kernels.pa_mqa_logits_fp4 import ( flydsl_pa_mqa_logits_fp4, ) @@ -70,6 +74,8 @@ "flydsl_hgemm", "flydsl_moe_stage1", "flydsl_moe_stage2", + "flydsl_kimi_k3_kda_decode", + "is_flydsl_kimi_k3_kda_decode_supported", "flydsl_pa_mqa_logits_fp4", "flydsl_pa_mqa_logits_fp4_prefill", "flydsl_pa_mqa_logits_fp4_varqlen", 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..56a6c51b9a2 --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py @@ -0,0 +1,564 @@ +# 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, vector +from flydsl.expr.typing import T + +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.lru_cache(maxsize=None) +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), + ): + 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/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kimi_k3_kda_decode.py new file mode 100644 index 00000000000..2c5d51a96f7 --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_kda_decode.py @@ -0,0 +1,305 @@ +# 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.tensor_shim import _run_compiled + +_HEADS = 12 +_DIM = 128 +_CONV_CHANNELS = 3 * _HEADS * _DIM +_CONV_WIDTH = 4 + + +@functools.lru_cache(maxsize=None) +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)}, " f"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}, " + f"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 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 not x.is_cuda: + raise ValueError("`x` must be a CUDA tensor.") + device = x.device + if not is_flydsl_kimi_k3_kda_decode_supported(device): + raise RuntimeError("`flydsl_kimi_k3_kda_decode` requires a gfx950 GPU.") + batch = x.shape[0] if x.ndim == 2 else -1 + if batch <= 0: + raise ValueError("`x` 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.") + + tensors = ( + ("conv_weight", conv_weight), + ("conv_state", conv_state), + ("raw_g", raw_g), + ("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), + ) + _check_same_device(tensors, 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( + "`state` must have shape [cache, 12, 128, 128], " + f"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_g", + raw_g, + shape=(1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(_DIM, 1), + ) + _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: + out = torch.empty( + (1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + ) + else: + _check_same_device((("out", out),), device) + _check_tensor( + "out", + out, + shape=(1, batch, _HEADS, _DIM), + dtype=torch.bfloat16, + device=device, + inner_strides=(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 + + +__all__ = [ + "flydsl_kimi_k3_kda_decode", + "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..06395c33937 --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py @@ -0,0 +1,340 @@ +# 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 + +from aiter.jit.utils.chip_info import get_gfx +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", +) + +from aiter.ops.flydsl.kimi_k3_kda_decode import ( # noqa: E402 + flydsl_kimi_k3_kda_decode, + is_flydsl_kimi_k3_kda_decode_supported, +) + +_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 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") + + +@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) From 63c8cb38c1e36ad69f2aaa09632d90c0cd11117b Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 09:41:18 +0000 Subject: [PATCH 2/6] perf(flydsl): fuse Kimi-K3 f_b projection into KDA decode Signed-off-by: Yanyuan Qin Assisted-by: OpenAI Codex --- aiter/ops/flydsl/__init__.py | 2 + .../flydsl/kernels/kimi_k3_kda_decode_fb.py | 623 ++++++++++++++++++ aiter/ops/flydsl/kimi_k3_kda_decode.py | 272 ++++++-- .../flydsl_tests/test_kimi_k3_kda_decode.py | 105 +++ 4 files changed, 943 insertions(+), 59 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index d148e18c41c..ba99353d996 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -50,6 +50,7 @@ ) 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 .kernels.pa_mqa_logits_fp4 import ( @@ -75,6 +76,7 @@ "flydsl_moe_stage1", "flydsl_moe_stage2", "flydsl_kimi_k3_kda_decode", + "flydsl_kimi_k3_kda_decode_with_f_b", "is_flydsl_kimi_k3_kda_decode_supported", "flydsl_pa_mqa_logits_fp4", "flydsl_pa_mqa_logits_fp4_prefill", 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..76e4fec2955 --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py @@ -0,0 +1,623 @@ +# 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, vector +from flydsl.expr.typing import T + +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.lru_cache(maxsize=None) +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), + ): + 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 index 2c5d51a96f7..0c449498738 100644 --- a/aiter/ops/flydsl/kimi_k3_kda_decode.py +++ b/aiter/ops/flydsl/kimi_k3_kda_decode.py @@ -13,6 +13,9 @@ 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 @@ -68,7 +71,7 @@ def _check_tensor( ) -> None: if tensor.shape != shape: raise ValueError( - f"`{name}` must have shape {list(shape)}, " f"got {list(tensor.shape)}." + 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}.") @@ -76,8 +79,7 @@ def _check_tensor( 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}, " - f"got {tensor.stride()}." + f"`{name}` must have inner strides {inner_strides}, got {tensor.stride()}." ) @@ -92,12 +94,16 @@ def _check_same_device( raise ValueError(f"`{name}` must be on {device}, got {tensor.device}.") -def flydsl_kimi_k3_kda_decode( +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_g: torch.Tensor, raw_beta: torch.Tensor, A_log: torch.Tensor, dt_bias: torch.Tensor, @@ -106,48 +112,33 @@ def flydsl_kimi_k3_kda_decode( state_indices: torch.Tensor, output_gate: torch.Tensor, norm_weight: torch.Tensor, - norm_eps: float, - out: torch.Tensor | None = None, + out: torch.Tensor | 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 not x.is_cuda: - raise ValueError("`x` must be a CUDA tensor.") - device = x.device + """Validate operands shared by both explicit KDA specializations.""" if not is_flydsl_kimi_k3_kda_decode_supported(device): - raise RuntimeError("`flydsl_kimi_k3_kda_decode` requires a gfx950 GPU.") - batch = x.shape[0] if x.ndim == 2 else -1 + raise RuntimeError(f"`{api_name}` requires a gfx950 GPU.") if batch <= 0: - raise ValueError("`x` must have a non-empty batch dimension.") + 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.") - tensors = ( - ("conv_weight", conv_weight), - ("conv_state", conv_state), - ("raw_g", raw_g), - ("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), + _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_same_device(tensors, device) - _check_tensor( "x", x, @@ -179,21 +170,12 @@ def flydsl_kimi_k3_kda_decode( _DIM, ): raise ValueError( - "`state` must have shape [cache, 12, 128, 128], " - f"got {list(state.shape)}." + 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_g", - raw_g, - shape=(1, batch, _HEADS, _DIM), - dtype=torch.bfloat16, - device=device, - inner_strides=(_DIM, 1), - ) _check_tensor( "raw_beta", raw_beta, @@ -244,21 +226,84 @@ def flydsl_kimi_k3_kda_decode( ) if out is None: - out = torch.empty( + return torch.empty( (1, batch, _HEADS, _DIM), dtype=torch.bfloat16, device=device, ) - else: - _check_same_device((("out", out),), device) - _check_tensor( - "out", - out, - shape=(1, batch, _HEADS, _DIM), - dtype=torch.bfloat16, - device=device, - inner_strides=(1,), - ) + _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 not x.is_cuda: + raise ValueError("`x` must be a CUDA tensor.") + device = x.device + batch = x.shape[0] if x.ndim == 2 else -1 + 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), @@ -299,7 +344,116 @@ def flydsl_kimi_k3_kda_decode( 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 not f_a.is_cuda: + raise ValueError("`f_a` must be a CUDA tensor.") + device = f_a.device + batch = f_a.shape[0] if f_a.ndim == 2 else -1 + _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 index 06395c33937..1b6741ce0b1 100644 --- a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py @@ -31,6 +31,7 @@ def _gfx950_flydsl_available() -> bool: from aiter.ops.flydsl.kimi_k3_kda_decode import ( # noqa: E402 flydsl_kimi_k3_kda_decode, + flydsl_kimi_k3_kda_decode_with_f_b, is_flydsl_kimi_k3_kda_decode_supported, ) @@ -288,6 +289,60 @@ def _run(inputs: Inputs) -> torch.Tensor: ) +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 @@ -300,6 +355,15 @@ def test_public_api_and_support_predicate() -> None: 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) @@ -338,3 +402,44 @@ def test_non_positive_slots_do_not_modify_caches() -> None: 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 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) From 1a7f8260dd658a78077e92877fe7f8a6aa228e7f Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Fri, 31 Jul 2026 21:15:05 +0000 Subject: [PATCH 3/6] style(flydsl): satisfy current Ruff checks Sort public imports and exports, use functools.cache, and annotate intentional FlyDSL stream defaults for the repository-wide lint configuration. Assisted-by: OpenAI Codex Signed-off-by: Yanyuan Qin --- aiter/ops/flydsl/__init__.py | 16 ++++++++-------- aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py | 4 ++-- .../ops/flydsl/kernels/kimi_k3_kda_decode_fb.py | 4 ++-- aiter/ops/flydsl/kimi_k3_kda_decode.py | 2 +- op_tests/flydsl_tests/test_kimi_k3_kda_decode.py | 2 +- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index ba99353d996..2b63111b126 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -48,11 +48,6 @@ from .kernels.fp8_mqa_logits import ( flydsl_fp8_mqa_logits, ) - 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 .kernels.pa_mqa_logits_fp4 import ( flydsl_pa_mqa_logits_fp4, ) @@ -62,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 @@ -73,15 +73,15 @@ "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_hgemm", - "flydsl_moe_stage1", - "flydsl_moe_stage2", "flydsl_kimi_k3_kda_decode", "flydsl_kimi_k3_kda_decode_with_f_b", - "is_flydsl_kimi_k3_kda_decode_supported", + "flydsl_moe_stage1", + "flydsl_moe_stage2", "flydsl_pa_mqa_logits_fp4", "flydsl_pa_mqa_logits_fp4_prefill", "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 index 56a6c51b9a2..945536412f5 100644 --- a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py @@ -34,7 +34,7 @@ _WAVES_PER_EU = 3 -@functools.lru_cache(maxsize=None) +@functools.cache def create_kimi_k3_kda_decode_kernel(norm_eps: float, lower_bound: float): """Build the fixed gfx950 BF16 Kimi-K3 decode specialization.""" @@ -517,7 +517,7 @@ def launch( stride_gate_head: fx.Int32, stride_out_token: fx.Int32, stride_out_head: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = fx.Stream(None), # noqa: B008 ): kernel( x_mem, diff --git a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py index 76e4fec2955..e5d620a435b 100644 --- a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py @@ -36,7 +36,7 @@ _WAVES_PER_EU = 2 -@functools.lru_cache(maxsize=None) +@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.""" @@ -573,7 +573,7 @@ def launch( stride_gate_head: fx.Int32, stride_out_token: fx.Int32, stride_out_head: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = fx.Stream(None), # noqa: B008 ): kernel( f_a_mem, diff --git a/aiter/ops/flydsl/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kimi_k3_kda_decode.py index 0c449498738..2082e9d4429 100644 --- a/aiter/ops/flydsl/kimi_k3_kda_decode.py +++ b/aiter/ops/flydsl/kimi_k3_kda_decode.py @@ -24,7 +24,7 @@ _CONV_WIDTH = 4 -@functools.lru_cache(maxsize=None) +@functools.cache def _rocm_arch(device: torch.device) -> str | None: properties = torch.cuda.get_device_properties(device) arch = getattr(properties, "gcnArchName", None) diff --git a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py index 1b6741ce0b1..1e7906bf938 100644 --- a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py @@ -29,7 +29,7 @@ def _gfx950_flydsl_available() -> bool: reason="gfx950 FlyDSL required", ) -from aiter.ops.flydsl.kimi_k3_kda_decode import ( # noqa: E402 +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 636d7862f26131da3a37852c05b20d77300ba0e7 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Fri, 31 Jul 2026 23:29:35 +0000 Subject: [PATCH 4/6] fix(flydsl): address KDA review feedback Signed-off-by: Yanyuan Qin --- .../ops/flydsl/kernels/kimi_k3_kda_decode.py | 4 +++- .../flydsl/kernels/kimi_k3_kda_decode_fb.py | 4 +++- aiter/ops/flydsl/kimi_k3_kda_decode.py | 8 +++++-- .../flydsl_tests/test_kimi_k3_kda_decode.py | 22 ++++++++++++++----- 4 files changed, 28 insertions(+), 10 deletions(-) diff --git a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py index 945536412f5..bd03af0342e 100644 --- a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py @@ -12,9 +12,11 @@ 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, 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 diff --git a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py index e5d620a435b..b4a84bea4e5 100644 --- a/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py +++ b/aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py @@ -12,9 +12,11 @@ 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, 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 diff --git a/aiter/ops/flydsl/kimi_k3_kda_decode.py b/aiter/ops/flydsl/kimi_k3_kda_decode.py index 2082e9d4429..2acdabd799a 100644 --- a/aiter/ops/flydsl/kimi_k3_kda_decode.py +++ b/aiter/ops/flydsl/kimi_k3_kda_decode.py @@ -272,10 +272,12 @@ def flydsl_kimi_k3_kda_decode( :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] if x.ndim == 2 else -1 + batch = x.shape[0] out = _validate_kda_inputs( api_name="flydsl_kimi_k3_kda_decode", batch_source="x", @@ -369,10 +371,12 @@ def flydsl_kimi_k3_kda_decode_with_f_b( 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] if f_a.ndim == 2 else -1 + batch = f_a.shape[0] _check_same_device((("f_b_weight", f_b_weight),), device) _check_tensor( "f_a", diff --git a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py index 1e7906bf938..87e25cd78b2 100644 --- a/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_decode.py @@ -11,7 +11,13 @@ 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 @@ -29,12 +35,6 @@ def _gfx950_flydsl_available() -> bool: reason="gfx950 FlyDSL required", ) -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, -) - _DEVICE = torch.device("cuda") _HEADS = 12 _DIM = 128 @@ -437,9 +437,19 @@ def test_f_b_non_positive_slots_do_not_modify_caches() -> None: 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) From 762b3e9bbc15c8a384842da510dc95b16d6d6c9f Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Sun, 2 Aug 2026 01:38:45 +0000 Subject: [PATCH 5/6] test(perf): add Kimi-K3 KDA f_b benchmark Signed-off-by: Yanyuan Qin --- .../flydsl/bench_kimi_k3_kda_decode_fb.py | 288 ++++++++++++++++++ 1 file changed, 288 insertions(+) create mode 100644 op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py diff --git a/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py b/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py new file mode 100644 index 00000000000..6e5eb2d5e4d --- /dev/null +++ b/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmark separate versus fused f_b projection in Kimi-K3 KDA decode.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +import torch + +from aiter.ops.flydsl.kimi_k3_kda_decode import ( + flydsl_kimi_k3_kda_decode, + flydsl_kimi_k3_kda_decode_with_f_b, +) + +HEADS = 12 +DIM = 128 +CHANNELS = 3 * HEADS * DIM +CONV_WIDTH = 4 +LOWER_BOUND = -5.0 +NORM_EPS = 1.0e-5 +ROTATIONS = 8 + + +@dataclass +class Case: + f_a: torch.Tensor + f_b_weight: torch.Tensor + x: torch.Tensor + conv_weight: torch.Tensor + conv_state: 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 + out: torch.Tensor + + +Operation = Callable[[Case], torch.Tensor] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--seed", type=int, default=20260728) + parser.add_argument("--operations-per-graph", type=int, default=100) + parser.add_argument("--warmup-replays", type=int, default=100) + parser.add_argument("--replays-per-trial", type=int, default=10) + parser.add_argument("--trials", type=int, default=21) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def make_case(generator: torch.Generator) -> Case: + slots = 3 + return Case( + f_a=torch.randn( + (1, DIM), generator=generator, device="cuda", dtype=torch.bfloat16 + ), + f_b_weight=torch.randn( + (HEADS, DIM, DIM), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ).mul_(0.05), + x=torch.randn( + (1, CHANNELS), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ), + conv_weight=torch.randn( + (CHANNELS, CONV_WIDTH), + generator=generator, + device="cuda", + dtype=torch.float32, + ).mul_(0.1), + conv_state=torch.randn( + (slots, CHANNELS, CONV_WIDTH - 1), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ), + raw_beta=torch.randn( + (1, 1, HEADS), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ), + A_log=torch.randn( + (HEADS,), generator=generator, device="cuda", dtype=torch.float32 + ).mul_(0.5), + dt_bias=torch.randn( + (HEADS * DIM,), + generator=generator, + device="cuda", + dtype=torch.float32, + ).mul_(0.1), + state=torch.randn( + (slots, HEADS, DIM, DIM), + generator=generator, + device="cuda", + dtype=torch.float32, + ).mul_(0.01), + state_indices=torch.tensor([1], device="cuda", dtype=torch.int32), + output_gate=torch.randn( + (1, HEADS, DIM), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ), + norm_weight=torch.randn( + (DIM,), generator=generator, device="cuda", dtype=torch.bfloat16 + ), + out=torch.empty((1, 1, HEADS, DIM), device="cuda", dtype=torch.bfloat16), + ) + + +def clone_case(case: Case) -> Case: + return Case( + **{name: getattr(case, name).clone() for name in Case.__dataclass_fields__} + ) + + +def control(case: Case): + from vllm.model_executor.layers.utils import rocm_unquantized_gemm_impl + + raw_g = rocm_unquantized_gemm_impl(case.f_a, case.f_b_weight.view(HEADS * DIM, DIM)) + return flydsl_kimi_k3_kda_decode( + case.x, + case.conv_weight, + None, + case.conv_state, + raw_g.view(1, 1, HEADS, DIM), + case.raw_beta, + case.A_log, + case.dt_bias, + LOWER_BOUND, + case.state, + case.state_indices, + case.output_gate, + case.norm_weight, + NORM_EPS, + case.out, + ) + + +def candidate(case: Case): + return flydsl_kimi_k3_kda_decode_with_f_b( + case.f_a, + case.f_b_weight, + case.x, + case.conv_weight, + None, + case.conv_state, + case.raw_beta, + case.A_log, + case.dt_bias, + LOWER_BOUND, + case.state, + case.state_indices, + case.output_gate, + case.norm_weight, + NORM_EPS, + case.out, + ) + + +def relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: + error = (actual.float() - expected.float()).square().mean().sqrt() + scale = expected.float().square().mean().sqrt().clamp_min(1.0e-8) + return (error / scale).item() + + +def capture(operation: Operation, cases: list[Case], operations: int): + for case in cases: + operation(case) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for index in range(operations): + operation(cases[index % len(cases)]) + return graph + + +def elapsed_us(graph, replays: int, operations: int): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000.0 / (replays * operations) + + +def main() -> None: + args = parse_args() + if ( + min( + args.operations_per_graph, + args.warmup_replays, + args.replays_per_trial, + args.trials, + ) + <= 0 + ): + raise ValueError("all benchmark counts must be positive") + if not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires a ROCm GPU") + properties = torch.cuda.get_device_properties(0) + arch = str(getattr(properties, "gcnArchName", "")).split(":", 1)[0] + if not torch.version.hip or arch != "gfx950": + raise RuntimeError(f"this benchmark requires ROCm gfx950, got {arch!r}") + + generator = torch.Generator(device="cuda").manual_seed(args.seed) + seed_case = make_case(generator) + control_case = clone_case(seed_case) + candidate_case = clone_case(seed_case) + expected = control(control_case).clone() + actual = candidate(candidate_case).clone() + torch.cuda.synchronize() + output_error = relative_rmse(actual, expected) + state_error = relative_rmse(candidate_case.state, control_case.state) + if max(output_error, state_error) >= 1.0e-3: + raise AssertionError( + f"correctness gate failed: output={output_error}, state={state_error}" + ) + + control_cases = [make_case(generator) for _ in range(ROTATIONS)] + candidate_cases = [clone_case(case) for case in control_cases] + graphs = { + "separate": capture(control, control_cases, args.operations_per_graph), + "fused": capture(candidate, candidate_cases, args.operations_per_graph), + } + for index in range(args.warmup_replays * len(graphs)): + graphs[("separate", "fused")[index % 2]].replay() + torch.cuda.synchronize() + samples = {name: [] for name in graphs} + for trial in range(args.trials): + order = ("separate", "fused") if trial % 2 == 0 else ("fused", "separate") + for name in order: + samples[name].append( + elapsed_us( + graphs[name], args.replays_per_trial, args.operations_per_graph + ) + ) + medians = {name: statistics.median(values) for name, values in samples.items()} + rotating_weight_bytes = sum( + case.f_b_weight.numel() * case.f_b_weight.element_size() + for case in control_cases + ) + result = { + "shape": "B1 Kimi-K3 TP8-local KDA decode plus f_b", + "runtime": { + "torch": torch.__version__, + "hip": torch.version.hip, + "device": properties.name, + "arch": arch, + }, + "seed": args.seed, + "rotations": ROTATIONS, + "rotating_f_b_weight_bytes": rotating_weight_bytes, + "operations_per_graph": args.operations_per_graph, + "warmup_replays": args.warmup_replays, + "replays_per_trial": args.replays_per_trial, + "trials": args.trials, + "relative_rmse": {"output": output_error, "state": state_error}, + "p50_us": medians, + "speedup": medians["separate"] / medians["fused"], + "samples_us": samples, + } + rendered = json.dumps(result, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.write_text(rendered) + print(rendered, end="") + + +if __name__ == "__main__": + main() From cd0f2a757f6c79e4747577ac4bf7cc392bc918c3 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Sun, 2 Aug 2026 05:06:11 +0000 Subject: [PATCH 6/6] Drop the flydsl microbenchmark from this PR op_tests/op_benchmarks/flydsl/ does not exist on main, no CI job references it, and none of the 84 existing flydsl kernel modules ship a companion benchmark. The script is kept out of tree; measured numbers go in the PR description instead. Signed-off-by: Yanyuan Qin --- .../flydsl/bench_kimi_k3_kda_decode_fb.py | 288 ------------------ 1 file changed, 288 deletions(-) delete mode 100644 op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py diff --git a/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py b/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py deleted file mode 100644 index 6e5eb2d5e4d..00000000000 --- a/op_tests/op_benchmarks/flydsl/bench_kimi_k3_kda_decode_fb.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Benchmark separate versus fused f_b projection in Kimi-K3 KDA decode.""" - -from __future__ import annotations - -import argparse -import json -import statistics -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path - -import torch - -from aiter.ops.flydsl.kimi_k3_kda_decode import ( - flydsl_kimi_k3_kda_decode, - flydsl_kimi_k3_kda_decode_with_f_b, -) - -HEADS = 12 -DIM = 128 -CHANNELS = 3 * HEADS * DIM -CONV_WIDTH = 4 -LOWER_BOUND = -5.0 -NORM_EPS = 1.0e-5 -ROTATIONS = 8 - - -@dataclass -class Case: - f_a: torch.Tensor - f_b_weight: torch.Tensor - x: torch.Tensor - conv_weight: torch.Tensor - conv_state: 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 - out: torch.Tensor - - -Operation = Callable[[Case], torch.Tensor] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--seed", type=int, default=20260728) - parser.add_argument("--operations-per-graph", type=int, default=100) - parser.add_argument("--warmup-replays", type=int, default=100) - parser.add_argument("--replays-per-trial", type=int, default=10) - parser.add_argument("--trials", type=int, default=21) - parser.add_argument("--output", type=Path) - return parser.parse_args() - - -def make_case(generator: torch.Generator) -> Case: - slots = 3 - return Case( - f_a=torch.randn( - (1, DIM), generator=generator, device="cuda", dtype=torch.bfloat16 - ), - f_b_weight=torch.randn( - (HEADS, DIM, DIM), - generator=generator, - device="cuda", - dtype=torch.bfloat16, - ).mul_(0.05), - x=torch.randn( - (1, CHANNELS), - generator=generator, - device="cuda", - dtype=torch.bfloat16, - ), - conv_weight=torch.randn( - (CHANNELS, CONV_WIDTH), - generator=generator, - device="cuda", - dtype=torch.float32, - ).mul_(0.1), - conv_state=torch.randn( - (slots, CHANNELS, CONV_WIDTH - 1), - generator=generator, - device="cuda", - dtype=torch.bfloat16, - ), - raw_beta=torch.randn( - (1, 1, HEADS), - generator=generator, - device="cuda", - dtype=torch.bfloat16, - ), - A_log=torch.randn( - (HEADS,), generator=generator, device="cuda", dtype=torch.float32 - ).mul_(0.5), - dt_bias=torch.randn( - (HEADS * DIM,), - generator=generator, - device="cuda", - dtype=torch.float32, - ).mul_(0.1), - state=torch.randn( - (slots, HEADS, DIM, DIM), - generator=generator, - device="cuda", - dtype=torch.float32, - ).mul_(0.01), - state_indices=torch.tensor([1], device="cuda", dtype=torch.int32), - output_gate=torch.randn( - (1, HEADS, DIM), - generator=generator, - device="cuda", - dtype=torch.bfloat16, - ), - norm_weight=torch.randn( - (DIM,), generator=generator, device="cuda", dtype=torch.bfloat16 - ), - out=torch.empty((1, 1, HEADS, DIM), device="cuda", dtype=torch.bfloat16), - ) - - -def clone_case(case: Case) -> Case: - return Case( - **{name: getattr(case, name).clone() for name in Case.__dataclass_fields__} - ) - - -def control(case: Case): - from vllm.model_executor.layers.utils import rocm_unquantized_gemm_impl - - raw_g = rocm_unquantized_gemm_impl(case.f_a, case.f_b_weight.view(HEADS * DIM, DIM)) - return flydsl_kimi_k3_kda_decode( - case.x, - case.conv_weight, - None, - case.conv_state, - raw_g.view(1, 1, HEADS, DIM), - case.raw_beta, - case.A_log, - case.dt_bias, - LOWER_BOUND, - case.state, - case.state_indices, - case.output_gate, - case.norm_weight, - NORM_EPS, - case.out, - ) - - -def candidate(case: Case): - return flydsl_kimi_k3_kda_decode_with_f_b( - case.f_a, - case.f_b_weight, - case.x, - case.conv_weight, - None, - case.conv_state, - case.raw_beta, - case.A_log, - case.dt_bias, - LOWER_BOUND, - case.state, - case.state_indices, - case.output_gate, - case.norm_weight, - NORM_EPS, - case.out, - ) - - -def relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: - error = (actual.float() - expected.float()).square().mean().sqrt() - scale = expected.float().square().mean().sqrt().clamp_min(1.0e-8) - return (error / scale).item() - - -def capture(operation: Operation, cases: list[Case], operations: int): - for case in cases: - operation(case) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for index in range(operations): - operation(cases[index % len(cases)]) - return graph - - -def elapsed_us(graph, replays: int, operations: int): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(replays): - graph.replay() - end.record() - end.synchronize() - return start.elapsed_time(end) * 1000.0 / (replays * operations) - - -def main() -> None: - args = parse_args() - if ( - min( - args.operations_per_graph, - args.warmup_replays, - args.replays_per_trial, - args.trials, - ) - <= 0 - ): - raise ValueError("all benchmark counts must be positive") - if not torch.cuda.is_available(): - raise RuntimeError("this benchmark requires a ROCm GPU") - properties = torch.cuda.get_device_properties(0) - arch = str(getattr(properties, "gcnArchName", "")).split(":", 1)[0] - if not torch.version.hip or arch != "gfx950": - raise RuntimeError(f"this benchmark requires ROCm gfx950, got {arch!r}") - - generator = torch.Generator(device="cuda").manual_seed(args.seed) - seed_case = make_case(generator) - control_case = clone_case(seed_case) - candidate_case = clone_case(seed_case) - expected = control(control_case).clone() - actual = candidate(candidate_case).clone() - torch.cuda.synchronize() - output_error = relative_rmse(actual, expected) - state_error = relative_rmse(candidate_case.state, control_case.state) - if max(output_error, state_error) >= 1.0e-3: - raise AssertionError( - f"correctness gate failed: output={output_error}, state={state_error}" - ) - - control_cases = [make_case(generator) for _ in range(ROTATIONS)] - candidate_cases = [clone_case(case) for case in control_cases] - graphs = { - "separate": capture(control, control_cases, args.operations_per_graph), - "fused": capture(candidate, candidate_cases, args.operations_per_graph), - } - for index in range(args.warmup_replays * len(graphs)): - graphs[("separate", "fused")[index % 2]].replay() - torch.cuda.synchronize() - samples = {name: [] for name in graphs} - for trial in range(args.trials): - order = ("separate", "fused") if trial % 2 == 0 else ("fused", "separate") - for name in order: - samples[name].append( - elapsed_us( - graphs[name], args.replays_per_trial, args.operations_per_graph - ) - ) - medians = {name: statistics.median(values) for name, values in samples.items()} - rotating_weight_bytes = sum( - case.f_b_weight.numel() * case.f_b_weight.element_size() - for case in control_cases - ) - result = { - "shape": "B1 Kimi-K3 TP8-local KDA decode plus f_b", - "runtime": { - "torch": torch.__version__, - "hip": torch.version.hip, - "device": properties.name, - "arch": arch, - }, - "seed": args.seed, - "rotations": ROTATIONS, - "rotating_f_b_weight_bytes": rotating_weight_bytes, - "operations_per_graph": args.operations_per_graph, - "warmup_replays": args.warmup_replays, - "replays_per_trial": args.replays_per_trial, - "trials": args.trials, - "relative_rmse": {"output": output_error, "state": state_error}, - "p50_us": medians, - "speedup": medians["separate"] / medians["fused"], - "samples_us": samples, - } - rendered = json.dumps(result, indent=2, sort_keys=True) + "\n" - if args.output is not None: - args.output.write_text(rendered) - print(rendered, end="") - - -if __name__ == "__main__": - main()