perf(flydsl): Kimi-K3 B1 route-sort + prepared-MoE handoff (stack)

.github/scripts/aiter_test.sh

@@ -130,6 +130,7 @@ if [[ "$mla_in_shard" == "true" && "$MULTIGPU" != "TRUE" ]]; then
         "-c 98304 -b 1 -n 16,1 -kvd fp8" \
         "-c 10000 100000 -b 1 3 4 -n 12,1 16,1 -kvd bf16 -lse" \
         "-c 1 21 63 64 65 256 -b 1 -n 16,1 -kvd bf16 -lse" \
+        "-c 1 31 32 63 64 65 8192 -b 1 -n 12,1 -d bf16 -kvd bf16 -k 512 -qn 512 -qr 64 -vh 512 -blk 1 --varlen --gluon-num-kv-splits 32" \
         "-c 16384 -b 4 -n 16,8 16,17 -kvd bf16"; do
         echo "=== extra: test_mla.py $args ===" | tee -a latest_test.log
         if ! timeout 10m python3 op_tests/test_mla.py $args 2>&1 | tee -a latest_test.log; then

aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv

@@ -15,7 +15,7 @@ gfx950,256,2048,3584,512,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloa
 gfx950,256,4096,3584,512,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,760.7278,flydsl_moe1_abf16_wfp4_bf16_t32x128x256_w3,35.5%,454.8517,flydsl_moe2_abf16_wfp4_bf16_t32x256x128_atomic_xcd4_persist,0.2%,1215.5795,0,0,0,593.59,4106.04,
 gfx950,256,8192,3584,512,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,1223.59,flydsl_moe1_abf16_wfp4_bf16_t32x128x256_w4_bnt0_xcd4,35.4%,802.0668,flydsl_moe2_abf16_wfp4_bf16_t32x256x128_atomic_xcd4,0.2%,2025.6568,0,0,0,712.42,2492.99,
 gfx950,256,16384,3584,512,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,2184.0333,flydsl_moe1_abf16_wfp4_bf16_t32x128x256_w4_bnt0_xcd4,35.4%,1510.1664,flydsl_moe2_abf16_wfp4_bf16_t32x256x128_atomic_xcd4_persist,0.2%,3694.1997,0,0,0,781.28,1398.78,
-gfx950,256,1,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,18.8559,flydsl_moe1_abf16_wfp4_bf16_t32x64x256_w3_kw4,34.6%,9.4261,flydsl_moe2_abf16_wfp4_bf16_t32x128x128_reduce_bnt2_persist,0.0%,28.282,0,0,0,4.67,130803.71,
+gfx950,256,1,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,10.5544,flydsl_moe1_abf16_wfp4_bf16_t16x64x256_w3_xcd4,34.6%,3.8114,flydsl_moe2_abf16_wfp4_bf16_t32x128x128_atomic_bnt2,0.0%,14.3658,0,0,0,9.19,257513.71,
 gfx950,256,2,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,19.8871,flydsl_moe1_abf16_wfp4_bf16_t32x64x256_w4_bnt0_xcd4,35.5%,11.2795,flydsl_moe2_abf16_wfp4_bf16_t32x256x256_reduce_bnt2,0.0%,31.1666,0,0,0,8.48,118697.73,
 gfx950,256,3,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,25.2668,flydsl_moe1_abf16_wfp4_bf16_t32x64x256_w4_xcd4_kw2,35.7%,13.6755,flydsl_moe2_abf16_wfp4_bf16_t32x256x256_reduce_bnt2,0.0%,38.9423,0,0,0,10.18,94997.45,
 gfx950,256,4,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.bfloat16,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,26.6074,flydsl_moe1_abf16_wfp4_bf16_t32x64x256_w2_bnt0_xcd4_kw2,35.6%,16.1406,flydsl_moe2_abf16_wfp4_bf16_t32x256x128_reduce,0.0%,42.748,0,0,0,12.36,86540.5,

aiter/ops/flydsl/__init__.py

@@ -38,8 +38,14 @@
         )
 
     from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8
+    from .latent_moe_tail import latent_moe_tail, supports_latent_moe_tail
     from .moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2
     from .fmha_kernels import flydsl_flash_attn_func
