diff --git a/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort_parallel.py b/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort_parallel.py new file mode 100644 index 0000000000..7ef4535b0f --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort_parallel.py @@ -0,0 +1,724 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Parallel Kimi-K3 B1 route, metadata, and stage-1 MXFP8 prep for gfx950.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, vector +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate +from flydsl.expr.typing import T + +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from aiter.ops.flydsl.kernels.quant_utils import emit_mx_e8m0_scale +from aiter.utility.mx_types import MxDtypeInt, MxScaleRoundModeInt + +_ROUTE_THREADS = 64 +_BLOCK_THREADS = 256 +_QUANT_THREADS = _BLOCK_THREADS - _ROUTE_THREADS +_EXPERTS = 896 +_TOPK = 16 +_LOCAL_ROUTE_CANDIDATES = _EXPERTS // _ROUTE_THREADS +_DPP_PRIORITY_HIGH_MASK = 0x38 +_SHORT_PRIORITY_LANES = _ROUTE_THREADS // 2 +_SHORT_LANE_CANDIDATES = 12 +_FULL_LANE_CANDIDATES = 16 +_SHORT_PRIORITY_EXPERTS = _SHORT_PRIORITY_LANES * _SHORT_LANE_CANDIDATES +_BLOCK_M = 32 +_SORTED_ROWS = _TOPK * _BLOCK_M +_MODEL_DIM = 3584 +_MX_GROUP_SIZE = 32 +_MX_SCALE_COLS = _MODEL_DIM // _MX_GROUP_SIZE +_MX_GROUPS_PER_ITERATION = _QUANT_THREADS // 2 +_MX_VALUES_PER_THREAD = _MX_GROUP_SIZE // 2 +_LOG2E = 1.4426950408889634 + + +@fx.struct +class _RouteSortStorage: + biased_scores: fx.Array[fx.Float32, _EXPERTS, 16] + priority_expert_ids: fx.Array[fx.Int32, _EXPERTS, 16] + route_scores: fx.Array[fx.Float32, _EXPERTS, 16] + selected_ids: fx.Array[fx.Int32, _TOPK, 16] + selected_scores: fx.Array[fx.Float32, _TOPK, 16] + route_sum: fx.Array[fx.Float32, 1, 16] + local_route_ids: fx.Array[ + fx.Int32, + _ROUTE_THREADS * _LOCAL_ROUTE_CANDIDATES, + 16, + ] + local_route_scores: fx.Array[ + fx.Float32, + _ROUTE_THREADS * _LOCAL_ROUTE_CANDIDATES, + 16, + ] + + +def _lds_load(ptr, idx): + return fx.ptr_load(ptr + fx.Int64(idx)) + + +def _lds_store(ptr, value, idx): + fx.ptr_store(value, ptr + fx.Int64(idx)) + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_b1_route_sort_parallel_module(): + """Build the fixed-shape overlapped route+MXFP8 preparation launcher.""" + + @flyc.kernel( + name="kimi_k3_b1_route_quant_parallel_gfx950", + known_block_size=[_BLOCK_THREADS, 1, 1], + ) + def route_sort_kernel( + logits: fx.Pointer, + correction_bias: fx.Pointer, + topk_weights: fx.Pointer, + topk_ids: fx.Pointer, + sorted_ids: fx.Pointer, + sorted_weights: fx.Pointer, + sorted_expert_ids: fx.Pointer, + num_valid_ids: fx.Pointer, + moe_buf: fx.Pointer, + moe_buf_i32_elements: fx.Int32, + hidden_states: fx.Pointer, + quantized_hidden_states: fx.Pointer, + quantized_scales: fx.Pointer, + ): + i32 = T.i32 + f32 = T.f32 + tid = ArithValue(gpu.thread_idx.x) + c_zero_i32 = arith.constant(0, type=i32) + c_one_i32 = arith.constant(1, type=i32) + c_zero_f32 = arith.constant(0.0, type=f32) + c_one_f32 = arith.constant(1.0, type=f32) + c_neg_inf = arith.constant(float("-inf"), type=f32) + vec4_bf16 = T.vec(4, T.bf16) + vec2_i32 = T.vec(2, i32) + vec2_f32 = T.vec(2, f32) + vec4_i32 = T.vec(4, i32) + vec4_f32 = T.vec(4, f32) + + logits_rsrc = ptr_rsrc(logits) + bias_rsrc = ptr_rsrc(correction_bias) + topk_weights_rsrc = ptr_rsrc(topk_weights) + topk_ids_rsrc = ptr_rsrc(topk_ids) + sorted_ids_rsrc = ptr_rsrc(sorted_ids) + sorted_weights_rsrc = ptr_rsrc(sorted_weights) + sorted_experts_rsrc = ptr_rsrc(sorted_expert_ids) + nvalid_rsrc = ptr_rsrc(num_valid_ids) + moe_buf_rsrc = ptr_rsrc(moe_buf) + hidden_rsrc = ptr_rsrc(hidden_states) + quantized_hidden_rsrc = ptr_rsrc(quantized_hidden_states) + quantized_scale_rsrc = ptr_rsrc(quantized_scales) + + lds = fx.SharedAllocator().allocate(_RouteSortStorage).peek() + biased_lds = lds.biased_scores.ptr + priority_expert_ids_lds = lds.priority_expert_ids.ptr + route_lds = lds.route_scores.ptr + selected_ids_lds = lds.selected_ids.ptr + selected_scores_lds = lds.selected_scores.ptr + route_sum_lds = lds.route_sum.ptr + local_route_ids_lds = lds.local_route_ids.ptr + local_route_scores_lds = lds.local_route_scores.ptr + + def lane_priority_rank(lane): + """Return the accepted DPP equal-score rank for a wave64 lane.""" + lane_bit2 = (lane >> arith.constant(2, type=i32)) & c_one_i32 + lane_bit3 = (lane >> arith.constant(3, type=i32)) & c_one_i32 + rank_hi = ( + lane ^ arith.constant(_DPP_PRIORITY_HIGH_MASK, type=i32) + ) & arith.constant(_DPP_PRIORITY_HIGH_MASK, type=i32) + rank_lo = (lane & arith.constant(7, type=i32)) ^ ( + (lane_bit3 << arith.constant(2, type=i32)) + | (lane_bit2 * arith.constant(3, type=i32)) + ) + return rank_hi | rank_lo + + def expert_priority_index(expert): + """Map an expert ID to the accepted global stable-tie order.""" + original_lane = (expert >> arith.constant(2, type=i32)) & arith.constant( + 63, type=i32 + ) + rank = lane_priority_rank(original_lane) + local_order = ( + (expert >> arith.constant(8, type=i32)) << arith.constant(2, type=i32) + ) | (expert & arith.constant(3, type=i32)) + short_lane = arith.cmpi( + CmpIPredicate.uge, + original_lane, + arith.constant(_SHORT_PRIORITY_LANES, type=i32), + ) + short_index = ( + rank * arith.constant(_SHORT_LANE_CANDIDATES, type=i32) + local_order + ) + full_index = ( + arith.constant(_SHORT_PRIORITY_EXPERTS, type=i32) + + (rank - arith.constant(_SHORT_PRIORITY_LANES, type=i32)) + * arith.constant(_FULL_LANE_CANDIDATES, type=i32) + + local_order + ) + return arith.select(short_lane, short_index, full_index) + + # Four waves compute sigmoid+bias in parallel. Biased scores are + # staged in the accepted global tie order so the route wave can split + # all 896 experts into 64 balanced, contiguous 14-score partitions. + for vec_base in range_constexpr(0, _EXPERTS, _BLOCK_THREADS * 4): + expert_base = tid * arith.constant(4, type=i32) + arith.constant( + vec_base, type=i32 + ) + vector_in_range = arith.cmpi( + CmpIPredicate.ult, + expert_base, + arith.constant(_EXPERTS, type=i32), + ) + load_if = scf.IfOp(vector_in_range) + with ir.InsertionPoint(load_if.then_block): + logits_vec = buffer_ops.buffer_load( + logits_rsrc, + expert_base, + vec_width=4, + dtype=f32, + ) + bias_i32 = buffer_ops.buffer_load( + bias_rsrc, + expert_base // arith.constant(2, type=i32), + vec_width=2, + dtype=i32, + ) + bias_vec = vector.bitcast(vec4_bf16, bias_i32) + priority_base = expert_priority_index(expert_base) + sigmoid_values = [] + biased_values = [] + expert_ids = [] + for lane_in_vec in range_constexpr(4): + expert = expert_base + arith.constant( + lane_in_vec, + type=i32, + ) + x = vector.extract( + logits_vec, + static_position=[lane_in_vec], + dynamic_position=[], + ) + bias_bf16 = vector.extract( + bias_vec, + static_position=[lane_in_vec], + dynamic_position=[], + ) + bias_f32 = arith.extf(f32, bias_bf16) + exp_value = llvm.call_intrinsic( + f32, + "llvm.amdgcn.exp2.f32", + [ArithValue(x) * arith.constant(-_LOG2E, type=f32)], + [], + [], + ) + sigmoid = llvm.call_intrinsic( + f32, + "llvm.amdgcn.rcp.f32", + [c_one_f32 + exp_value], + [], + [], + ) + sigmoid_values.append(sigmoid) + biased_values.append(ArithValue(sigmoid) + bias_f32) + expert_ids.append(expert) + fx.ptr_store( + vector.from_elements(vec4_f32, sigmoid_values), + route_lds + fx.Int64(expert_base), + ) + fx.ptr_store( + vector.from_elements(vec4_f32, biased_values), + biased_lds + fx.Int64(priority_base), + ) + fx.ptr_store( + vector.from_elements(vec4_i32, expert_ids), + priority_expert_ids_lds + fx.Int64(priority_base), + ) + scf.YieldOp([]) + gpu.barrier() + + route_active = arith.cmpi( + CmpIPredicate.ult, + tid, + arith.constant(_ROUTE_THREADS, type=i32), + ) + route_if = scf.IfOp(route_active) + with ir.InsertionPoint(route_if.then_block): + local_scores = [ + ArithValue(c_neg_inf) for _ in range(_LOCAL_ROUTE_CANDIDATES) + ] + local_ids = [ArithValue(c_zero_i32) for _ in range(_LOCAL_ROUTE_CANDIDATES)] + + def insert_candidate(score, expert, position): + local_scores[position] = ArithValue(score) + local_ids[position] = ArithValue(expert) + for sort_offset in range_constexpr(position): + right_position = position - sort_offset + left_position = right_position - 1 + right_score = local_scores[right_position] + left_score = local_scores[left_position] + right_id = local_ids[right_position] + left_id = local_ids[left_position] + swap = arith.cmpf( + CmpFPredicate.OGT, + right_score, + left_score, + ) + local_scores[left_position] = ArithValue( + arith.select( + swap, + _raw(right_score), + _raw(left_score), + ) + ) + local_scores[right_position] = ArithValue( + arith.select( + swap, + _raw(left_score), + _raw(right_score), + ) + ) + local_ids[left_position] = ArithValue( + arith.select(swap, _raw(right_id), _raw(left_id)) + ) + local_ids[right_position] = ArithValue( + arith.select(swap, _raw(left_id), _raw(right_id)) + ) + + local_position = 0 + priority_lane_base = lane_priority_rank(tid) * arith.constant( + _LOCAL_ROUTE_CANDIDATES, type=i32 + ) + for vec_base in range_constexpr(0, 12, 4): + priority_base = priority_lane_base + arith.constant( + vec_base, + type=i32, + ) + score_vector = fx.ptr_load( + biased_lds + fx.Int64(priority_base), + result_type=vec4_f32, + ) + expert_vector = fx.ptr_load( + priority_expert_ids_lds + fx.Int64(priority_base), + result_type=vec4_i32, + ) + for lane_in_vec in range_constexpr(4): + expert = vector.extract( + expert_vector, + static_position=[lane_in_vec], + dynamic_position=[], + ) + score = vector.extract( + score_vector, + static_position=[lane_in_vec], + dynamic_position=[], + ) + insert_candidate(score, expert, local_position) + local_position += 1 + + priority_base = priority_lane_base + arith.constant(12, type=i32) + score_vector = fx.ptr_load( + biased_lds + fx.Int64(priority_base), + result_type=vec2_f32, + ) + expert_vector = fx.ptr_load( + priority_expert_ids_lds + fx.Int64(priority_base), + result_type=vec2_i32, + ) + for lane_in_vec in range_constexpr(2): + expert = vector.extract( + expert_vector, + static_position=[lane_in_vec], + dynamic_position=[], + ) + score = vector.extract( + score_vector, + static_position=[lane_in_vec], + dynamic_position=[], + ) + insert_candidate(score, expert, local_position) + local_position += 1 + + local_route_base = tid * arith.constant( + _LOCAL_ROUTE_CANDIDATES, + type=i32, + ) + for position in range_constexpr(_LOCAL_ROUTE_CANDIDATES): + local_route_offset = local_route_base + arith.constant( + position, + type=i32, + ) + _lds_store( + local_route_scores_lds, + local_scores[position], + local_route_offset, + ) + _lds_store( + local_route_ids_lds, + local_ids[position], + local_route_offset, + ) + + route_sum = ArithValue(c_zero_f32) + local_rank = ArithValue(c_zero_i32) + for k in range_constexpr(_TOPK): + local_route_offset = local_route_base + local_rank + local_max = ArithValue( + _lds_load(local_route_scores_lds, local_route_offset) + ) + local_id = ArithValue( + _lds_load(local_route_ids_lds, local_route_offset) + ) + lane_candidate_id = local_id + + for dpp_control in (0xB1, 0x4E, 0x141, 0x140, 0x142, 0x143): + local_max_i32 = arith.bitcast(i32, _raw(local_max)) + remote_max_i32 = rocdl.update_dpp( + i32, + c_zero_i32, + local_max_i32, + dpp_control, + 0xF, + 0xF, + True, + ) + remote_max = ArithValue(arith.bitcast(f32, remote_max_i32)) + remote_id = ArithValue( + rocdl.update_dpp( + i32, + c_zero_i32, + _raw(local_id), + dpp_control, + 0xF, + 0xF, + True, + ) + ) + take_remote = arith.cmpf( + CmpFPredicate.OGT, + remote_max, + local_max, + ) + local_max = ArithValue( + arith.select( + take_remote, + _raw(remote_max), + _raw(local_max), + ) + ) + local_id = ArithValue( + arith.select( + take_remote, + _raw(remote_id), + _raw(local_id), + ) + ) + + selected_id = ArithValue( + rocdl.readlane( + i32, + _raw(local_id), + arith.constant(_ROUTE_THREADS - 1, type=i32), + ) + ) + + selected_score = ArithValue(_lds_load(route_lds, selected_id)) + is_local_winner = arith.cmpi( + CmpIPredicate.eq, + lane_candidate_id, + selected_id, + ) + local_rank = local_rank + ArithValue( + arith.select( + is_local_winner, + c_one_i32, + c_zero_i32, + ) + ) + route_sum = route_sum + selected_score + _lds_store( + selected_ids_lds, + selected_id, + arith.constant(k, type=i32), + ) + _lds_store( + selected_scores_lds, + selected_score, + arith.constant(k, type=i32), + ) + + is_route_writer = arith.cmpi( + CmpIPredicate.eq, + tid, + c_zero_i32, + ) + route_writer_if = scf.IfOp(is_route_writer) + with ir.InsertionPoint(route_writer_if.then_block): + _lds_store(route_sum_lds, route_sum, c_zero_i32) + scf.YieldOp([]) + + scf.YieldOp([]) + + # Waves 1-3 prepare the one B1 activation row for A8W4 stage 1 while + # wave 0 performs the exact 16-step route selection above. Two adjacent + # lanes own one 32-value MX group. The 192 quant lanes cover 96 groups + # in the first iteration and the remaining 16 groups in the second. + # This preserves the accepted RoundUp E8M0 and OCP FP8 E4M3 bytes while + # removing the standalone activation-quant launch. + quant_active = arith.cmpi( + CmpIPredicate.uge, + tid, + arith.constant(_ROUTE_THREADS, type=i32), + ) + quant_if = scf.IfOp(quant_active) + with ir.InsertionPoint(quant_if.then_block): + quant_tid = tid - arith.constant(_ROUTE_THREADS, type=i32) + lane_in_group = quant_tid % arith.constant(2, type=i32) + group_in_iteration = quant_tid // arith.constant(2, type=i32) + c_scale_exp = arith.constant(254, type=i32) + c_exp_shift = arith.constant(23, type=i32) + c_amax_floor = arith.constant(1.0e-10, type=f32) + for group_iteration in range_constexpr( + (_MX_SCALE_COLS + _MX_GROUPS_PER_ITERATION - 1) + // _MX_GROUPS_PER_ITERATION + ): + group = group_in_iteration + arith.constant( + group_iteration * _MX_GROUPS_PER_ITERATION, + type=i32, + ) + group_in_range = arith.cmpi( + CmpIPredicate.ult, + group, + arith.constant(_MX_SCALE_COLS, type=i32), + ) + group_if = scf.IfOp(group_in_range) + with ir.InsertionPoint(group_if.then_block): + element_base = group * arith.constant( + _MX_GROUP_SIZE, + type=i32, + ) + lane_in_group * arith.constant( + _MX_VALUES_PER_THREAD, + type=i32, + ) + values = [] + local_amax = c_amax_floor + for element_offset in range_constexpr(_MX_VALUES_PER_THREAD): + value_bf16 = buffer_ops.buffer_load( + hidden_rsrc, + element_base + arith.constant(element_offset, type=i32), + vec_width=1, + dtype=T.bf16, + ) + value = arith.extf(f32, value_bf16) + values.append(value) + local_amax = arith.maximumf( + local_amax, + fmath.absf(value), + ) + + peer_amax = ArithValue(local_amax).shuffle_xor( + arith.constant(1, type=i32), + arith.constant(_ROUTE_THREADS, type=i32), + ) + group_amax = arith.maximumf(local_amax, _raw(peer_amax)) + e8m0 = emit_mx_e8m0_scale( + group_amax, + mode=MxScaleRoundModeInt.RoundUp, + dtype=MxDtypeInt.FP8_E4M3, + ) + quant_scale = ((c_scale_exp - e8m0) << c_exp_shift).bitcast(f32) + + for pack_index in range_constexpr(_MX_VALUES_PER_THREAD // 4): + value_base = pack_index * 4 + packed = arith.constant(0, type=i32) + packed = rocdl.cvt_pk_fp8_f32( + i32, + arith.mulf( + values[value_base], + quant_scale, + ), + arith.mulf( + values[value_base + 1], + quant_scale, + ), + packed, + 0, + ) + packed = rocdl.cvt_pk_fp8_f32( + i32, + arith.mulf( + values[value_base + 2], + quant_scale, + ), + arith.mulf( + values[value_base + 3], + quant_scale, + ), + packed, + 1, + ) + output_byte = element_base + arith.constant( + pack_index * 4, + type=i32, + ) + buffer_ops.buffer_store( + packed, + quantized_hidden_rsrc, + output_byte, + offset_is_bytes=True, + ) + + is_scale_writer = arith.cmpi( + CmpIPredicate.eq, + lane_in_group, + c_zero_i32, + ) + scale_if = scf.IfOp(is_scale_writer) + with ir.InsertionPoint(scale_if.then_block): + scale_tile = group // arith.constant(8, type=i32) + scale_lane4 = group % arith.constant(4, type=i32) + scale_half8 = ( + group % arith.constant(8, type=i32) + ) // arith.constant(4, type=i32) + within_rank = ( + scale_tile * arith.constant(256, type=i32) + + scale_lane4 * arith.constant(64, type=i32) + + scale_half8 * arith.constant(2, type=i32) + ) + e8m0_i8 = arith.trunci(T.i8, _raw(e8m0)) + for rank in range_constexpr(_TOPK): + scale_offset = within_rank + arith.constant( + rank * _MX_SCALE_COLS * _BLOCK_M, + type=i32, + ) + buffer_ops.buffer_store( + e8m0_i8, + quantized_scale_rsrc, + scale_offset, + offset_is_bytes=True, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + gpu.barrier() + + # Initialize every padded metadata row and zero the atomic-output + # buffer while the selected routes remain resident. + sentinel = arith.constant((_TOPK << 24) | 1, type=i32) + for row_base in range_constexpr(0, _SORTED_ROWS, _BLOCK_THREADS): + row = tid + arith.constant(row_base, type=i32) + buffer_ops.buffer_store(sentinel, sorted_ids_rsrc, row) + buffer_ops.buffer_store(c_zero_f32, sorted_weights_rsrc, row) + + moe_count = ArithValue(moe_buf_i32_elements) + loop_lower = arith.index_cast(T.index, tid) + loop_upper = arith.index_cast(T.index, moe_count) + loop_step = arith.index(_BLOCK_THREADS) + zero_loop = scf.ForOp(loop_lower, loop_upper, loop_step) + with ir.InsertionPoint(zero_loop.body): + zero_idx = arith.index_cast(i32, zero_loop.induction_variable) + buffer_ops.buffer_store(c_zero_i32, moe_buf_rsrc, zero_idx) + scf.YieldOp([]) + gpu.barrier() + + active = arith.cmpi( + CmpIPredicate.ult, + tid, + arith.constant(_TOPK, type=i32), + ) + active_if = scf.IfOp(active) + with ir.InsertionPoint(active_if.then_block): + route_id = _lds_load(selected_ids_lds, tid) + route_score = _lds_load(selected_scores_lds, tid) + route_sum = _lds_load(route_sum_lds, c_zero_i32) + normalized = arith.divf(_raw(route_score), _raw(route_sum)) + buffer_ops.buffer_store(route_id, topk_ids_rsrc, tid) + buffer_ops.buffer_store(normalized, topk_weights_rsrc, tid) + + rank = ArithValue(c_zero_i32) + for other_slot in range_constexpr(_TOPK): + other_id = _lds_load( + selected_ids_lds, + arith.constant(other_slot, type=i32), + ) + is_before = arith.cmpi(CmpIPredicate.slt, other_id, route_id) + rank = rank + ArithValue(arith.select(is_before, c_one_i32, c_zero_i32)) + + sorted_base = rank * arith.constant(_BLOCK_M, type=i32) + packed_route = tid << arith.constant(24, type=i32) + buffer_ops.buffer_store(route_id, sorted_experts_rsrc, rank) + buffer_ops.buffer_store(packed_route, sorted_ids_rsrc, sorted_base) + buffer_ops.buffer_store(normalized, sorted_weights_rsrc, sorted_base) + scf.YieldOp([]) + + is_first = arith.cmpi(CmpIPredicate.eq, tid, c_zero_i32) + first_if = scf.IfOp(is_first) + with ir.InsertionPoint(first_if.then_block): + buffer_ops.buffer_store( + arith.constant(_SORTED_ROWS, type=i32), + nvalid_rsrc, + c_zero_i32, + ) + buffer_ops.buffer_store(c_one_i32, nvalid_rsrc, c_one_i32) + scf.YieldOp([]) + + @flyc.jit + def launch_route_sort( + logits: fx.Pointer, + correction_bias: fx.Pointer, + topk_weights: fx.Pointer, + topk_ids: fx.Pointer, + sorted_ids: fx.Pointer, + sorted_weights: fx.Pointer, + sorted_expert_ids: fx.Pointer, + num_valid_ids: fx.Pointer, + moe_buf: fx.Pointer, + moe_buf_i32_elements: fx.Int32, + hidden_states: fx.Pointer, + quantized_hidden_states: fx.Pointer, + quantized_scales: fx.Pointer, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + pass + route_sort_kernel( + logits, + correction_bias, + topk_weights, + topk_ids, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + moe_buf_i32_elements, + hidden_states, + quantized_hidden_states, + quantized_scales, + ).launch( + grid=(arith.index(1), 1, 1), + block=(_BLOCK_THREADS, 1, 1), + stream=stream, + ) + + launch_route_sort.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_route_sort diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 896ec48423..44fbb8a997 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -248,12 +248,12 @@ def load_bias_scalar(bias_rsrc, offset): raise ValueError("mock_gate_only requires k_batch > 1 (split-K)") if is_splitk: k_per_batch = model_dim // k_batch - assert ( - model_dim % k_batch == 0 - ), f"model_dim={model_dim} not divisible by k_batch={k_batch}" - assert ( - k_per_batch % tile_k == 0 - ), f"K_per_batch={k_per_batch} not divisible by tile_k={tile_k}" + assert model_dim % k_batch == 0, ( + f"model_dim={model_dim} not divisible by k_batch={k_batch}" + ) + assert k_per_batch % tile_k == 0, ( + f"K_per_batch={k_per_batch} not divisible by tile_k={tile_k}" + ) out_dtype = "bf16" else: @@ -442,7 +442,7 @@ def x_lds_elem(): pipe_phases.append(phase) bi = 0 - for _p in range(1, pipe_n_phases): + for _p in range(pipe_n_phases): rem_b = len(pipe_b_loads) - bi rem_p = pipe_n_phases - _p n_b = (rem_b + rem_p - 1) // rem_p if rem_p > 0 else 0 @@ -1579,9 +1579,10 @@ def interleaved_half( lds_read (already DMA'd in previous half). Interleaving schedule (per half): - Phase 0: scale VMEM + 2 ds_read(A) -> 4 MFMA(prev) - Phase 1..N: B VMEM(distributed) + 2 ds_read(A, if avail) -> 4 MFMA(prev) - Phase N+1..: remaining B VMEM -> 4 MFMA(prev) + Phase 0: scale VMEM + early B VMEM + ds_read(A) + -> MFMA(prev) + Phase 1..N: remaining B VMEM + ds_read(A, if available) + -> MFMA(prev) """ abs_k = k_base_idx + arith.constant(next_k_load, index=True) bk = abs_k // arith.constant(b_byte_div, index=True) @@ -2161,13 +2162,21 @@ def act_elem(g, u): u = _clamp_lin(u) return silu_elem(g) * u - kwave_fused = const_expr( + kwave_separated_fused = const_expr( k_wave > 1 and not enable_bias and not is_splitk and not gate_up_interleave and need_quant ) + kwave_gui_fused = const_expr( + k_wave > 1 + and not enable_bias + and not is_splitk + and gate_up_interleave + and need_quant + ) + kwave_fused = const_expr(kwave_separated_fused or kwave_gui_fused) if const_expr(k_wave > 1 and not kwave_fused): has_up = const_expr(acc_up is not None) @@ -2287,7 +2296,9 @@ def act_elem(g, u): ) acc_up[aidx] = arith.addf(acc_up[aidx], bsplat) - if const_expr(gate_up_interleave and not is_splitk): + if const_expr( + gate_up_interleave and not is_splitk and not kwave_gui_fused + ): gui_out_n = num_acc_n // pack_N acc = [None] * (gui_out_n * m_repeat) for mi in range_constexpr(m_repeat): @@ -2681,7 +2692,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): else (ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get()) ) - if const_expr(kwave_fused): + if const_expr(kwave_separated_fused): slab_n = tile_m * tile_n slab_ty = _mT.memref( k_wave * slab_n, f32, memory_space=_lds_space() @@ -2810,6 +2821,170 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): with ir.InsertionPoint(ifr.then_block): fused_read() scf.YieldOp([]) + elif const_expr(kwave_gui_fused): + # Interleaved gate/up uses 16-column chunks: + # [gate0:16, up0:16, gate1:16, up1:16, ...]. Keep every + # K-wave partial in one LDS slab, then combine K-wave and + # gate/up ownership in the quantizing reader. This avoids + # the generic K-wave reduction slab followed by a second + # CShuffle write/read phase. + slab_n = tile_m * tile_n + slab_ty = _mT.memref( + k_wave * slab_n, + f32, + memory_space=_lds_space(), + ) + gui_slab = memref.view( + slab_ty, + base_ptr_pong, + arith.constant(lds_pong_offset, index=True), + sizes=[], + ) + c_tn = arith.constant(tile_n, index=True) + c_slabn = arith.constant(slab_n, index=True) + kg_base = wave_k_id * c_slabn + vec1_f32 = T.vec(1, f32) + vecev_f32 = T.vec(e_vec, f32) + + gpu.barrier() + + def gui_fused_write(mi, ii, row_in_tile, row): + rb = row_in_tile * c_tn + for ni in range_constexpr(num_acc_n): + col = ( + n_tile_base + + lane_mod_16 + + arith.constant(ni * 16, index=True) + ) + aidx = mi * num_acc_n + ni + value = vector.extract( + acc_gate[aidx], + static_position=[ii], + dynamic_position=[], + ) + vector.store( + vector.from_elements(vec1_f32, [value]), + gui_slab, + [kg_base + rb + col], + alignment=4, + ) + + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=gui_fused_write, + ) + gpu.barrier() + + gui_tile_n = tile_n // 2 + gui_cshuffle_nlane = min(32, gui_tile_n // e_vec) + gui_by_n = by_n // arith.constant(2, index=True) + cn = int(gui_cshuffle_nlane) + gui_cshuffle_threads = min( + int(total_threads), + int(tile_m) * cn, + ) + cm = gui_cshuffle_threads // cn + mreps = int(tile_m) // cm + nreps = int(gui_tile_n) // (cn * int(e_vec)) + c_cn = arith.constant(cn, index=True) + c_ev = arith.constant(e_vec, index=True) + c16_idx = arith.constant(16, index=True) + c32_idx = arith.constant(32, index=True) + gui_reader_active = arith.cmpi( + CmpIPredicate.ult, + tx, + arith.constant(gui_cshuffle_threads, index=True), + ) + m_lane = tx / c_cn + n_lane = tx % c_cn + for mr in range_constexpr(mreps): + row_local = arith.constant(mr * cm, index=True) + m_lane + row = bx_m + row_local + row_context, row_valid = precompute_row( + row_local=row_local, + row=row, + ) + row_valid = arith.andi(row_valid, gui_reader_active) + + def gui_fused_read( + row_local=row_local, + row=row, + row_context=row_context, + ): + row_base = row_local * c_tn + for nr in range_constexpr(nreps): + output_col = ( + arith.constant( + nr * (cn * int(e_vec)), + index=True, + ) + + n_lane * c_ev + ) + chunk = output_col / c16_idx + within_chunk = output_col % c16_idx + gate_col = chunk * c32_idx + within_chunk + gate_sum = None + up_sum = None + for kg in range_constexpr(k_wave): + wave_base = ( + arith.constant(kg, index=True) * c_slabn + + row_base + ) + gate_values = vector.load_op( + vecev_f32, + gui_slab, + [wave_base + gate_col], + ) + up_values = vector.load_op( + vecev_f32, + gui_slab, + [wave_base + gate_col + c16_idx], + ) + if kg == 0: + gate_sum = gate_values + up_sum = up_values + else: + gate_sum = arith.addf( + gate_sum, + gate_values, + ) + up_sum = arith.addf( + up_sum, + up_values, + ) + activated = [] + for element in range_constexpr(int(e_vec)): + gate_value = vector.extract( + gate_sum, + static_position=[element], + dynamic_position=[], + ) + up_value = vector.extract( + up_sum, + static_position=[element], + dynamic_position=[], + ) + activated.append(act_elem(gate_value, up_value)) + store_pair( + row_local=row_local, + row=row, + row_ctx=row_context, + col_pair0=output_col, + col_g0=gui_by_n + output_col, + frag=vector.from_elements( + vecev_f32, + activated, + ), + ) + + row_if = scf.IfOp(row_valid) + with ir.InsertionPoint(row_if.then_block): + gui_fused_read() + scf.YieldOp([]) elif const_expr(gate_up_interleave and not is_splitk): gui_eff_n = gui_out_n gui_tile_n = tile_n // 2 @@ -3103,7 +3278,8 @@ def compile_mixed_moe_gemm2( xcd_swizzle: int = 0, ): """Compile stage2 kernel (moe_gemm2): A2 @ W2.T -> [tokens, model_dim], atomic-add.""" - del b_nt + if not isinstance(b_nt, int) or b_nt not in (0, 1, 2, 3): + raise ValueError(f"b_nt must be one of 0, 1, 2, 3, got {b_nt!r}") _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m if const_expr(_sort_block_m != tile_m and _sort_block_m % tile_m != 0): raise ValueError( @@ -3803,6 +3979,7 @@ def load_cell(k0): vec_elems=vec_elems, elem_bytes=b_elem_bytes, offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, ) b_i64x2 = vector.bitcast(vec2_i64, b16) return ( @@ -4179,14 +4356,14 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): if const_expr(xdl_arb_hint): rocdl.disable_xdl_arb_stall() - if const_expr(b_hi_loader is not None): - b_hi = b_hi_loader() - for bhi_i in range_constexpr(len(b_hi)): - b_tile_full[b_split_ku + bhi_i] = b_hi[bhi_i] - rocdl.s_setprio(1) for k_idx in range_constexpr(ku_loop): + if const_expr(b_hi_loader is not None and k_idx == b_split_ku): + b_hi = b_hi_loader() + for bhi_i in range_constexpr(len(b_hi)): + b_tile_full[b_split_ku + bhi_i] = b_hi[bhi_i] + ku128 = k_idx >> pack_K_shift ikxdl = k_idx & pack_K_mask @@ -5125,12 +5302,12 @@ def x_lds_elem(): else: if _is_splitk: _k_per_batch = model_dim // k_batch - assert ( - model_dim % k_batch == 0 - ), f"model_dim={model_dim} not divisible by k_batch={k_batch}" - assert ( - _k_per_batch % tile_k == 0 - ), f"K_per_batch={_k_per_batch} not divisible by tile_k={tile_k}" + assert model_dim % k_batch == 0, ( + f"model_dim={model_dim} not divisible by k_batch={k_batch}" + ) + assert _k_per_batch % tile_k == 0, ( + f"K_per_batch={_k_per_batch} not divisible by tile_k={tile_k}" + ) out_dtype = "bf16" else: _k_per_batch = model_dim diff --git a/aiter/ops/flydsl/kimi_k3_moe_route_parallel.py b/aiter/ops/flydsl/kimi_k3_moe_route_parallel.py new file mode 100644 index 0000000000..9fbf3ea44c --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_moe_route_parallel.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Narrow gfx950 wrapper for overlapped Kimi-K3 B1 route and MXFP8 prep.""" + +import functools + +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +KimiK3RouteSortResult = tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +] + + +@functools.cache +def _compiled_kimi_k3_b1_route_sort_parallel(): + from aiter.ops.flydsl.kernels.kimi_k3_b1_route_sort_parallel import ( + build_kimi_k3_b1_route_sort_parallel_module, + ) + + return build_kimi_k3_b1_route_sort_parallel_module() + + +def _supports_route_contract( + logits: torch.Tensor, + correction_bias: torch.Tensor, +) -> bool: + return ( + logits.is_cuda + and correction_bias.is_cuda + and logits.device == correction_bias.device + and logits.dtype == torch.float32 + and correction_bias.dtype == torch.bfloat16 + and logits.is_contiguous() + and correction_bias.is_contiguous() + and tuple(logits.shape) == (1, 896) + and tuple(correction_bias.shape) == (896,) + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +def _supports_hidden_contract( + hidden_states: torch.Tensor, + logits: torch.Tensor, + *, + model_dim: int, +) -> bool: + return ( + hidden_states.is_cuda + and hidden_states.device == logits.device + and hidden_states.dtype == torch.bfloat16 + and hidden_states.is_contiguous() + and tuple(hidden_states.shape) == (1, 3584) + and model_dim == 3584 + ) + + +def supports_kimi_k3_b1_route_sort_parallel( + hidden_states: torch.Tensor, + logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + model_dim: int, +) -> bool: + """Return whether the fixed B1 route/quant specialization is safe.""" + + return _supports_route_contract( + logits, correction_bias + ) and _supports_hidden_contract( + hidden_states, + logits, + model_dim=model_dim, + ) + + +def kimi_k3_b1_route_sort_parallel( + hidden_states: torch.Tensor, + logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + model_dim: int, +) -> KimiK3RouteSortResult: + """Route one token and emit exact prequantized A8W4 stage-1 inputs.""" + + if not _supports_route_contract(logits, correction_bias): + raise NotImplementedError( + "parallel route-sort only supports contiguous gfx950 FP32/BF16 " + "B1x896, topk=16, group=1/1, block_size_m=32" + ) + if not _supports_hidden_contract( + hidden_states, + logits, + model_dim=model_dim, + ): + raise ValueError( + "hidden_states must be contiguous gfx950 BF16 [1, 3584] and " + f"model_dim must be 3584; got {hidden_states.shape=}, " + f"{hidden_states.dtype=}, {model_dim=}" + ) + + from aiter import dtypes + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + device = logits.device + topk_weights = torch.empty((1, 16), dtype=torch.float32, device=device) + topk_ids = torch.empty((1, 16), dtype=torch.int32, device=device) + sorted_ids = torch.empty(16 * 32, dtype=torch.int32, device=device) + sorted_weights = torch.empty(16 * 32, dtype=torch.float32, device=device) + sorted_expert_ids = torch.empty(16, dtype=torch.int32, device=device) + num_valid_ids = torch.empty(2, dtype=torch.int32, device=device) + moe_buf = torch.empty((1, model_dim), dtype=torch.bfloat16, device=device) + quantized_hidden = torch.empty( + (1, model_dim), + dtype=dtypes.fp8, + device=device, + ) + quantized_scales = torch.empty( + (16 * 32, model_dim // 32), + dtype=dtypes.fp8_e8m0, + device=device, + ) + + _compiled_kimi_k3_b1_route_sort_parallel()( + ptr_arg(logits), + ptr_arg(correction_bias), + ptr_arg(topk_weights), + ptr_arg(topk_ids), + ptr_arg(sorted_ids), + ptr_arg(sorted_weights), + ptr_arg(sorted_expert_ids), + ptr_arg(num_valid_ids), + ptr_arg(moe_buf), + moe_buf.numel() // 2, + ptr_arg(hidden_states), + ptr_arg(quantized_hidden), + ptr_arg(quantized_scales), + stream=torch.cuda.current_stream(logits.device), + ) + return ( + topk_weights, + topk_ids, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + quantized_hidden, + quantized_scales, + ) diff --git a/aiter/ops/flydsl/kimi_k3_persistent_moe.py b/aiter/ops/flydsl/kimi_k3_persistent_moe.py new file mode 100644 index 0000000000..7772d5d5ef --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_persistent_moe.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Typed Kimi-K3 B1 route-to-expert specialization for gfx950.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +import torch + +from aiter import dtypes +from aiter.ops.flydsl.kimi_k3_moe_route_parallel import ( + kimi_k3_b1_route_sort_parallel, + supports_kimi_k3_b1_route_sort_parallel, +) +from aiter.ops.flydsl.moe_kernels import ( + flydsl_moe_stage1, + flydsl_moe_stage2, + get_flydsl_kernel_params, +) + +_EXPERTS = 896 +_TOPK = 16 +_MODEL_DIM = 3584 +_INTERMEDIATE_DIM = 384 +_SORT_BLOCK_M = 32 +_STAGE1_KERNEL = "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_gui_kw7_fp8" +_STAGE1_EXPECTED = { + "stage": 1, + "a_dtype": "fp8", + "b_dtype": "fp4", + "out_dtype": "fp8", + "tile_m": 32, + "tile_n": 64, + "tile_k": 256, + "waves_per_eu": 1, + "b_nt": 2, + "gate_mode": "interleave", + "k_wave": 7, +} +_W1_SCALE_SHAPE = ( + _EXPERTS * 2 * _INTERMEDIATE_DIM, + _MODEL_DIM // 32, +) +_W2_SCALE_SHAPE = ( + ((_EXPERTS * _MODEL_DIM + 255) // 256) * 256, + ((_INTERMEDIATE_DIM // 32 + 7) // 8) * 8, +) + + +@dataclass(frozen=True) +class KimiK3PersistentMoERequest: + """Complete immutable contract for one specialized MoE invocation.""" + + hidden_states: torch.Tensor + router_logits: torch.Tensor + correction_bias: torch.Tensor | None + w1: torch.Tensor | None + w2: torch.Tensor | None + w1_scale: torch.Tensor | None + w2_scale: torch.Tensor | None + situ_beta: float + situ_linear_beta: float + w13_layout: str | None + weights_shuffled: bool + quantization_supported: bool + activation: str + num_experts: int + topk: int + num_expert_group: int + topk_group: int + renormalize: bool + scoring_func: str + routed_scaling_factor: float + expert_parallel: bool + eplb_enabled: bool + lora_enabled: bool + has_expert_bias: bool + apply_router_weight_on_input: bool + expert_map_active: bool + routing_capture_enabled: bool + custom_routing_active: bool + input_ids_active: bool + routing_method: str + + +@dataclass(frozen=True) +class KimiK3PersistentMoEMetadata: + """Route metadata and prequantized stage-1 activation owned together.""" + + routing_weights: torch.Tensor + expert_ids: torch.Tensor + sorted_token_ids: torch.Tensor + sorted_weights: torch.Tensor + sorted_expert_ids: torch.Tensor + num_valid_ids: torch.Tensor + moe_buf: torch.Tensor + quantized_hidden_states: torch.Tensor + quantized_scales: torch.Tensor + + +def _is_packed_tensor( + tensor: torch.Tensor | None, + *, + device: torch.device, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> bool: + return ( + isinstance(tensor, torch.Tensor) + and tensor.is_cuda + and tensor.device == device + and tensor.dtype == dtype + and tensor.is_contiguous() + and tuple(tensor.shape) == shape + ) + + +def _has_validated_stage1_kernel() -> bool: + parameters = get_flydsl_kernel_params(_STAGE1_KERNEL) + return parameters is not None and all( + parameters.get(name) == expected for name, expected in _STAGE1_EXPECTED.items() + ) + + +def _has_native_weight_layout( + request: KimiK3PersistentMoERequest, + device: torch.device, +) -> bool: + return ( + _is_packed_tensor( + request.w1, + device=device, + dtype=dtypes.fp4x2, + shape=(_EXPERTS, 2 * _INTERMEDIATE_DIM, _MODEL_DIM // 2), + ) + and _is_packed_tensor( + request.w2, + device=device, + dtype=dtypes.fp4x2, + shape=(_EXPERTS, _MODEL_DIM, _INTERMEDIATE_DIM // 2), + ) + and _is_packed_tensor( + request.w1_scale, + device=device, + dtype=dtypes.fp8_e8m0, + shape=_W1_SCALE_SHAPE, + ) + and _is_packed_tensor( + request.w2_scale, + device=device, + dtype=dtypes.fp8_e8m0, + shape=_W2_SCALE_SHAPE, + ) + ) + + +def _has_supported_quantized_activation( + request: KimiK3PersistentMoERequest, +) -> bool: + return ( + request.activation == "situ" + and request.situ_beta == 4.0 + and request.situ_linear_beta == 25.0 + and request.quantization_supported + and request.weights_shuffled + and request.w13_layout == "gate_up_interleaved_preshuffled" + ) + + +def _has_supported_routing(request: KimiK3PersistentMoERequest) -> bool: + return ( + request.num_experts == _EXPERTS + and request.topk == _TOPK + and request.num_expert_group == 1 + and request.topk_group == 1 + and request.renormalize + and request.scoring_func == "sigmoid" + and request.routed_scaling_factor == 1.0 + and request.routing_method == "DeepSeekV3" + ) + + +def _has_unsupported_runtime_features( + request: KimiK3PersistentMoERequest, +) -> bool: + return any( + ( + request.expert_parallel, + request.eplb_enabled, + request.lora_enabled, + request.has_expert_bias, + request.apply_router_weight_on_input, + request.expert_map_active, + request.routing_capture_enabled, + request.custom_routing_active, + request.input_ids_active, + ) + ) + + +def supports_kimi_k3_b1_persistent_moe( + request: KimiK3PersistentMoERequest, +) -> bool: + """Return whether the complete fixed-shape specialization is safe.""" + + hidden_states = request.hidden_states + correction_bias = request.correction_bias + if not isinstance(correction_bias, torch.Tensor): + return False + device = hidden_states.device + return ( + os.environ.get("AITER_DISABLE", "0") != "1" + and os.environ.get("AITER_SITUV2_A8W4", "0") == "1" + and supports_kimi_k3_b1_route_sort_parallel( + hidden_states, + request.router_logits, + correction_bias, + model_dim=_MODEL_DIM, + ) + and _has_validated_stage1_kernel() + and _has_native_weight_layout(request, device) + and _has_supported_quantized_activation(request) + and _has_supported_routing(request) + and not _has_unsupported_runtime_features(request) + ) + + +def prepare_kimi_k3_b1_persistent_moe( + request: KimiK3PersistentMoERequest, +) -> KimiK3PersistentMoEMetadata | None: + """Prepare route and stage-1 activation, or select the generic fallback.""" + + if not supports_kimi_k3_b1_persistent_moe(request): + return None + assert request.correction_bias is not None + return KimiK3PersistentMoEMetadata( + *kimi_k3_b1_route_sort_parallel( + request.hidden_states, + request.router_logits, + request.correction_bias, + model_dim=_MODEL_DIM, + ) + ) + + +def consume_kimi_k3_b1_persistent_moe( + request: KimiK3PersistentMoERequest, + metadata: KimiK3PersistentMoEMetadata, +) -> torch.Tensor: + """Consume prepared metadata exactly once through both expert GEMMs.""" + + if not supports_kimi_k3_b1_persistent_moe(request): + raise NotImplementedError("unsupported Kimi-K3 persistent-MoE contract") + assert request.w1 is not None + assert request.w2 is not None + assert request.w1_scale is not None + assert request.w2_scale is not None + + intermediate, intermediate_scale = flydsl_moe_stage1( + a=metadata.quantized_hidden_states, + w1=request.w1, + sorted_token_ids=metadata.sorted_token_ids, + sorted_expert_ids=metadata.sorted_expert_ids, + num_valid_ids=metadata.num_valid_ids, + out=None, + topk=_TOPK, + tile_m=32, + tile_n=64, + tile_k=256, + a_dtype="fp8", + b_dtype="fp4", + out_dtype="fp8", + act="situv2", + situ_beta=request.situ_beta, + situ_linear_beta=request.situ_linear_beta, + w1_scale=request.w1_scale, + a1_scale=metadata.quantized_scales, + sorted_weights=None, + use_async_copy=True, + k_batch=1, + waves_per_eu=1, + b_nt=2, + gate_mode="interleave", + model_dim_pad=0, + inter_dim_pad=0, + xcd_swizzle=0, + k_wave=7, + ) + return flydsl_moe_stage2( + inter_states=intermediate.view( + 1, + _TOPK, + _INTERMEDIATE_DIM, + ), + w2=request.w2, + sorted_token_ids=metadata.sorted_token_ids, + sorted_expert_ids=metadata.sorted_expert_ids, + num_valid_ids=metadata.num_valid_ids, + out=metadata.moe_buf, + topk=_TOPK, + tile_m=32, + tile_n=128, + tile_k=128, + a_dtype="fp8", + b_dtype="fp4", + out_dtype="bf16", + mode="atomic", + w2_scale=request.w2_scale, + a2_scale=intermediate_scale, + sorted_weights=metadata.sorted_weights, + sort_block_m=_SORT_BLOCK_M, + persist=False, + waves_per_eu=1, + use_async_copy=False, + cu_num_mul=1, + b_nt=2, + model_dim_pad=0, + inter_dim_pad=0, + xcd_swizzle=0, + ) + + +__all__ = [ + "KimiK3PersistentMoEMetadata", + "KimiK3PersistentMoERequest", + "consume_kimi_k3_b1_persistent_moe", + "prepare_kimi_k3_b1_persistent_moe", + "supports_kimi_k3_b1_persistent_moe", +] diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index c0881aa61f..086ef53d17 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -277,9 +277,24 @@ def get_flydsl_stage1_kernels( "xcd_swizzle": xcd, "k_wave": kw, } + _register_production_variants_stage1(kernels, a_dtype, b_dtype, out_dtype) return kernels +def _register_production_variants_stage1( + kernels: dict[str, dict], a_dtype: str, b_dtype: str, out_dtype: str +) -> None: + """Append independently validated stage1 variants to ``kernels`` in-place.""" + + if (a_dtype, b_dtype, out_dtype) != ("fp8", "fp4", "bf16"): + return + + base = flydsl_kernel_name(1, a_dtype, b_dtype, out_dtype, 32, 64, 256) + "_gui" + if base not in kernels: + return + kernels[base + "_kw7"] = {**kernels[base], "k_wave": 7} + + def get_flydsl_stage2_kernels( a_dtype: str, b_dtype: str, out_dtype: str ) -> dict[str, dict]: @@ -2152,15 +2167,14 @@ def flydsl_moe_fused_route_quant_scatter( numel = token_num * topk model_dim = hidden_states.shape[-1] rows_per_tile = wmma_rep * 16 - assert ( - max_m % rows_per_tile == 0 - ), f"max_m ({max_m}) must be a multiple of wmma_rep*16 ({rows_per_tile})" + assert max_m % rows_per_tile == 0, ( + f"max_m ({max_m}) must be a multiple of wmma_rep*16 ({rows_per_tile})" + ) out_E = E if out_E is None else int(out_E) out_max_m = max_m if out_max_m is None else int(out_max_m) assert out_max_m % rows_per_tile == 0, ( - f"out_max_m ({out_max_m}) must be a multiple of wmma_rep*16 " - f"({rows_per_tile})" + f"out_max_m ({out_max_m}) must be a multiple of wmma_rep*16 ({rows_per_tile})" ) payload_bytes_per_row = model_dim if quant_mode == "fp8" else model_dim // 2 @@ -2522,15 +2536,14 @@ def flydsl_moe_fused_quant_preshuffle( "unsupported (expected 'fp4' or 'fp8')." ) assert grouped_in.dtype == torch.bfloat16, ( - "fused grouped quant+preshuffle requires bf16 input " - f"(got {grouped_in.dtype})" + f"fused grouped quant+preshuffle requires bf16 input (got {grouped_in.dtype})" ) device = grouped_in.device feat_dim = grouped_in.shape[-1] rows_per_tile = wmma_rep * 16 - assert ( - max_m % rows_per_tile == 0 - ), f"max_m ({max_m}) must be a multiple of wmma_rep*16 ({rows_per_tile})" + assert max_m % rows_per_tile == 0, ( + f"max_m ({max_m}) must be a multiple of wmma_rep*16 ({rows_per_tile})" + ) n_rows = E * max_m Pb = feat_dim if quant_mode == "fp8" else feat_dim // 2 diff --git a/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort_parallel.py b/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort_parallel.py new file mode 100644 index 0000000000..fd0a046f2c --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort_parallel.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch + +import aiter +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.kimi_k3_moe_route_parallel import ( + kimi_k3_b1_route_sort_parallel, + supports_kimi_k3_b1_route_sort_parallel, +) +from aiter.ops.flydsl.utils import is_flydsl_available + + +def _gfx950_flydsl_available() -> bool: + return ( + torch.cuda.is_available() + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +def test_parallel_route_support_predicate_fails_closed_on_cpu(): + hidden = torch.empty((1, 3584), dtype=torch.bfloat16) + logits = torch.empty((1, 896), dtype=torch.float32) + bias = torch.empty(896, dtype=torch.bfloat16) + + assert not supports_kimi_k3_b1_route_sort_parallel( + hidden, + logits, + bias, + model_dim=3584, + ) + + +def _make_case(case: str) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(20260729) + logits = torch.randn((1, 896), generator=generator) + bias = ( + torch.empty(896, dtype=torch.float32) + .uniform_(-0.125, 0.125, generator=generator) + .to(torch.bfloat16) + ) + if case == "all_equal": + logits.zero_() + bias.zero_() + elif case == "repeating_bias": + logits.zero_() + bias = ( + torch.arange(896, dtype=torch.int32) + .remainder_(8) + .to(torch.float32) + .sub_(4.0) + .mul_(0.03125) + .to(torch.bfloat16) + ) + elif case == "partition_boundary_ties": + logits.fill_(-5.0) + tied = torch.tensor( + [ + 0, + 13, + 14, + 63, + 64, + 127, + 128, + 255, + 256, + 447, + 448, + 511, + 512, + 767, + 768, + 895, + ], + dtype=torch.long, + ) + logits[0, tied] = 5.0 + bias.zero_() + return logits.cuda(), bias.cuda() + + +def _expected_metadata( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sentinel = (16 << 24) | 1 + sorted_ids = torch.full( + (16 * 32,), + sentinel, + dtype=torch.int32, + device=topk_ids.device, + ) + sorted_weights = torch.zeros( + 16 * 32, + dtype=torch.float32, + device=topk_ids.device, + ) + sorted_expert_ids, slots = torch.sort(topk_ids.flatten()) + rows = torch.arange(16, device=topk_ids.device) * 32 + sorted_ids[rows] = slots.to(torch.int32) << 24 + sorted_weights[rows] = topk_weights.flatten()[slots] + return sorted_ids, sorted_weights, sorted_expert_ids + + +@pytest.mark.skipif( + not _gfx950_flydsl_available(), + reason="requires FlyDSL on gfx950", +) +@pytest.mark.parametrize( + "case", + [ + "random_nonzero_bias", + "all_equal", + "repeating_bias", + "partition_boundary_ties", + ], +) +def test_parallel_route_matches_exact_metadata_and_quantization(case: str): + logits, bias = _make_case(case) + hidden = torch.randn( + (1, 3584), + dtype=torch.bfloat16, + device="cuda", + ) + expected_weights = torch.empty( + (1, 16), + dtype=torch.float32, + device="cuda", + ) + expected_ids = torch.empty( + (1, 16), + dtype=torch.int32, + device="cuda", + ) + aiter.biased_grouped_topk_hip( + logits, + bias.float(), + expected_weights, + expected_ids, + 1, + 1, + True, + 1.0, + ) + ( + expected_sorted_ids, + expected_sorted_weights, + expected_sorted_expert_ids, + ) = _expected_metadata(expected_weights, expected_ids) + expected_quantized, expected_scales = aiter.fused_dynamic_mxfp8_quant_moe_sort( + hidden, + sorted_ids=expected_sorted_ids, + num_valid_ids=torch.tensor( + [512, 1], + dtype=torch.int32, + device="cuda", + ), + token_num=1, + topk=16, + block_size=32, + sorted_weights=expected_sorted_weights, + ) + + ( + topk_weights, + topk_ids, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + quantized_hidden, + quantized_scales, + ) = kimi_k3_b1_route_sort_parallel( + hidden, + logits, + bias, + model_dim=3584, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(topk_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close( + topk_weights, + expected_weights, + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close(sorted_ids, expected_sorted_ids, rtol=0, atol=0) + torch.testing.assert_close( + sorted_weights, + expected_sorted_weights, + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close( + sorted_expert_ids, + expected_sorted_expert_ids, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + num_valid_ids, + torch.tensor([512, 1], dtype=torch.int32, device="cuda"), + ) + assert torch.count_nonzero(moe_buf).item() == 0 + torch.testing.assert_close( + quantized_hidden.view(torch.uint8), + expected_quantized.view(torch.uint8), + rtol=0, + atol=0, + ) + + active_scale_offsets = torch.tensor( + [ + rank * 112 * 32 + + (group // 8) * 256 + + (group % 4) * 64 + + ((group % 8) // 4) * 2 + for rank in range(16) + for group in range(112) + ], + dtype=torch.int64, + device="cuda", + ) + torch.testing.assert_close( + quantized_scales.view(torch.uint8).flatten()[active_scale_offsets], + expected_scales.view(torch.uint8).flatten()[active_scale_offsets], + rtol=0, + atol=0, + ) diff --git a/op_tests/flydsl_tests/test_kimi_k3_persistent_moe.py b/op_tests/flydsl_tests/test_kimi_k3_persistent_moe.py new file mode 100644 index 0000000000..dbd522083b --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_persistent_moe.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +from dataclasses import replace + +import pytest +import torch + +import aiter.ops.flydsl.kimi_k3_persistent_moe as persistent_moe +from aiter.ops.flydsl.kimi_k3_persistent_moe import ( + KimiK3PersistentMoERequest, + prepare_kimi_k3_b1_persistent_moe, + supports_kimi_k3_b1_persistent_moe, +) +from aiter.ops.flydsl.moe_kernels import get_flydsl_kernel_params + + +def _request() -> KimiK3PersistentMoERequest: + hidden = torch.empty((1, 3584), dtype=torch.bfloat16) + return KimiK3PersistentMoERequest( + hidden_states=hidden, + router_logits=torch.empty((1, 896), dtype=torch.float32), + correction_bias=torch.empty(896, dtype=torch.bfloat16), + w1=torch.empty(1), + w2=torch.empty(1), + w1_scale=torch.empty(1), + w2_scale=torch.empty(1), + situ_beta=4.0, + situ_linear_beta=25.0, + w13_layout="gate_up_interleaved_preshuffled", + weights_shuffled=True, + quantization_supported=True, + activation="situ", + num_experts=896, + topk=16, + num_expert_group=1, + topk_group=1, + renormalize=True, + scoring_func="sigmoid", + routed_scaling_factor=1.0, + expert_parallel=False, + eplb_enabled=False, + lora_enabled=False, + has_expert_bias=False, + apply_router_weight_on_input=False, + expert_map_active=False, + routing_capture_enabled=False, + custom_routing_active=False, + input_ids_active=False, + routing_method="DeepSeekV3", + ) + + +def test_persistent_moe_fails_closed_on_cpu(): + request = _request() + + assert not supports_kimi_k3_b1_persistent_moe(request) + assert prepare_kimi_k3_b1_persistent_moe(request) is None + + +def test_persistent_moe_accepts_native_kimi_k3_weight_layout( + monkeypatch: pytest.MonkeyPatch, +): + packed_shapes = [] + monkeypatch.setenv("AITER_SITUV2_A8W4", "1") + monkeypatch.setattr( + persistent_moe, + "supports_kimi_k3_b1_route_sort_parallel", + lambda *args, **kwargs: True, + ) + monkeypatch.setattr( + persistent_moe, + "_has_validated_stage1_kernel", + lambda: True, + ) + + def record_packed_shape(*_args, shape, **_kwargs): + packed_shapes.append(shape) + return True + + monkeypatch.setattr( + persistent_moe, + "_is_packed_tensor", + record_packed_shape, + ) + + assert supports_kimi_k3_b1_persistent_moe(_request()) + assert packed_shapes == [ + (896, 768, 1792), + (896, 3584, 192), + (688128, 112), + (3211264, 16), + ] + + +def test_persistent_moe_stage1_variant_is_registered_explicitly(): + parameters = get_flydsl_kernel_params( + "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_gui_kw7_fp8" + ) + + assert parameters is not None + assert { + "stage": parameters["stage"], + "a_dtype": parameters["a_dtype"], + "b_dtype": parameters["b_dtype"], + "out_dtype": parameters["out_dtype"], + "tile_m": parameters["tile_m"], + "tile_n": parameters["tile_n"], + "tile_k": parameters["tile_k"], + "waves_per_eu": parameters["waves_per_eu"], + "b_nt": parameters["b_nt"], + "gate_mode": parameters["gate_mode"], + "k_wave": parameters["k_wave"], + } == persistent_moe._STAGE1_EXPECTED + + +@pytest.mark.parametrize( + ("field", "unsupported"), + [ + ("situ_beta", 1.0), + ("situ_linear_beta", 1.0), + ("w13_layout", "gate_up_separated_preshuffled"), + ("weights_shuffled", False), + ("quantization_supported", False), + ("activation", "silu"), + ("num_experts", 128), + ("topk", 8), + ("num_expert_group", 8), + ("topk_group", 4), + ("renormalize", False), + ("scoring_func", "softmax"), + ("routed_scaling_factor", 2.5), + ("expert_parallel", True), + ("eplb_enabled", True), + ("lora_enabled", True), + ("has_expert_bias", True), + ("apply_router_weight_on_input", True), + ("expert_map_active", True), + ("routing_capture_enabled", True), + ("custom_routing_active", True), + ("input_ids_active", True), + ("routing_method", "Renormalize"), + ], +) +def test_persistent_moe_contract_fails_closed( + monkeypatch: pytest.MonkeyPatch, + field: str, + unsupported: object, +): + monkeypatch.setenv("AITER_SITUV2_A8W4", "1") + monkeypatch.setattr( + persistent_moe, + "supports_kimi_k3_b1_route_sort_parallel", + lambda *args, **kwargs: True, + ) + monkeypatch.setattr( + persistent_moe, + "_has_validated_stage1_kernel", + lambda: True, + ) + monkeypatch.setattr( + persistent_moe, + "_is_packed_tensor", + lambda *args, **kwargs: True, + ) + + assert supports_kimi_k3_b1_persistent_moe(_request()) + assert not supports_kimi_k3_b1_persistent_moe( + replace(_request(), **{field: unsupported}) + )