+    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.qk_norm_rope_quant import flydsl_qk_norm_rope_quant
     from .kernels.pa_mqa_logits_fp4 import (
         flydsl_pa_mqa_logits_fp4,
@@ -60,7 +66,12 @@
         "flydsl_moe_stage1",
         "flydsl_moe_stage2",
         "flydsl_hgemm",
+        "latent_moe_tail",
+        "supports_latent_moe_tail",
         "flydsl_flash_attn_func",
+        "flydsl_kimi_k3_kda_decode",
+        "flydsl_kimi_k3_kda_decode_with_f_b",
+        "is_flydsl_kimi_k3_kda_decode_supported",
         "flydsl_qk_norm_rope_quant",
         "flydsl_pa_mqa_logits_fp4",
         "flydsl_pa_mqa_logits_fp4_prefill",

aiter/ops/flydsl/kernels/kimi_k3_b1_gate_projection.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Fixed-shape Kimi-K3 BF16 router projection for gfx950."""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl
+from flydsl.expr.arith import ArithValue, 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,
+)
+
+_HIDDEN_SIZE = 7168
+_EXPERTS = 896
+_WAVE_SIZE = 64
+_BLOCK_THREADS = 64
+_WORKGROUPS = _EXPERTS
+
+
+def _raw(value):
+    return value.ir_value() if hasattr(value, "ir_value") else value
+
+
+def build_kimi_k3_b1_gate_projection_module():
+    """Build the fixed B1x7168 by 896x7168 projection launcher."""
+
+    @flyc.kernel(
+        name="kimi_k3_b1_gate_projection_gfx950",
+        known_block_size=[_BLOCK_THREADS, 1, 1],
+    )
+    def projection_kernel(
+        hidden: fx.Pointer,
+        router_weight: fx.Pointer,
+        logits: fx.Pointer,
+    ):
+        i32 = T.i32
+        f32 = T.f32
+        lane = ArithValue(gpu.thread_idx.x)
+        wave = ArithValue(gpu.block_idx.x)
+        expert = wave
+
+        hidden_rsrc = ptr_rsrc(hidden)
+        weight_rsrc = ptr_rsrc(router_weight)
+        logits_rsrc = ptr_rsrc(logits)
+        zero_i32 = arith.constant(0, type=i32)
+        zero_f32 = arith.constant(0.0, type=f32)
+
+        accumulator = ArithValue(zero_f32)
+        for k_base in range_constexpr(0, _HIDDEN_SIZE, _WAVE_SIZE):
+            k = lane + arith.constant(k_base, type=i32)
+            hidden_bf16 = buffer_ops.buffer_load(
+                hidden_rsrc,
+                k,
+                vec_width=1,
+                dtype=T.bf16,
+            )
+            hidden_f32 = ArithValue(arith.extf(f32, hidden_bf16))
+            weight_index = expert * arith.constant(_HIDDEN_SIZE, type=i32) + k
+            weight_bf16 = buffer_ops.buffer_load(
+                weight_rsrc,
+                weight_index,
+                vec_width=1,
+                dtype=T.bf16,
+            )
+            weight_f32 = ArithValue(arith.extf(f32, weight_bf16))
+            accumulator = accumulator + hidden_f32 * weight_f32
+
+        # The same gfx9 DPP tree used by the accepted route kernel reduces a
+        # full 64-lane wave and leaves the total in lane 63.
+        for dpp_control in (0xB1, 0x4E, 0x141, 0x140, 0x142, 0x143):
+            remote_i32 = rocdl.update_dpp(
+                i32,
+                zero_i32,
+                arith.bitcast(i32, _raw(accumulator)),
+                dpp_control,
+                0xF,
+                0xF,
+                True,
+            )
+            remote = ArithValue(arith.bitcast(f32, remote_i32))
+            accumulator = accumulator + remote
+
+        is_last_lane = arith.cmpi(
+            CmpIPredicate.eq,
+            lane,
+            arith.constant(_WAVE_SIZE - 1, type=i32),
+        )
+        store_if = scf.IfOp(is_last_lane)
+        with ir.InsertionPoint(store_if.then_block):
+            # GateLinear's production contract is BF16 linear followed by an
+            # FP32 cast. Round once to BF16 before materializing FP32.
+            rounded = arith.trunc_f(T.bf16, _raw(accumulator))
+            projected = arith.extf(f32, rounded)
+            buffer_ops.buffer_store(projected, logits_rsrc, expert)
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_projection(
+        hidden: fx.Pointer,
+        router_weight: fx.Pointer,
+        logits: fx.Pointer,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            pass
+        projection_kernel(hidden, router_weight, logits).launch(
+            grid=(arith.index(_WORKGROUPS), 1, 1),
+            block=(_BLOCK_THREADS, 1, 1),
+            stream=stream,
+        )
+
+    launch_projection.compile_hints = {
+        "llvm_options": {
+            "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD,
+            "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT,
+        },
+    }
+    return launch_projection