From ac4ba2dfe3efef1b482d1f3f472e7de6273f6098 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Mon, 27 Jul 2026 22:15:03 +0000 Subject: [PATCH 1/9] perf(mla): expose split override for graph decode Graph-captured callers may not have a useful runtime sequence length for automatic split-KV scheduling. Add a validated num_kv_splits override while preserving the existing automatic default and positional API compatibility. Cover policy validation and the Kimi-K3 TP8 bh16/BF16 shape across short-sequence split boundaries and 8K context. Assisted-by: OpenAI Codex Signed-off-by: Yanyuan Qin --- .github/scripts/aiter_test.sh | 1 + aiter/ops/triton/gluon/README.md | 28 +++--- aiter/ops/triton/gluon/mla_gluon.py | 85 ++++++++++++++----- op_tests/test_mla.py | 16 +++- .../attention/test_mla_gluon_split_policy.py | 28 ++++++ 5 files changed, 125 insertions(+), 33 deletions(-) create mode 100644 op_tests/triton_tests/attention/test_mla_gluon_split_policy.py diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index c8e0ec6970..51db488829 100755 --- a/.github/scripts/aiter_test.sh +++ b/.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 diff --git a/aiter/ops/triton/gluon/README.md b/aiter/ops/triton/gluon/README.md index d6477dc63a..fe8a76dcbc 100644 --- a/aiter/ops/triton/gluon/README.md +++ b/aiter/ops/triton/gluon/README.md @@ -157,17 +157,25 @@ python op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py [-gluon] ### `mla_gluon.py` — MLA Decode + DeepSeek V4 Sparse Prefill -**Function:** `mla_gluon(q_nope, q_pe, kv_c, o, page_table, seq_info, sm_scale, k_pe=None, kv_pe_offset=512, use_2d_view=True, kv_scale=1.0, min_kv_seq_len=1, return_lse=False)` +**Function:** `mla_gluon(q_nope, q_pe, kv_c, o, page_table, seq_info, sm_scale, k_pe=None, kv_pe_offset=512, use_2d_view=True, kv_scale=1.0, min_kv_seq_len=1, return_lse=False, has_pe=True, attn_sink=None, num_kv_splits=None)` **Description:** Multi-head Latent Attention (DeepSeek MLA) kernel with split-KV. For MLA Decode, Q is split into compressed latent (`q_nope`, dim=kv_lora_rank) and rope positional encoding (`q_pe`, dim=qk_rope_head_dim). KV cache is a flat `[N, 576]` buffer (`kv_c`). For DSv4 Sparse Prefill, Q packs compressed latent and positional encoding into one contiguous row (448 NoPE + 64 RoPE, `q_nope` with shape `[nquery, nhead, 512]`), KV cache has aligned `head_dim=512`, `q_pe` and `k_pe` can be left as placeholders. Uses 3-stage async copy pipeline with double-buffered page numbers and KV tiles. The wrapper dispatches by `(nhead, kv_c.dtype)` to one of three compile-time regimes (single `@gluon.jit` kernel, REGIME constexpr gates layouts and grid mapping): - **`bh64`** (`nhead in {64, 128}`): bf16 KV, BLOCK_H=64, BLOCK_N=64, multi-batch + XCD-aware 3-D grid. `NUM_KV_SPLITS` auto-picked ∈ {1, 2, 4} so the launch fills ~256 workgroups (one wave on MI350). When `NUM_KV_SPLITS == 1`, stage-1 writes the final attention output directly to `o` (no temp buffer, no reduce). When `NUM_KV_SPLITS > 1`, stage-1 writes per-split `(acc, fp32 lse)` and stage-2 (`_mla_softmax_reducev_kernel`) reduces them into `o`. -- **`bh16bn128`** (`nhead ≤ 16`, `batch_size == 1`, fp8 KV): BLOCK_H=16, BLOCK_N=128, 2-D grid `(1, NUM_KV_SPLITS)` with token-bound `NUM_KV_SPLITS = max(1, min(256, min_kv_seq_len))` — 256 for the normal long-context path, reduced only for small kv (`min_kv_seq_len < 256`) so every split stays non-empty. Optional `kv_scale` dequant. Stage-2 reduce runs whenever `NUM_KV_SPLITS > 1` (skipped via the fast path only at `min_kv_seq_len == 1`). Supports the general case `num_iter ∈ {1, 2, ...}` (no `gl.assume(num_iter >= 3)`). `NHEAD < BLOCK_H` masks OOB heads on Q load and O store (wasted MFMA lanes are free; this regime is memory-bound). -- **`bh16bn64`** (`nhead ≤ 16`, bf16 KV): BLOCK_H=16, BLOCK_N=64, 2-D grid `(batch_size, NUM_KV_SPLITS)` with block-bound `NUM_KV_SPLITS = max(1, min(256 // batch_size, cdiv(min_kv_seq_len, BLOCK_N)))` — fills ~256 WGs but never splits a sequence into more than its 64-token block count, so small kv is supported and it collapses to 1 (one WG per batch over the whole sequence) when `min_kv_seq_len <= 64`. Use when KV is kept in bf16 (no fp8 quant). Same `NHEAD < BLOCK_H` masking. Full decode (stage-1, plus stage-2 reduce into `o` when `NUM_KV_SPLITS > 1`). +- **`bh16bn128`** (`nhead ≤ 16`, `batch_size == 1`, fp8 KV): BLOCK_H=16, BLOCK_N=128, 3-D grid `(1, NUM_KV_SPLITS, qlen)` with token-bound `NUM_KV_SPLITS = max(1, min(256 // qlen, min_kv_seq_len))` — 256 for single-query long-context decode, reduced for larger query lengths or small KV so the automatic policy keeps every split non-empty. Optional `kv_scale` dequant. Stage-2 reduce runs whenever `NUM_KV_SPLITS > 1`. Supports the general case `num_iter ∈ {1, 2, ...}` (no `gl.assume(num_iter >= 3)`). `NHEAD < BLOCK_H` masks OOB heads on Q load and O store (wasted MFMA lanes are free; this regime is memory-bound). +- **`bh16bn64`** (`nhead ≤ 16`, bf16 KV): BLOCK_H=16, BLOCK_N=64, 3-D grid `(batch_size, NUM_KV_SPLITS, qlen)` with block-bound `NUM_KV_SPLITS = max(1, min(256 // (batch_size * qlen), cdiv(min_kv_seq_len, BLOCK_N)))` — fills ~256 WGs but never splits a sequence into more than its 64-token block count, so small kv is supported and it collapses to 1 (one WG per batch and query position over the whole sequence) when `min_kv_seq_len <= 64`. Use when KV is kept in bf16 (no fp8 quant). Same `NHEAD < BLOCK_H` masking. Full decode (stage-1, plus stage-2 reduce into `o` when `NUM_KV_SPLITS > 1`). -All three regimes run the full decode and dsv4 prefill. `return_lse=True` also returns the merged fp32 lse `[batch, nhead]`, so `mla_gluon(...)` returns `(o, final_lse)` instead of `(o, None)`. +All three regimes run the full decode and dsv4 prefill. `return_lse=True` also returns the merged fp32 lse `[batch, qlen, nhead]` (`qlen=1` for plain decode), so `mla_gluon(...)` returns `(o, final_lse)` instead of `(o, None)`. + +By default, each regime derives `NUM_KV_SPLITS` from the static batch shape and +`min_kv_seq_len`. Graph-captured callers that cannot provide a useful runtime +minimum may pass an explicitly tuned `num_kv_splits` in `[1, 256]`. This is a +caller-owned scheduling policy: the generic wrapper does not infer a model from +its tensor shape. If the override creates empty leading splits for a short +sequence, stage 1 skips them and stage 2 uses the runtime sequence length to +reduce only initialized splits. Modified from [FlashMLA](https://github.com/deepseek-ai/FlashMLA/blob/main/benchmark/bench_flash_mla.py). @@ -183,17 +191,17 @@ Modified from [FlashMLA](https://github.com/deepseek-ai/FlashMLA/blob/main/bench | BLOCK_H | 64 | 16 | 16 | | BLOCK_N | 64 | 128 | 64 | | MFMA | 16×16×32, warps=[4,1] | 16×16×32, warps=[1,4] | 16×16×32, warps=[1,4] | -| Grid | 3-D XCD-aware | 2-D `(1, NUM_KV_SPLITS)` | 2-D `(batch, NUM_KV_SPLITS)` | -| NUM_KV_SPLITS | auto ∈ {1, 2, 4} from (batch, nhead) | `max(1, min(256, min_kv_seq_len))` (token-bound; 256 for ctx ≥ 256) | `max(1, min(256 // batch_size, cdiv(min_kv_seq_len, 64)))` (block-bound; collapses to 1 for ctx ≤ 64) | +| Grid | 3-D XCD-aware | 3-D `(1, NUM_KV_SPLITS, qlen)` | 3-D `(batch, NUM_KV_SPLITS, qlen)` | +| NUM_KV_SPLITS | auto ∈ {1, 2, 4} from (batch, nhead, qlen) | `max(1, min(256 // qlen, min_kv_seq_len))` (token-bound; 256 for qlen=1 and ctx ≥ 256) | `max(1, min(256 // (batch_size * qlen), cdiv(min_kv_seq_len, 64)))` (block-bound; collapses to 1 for ctx ≤ 64) | | `kv_scale` | unused (pass 1.0) | dequant scale folded into `qk_scale` (applied before softmax for fp8 correctness) | unused (pass 1.0) | -| Seq constraint | `min_kv_seq_len > NUM_KV_SPLITS * (3 * BLOCK_N + NUM_KV_SPLITS)` (the `3` matches the kernel's `gl.assume(num_iter > 3)`) | `min_kv_seq_len ≥ 1` (small kv 1..256 supported; token-bound clamp keeps splits non-empty) | `min_kv_seq_len ≥ 1` (small kv 1..256 supported; block-bound clamp keeps splits non-empty) | -| Stage-2 reduce | skipped when `NUM_KV_SPLITS == 1` | skipped when `NUM_KV_SPLITS == 1` (i.e. `min_kv_seq_len == 1`) | skipped when `NUM_KV_SPLITS == 1` | +| Seq constraint | `min_kv_seq_len > NUM_KV_SPLITS * (3 * BLOCK_N + NUM_KV_SPLITS)` (the `3` matches the kernel's `gl.assume(num_iter > 3)`) | `min_kv_seq_len ≥ 1` (automatic token-bound clamp keeps splits non-empty; explicit empty leading splits are supported) | `min_kv_seq_len ≥ 1` (automatic block-bound clamp keeps splits non-empty; explicit empty leading splits are supported) | +| Stage-2 reduce | skipped when `NUM_KV_SPLITS == 1` | skipped when `NUM_KV_SPLITS == 1` | skipped when `NUM_KV_SPLITS == 1` | -**Page table modes** (`use_2d_view`, both regimes): +**Page table modes** (`use_2d_view`, all regimes): - `True`: `page_table = block_table [batch, max_seqlen]`, `seq_info = cache_seqlens [batch]`. Use for fixed-length or pre-padded variable-length sequences. - `False`: `page_table = kv_indices [total_kv]`, `seq_info = kv_indptr [batch+1]`. Use for variable-length sequences without block_table construction. -**KV layout** (both regimes): By default `kv_c` is a flat `[N, 576]` buffer containing both the compressed latent (columns `[0, 512)`) and rope PE (columns `[512, 576)`). The kernel adds `kv_pe_offset` to k_pe column offsets — set to `kv_lora_rank` (512) when `k_pe` shares `kv_c` (default), or `0` when `k_pe` is a separate buffer. The kernel auto-selects the load instruction via `WITHIN_2GB`: `buffer_load_to_shared` (scalar base + 32-bit offsets) when KV caches ≤ 2 GB, or `global_load_to_shared` (64-bit pointer tensors) when KV caches > 2 GB. +**KV layout** (all regimes): By default `kv_c` is a flat `[N, 576]` buffer containing both the compressed latent (columns `[0, 512)`) and rope PE (columns `[512, 576)`). The kernel adds `kv_pe_offset` to k_pe column offsets — set to `kv_lora_rank` (512) when `k_pe` shares `kv_c` (default), or `0` when `k_pe` is a separate buffer. The kernel auto-selects the load instruction via `WITHIN_2GB`: `buffer_load_to_shared` (scalar base + 32-bit offsets) when KV caches ≤ 2 GB, or `global_load_to_shared` (64-bit pointer tensors) when KV caches > 2 GB. **`bh64` perf** (MI350, ctx=16384, bf16 Q + bf16 KV; compute-bound): diff --git a/aiter/ops/triton/gluon/mla_gluon.py b/aiter/ops/triton/gluon/mla_gluon.py index 9cd8b4b43a..12988da6a7 100644 --- a/aiter/ops/triton/gluon/mla_gluon.py +++ b/aiter/ops/triton/gluon/mla_gluon.py @@ -10,14 +10,15 @@ # Fast path: when NUM_KV_SPLITS==1, stage-1 writes the # final output directly to O and stage-2 reduce is skipped. # REGIME='bh16bn128' - bf16 Q + fp8 KV, BLOCK_H=16, BLOCK_N=128, -# nhead <= 16, batch_size=1, NUM_KV_SPLITS=256. -# 2-D (batch, split) grid. Always splits + always -# reduces. NHEAD < BLOCK_H masks OOB heads on Q load -# and O store. +# nhead <= 16, batch_size=1. 3-D +# (batch, split, q_pos) grid. Automatic NUM_KV_SPLITS +# is bounded by qlen and the minimum KV sequence length. +# NHEAD < BLOCK_H masks OOB heads on Q load and O store. # REGIME='bh16bn64' - bf16 Q + bf16 KV, BLOCK_H=16, BLOCK_N=64, -# nhead <= 16, batch_size >= 1, 2-D (batch, split) grid, -# NUM_KV_SPLITS = max(1, 256 // batch_size). Full decode -# (stage-1 + stage-2 reduce into the final O). +# nhead <= 16, batch_size >= 1. 3-D +# (batch, split, q_pos) grid. Automatic NUM_KV_SPLITS +# is bounded by batch_size, qlen, and KV block count. +# Full decode (stage-1 + stage-2 reduce into final O). # NHEAD < BLOCK_H masks OOB heads on Q load and O store. # # The bh16 regimes support num_iter in {1, 2, ...} (no gl.assume(num_iter>=3)); @@ -28,7 +29,7 @@ # # Full decode for all regimes. For NUM_KV_SPLITS>1 stage-1 writes per-split acc + # fp32 lse; stage-2 (_mla_softmax_reducev_kernel) reduces into O. RETURN_LSE also -# returns the merged fp32 lse [B, H] (stage-2 for splits>1, else stage-1). +# returns the merged fp32 lse [B, QLEN, H] (stage-2 for splits>1, else stage-1). # # 3-stage software pipeline (double-buffered, BLOCK_N with 2x(BLOCK_N/2) KV slices): # AC = async_copy (global->LDS), LL = load (LDS->reg), P = page, K = K-cache, V = V-cache @@ -49,9 +50,29 @@ from triton.experimental import gluon from triton.experimental.gluon import language as gl -import aiter.ops.triton.utils._triton.arch_info as arch_info +from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.device_info import get_num_xcds +_MAX_NUM_KV_SPLITS = 256 + + +def _resolve_num_kv_splits(auto_num_kv_splits, num_kv_splits): + """Return the automatic split count or a validated caller override.""" + if num_kv_splits is None: + return auto_num_kv_splits + if isinstance(num_kv_splits, bool) or not isinstance(num_kv_splits, int): + raise TypeError( + "num_kv_splits must be an int or None, " + f"got {type(num_kv_splits).__name__}" + ) + if not 1 <= num_kv_splits <= _MAX_NUM_KV_SPLITS: + raise ValueError( + f"num_kv_splits must be in [1, {_MAX_NUM_KV_SPLITS}], " + f"got {num_kv_splits}" + ) + return num_kv_splits + + # fmt: off @gluon.jit def _mla_gluon( @@ -105,10 +126,9 @@ def _mla_gluon( HAS_PE: gl.constexpr, HAS_ATTN_SINK: gl.constexpr, ): - # Grid mapping: bh64 uses 3-D XCD-aware multi-batch; bh16bn64 and bh16bn128 - # use 2-D (batch, split) — for batch_size=1 this is (1, NUM_KV_SPLITS). - # MTP: an extra q_pos axis carries the query position within QLEN. bh64 packs - # it into grid axis 1 (after the head-block index); bh16 uses grid axis 2. + # Grid mapping: all stage-1 launches are 3-D. bh64 uses an XCD-aware + # multi-batch mapping and packs q_pos into grid axis 1 after the head-block + # index. bh16 uses (batch, split, q_pos), with q_pos on grid axis 2. # When QLEN==1, q_pos is always 0 and the layout below is identical to before. if REGIME == 'bh64': NUM_M_BLOCKS: gl.constexpr = (NHEAD + BLOCK_H - 1) // BLOCK_H @@ -816,6 +836,7 @@ def mla_gluon( return_lse=False, has_pe=True, attn_sink=None, # [nhead] fp32 per-head sink bias, None means no sink + num_kv_splits=None, ): """Unified Gluon MLA entry (gfx950 / CDNA4) — decode and DeepSeek V4 sparse prefill. @@ -836,6 +857,10 @@ def mla_gluon( return_lse=True: additionally returns the merged log-sum-exp, a separate fp32 tensor [batch, qlen, nhead] + num_kv_splits=None uses the regime's automatic split policy. + Set num_kv_splits to an integer in [1, 256] when a graph-captured caller + cannot provide a useful min_kv_seq_len and has an externally tuned schedule. + DSv4 Sparse prefill packs NoPE and RoPE in to one contiguous row (448+64). To run DSv4 prefill, it requires has_pe=False, prepares valid Q / K in q_nope / kv_c, and attn_sink, q_pe / k_pe are unused placeholders. @@ -902,7 +927,11 @@ def mla_gluon( base_grid = ( NUM_XCDS * triton.cdiv(nhead, BLOCK_H) * qlen * (batch_size // NUM_XCDS) ) - NUM_KV_SPLITS = max(1, triton.next_power_of_2(triton.cdiv(256, base_grid))) + auto_num_kv_splits = max( + 1, + triton.next_power_of_2(triton.cdiv(_MAX_NUM_KV_SPLITS, base_grid)), + ) + NUM_KV_SPLITS = _resolve_num_kv_splits(auto_num_kv_splits, num_kv_splits) assert ( batch_size % 64 == 0 @@ -924,25 +953,37 @@ def mla_gluon( BLOCK_H = 16 BLOCK_N = 128 if REGIME == "bh16bn128" else 64 kv_dtype = torch.float8_e4m3fn if REGIME == "bh16bn128" else torch.bfloat16 - NUM_XCDS = 1 # unused by 2-D split grid mapping - # 2-D grid (batch, split). Both bh16 regimes support num_iter in {1, 2, ...} - # (no gl.assume(num_iter >= 3) in the kernel); the only correctness need is - # that every split is non-empty (floor split size = min_kv_seq_len // - # NUM_KV_SPLITS >= 1). Each clamp below keeps NUM_KV_SPLITS <= min_kv_seq_len, + NUM_XCDS = 1 # unused by the bh16 (batch, split, q_pos) grid mapping + # Both bh16 regimes support num_iter in {1, 2, ...} + # (no gl.assume(num_iter >= 3) in the kernel). The automatic policy keeps + # every split non-empty. An explicit override may create empty leading + # splits for short sequences; stage 1 skips them and stage 2 derives the + # valid split range from the runtime sequence length. if REGIME == "bh16bn128": assert ( batch_size == 1 ), f"mla_gluon[bh16bn128] requires batch_size=1, got {batch_size}" - NUM_KV_SPLITS = max(1, min(256 // (batch_size * qlen), min_kv_seq_len)) + auto_num_kv_splits = max( + 1, + min( + _MAX_NUM_KV_SPLITS // (batch_size * qlen), + min_kv_seq_len, + ), + ) else: # bh16bn64 # Fill ~256 WGs (total WGs = B * NUM_KV_SPLITS <= 256, one MI350 wave), # but never split a sequence into more blocks than it has: bound by the # shortest seq's block count so every split holds >= 1 block (no wasted # partial-block MFMA). For min_kv_seq_len <= BLOCK_N this collapses to # NUM_KV_SPLITS=1, i.e. one WG per batch computing the whole (short) seq. - NUM_KV_SPLITS = max( - 1, min(256 // (batch_size * qlen), triton.cdiv(min_kv_seq_len, BLOCK_N)) + auto_num_kv_splits = max( + 1, + min( + _MAX_NUM_KV_SPLITS // (batch_size * qlen), + triton.cdiv(min_kv_seq_len, BLOCK_N), + ), ) + NUM_KV_SPLITS = _resolve_num_kv_splits(auto_num_kv_splits, num_kv_splits) assert ( q_nope.dtype == torch.bfloat16 and q_pe.dtype == torch.bfloat16 ), f"q_nope/q_pe must be bf16, got {q_nope.dtype}/{q_pe.dtype}" diff --git a/op_tests/test_mla.py b/op_tests/test_mla.py index eaa7c6deb1..4b53d8f4f1 100644 --- a/op_tests/test_mla.py +++ b/op_tests/test_mla.py @@ -140,6 +140,7 @@ def test_mla( return_lse=False, is_causal=True, sequential_page_indices=False, + gluon_num_kv_splits=None, ): ret = {} @@ -232,11 +233,14 @@ def test_normal_prefill(): us_aiter = None prefill_ref_token_cap = 512 * 1024 + ck_prefill_max_qk_head_dim = 256 # Prefill ref builds [nhead, (batch*ctx)^2] fp32 attn weights; bound both # the lazy "tile area" gate and the per-call ctx so decode-scale ctx_lens - # (1M+) never trigger the O(N^2) ref. + # (1M+) never trigger the O(N^2) ref. CK prefill supports QK head + # dimensions up to 256; larger absorbed-MLA decode shapes remain enabled. if ( (dtype == torch.bfloat16 and kvtype == torch.bfloat16) + and qk_head_dim <= ck_prefill_max_qk_head_dim and batch_size * ctx_lens * nhead < 256 * 8192 * 16 and ctx_lens <= 16384 and total_qo <= prefill_ref_token_cap @@ -518,6 +522,7 @@ def test_absorb_decode_gluon(): sm_scale, use_2d_view=use_2d_view, min_kv_seq_len=ctx_lens, + num_kv_splits=gluon_num_kv_splits, return_lse=return_lse, ) @@ -580,6 +585,7 @@ def test_absorb_decode_gluon_bh16(name): use_2d_view=use_2d_view, kv_scale=1.0, min_kv_seq_len=ctx_lens, + num_kv_splits=gluon_num_kv_splits, return_lse=return_lse, ) @@ -856,6 +862,13 @@ def test_absorb_decode_gluon_bh16(name): help="""Enable/disable causal masking. Default: True. --causal / --no-causal""", ) +parser.add_argument( + "--gluon-num-kv-splits", + type=int, + default=None, + help="""Override Gluon MLA's automatic split count. + e.g.: --gluon-num-kv-splits 32""", +) args = parser.parse_args() @@ -883,6 +896,7 @@ def test_absorb_decode_gluon_bh16(name): return_lse=args.return_lse, is_causal=args.causal, sequential_page_indices=args.sequential_page_indices, + gluon_num_kv_splits=args.gluon_num_kv_splits, ) df.append(ret) df = pd.DataFrame(df) diff --git a/op_tests/triton_tests/attention/test_mla_gluon_split_policy.py b/op_tests/triton_tests/attention/test_mla_gluon_split_policy.py new file mode 100644 index 0000000000..0866ffbcee --- /dev/null +++ b/op_tests/triton_tests/attention/test_mla_gluon_split_policy.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest + +from aiter.ops.triton.gluon.mla_gluon import _resolve_num_kv_splits + + +@pytest.mark.parametrize("auto_num_kv_splits", [1, 4, 32, 256]) +def test_default_num_kv_splits_uses_auto_policy(auto_num_kv_splits): + assert _resolve_num_kv_splits(auto_num_kv_splits, None) == auto_num_kv_splits + + +@pytest.mark.parametrize("num_kv_splits", [1, 16, 32, 256]) +def test_num_kv_splits_override(num_kv_splits): + assert _resolve_num_kv_splits(1, num_kv_splits) == num_kv_splits + + +@pytest.mark.parametrize("num_kv_splits", [0, 257, -1]) +def test_num_kv_splits_override_rejects_out_of_range(num_kv_splits): + with pytest.raises(ValueError, match=r"must be in \[1, 256\]"): + _resolve_num_kv_splits(1, num_kv_splits) + + +@pytest.mark.parametrize("num_kv_splits", [True, 32.0, "32"]) +def test_num_kv_splits_override_rejects_non_integer(num_kv_splits): + with pytest.raises(TypeError, match="must be an int or None"): + _resolve_num_kv_splits(1, num_kv_splits) From 2790dbfe9a33a617e6888ed3b99d528db76b9009 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 11:28:25 +0000 Subject: [PATCH 2/9] [Kimi-K3] Fuse gfx950 B1 routing and MoE sort metadata Add a narrow FlyDSL specialization for the 896-expert, top-16, single-group decode shape. Preserve AITER tie ordering and emit Opus-compatible metadata directly from the routing launch. Assisted-by: Codex Signed-off-by: Yanyuan Qin --- .../flydsl/kernels/kimi_k3_b1_route_sort.py | 316 ++++++++++++++++++ aiter/ops/flydsl/kimi_k3_moe_route.py | 156 +++++++++ .../test_kimi_k3_b1_route_sort.py | 103 ++++++ 3 files changed, 575 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort.py create mode 100644 aiter/ops/flydsl/kimi_k3_moe_route.py create mode 100644 op_tests/flydsl_tests/test_kimi_k3_b1_route_sort.py diff --git a/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort.py b/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort.py new file mode 100644 index 0000000000..2170f560bf --- /dev/null +++ b/aiter/ops/flydsl/kernels/kimi_k3_b1_route_sort.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K3 B1 route and Opus-compatible sort metadata kernel for gfx950.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, 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, 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, +) + +_BLOCK_THREADS = 64 +_EXPERTS = 896 +_TOPK = 16 +_BLOCK_M = 32 +_SORTED_ROWS = _TOPK * _BLOCK_M +_LOG2E = 1.4426950408889634 + + +@fx.struct +class _RouteSortStorage: + biased_scores: fx.Array[fx.Float32, _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] + + +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_module(): + """Build the fixed-shape gfx950 route+metadata launcher.""" + + @flyc.kernel( + name="kimi_k3_b1_route_sort_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, + ): + 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) + + 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) + + lds = fx.SharedAllocator().allocate(_RouteSortStorage).peek() + biased_lds = lds.biased_scores.ptr + route_lds = lds.route_scores.ptr + selected_ids_lds = lds.selected_ids.ptr + selected_scores_lds = lds.selected_scores.ptr + + # Match biased_grouped_topk's vec4/thread traversal. The group-selection + # branch disappears because NUM_GRP == TOPK_GRP == 1. + 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 + ) + for lane_in_vec in range_constexpr(4): + expert = expert_base + arith.constant(lane_in_vec, type=i32) + in_range = arith.cmpi( + CmpIPredicate.ult, + expert, + arith.constant(_EXPERTS, type=i32), + ) + load_if = scf.IfOp(in_range) + with ir.InsertionPoint(load_if.then_block): + x = buffer_ops.buffer_load( + logits_rsrc, expert, vec_width=1, dtype=f32 + ) + bias_bf16 = buffer_ops.buffer_load( + bias_rsrc, expert, vec_width=1, dtype=T.bf16 + ) + 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], + [], + [], + ) + _lds_store(route_lds, sigmoid, expert) + _lds_store(biased_lds, ArithValue(sigmoid) + bias_f32, expert) + scf.YieldOp([]) + gpu.barrier() + + route_sum = ArithValue(c_zero_f32) + for k in range_constexpr(_TOPK): + local_max = ArithValue(c_neg_inf) + local_id = ArithValue(arith.constant(k, type=i32)) + 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 + ) + for lane_in_vec in range_constexpr(4): + expert = expert_base + arith.constant(lane_in_vec, type=i32) + in_range = arith.cmpi( + CmpIPredicate.ult, + expert, + arith.constant(_EXPERTS, type=i32), + ) + score_if = scf.IfOp(in_range, results_=[f32], has_else=True) + with ir.InsertionPoint(score_if.then_block): + score = _lds_load(biased_lds, expert) + scf.YieldOp([_raw(score)]) + with ir.InsertionPoint(score_if.else_block): + scf.YieldOp([c_neg_inf]) + score = ArithValue(score_if.results[0]) + take = arith.cmpf(CmpFPredicate.OGT, score, local_max) + local_max = ArithValue( + arith.select(take, _raw(score), _raw(local_max)) + ) + local_id = ArithValue( + arith.select(take, _raw(expert), _raw(local_id)) + ) + + # Reproduce hip_reduce.h's gfx9 DPP tree exactly. Its asymmetric + # tie behavior is observable: a fully uniform row starts at expert + # 252, while a tied prefix starts at expert 0. + 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(_BLOCK_THREADS - 1, type=i32), + ) + ) + selected_score = ArithValue(_lds_load(route_lds, selected_id)) + _lds_store(biased_lds, c_neg_inf, selected_id) + 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), + ) + + # Initialize every padded metadata row and zero the accumulation buffer. + 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) + 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, + 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, + ).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/kimi_k3_moe_route.py b/aiter/ops/flydsl/kimi_k3_moe_route.py new file mode 100644 index 0000000000..57fe7dc433 --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_moe_route.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Narrow gfx950 dispatch for Kimi-K3 B1 route+sort metadata fusion.""" + +import enum +import functools + +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + + +class KimiK3RouteSortDispatch(enum.Enum): + """Dispatch decision for the fixed production-shape specialization.""" + + FLYDSL_GFX950_B1 = "flydsl_gfx950_b1" + UNSUPPORTED = "unsupported" + + +KimiK3RouteSortResult = tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +] + + +def supports_kimi_k3_b1_route_sort( + logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + num_experts: int, + topk: int, + num_expert_group: int, + topk_group: int, + block_size_m: int, +) -> bool: + """Return whether the exact Kimi-K3 B1 FlyDSL specialization is safe.""" + + 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 num_experts == 896 + and topk == 16 + and num_expert_group == 1 + and topk_group == 1 + and block_size_m == 32 + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +def kimi_k3_route_sort_dispatch( + logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + num_experts: int, + topk: int, + num_expert_group: int, + topk_group: int, + block_size_m: int, +) -> KimiK3RouteSortDispatch: + """Choose the named route+sort implementation without changing fallbacks.""" + + if supports_kimi_k3_b1_route_sort( + logits, + correction_bias, + num_experts=num_experts, + topk=topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + block_size_m=block_size_m, + ): + return KimiK3RouteSortDispatch.FLYDSL_GFX950_B1 + return KimiK3RouteSortDispatch.UNSUPPORTED + + +@functools.cache +def _compiled_kimi_k3_b1_route_sort(): + from aiter.ops.flydsl.kernels.kimi_k3_b1_route_sort import ( + build_kimi_k3_b1_route_sort_module, + ) + + return build_kimi_k3_b1_route_sort_module() + + +def kimi_k3_b1_route_sort( + logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + model_dim: int, +) -> KimiK3RouteSortResult: + """Route one token and emit the standard AITER/Opus metadata tuple.""" + + decision = kimi_k3_route_sort_dispatch( + logits, + correction_bias, + num_experts=896, + topk=16, + num_expert_group=1, + topk_group=1, + block_size_m=32, + ) + if decision is not KimiK3RouteSortDispatch.FLYDSL_GFX950_B1: + raise NotImplementedError( + "kimi_k3_b1_route_sort only supports contiguous gfx950 FP32/BF16 " + "B1x896, topk=16, group=1/1, block_size_m=32" + ) + if model_dim <= 0 or model_dim % 2: + raise ValueError(f"model_dim must be positive and even, got {model_dim}") + + 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) + + _compiled_kimi_k3_b1_route_sort()( + 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, + stream=torch.cuda.current_stream(logits.device), + ) + return ( + topk_weights, + topk_ids, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + ) diff --git a/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort.py b/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort.py new file mode 100644 index 0000000000..8a94ba7ecf --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_b1_route_sort.py @@ -0,0 +1,103 @@ +# 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 import ( + KimiK3RouteSortDispatch, + kimi_k3_b1_route_sort, + kimi_k3_route_sort_dispatch, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_gfx_runtime() != "gfx950", + reason="Kimi-K3 route+sort specialization requires gfx950", +) + + +def _incumbent_route(logits, bias): + weights = torch.empty((1, 16), dtype=torch.float32, device=logits.device) + expert_ids = torch.empty((1, 16), dtype=torch.int32, device=logits.device) + aiter.biased_grouped_topk_hip( + logits, + bias.float(), + weights, + expert_ids, + 1, + 1, + True, + 1.0, + ) + return weights, expert_ids + + +@pytest.mark.parametrize("case", ["random_bias", "uniform", "tied_prefix"]) +def test_kimi_k3_b1_route_sort_matches_incumbent(case): + generator = torch.Generator(device="cpu").manual_seed(20260728) + logits = torch.randn((1, 896), generator=generator).cuda() + bias = (torch.randn(896, generator=generator) * 0.01).bfloat16().cuda() + if case == "uniform": + logits.zero_() + bias.zero_() + elif case == "tied_prefix": + logits.fill_(-16.0) + logits[:, :18] = 2.0 + bias.zero_() + + expected_weights, expected_ids = _incumbent_route(logits, bias) + ( + weights, + expert_ids, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + ) = kimi_k3_b1_route_sort(logits, bias, model_dim=128) + torch.cuda.synchronize() + + torch.testing.assert_close(expert_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(weights, expected_weights, rtol=1e-6, atol=1e-7) + torch.testing.assert_close( + num_valid_ids, torch.tensor([512, 1], dtype=torch.int32, device="cuda") + ) + torch.testing.assert_close( + sorted_expert_ids, torch.sort(expert_ids.flatten()).values + ) + assert torch.count_nonzero(moe_buf).item() == 0 + + sentinel = (16 << 24) | 1 + for rank, expert in enumerate(sorted_expert_ids.tolist()): + slot = torch.nonzero(expert_ids[0] == expert).item() + base = rank * 32 + assert sorted_ids[base].item() == slot << 24 + assert torch.all(sorted_ids[base + 1 : base + 32] == sentinel) + torch.testing.assert_close(sorted_weights[base], weights[0, slot]) + assert torch.count_nonzero(sorted_weights[base + 1 : base + 32]).item() == 0 + + +def test_kimi_k3_route_sort_dispatch_is_narrow(): + logits = torch.empty((1, 896), dtype=torch.float32, device="cuda") + bias = torch.empty(896, dtype=torch.bfloat16, device="cuda") + common = { + "num_experts": 896, + "topk": 16, + "num_expert_group": 1, + "topk_group": 1, + "block_size_m": 32, + } + assert ( + kimi_k3_route_sort_dispatch(logits, bias, **common) + is KimiK3RouteSortDispatch.FLYDSL_GFX950_B1 + ) + assert ( + kimi_k3_route_sort_dispatch(logits.expand(2, -1), bias, **common) + is KimiK3RouteSortDispatch.UNSUPPORTED + ) + assert ( + kimi_k3_route_sort_dispatch(logits, bias.float(), **common) + is KimiK3RouteSortDispatch.UNSUPPORTED + ) From d1a3ef562de5c4e13d944b1a54e4f8fa103b9fd3 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 11:54:11 +0000 Subject: [PATCH 3/9] perf(moe): specialize Kimi-K3 B1 gate projection Add a narrow gfx950 FlyDSL projection for the production 1x7168 by 896x7168 BF16 router shape. Preserve the GateLinear BF16 output contract, isolate dispatch behind a named capability predicate, and leave unsupported shapes on their existing paths. Assisted-by: Codex Signed-off-by: Yanyuan Qin --- .../kernels/kimi_k3_b1_gate_projection.py | 127 ++++++++++++++++++ aiter/ops/flydsl/kimi_k3_gate.py | 97 +++++++++++++ .../test_kimi_k3_b1_gate_projection.py | 110 +++++++++++++++ 3 files changed, 334 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/kimi_k3_b1_gate_projection.py create mode 100644 aiter/ops/flydsl/kimi_k3_gate.py create mode 100644 op_tests/flydsl_tests/test_kimi_k3_b1_gate_projection.py diff --git a/aiter/ops/flydsl/kernels/kimi_k3_b1_gate_projection.py b/aiter/ops/flydsl/kernels/kimi_k3_b1_gate_projection.py new file mode 100644 index 0000000000..437dd8862d --- /dev/null +++ b/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 diff --git a/aiter/ops/flydsl/kimi_k3_gate.py b/aiter/ops/flydsl/kimi_k3_gate.py new file mode 100644 index 0000000000..e90a1f81da --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_gate.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Narrow gfx950 dispatch for the Kimi-K3 B1 router projection.""" + +import enum +import functools + +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + + +class KimiK3GateProjectionDispatch(enum.Enum): + """Dispatch decision for the fixed production-shape projection.""" + + FLYDSL_GFX950_B1 = "flydsl_gfx950_b1" + UNSUPPORTED = "unsupported" + + +def supports_kimi_k3_b1_gate_projection( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, +) -> bool: + """Return whether the exact Kimi-K3 B1 projection is safe.""" + + return ( + hidden_states.is_cuda + and router_weight.is_cuda + and hidden_states.device == router_weight.device + and hidden_states.dtype == torch.bfloat16 + and router_weight.dtype == torch.bfloat16 + and hidden_states.is_contiguous() + and router_weight.is_contiguous() + and tuple(hidden_states.shape) == (1, 7168) + and tuple(router_weight.shape) == (896, 7168) + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +def kimi_k3_gate_projection_dispatch( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, +) -> KimiK3GateProjectionDispatch: + """Choose the named projection implementation without changing fallbacks.""" + + if supports_kimi_k3_b1_gate_projection(hidden_states, router_weight): + return KimiK3GateProjectionDispatch.FLYDSL_GFX950_B1 + return KimiK3GateProjectionDispatch.UNSUPPORTED + + +@functools.cache +def _compiled_kimi_k3_b1_gate_projection(): + from aiter.ops.flydsl.kernels.kimi_k3_b1_gate_projection import ( + build_kimi_k3_b1_gate_projection_module, + ) + + return build_kimi_k3_b1_gate_projection_module() + + +def kimi_k3_b1_gate_projection( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Project one hidden-state row to FP32 Kimi-K3 router logits.""" + + decision = kimi_k3_gate_projection_dispatch(hidden_states, router_weight) + if decision is not KimiK3GateProjectionDispatch.FLYDSL_GFX950_B1: + raise NotImplementedError( + "kimi_k3_b1_gate_projection only supports contiguous gfx950 BF16 " + "B1x7168 hidden states and 896x7168 router weights" + ) + if out is None: + out = torch.empty((1, 896), dtype=torch.float32, device=hidden_states.device) + elif ( + out.device != hidden_states.device + or out.dtype != torch.float32 + or not out.is_contiguous() + or tuple(out.shape) != (1, 896) + ): + raise ValueError( + "out must be contiguous FP32 shape (1, 896) on the input device" + ) + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + _compiled_kimi_k3_b1_gate_projection()( + ptr_arg(hidden_states), + ptr_arg(router_weight), + ptr_arg(out), + stream=torch.cuda.current_stream(hidden_states.device), + ) + return out diff --git a/op_tests/flydsl_tests/test_kimi_k3_b1_gate_projection.py b/op_tests/flydsl_tests/test_kimi_k3_b1_gate_projection.py new file mode 100644 index 0000000000..c3ee4dbea0 --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_b1_gate_projection.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch +import torch.nn.functional as F + +import aiter +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.kimi_k3_gate import ( + KimiK3GateProjectionDispatch, + kimi_k3_b1_gate_projection, + kimi_k3_gate_projection_dispatch, + supports_kimi_k3_b1_gate_projection, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_gfx_runtime() != "gfx950", + reason="Kimi-K3 gate projection specialization requires gfx950", +) + + +def _route(logits, bias): + weights = torch.empty((1, 16), dtype=torch.float32, device=logits.device) + expert_ids = torch.empty((1, 16), dtype=torch.int32, device=logits.device) + aiter.biased_grouped_topk_hip( + logits, + bias.float(), + weights, + expert_ids, + 1, + 1, + True, + 1.0, + ) + return weights, expert_ids + + +@pytest.mark.parametrize("seed", [1, 17, 20260728]) +def test_kimi_k3_b1_gate_projection_matches_gate_linear(seed): + generator = torch.Generator(device="cpu").manual_seed(seed) + hidden = torch.randn((1, 7168), generator=generator).bfloat16().cuda() + weight = ( + torch.randn((896, 7168), generator=generator).mul_(7168**-0.5).bfloat16().cuda() + ) + bias = torch.randn(896, generator=generator).mul_(0.01).bfloat16().cuda() + + expected_logits = F.linear(hidden, weight).float() + expected_weights, expected_ids = _route(expected_logits, bias) + actual_logits = kimi_k3_b1_gate_projection(hidden, weight) + actual_weights, actual_ids = _route(actual_logits, bias) + torch.cuda.synchronize() + + torch.testing.assert_close(actual_logits, expected_logits, rtol=0, atol=0) + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_weights, expected_weights, rtol=1e-5, atol=0) + + +def test_kimi_k3_gate_projection_dispatch_is_narrow(): + hidden = torch.empty((1, 7168), dtype=torch.bfloat16, device="cuda") + weight = torch.empty((896, 7168), dtype=torch.bfloat16, device="cuda") + noncontiguous_weight = torch.empty( + (7168, 896), dtype=torch.bfloat16, device="cuda" + ).T + + assert supports_kimi_k3_b1_gate_projection(hidden, weight) + assert ( + kimi_k3_gate_projection_dispatch(hidden, weight) + is KimiK3GateProjectionDispatch.FLYDSL_GFX950_B1 + ) + assert ( + kimi_k3_gate_projection_dispatch(hidden.expand(2, -1), weight) + is KimiK3GateProjectionDispatch.UNSUPPORTED + ) + assert ( + kimi_k3_gate_projection_dispatch(hidden, weight.float()) + is KimiK3GateProjectionDispatch.UNSUPPORTED + ) + assert ( + kimi_k3_gate_projection_dispatch(hidden, noncontiguous_weight) + is KimiK3GateProjectionDispatch.UNSUPPORTED + ) + + +def test_kimi_k3_gate_projection_reuses_valid_output(): + hidden = torch.randn((1, 7168), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((896, 7168), dtype=torch.bfloat16, device="cuda") + out = torch.empty((1, 896), dtype=torch.float32, device="cuda") + + actual = kimi_k3_b1_gate_projection(hidden, weight, out=out) + torch.cuda.synchronize() + + assert actual is out + torch.testing.assert_close(actual, F.linear(hidden, weight).float(), rtol=0, atol=0) + + +@pytest.mark.parametrize( + "shape,dtype", + [ + ((896,), torch.float32), + ((1, 896), torch.bfloat16), + ], +) +def test_kimi_k3_gate_projection_rejects_invalid_output(shape, dtype): + hidden = torch.empty((1, 7168), dtype=torch.bfloat16, device="cuda") + weight = torch.empty((896, 7168), dtype=torch.bfloat16, device="cuda") + out = torch.empty(shape, dtype=dtype, device="cuda") + + with pytest.raises(ValueError, match="out must be contiguous FP32"): + kimi_k3_b1_gate_projection(hidden, weight, out=out) From c257ea65092652698386cd8736ff1b63f4505c1f Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 13:27:43 +0000 Subject: [PATCH 4/9] perf(moe): add Kimi-K3 B1 stage1 direct tile Register one validated BF16 x MXFP4 16-row FlyDSL variant and advance it with the producer 32-row sort-bucket stride. Select it only through the existing exact gfx950 Kimi-K3 B1 tuning row. The complete expert boundary improves from 20.70 us to 18.20 us across two independent processes, with 90/90 adversarial correctness cases passing. Signed-off-by: Yanyuan Qin --- .../model_configs/kimik3_fp4_tuned_fmoe.csv | 2 +- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 5 +- aiter/ops/flydsl/moe_kernels.py | 27 +++++++++ .../test_kimi_k3_b1_moe_stage1.py | 58 +++++++++++++++++++ 4 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage1.py diff --git a/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv b/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv index 7f4937fc79..8e4bf05ede 100644 --- a/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv +++ b/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%,9.4261,flydsl_moe2_abf16_wfp4_bf16_t32x128x128_reduce_bnt2_persist,0.0%,19.9805,0,0,0,6.61,185150.05, 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, diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 3aa2809d30..11a99ee0fb 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -5508,7 +5508,10 @@ def _act_elem(g_e, u_e): by = gpu.block_id("x") bx = gpu.block_id("y") - bx_m = bx * arith.index(tile_m) + # Route metadata remains in the standard 32-row expert buckets + # even when the B1 direct path issues a native 16-row MFMA tile. + # Advance by the producer's bucket stride, not the compute tile. + bx_m = bx * arith.index(sort_block_m) numids_rsrc = ptr_buffer_resource( arg_num_valid_ids, arith.constant(4, type=i32) ) diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 56a5fac5e1..d1a00417f5 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -256,9 +256,36 @@ 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 hand-tuned stage1 variants to ``kernels`` in-place.""" + if (a_dtype, b_dtype, out_dtype) != ("bf16", "fp4", "bf16"): + return + + name = "flydsl_moe1_abf16_wfp4_bf16_t16x64x256_w3_xcd4" + kernels[name] = { + "stage": 1, + "a_dtype": a_dtype, + "b_dtype": b_dtype, + "out_dtype": out_dtype, + "tile_m": 16, + "tile_n": 64, + "tile_k": 256, + "MPerBlock": 16, + "waves_per_eu": 3, + "k_batch": 1, + "b_nt": 2, + "gate_mode": "separated", + "xcd_swizzle": 4, + "k_wave": 1, + } + + def get_flydsl_stage2_kernels( a_dtype: str, b_dtype: str, out_dtype: str ) -> Dict[str, Dict]: diff --git a/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage1.py b/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage1.py new file mode 100644 index 0000000000..ca657f0b0c --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage1.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Regression coverage for the Kimi-K3 gfx950 B1 stage-1 configuration.""" + +import csv +from pathlib import Path + +from aiter.ops.flydsl.moe_kernels import get_flydsl_kernel_params + +STAGE1_KERNEL = "flydsl_moe1_abf16_wfp4_bf16_t16x64x256_w3_xcd4" +CONFIG_PATH = ( + Path(__file__).parents[2] / "aiter" / "configs" / "model_configs" / "kimik3_fp4_tuned_fmoe.csv" +) + + +def test_kimi_k3_b1_stage1_variant_is_narrowly_registered(): + params = get_flydsl_kernel_params(STAGE1_KERNEL) + assert params is not None + assert params["stage"] == 1 + assert params["a_dtype"] == "bf16" + assert params["b_dtype"] == "fp4" + assert params["out_dtype"] == "bf16" + assert params["tile_m"] == 16 + assert params["tile_n"] == 64 + assert params["tile_k"] == 256 + assert params["gate_mode"] == "separated" + + +def test_kimi_k3_b1_selects_stage1_variant_through_tuned_config(): + with CONFIG_PATH.open(newline="", encoding="utf-8") as config_file: + rows = list(csv.DictReader(config_file)) + + matching = [ + row + for row in rows + if row["gfx"] == "gfx950" + and row["cu_num"] == "256" + and row["token"] == "1" + and row["model_dim"] == "3584" + and row["inter_dim"] == "384" + and row["expert"] == "896" + and row["topk"] == "16" + ] + assert len(matching) == 1 + assert matching[0]["kernelName1"] == STAGE1_KERNEL + + other_batches = [ + row + for row in rows + if row["gfx"] == "gfx950" + and row["model_dim"] == "3584" + and row["inter_dim"] == "384" + and row["expert"] == "896" + and row["topk"] == "16" + and row["token"] != "1" + ] + assert all(row["kernelName1"] != STAGE1_KERNEL for row in other_batches) From 8687dca5459d1645b2915b4d1ab4ef1f52fc8aa3 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 13:27:51 +0000 Subject: [PATCH 5/9] perf(moe): use atomic Kimi-K3 B1 stage2 reduction Select the existing direct atomic weighted-reduction kernel through the exact gfx950 Kimi-K3 B1 tuning row. Larger batch rows and generic dispatch remain unchanged. The complete expert boundary improves from 20.70 us to 18.62 us across two independent processes, with 90/90 adversarial correctness cases passing. Signed-off-by: Yanyuan Qin --- .../model_configs/kimik3_fp4_tuned_fmoe.csv | 2 +- .../test_kimi_k3_b1_moe_stage2.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage2.py diff --git a/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv b/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv index 8e4bf05ede..bab2d39acf 100644 --- a/aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv +++ b/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,10.5544,flydsl_moe1_abf16_wfp4_bf16_t16x64x256_w3_xcd4,34.6%,9.4261,flydsl_moe2_abf16_wfp4_bf16_t32x128x128_reduce_bnt2_persist,0.0%,19.9805,0,0,0,6.61,185150.05, +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, diff --git a/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage2.py b/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage2.py new file mode 100644 index 0000000000..bcec81935f --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_b1_moe_stage2.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Regression coverage for the Kimi-K3 gfx950 B1 stage-2 configuration.""" + +import csv +from pathlib import Path + +STAGE2_KERNEL = "flydsl_moe2_abf16_wfp4_bf16_t32x128x128_atomic_bnt2" +CONFIG_PATH = ( + Path(__file__).parents[2] / "aiter" / "configs" / "model_configs" / "kimik3_fp4_tuned_fmoe.csv" +) + + +def test_kimi_k3_b1_selects_atomic_stage2_through_tuned_config(): + with CONFIG_PATH.open(newline="", encoding="utf-8") as config_file: + rows = list(csv.DictReader(config_file)) + + matching = [ + row + for row in rows + if row["gfx"] == "gfx950" + and row["cu_num"] == "256" + and row["token"] == "1" + and row["model_dim"] == "3584" + and row["inter_dim"] == "384" + and row["expert"] == "896" + and row["topk"] == "16" + ] + assert len(matching) == 1 + assert matching[0]["kernelName2"] == STAGE2_KERNEL + + other_batches = [ + row + for row in rows + if row["gfx"] == "gfx950" + and row["model_dim"] == "3584" + and row["inter_dim"] == "384" + and row["expert"] == "896" + and row["topk"] == "16" + and row["token"] != "1" + ] + assert all(row["kernelName2"] != STAGE2_KERNEL for row in other_batches) From 15d1d1744f6baf218992573221199682789725cd Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 03:32:06 +0000 Subject: [PATCH 6/9] perf(flydsl): fuse Kimi-K3 KDA decode on gfx950 --- 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 b91158f5ad..087aa0ab69 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -40,6 +40,10 @@ from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8 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, + 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, @@ -61,6 +65,8 @@ "flydsl_moe_stage2", "flydsl_hgemm", "flydsl_flash_attn_func", + "flydsl_kimi_k3_kda_decode", + "is_flydsl_kimi_k3_kda_decode_supported", "flydsl_qk_norm_rope_quant", "flydsl_pa_mqa_logits_fp4", "flydsl_pa_mqa_logits_fp4_prefill", 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 0000000000..56a6c51b9a --- /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 0000000000..2c5d51a96f --- /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 0000000000..06395c3393 --- /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 95dfd13ee87ce698edfdcbc75f16d92b7feebb9a Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 09:41:18 +0000 Subject: [PATCH 7/9] perf(flydsl): fuse Kimi-K3 f_b projection into KDA decode Signed-off-by: Yanyuan Qin --- 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 087aa0ab69..529de9d7fb 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -42,6 +42,7 @@ 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 @@ -66,6 +67,7 @@ "flydsl_hgemm", "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", 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 0000000000..76e4fec295 --- /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 2c5d51a96f..0c44949873 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 06395c3393..1b6741ce0b 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 d4fb7ab2b4bdf4328094a6b4e9edfa4d903b29bb Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 14:46:16 +0000 Subject: [PATCH 8/9] perf(flydsl): fuse Kimi-K3 B1 latent MoE tail --- aiter/ops/flydsl/__init__.py | 3 + .../flydsl/kernels/latent_moe_tail_gfx950.py | 374 ++++++++++++++++++ aiter/ops/flydsl/latent_moe_tail.py | 144 +++++++ op_tests/flydsl_tests/test_latent_moe_tail.py | 118 ++++++ 4 files changed, 639 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/latent_moe_tail_gfx950.py create mode 100644 aiter/ops/flydsl/latent_moe_tail.py create mode 100644 op_tests/flydsl_tests/test_latent_moe_tail.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 529de9d7fb..4dae2ba3ae 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -38,6 +38,7 @@ ) 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 ( @@ -65,6 +66,8 @@ "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", diff --git a/aiter/ops/flydsl/kernels/latent_moe_tail_gfx950.py b/aiter/ops/flydsl/kernels/latent_moe_tail_gfx950.py new file mode 100644 index 0000000000..38dab0aa28 --- /dev/null +++ b/aiter/ops/flydsl/kernels/latent_moe_tail_gfx950.py @@ -0,0 +1,374 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Fixed-shape gfx950 BF16 RMSNorm, GEMV, and add kernel.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, vector +from flydsl.expr import math as fmath +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.typing import T +from flydsl.expr.vector import ReductionOp + +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) + +_LATENT_DIM = 3584 +_HIDDEN_DIM = 7168 +_WAVE_SIZE = 64 + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def _lds_load(ptr, index): + return fx.ptr_load(ptr + fx.Int64(index)) + + +def _lds_store(ptr, value, index): + fx.ptr_store(value, ptr + fx.Int64(index)) + + +def build_b1_latent_moe_tail_module( + rows_per_block: int = 4, + waves_per_eu: int = 0, + normalize_in_kernel: bool = True, + elements_per_thread: int = 8, + use_dot2: bool = True, + weight_cache_modifier: int = 0, +): + """Build a B1 launcher with one multi-row workgroup per output tile.""" + + if not 2 <= rows_per_block <= 64: + raise ValueError("rows_per_block must be between 2 and 64") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if elements_per_thread not in (8, 16, 32): + raise ValueError("elements_per_thread must be 8, 16, or 32") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + block_threads = ( + ((_LATENT_DIM + elements_per_thread - 1) // elements_per_thread + _WAVE_SIZE - 1) + // _WAVE_SIZE + * _WAVE_SIZE + ) + waves = block_threads // _WAVE_SIZE + vectors_per_thread = elements_per_thread // 8 + + @fx.struct + class TailStorage: + rms_sums: fx.Array[fx.Float32, waves, 16] + inverse_rms: fx.Array[fx.Float32, 1, 16] + dot_sums: fx.Array[fx.Float32, rows_per_block * waves, 16] + + kernel_name = ( + f"latent_moe_tail_b1_bf16_gfx950_r{rows_per_block}" + f"_wpe{waves_per_eu}_norm{int(normalize_in_kernel)}" + f"_ept{elements_per_thread}" + f"_dot2{int(use_dot2)}" + f"_wcm{weight_cache_modifier}" + ) + + @flyc.kernel( + name=kernel_name, + known_block_size=[block_threads, 1, 1], + ) + def tail_kernel( + routed: fx.Pointer, + shared: fx.Pointer, + rms_weight: fx.Pointer, + up_weight: fx.Pointer, + output: fx.Pointer, + epsilon: fx.Float32, + ): + i32 = T.i32 + f32 = T.f32 + i1 = ir.IntegerType.get_signless(1) + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + output_base = ArithValue(gpu.block_idx.x) * arith.constant(rows_per_block, type=i32) + k_base = tid * arith.constant(elements_per_thread, type=i32) + + routed_rsrc = ptr_rsrc(routed) + shared_rsrc = ptr_rsrc(shared) + rms_weight_rsrc = ptr_rsrc(rms_weight) + up_weight_rsrc = ptr_rsrc(up_weight) + output_rsrc = ptr_rsrc(output) + lds = fx.SharedAllocator().allocate(TailStorage).peek() + rms_sums = lds.rms_sums.ptr + inverse_rms = lds.inverse_rms.ptr + dot_sums = lds.dot_sums.ptr + + zero_f32 = arith.constant(0.0, type=f32) + one_over_dim = arith.constant(1.0 / _LATENT_DIM, type=f32) + vec8_bf16 = T.vec(8, T.bf16) + vec8_f32 = T.vec(8, f32) + zero_bf16_vec = vector.broadcast(vec8_bf16, arith.constant(0.0, type=T.bf16)) + + def load_bf16x8(resource, dword_index, cache_modifier=0): + dwords = buffer_ops.buffer_load( + resource, + dword_index, + vec_width=4, + dtype=i32, + cache_modifier=cache_modifier, + ) + return vector.bitcast(vec8_bf16, dwords) + + def load_bf16x8_masked(resource, element_index): + if const_expr(block_threads * elements_per_thread == _LATENT_DIM): + return load_bf16x8(resource, element_index // arith.constant(2, type=i32)) + valid = arith.cmpi( + CmpIPredicate.ult, + element_index, + arith.constant(_LATENT_DIM, type=i32), + ) + load_if = scf.IfOp(valid, results_=[vec8_bf16], has_else=True) + with ir.InsertionPoint(load_if.then_block): + loaded = load_bf16x8(resource, element_index // arith.constant(2, type=i32)) + scf.YieldOp([_raw(loaded)]) + with ir.InsertionPoint(load_if.else_block): + scf.YieldOp([_raw(zero_bf16_vec)]) + return load_if.results[0] + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith.AddFOp(reduced, peer, fastmath=fm_fast).result + return reduced + + def dot_bf16x8(left, right): + dot = zero_f32 + for pair_index in range_constexpr(4): + left_pair = vector.from_elements( + T.vec(2, T.bf16), + [ + vector.extract( + left, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + left, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + right_pair = vector.from_elements( + T.vec(2, T.bf16), + [ + vector.extract( + right, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + right, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + dot = llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + left_pair, + right_pair, + dot, + arith.constant(False, type=i1), + ], + [], + [], + ) + return dot + + routed_f32_vectors = [] + for vector_index in range_constexpr(vectors_per_thread): + element_index = k_base + arith.constant(vector_index * 8, type=i32) + routed_bf16 = load_bf16x8_masked(routed_rsrc, element_index) + routed_f32_vectors.append(ArithValue(routed_bf16).extf(vec8_f32)) + is_lane_zero = arith.cmpi(CmpIPredicate.eq, lane, arith.constant(0, type=i32)) + if const_expr(normalize_in_kernel): + local_square_sum = ArithValue(zero_f32) + for routed_f32 in routed_f32_vectors: + local_square_sum = local_square_sum + (routed_f32 * routed_f32).reduce( + ReductionOp.ADD, fastmath=fm_fast + ) + wave_square_sum = wave_reduce_add(local_square_sum) + lane_zero_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(lane_zero_if.then_block): + _lds_store(rms_sums, wave_square_sum, wave) + scf.YieldOp([]) + gpu.barrier() + + is_thread_zero = arith.cmpi(CmpIPredicate.eq, tid, arith.constant(0, type=i32)) + thread_zero_if = scf.IfOp(is_thread_zero) + with ir.InsertionPoint(thread_zero_if.then_block): + total_square_sum = ArithValue(zero_f32) + for wave_index in range_constexpr(waves): + total_square_sum = total_square_sum + _lds_load( + rms_sums, arith.constant(wave_index, type=i32) + ) + variance = total_square_sum * one_over_dim + inverse = fmath.rsqrt(variance + ArithValue(epsilon), fastmath=fm_fast) + _lds_store(inverse_rms, _raw(inverse), arith.constant(0, type=i32)) + scf.YieldOp([]) + gpu.barrier() + + inverse = ArithValue(_lds_load(inverse_rms, arith.constant(0, type=i32))) + normalized_dot_vectors = [] + normalized_bf16_vectors = [] + for vector_index in range_constexpr(vectors_per_thread): + element_index = k_base + arith.constant(vector_index * 8, type=i32) + gamma_bf16 = load_bf16x8_masked(rms_weight_rsrc, element_index) + gamma_f32 = ArithValue(gamma_bf16).extf(vec8_f32) + normalized_f32 = routed_f32_vectors[vector_index] * gamma_f32 * inverse + normalized_bf16 = normalized_f32.truncf(vec8_bf16) + normalized_bf16_vectors.append(normalized_bf16) + normalized_dot_vectors.append(normalized_bf16.extf(vec8_f32)) + else: + normalized_dot_vectors = routed_f32_vectors + + accumulators = [] + for row_index in range_constexpr(rows_per_block): + row = output_base + arith.constant(row_index, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_HIDDEN_DIM, type=i32), + ) + safe_row = arith.select(row_in_range, row, arith.constant(0, type=i32)) + local_dot = ArithValue(zero_f32) + for vector_index in range_constexpr(vectors_per_thread): + row_element = k_base + arith.constant(vector_index * 8, type=i32) + weight_element = safe_row * arith.constant(_LATENT_DIM, type=i32) + row_element + if const_expr(block_threads * elements_per_thread == _LATENT_DIM): + weight_bf16 = load_bf16x8( + up_weight_rsrc, + weight_element // arith.constant(2, type=i32), + weight_cache_modifier, + ) + else: + valid = arith.cmpi( + CmpIPredicate.ult, + row_element, + arith.constant(_LATENT_DIM, type=i32), + ) + weight_if = scf.IfOp(valid, results_=[vec8_bf16], has_else=True) + with ir.InsertionPoint(weight_if.then_block): + loaded_weight = load_bf16x8( + up_weight_rsrc, + weight_element // arith.constant(2, type=i32), + weight_cache_modifier, + ) + scf.YieldOp([_raw(loaded_weight)]) + with ir.InsertionPoint(weight_if.else_block): + scf.YieldOp([_raw(zero_bf16_vec)]) + weight_bf16 = weight_if.results[0] + if const_expr(use_dot2 and normalize_in_kernel): + local_dot = local_dot + dot_bf16x8( + normalized_bf16_vectors[vector_index], + weight_bf16, + ) + else: + weight_f32 = ArithValue(weight_bf16).extf(vec8_f32) + local_dot = local_dot + ( + normalized_dot_vectors[vector_index] * weight_f32 + ).reduce(ReductionOp.ADD, fastmath=fm_fast) + accumulators.append(wave_reduce_add(local_dot)) + + lane_zero_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(lane_zero_if.then_block): + for row_index in range_constexpr(rows_per_block): + index = arith.constant(row_index * waves, type=i32) + wave + _lds_store(dot_sums, accumulators[row_index], index) + scf.YieldOp([]) + gpu.barrier() + + writes_output = arith.cmpi( + CmpIPredicate.ult, + tid, + arith.constant(rows_per_block, type=i32), + ) + output_index = output_base + tid + output_in_range = arith.cmpi( + CmpIPredicate.ult, + output_index, + arith.constant(_HIDDEN_DIM, type=i32), + ) + writes_output = arith.andi(writes_output, output_in_range) + write_if = scf.IfOp(writes_output) + with ir.InsertionPoint(write_if.then_block): + dot = ArithValue(zero_f32) + for wave_index in range_constexpr(waves): + index = tid * arith.constant(waves, type=i32) + arith.constant(wave_index, type=i32) + dot = dot + _lds_load(dot_sums, index) + projected_bf16 = arith.trunc_f(T.bf16, _raw(dot)) + projected_f32 = ArithValue(arith.extf(f32, projected_bf16)) + shared_bf16 = buffer_ops.buffer_load( + shared_rsrc, output_index, vec_width=1, dtype=T.bf16 + ) + shared_f32 = ArithValue(arith.extf(f32, shared_bf16)) + result = arith.trunc_f(T.bf16, _raw(projected_f32 + shared_f32)) + buffer_ops.buffer_store(result, output_rsrc, output_index) + scf.YieldOp([]) + + @flyc.jit + def launch_tail( + routed: fx.Pointer, + shared: fx.Pointer, + rms_weight: fx.Pointer, + up_weight: fx.Pointer, + output: fx.Pointer, + epsilon: fx.Float32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in ctx.gpu_module_body.operations: + if hasattr(operation, "attributes") and operation.OPERATION_NAME == "gpu.func": + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, int(waves_per_eu) + ) + tail_kernel( + routed, + shared, + rms_weight, + up_weight, + output, + epsilon, + ).launch( + grid=( + (_HIDDEN_DIM + rows_per_block - 1) // rows_per_block, + 1, + 1, + ), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_tail.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_tail diff --git a/aiter/ops/flydsl/latent_moe_tail.py b/aiter/ops/flydsl/latent_moe_tail.py new file mode 100644 index 0000000000..d9e8987e9c --- /dev/null +++ b/aiter/ops/flydsl/latent_moe_tail.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Narrow BF16 latent-MoE local-tail primitive.""" + +import functools +import math + +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +_LATENT_DIM = 3584 +_HIDDEN_DIM = 7168 +_B1_ROWS_PER_BLOCK = 14 +_B1_WAVES_PER_EU = 4 +# Policy 2 bypasses cache levels that would otherwise retain the one-use +# 49 MiB projection matrix. It is the same FlyDSL cache modifier used by +# existing streamed mixed-MoE weight loads. +_B1_WEIGHT_CACHE_MODIFIER = 2 + + +def supports_latent_moe_tail( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + epsilon: float, +) -> bool: + """Return whether the fixed gfx950 BF16 primitive supports these tensors.""" + + tensors = (routed, shared, rms_weight, up_weight) + return ( + all(tensor.is_cuda for tensor in tensors) + and len({tensor.device for tensor in tensors}) == 1 + and all(tensor.dtype == torch.bfloat16 for tensor in tensors) + and all(tensor.is_contiguous() for tensor in tensors) + and tuple(routed.shape) == (1, _LATENT_DIM) + and tuple(shared.shape) == (1, _HIDDEN_DIM) + and tuple(rms_weight.shape) == (_LATENT_DIM,) + and tuple(up_weight.shape) == (_HIDDEN_DIM, _LATENT_DIM) + and math.isfinite(epsilon) + and epsilon > 0.0 + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +@functools.cache +def _compiled_b1_latent_moe_tail( + rows_per_block: int, + waves_per_eu: int, + normalize_in_kernel: bool, + elements_per_thread: int, + use_dot2: bool, + weight_cache_modifier: int, +): + from aiter.ops.flydsl.kernels.latent_moe_tail_gfx950 import ( + build_b1_latent_moe_tail_module, + ) + + return build_b1_latent_moe_tail_module( + rows_per_block, + waves_per_eu, + normalize_in_kernel, + elements_per_thread, + use_dot2, + weight_cache_modifier, + ) + + +def _launch_b1_latent_moe_tail( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + epsilon: float, + *, + out: torch.Tensor, + rows_per_block: int, + waves_per_eu: int, + normalize_in_kernel: bool = True, + elements_per_thread: int = 8, + use_dot2: bool = True, + weight_cache_modifier: int = 0, +) -> torch.Tensor: + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + _compiled_b1_latent_moe_tail( + rows_per_block, + waves_per_eu, + normalize_in_kernel, + elements_per_thread, + use_dot2, + weight_cache_modifier, + )( + ptr_arg(routed), + ptr_arg(shared), + ptr_arg(rms_weight), + ptr_arg(up_weight), + ptr_arg(out), + float(epsilon), + stream=torch.cuda.current_stream(routed.device), + ) + return out + + +def latent_moe_tail( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + epsilon: float, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Fuse BF16 RMSNorm, FP32-accumulated projection, and BF16 shared add.""" + + if not supports_latent_moe_tail(routed, shared, rms_weight, up_weight, epsilon): + raise NotImplementedError( + "latent_moe_tail requires contiguous gfx950 BF16 tensors with " + "shapes (1,3584), (1,7168), (3584,), and (7168,3584)" + ) + if out is None: + out = torch.empty_like(shared) + elif ( + out.device != routed.device + or out.dtype != torch.bfloat16 + or not out.is_contiguous() + or tuple(out.shape) != (1, _HIDDEN_DIM) + ): + raise ValueError("out must be contiguous BF16 shape (1, 7168) on the input device") + return _launch_b1_latent_moe_tail( + routed, + shared, + rms_weight, + up_weight, + epsilon, + out=out, + rows_per_block=_B1_ROWS_PER_BLOCK, + waves_per_eu=_B1_WAVES_PER_EU, + weight_cache_modifier=_B1_WEIGHT_CACHE_MODIFIER, + ) diff --git a/op_tests/flydsl_tests/test_latent_moe_tail.py b/op_tests/flydsl_tests/test_latent_moe_tail.py new file mode 100644 index 0000000000..ab3be9535d --- /dev/null +++ b/op_tests/flydsl_tests/test_latent_moe_tail.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch + +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.latent_moe_tail import ( + latent_moe_tail, + supports_latent_moe_tail, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_gfx_runtime() != "gfx950", + reason="Kimi-K3 latent-MoE local-tail specialization requires gfx950", +) + +LATENT_DIM = 3584 +HIDDEN_DIM = 7168 +EPSILON = 1.0e-6 + + +def _inputs(seed: int = 20260728): + generator = torch.Generator(device="cpu").manual_seed(seed) + routed = torch.randn((1, LATENT_DIM), generator=generator).bfloat16().cuda() + shared = torch.randn((1, HIDDEN_DIM), generator=generator).bfloat16().cuda() + rms_weight = torch.randn(LATENT_DIM, generator=generator).bfloat16().cuda() + up_weight = ( + torch.randn((HIDDEN_DIM, LATENT_DIM), generator=generator, dtype=torch.float32) + .mul_(LATENT_DIM**-0.5) + .bfloat16() + .cuda() + ) + return routed, shared, rms_weight, up_weight + + +def _oracle(routed, shared, rms_weight, up_weight): + inverse_rms = torch.rsqrt(routed.float().square().mean(dim=-1, keepdim=True) + EPSILON) + normalized = (routed.float() * inverse_rms * rms_weight.float()).bfloat16() + projected = torch.mm(normalized.float(), up_weight.float().T).bfloat16() + return (projected.float() + shared.float()).bfloat16() + + +@pytest.mark.parametrize("seed", [1, 17, 20260728]) +def test_latent_moe_tail_matches_explicit_fp32_oracle(seed): + routed, shared, rms_weight, up_weight = _inputs(seed) + routed_before = routed.clone() + shared_before = shared.clone() + + actual = latent_moe_tail(routed, shared, rms_weight, up_weight, EPSILON) + expected = _oracle(routed, shared, rms_weight, up_weight) + torch.cuda.synchronize() + + torch.testing.assert_close(actual, expected, rtol=0.01, atol=0.015625) + torch.testing.assert_close(routed, routed_before, rtol=0, atol=0) + torch.testing.assert_close(shared, shared_before, rtol=0, atol=0) + + +def test_latent_moe_tail_support_predicate_is_narrow(): + routed, shared, rms_weight, up_weight = _inputs() + noncontiguous = torch.empty((LATENT_DIM, 2), dtype=torch.bfloat16, device="cuda")[ + :, 0 + ].unsqueeze(0) + + assert supports_latent_moe_tail(routed, shared, rms_weight, up_weight, EPSILON) + assert not supports_latent_moe_tail( + routed.expand(8, -1), + shared.expand(8, -1), + rms_weight, + up_weight, + EPSILON, + ) + assert not supports_latent_moe_tail( + routed.expand(16, -1), + shared.expand(16, -1), + rms_weight, + up_weight, + EPSILON, + ) + assert not supports_latent_moe_tail(noncontiguous, shared, rms_weight, up_weight, EPSILON) + assert not supports_latent_moe_tail(routed, shared, rms_weight, up_weight.float(), EPSILON) + assert not supports_latent_moe_tail(routed, shared, rms_weight, up_weight, float("nan")) + assert not supports_latent_moe_tail(routed, shared, rms_weight, up_weight, float("inf")) + assert not supports_latent_moe_tail(routed, shared, rms_weight, up_weight, 0.0) + + +def test_latent_moe_tail_rejects_noncontiguous_input(): + _, shared, rms_weight, up_weight = _inputs() + routed = torch.empty((LATENT_DIM, 2), dtype=torch.bfloat16, device="cuda")[:, 0].unsqueeze(0) + + with pytest.raises(NotImplementedError, match="requires contiguous gfx950"): + latent_moe_tail(routed, shared, rms_weight, up_weight, EPSILON) + + +def test_latent_moe_tail_graph_capture_and_output_reuse(): + routed, shared, rms_weight, up_weight = _inputs() + out = torch.empty_like(shared) + graph = torch.cuda.CUDAGraph() + + with torch.cuda.graph(graph): + actual = latent_moe_tail( + routed, + shared, + rms_weight, + up_weight, + EPSILON, + out=out, + ) + graph.replay() + torch.cuda.synchronize() + + assert actual is out + torch.testing.assert_close( + actual, + _oracle(routed, shared, rms_weight, up_weight), + rtol=0.01, + atol=0.015625, + ) From 459601b680a3b23e13dc3203fae5d772b37022ba Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 16:42:32 +0000 Subject: [PATCH 9/9] perf(flydsl): consume Kimi-K3 prepared MoE routing Add a typed, fail-closed gfx950 B1 handoff that owns Kimi-K3 routing and feeds its prepared metadata directly into the existing MXFP4 expert stages. Support both production A8W4/interleaved and A16W4/separated layouts while preserving fallback outside the exact shape contract. The live A8W4 route-to-expert boundary improves from 41.83-42.01 us to 25.40-25.41 us (1.647-1.653x), with exact route IDs, 6.66e-8 routing-weight RRMSE, zero output RRMSE, and no duplicate route/sort launch. Assisted-by: Codex Signed-off-by: Yanyuan Qin --- aiter/ops/flydsl/kimi_k3_moe_handoff.py | 254 ++++++++++++++++++ .../flydsl_tests/test_kimi_k3_moe_handoff.py | 250 +++++++++++++++++ 2 files changed, 504 insertions(+) create mode 100644 aiter/ops/flydsl/kimi_k3_moe_handoff.py create mode 100644 op_tests/flydsl_tests/test_kimi_k3_moe_handoff.py diff --git a/aiter/ops/flydsl/kimi_k3_moe_handoff.py b/aiter/ops/flydsl/kimi_k3_moe_handoff.py new file mode 100644 index 0000000000..e28ba62247 --- /dev/null +++ b/aiter/ops/flydsl/kimi_k3_moe_handoff.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Typed Kimi-K3 route-to-MXFP4-expert ownership boundary.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from enum import Enum + +import torch + +from aiter import ActivationType, QuantType, dtypes +from aiter.ops.flydsl.kimi_k3_moe_route import ( + kimi_k3_b1_route_sort, + supports_kimi_k3_b1_route_sort, +) +from aiter.ops.flydsl.moe_common import GateMode + +_EXPERTS = 896 +_TOPK = 16 +_MODEL_DIM = 3584 +_INTER_DIM = 384 +_BLOCK_M = 32 +_W1_SCALE_SHAPE = (_EXPERTS * 2 * _INTER_DIM, _MODEL_DIM // 32) +_W2_SCALE_SHAPE = ( + _EXPERTS * _MODEL_DIM, + ((_INTER_DIM // 32 + 7) // 8) * 8, +) + + +class KimiK3MXFP4W13Layout(str, Enum): + """Stable stage-1 weight layouts selected during vLLM weight loading.""" + + GATE_UP_SEPARATED_PRESHUFFLED = "gate_up_separated_preshuffled" + GATE_UP_INTERLEAVED_PRESHUFFLED = "gate_up_interleaved_preshuffled" + + +@dataclass(frozen=True) +class KimiK3MXFP4ExpertMode: + """Activation and W13 layout contract for one Kimi-K3 expert mode.""" + + activation_dtype: torch.dtype + gate_mode: GateMode + w13_layout: KimiK3MXFP4W13Layout + + +def kimi_k3_mxfp4_expert_mode() -> KimiK3MXFP4ExpertMode: + """Own the live A8W4/A16W4 activation and weight-layout selection.""" + + if os.environ.get("AITER_SITUV2_A8W4", "0") == "1": + return KimiK3MXFP4ExpertMode( + activation_dtype=dtypes.fp8, + gate_mode=GateMode.INTERLEAVE, + w13_layout=KimiK3MXFP4W13Layout.GATE_UP_INTERLEAVED_PRESHUFFLED, + ) + return KimiK3MXFP4ExpertMode( + activation_dtype=dtypes.bf16, + gate_mode=GateMode.SEPARATED, + w13_layout=KimiK3MXFP4W13Layout.GATE_UP_SEPARATED_PRESHUFFLED, + ) + + +@dataclass(frozen=True) +class KimiK3MXFP4ExpertRequest: + """All tensors owned by the synchronous Kimi-K3 expert handoff.""" + + hidden_states: torch.Tensor + router_logits: torch.Tensor + correction_bias: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w1_scale: torch.Tensor + w2_scale: torch.Tensor + situ_beta: float + situ_linear_beta: float + + +@dataclass(frozen=True) +class KimiK3MXFP4ExpertResult: + """Routing, Opus stage metadata, and the routed BF16 expert output.""" + + expert_ids: torch.Tensor + routing_weights: torch.Tensor + stage1_sorted_token_ids: torch.Tensor + stage1_sorted_weights: torch.Tensor + sorted_expert_ids: torch.Tensor + num_valid_ids: torch.Tensor + expert_output: torch.Tensor + activation_dtype: torch.dtype + w13_layout: KimiK3MXFP4W13Layout + + +def _is_packed_scale( + tensor: torch.Tensor, + *, + device: torch.device, + shape: tuple[int, int], +) -> bool: + return ( + tensor.is_cuda + and tensor.device == device + and tensor.dtype == dtypes.fp8_e8m0 + and tensor.is_contiguous() + and tuple(tensor.shape) == shape + ) + + +def _has_expected_w13_layout( + tensor: torch.Tensor, + mode: KimiK3MXFP4ExpertMode, +) -> bool: + layout = getattr(tensor, "kimi_k3_w13_layout", None) + if isinstance(layout, KimiK3MXFP4W13Layout): + layout = layout.value + if mode.gate_mode == GateMode.INTERLEAVE: + return layout == mode.w13_layout.value + # Preserve the accepted A16W4 API for preshuffled tensors created before + # the explicit layout tag existed. + return layout in (None, mode.w13_layout.value) + + +def supports_kimi_k3_mxfp4_expert_handoff( + request: KimiK3MXFP4ExpertRequest, +) -> bool: + """Return whether the exact gfx950 Kimi-K3 TP8 specialization is safe.""" + + hidden_states = request.hidden_states + device = hidden_states.device + mode = kimi_k3_mxfp4_expert_mode() + return ( + os.environ.get("AITER_DISABLE", "0") != "1" + and supports_kimi_k3_b1_route_sort( + request.router_logits, + request.correction_bias, + num_experts=_EXPERTS, + topk=_TOPK, + num_expert_group=1, + topk_group=1, + block_size_m=_BLOCK_M, + ) + and hidden_states.is_cuda + and hidden_states.dtype == torch.bfloat16 + and hidden_states.is_contiguous() + and tuple(hidden_states.shape) == (1, _MODEL_DIM) + and request.router_logits.device == device + and request.w1.is_cuda + and request.w1.device == device + and request.w1.dtype == dtypes.fp4x2 + and request.w1.is_contiguous() + and tuple(request.w1.shape) == (_EXPERTS, 2 * _INTER_DIM, _MODEL_DIM // 2) + and bool(getattr(request.w1, "is_shuffled", False)) + and _has_expected_w13_layout(request.w1, mode) + and request.w2.is_cuda + and request.w2.device == device + and request.w2.dtype == dtypes.fp4x2 + and request.w2.is_contiguous() + and tuple(request.w2.shape) == (_EXPERTS, _MODEL_DIM, _INTER_DIM // 2) + and bool(getattr(request.w2, "is_shuffled", False)) + and _is_packed_scale( + request.w1_scale, + device=device, + shape=_W1_SCALE_SHAPE, + ) + and _is_packed_scale( + request.w2_scale, + device=device, + shape=_W2_SCALE_SHAPE, + ) + and request.situ_beta == 4.0 + and request.situ_linear_beta == 25.0 + ) + + +def kimi_k3_mxfp4_expert_handoff( + request: KimiK3MXFP4ExpertRequest, +) -> KimiK3MXFP4ExpertResult: + """Consume prepared route metadata exactly once in stage1/stage2.""" + + if not supports_kimi_k3_mxfp4_expert_handoff(request): + mode = kimi_k3_mxfp4_expert_mode() + raise NotImplementedError( + "Kimi-K3 prepared-route expert handoff requires contiguous gfx950 " + "B1 BF16 activations, FP32 1x896 logits, BF16 correction bias, " + "and preshuffled TP8 896x3584x384 MXFP4 weights matching " + f"{mode.w13_layout.value}" + ) + + mode = kimi_k3_mxfp4_expert_mode() + ( + routing_weights, + expert_ids, + sorted_token_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + ) = kimi_k3_b1_route_sort( + request.router_logits, + request.correction_bias, + model_dim=_MODEL_DIM, + ) + + # Import locally to avoid a package cycle while aiter.fused_moe imports + # FlyDSL dispatch helpers. + from aiter.fused_moe import fused_moe_2stages + + expert_output = fused_moe_2stages( + request.hidden_states, + request.w1, + request.w2, + _TOPK, + sorted_token_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + moe_buf, + True, + _BLOCK_M, + activation=ActivationType.Situv2, + quant_type=QuantType.per_1x32, + q_dtype_a=mode.activation_dtype, + q_dtype_w=dtypes.fp4x2, + w1_scale=request.w1_scale, + w2_scale=request.w2_scale, + topk_ids=expert_ids, + topk_weights=routing_weights, + beta=request.situ_beta, + linear_beta=request.situ_linear_beta, + gate_mode=mode.gate_mode.value, + ) + return KimiK3MXFP4ExpertResult( + expert_ids=expert_ids, + routing_weights=routing_weights, + stage1_sorted_token_ids=sorted_token_ids, + stage1_sorted_weights=sorted_weights, + sorted_expert_ids=sorted_expert_ids, + num_valid_ids=num_valid_ids, + expert_output=expert_output, + activation_dtype=mode.activation_dtype, + w13_layout=mode.w13_layout, + ) + + +__all__ = [ + "KimiK3MXFP4ExpertMode", + "KimiK3MXFP4ExpertRequest", + "KimiK3MXFP4ExpertResult", + "KimiK3MXFP4W13Layout", + "kimi_k3_mxfp4_expert_handoff", + "kimi_k3_mxfp4_expert_mode", + "supports_kimi_k3_mxfp4_expert_handoff", +] diff --git a/op_tests/flydsl_tests/test_kimi_k3_moe_handoff.py b/op_tests/flydsl_tests/test_kimi_k3_moe_handoff.py new file mode 100644 index 0000000000..a7a15c7747 --- /dev/null +++ b/op_tests/flydsl_tests/test_kimi_k3_moe_handoff.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import importlib +import math +from dataclasses import replace + +import aiter +import pytest +import torch +from aiter import dtypes +from aiter.fused_moe import fused_moe_2stages, moe_sorting +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.kimi_k3_moe_handoff import ( + KimiK3MXFP4ExpertRequest, + KimiK3MXFP4W13Layout, + kimi_k3_mxfp4_expert_handoff, + kimi_k3_mxfp4_expert_mode, + supports_kimi_k3_mxfp4_expert_handoff, +) +from aiter.ops.flydsl.moe_common import GateMode + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_gfx_runtime() != "gfx950", + reason="Kimi-K3 prepared-route expert handoff requires gfx950", +) + +EXPERTS = 896 +TOPK = 16 +MODEL_DIM = 3584 +INTER_DIM = 384 + + +def _request() -> KimiK3MXFP4ExpertRequest: + generator = torch.Generator(device="cuda").manual_seed(20260728) + + def packed(shape): + tensor = torch.full(shape, 0x11, dtype=torch.uint8, device="cuda") + tensor = tensor.view(dtypes.fp4x2) + tensor.is_shuffled = True + return tensor + + w1 = packed((EXPERTS, 2 * INTER_DIM, MODEL_DIM // 2)) + w2 = packed((EXPERTS, MODEL_DIM, INTER_DIM // 2)) + w1.kimi_k3_w13_layout = kimi_k3_mxfp4_expert_mode().w13_layout.value + w1_scale = torch.full( + (EXPERTS * 2 * INTER_DIM, MODEL_DIM // 32), + 0x7F, + dtype=torch.uint8, + device="cuda", + ).view(dtypes.fp8_e8m0) + w2_scale = torch.full( + ( + math.ceil(EXPERTS * MODEL_DIM / 256) * 256, + math.ceil((INTER_DIM // 32) / 8) * 8, + ), + 0x7F, + dtype=torch.uint8, + device="cuda", + ).view(dtypes.fp8_e8m0) + router_logits = torch.full((1, EXPERTS), -8.0, dtype=torch.float32, device="cuda") + router_logits[:, :32] = 2.0 + return KimiK3MXFP4ExpertRequest( + hidden_states=torch.randn( + (1, MODEL_DIM), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ), + router_logits=router_logits, + correction_bias=torch.zeros((EXPERTS,), dtype=torch.bfloat16, device="cuda"), + w1=w1, + w2=w2, + w1_scale=w1_scale, + w2_scale=w2_scale, + situ_beta=4.0, + situ_linear_beta=25.0, + ) + + +def test_kimi_k3_handoff_consumes_prepared_sort_once(monkeypatch): + request = _request() + mode = kimi_k3_mxfp4_expert_mode() + assert supports_kimi_k3_mxfp4_expert_handoff(request) + + expected_weights = torch.empty((1, TOPK), dtype=torch.float32, device="cuda") + expected_ids = torch.empty((1, TOPK), dtype=torch.int32, device="cuda") + aiter.biased_grouped_topk_hip( + request.router_logits, + request.correction_bias.float(), + expected_weights, + expected_ids, + 1, + 1, + True, + 1.0, + ) + ( + expected_sorted_ids, + expected_sorted_weights, + expected_sorted_experts, + expected_num_valid, + expected_moe_buf, + ) = moe_sorting( + expected_ids, + expected_weights, + EXPERTS, + MODEL_DIM, + torch.bfloat16, + 32, + ) + expected_output = fused_moe_2stages( + request.hidden_states, + request.w1, + request.w2, + TOPK, + expected_sorted_ids, + expected_sorted_weights, + expected_sorted_experts, + expected_num_valid, + expected_moe_buf, + True, + 32, + activation=aiter.ActivationType.Situv2, + quant_type=aiter.QuantType.per_1x32, + q_dtype_a=mode.activation_dtype, + q_dtype_w=dtypes.fp4x2, + w1_scale=request.w1_scale, + w2_scale=request.w2_scale, + topk_ids=expected_ids, + topk_weights=expected_weights, + beta=request.situ_beta, + linear_beta=request.situ_linear_beta, + gate_mode=mode.gate_mode.value, + ) + torch.cuda.synchronize() + + def duplicate_sort_is_a_failure(*_args, **_kwargs): + raise AssertionError("prepared-route handoff repeated routing or sorting") + + monkeypatch.setattr("aiter.fused_moe.moe_sorting", duplicate_sort_is_a_failure) + monkeypatch.setattr("aiter.biased_grouped_topk_hip", duplicate_sort_is_a_failure) + quant_sorted_ids = [] + if mode.activation_dtype == dtypes.bf16: + monkeypatch.setattr( + "aiter.ops.quant.fused_dynamic_mx_quant_moe_sort", + duplicate_sort_is_a_failure, + ) + else: + fused_moe_module = importlib.import_module("aiter.fused_moe") + original_quant = fused_moe_module.fused_dynamic_mxfp8_quant_moe_sort + + def audited_quant(*args, **kwargs): + quant_sorted_ids.append(kwargs["sorted_ids"].data_ptr()) + return original_quant(*args, **kwargs) + + monkeypatch.setattr( + "aiter.fused_moe.fused_dynamic_mxfp8_quant_moe_sort", + audited_quant, + ) + actual = kimi_k3_mxfp4_expert_handoff(request) + torch.cuda.synchronize() + if mode.activation_dtype == dtypes.fp8: + # Stage1 consumes the prepared sort for activation quantization. Its + # tuned `_fp8` epilogue fuses the inter-stage quantization. + assert quant_sorted_ids == [actual.stage1_sorted_token_ids.data_ptr()] + + torch.testing.assert_close(actual.expert_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual.routing_weights, expected_weights, rtol=1e-6, atol=1e-7) + torch.testing.assert_close( + actual.sorted_expert_ids, expected_sorted_experts[:TOPK], rtol=0, atol=0 + ) + torch.testing.assert_close( + actual.stage1_sorted_token_ids, + expected_sorted_ids[: actual.stage1_sorted_token_ids.numel()], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual.stage1_sorted_weights, + expected_sorted_weights[: actual.stage1_sorted_weights.numel()], + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close(actual.num_valid_ids, expected_num_valid, rtol=0, atol=0) + expected_f32 = expected_output.float() + actual_f32 = actual.expert_output.float() + relative_rmse = torch.sqrt(torch.mean((actual_f32 - expected_f32).square())) + relative_rmse /= torch.sqrt(torch.mean(expected_f32.square())).clamp_min(1e-12) + cosine = torch.nn.functional.cosine_similarity( + expected_f32.flatten(), actual_f32.flatten(), dim=0 + ) + assert float(relative_rmse) <= 1e-3 + assert float(cosine) >= 0.999 + assert torch.isfinite(actual.expert_output).all() + assert actual.activation_dtype == mode.activation_dtype + assert actual.w13_layout == mode.w13_layout + + +@pytest.mark.parametrize( + ("a8w4", "activation_dtype", "gate_mode", "w13_layout"), + [ + ( + "0", + dtypes.bf16, + GateMode.SEPARATED, + KimiK3MXFP4W13Layout.GATE_UP_SEPARATED_PRESHUFFLED, + ), + ( + "1", + dtypes.fp8, + GateMode.INTERLEAVE, + KimiK3MXFP4W13Layout.GATE_UP_INTERLEAVED_PRESHUFFLED, + ), + ], +) +def test_kimi_k3_handoff_mode_owns_activation_and_layout( + monkeypatch, a8w4, activation_dtype, gate_mode, w13_layout +): + monkeypatch.setenv("AITER_SITUV2_A8W4", a8w4) + mode = kimi_k3_mxfp4_expert_mode() + assert mode.activation_dtype == activation_dtype + assert mode.gate_mode == gate_mode + assert mode.w13_layout == w13_layout + assert supports_kimi_k3_mxfp4_expert_handoff(_request()) + + +def test_kimi_k3_handoff_support_is_narrow(): + request = _request() + assert supports_kimi_k3_mxfp4_expert_handoff(request) + assert not supports_kimi_k3_mxfp4_expert_handoff( + replace(request, hidden_states=request.hidden_states.expand(2, -1)) + ) + assert not supports_kimi_k3_mxfp4_expert_handoff( + replace(request, correction_bias=request.correction_bias.float()) + ) + assert not supports_kimi_k3_mxfp4_expert_handoff(replace(request, situ_beta=1.0)) + + +def test_kimi_k3_a8w4_requires_explicit_interleaved_weight_owner(monkeypatch): + monkeypatch.setenv("AITER_SITUV2_A8W4", "1") + request = _request() + request.w1.kimi_k3_w13_layout = KimiK3MXFP4W13Layout.GATE_UP_SEPARATED_PRESHUFFLED.value + assert not supports_kimi_k3_mxfp4_expert_handoff(request) + + +def test_kimi_k3_handoff_honors_aiter_opt_out(monkeypatch): + request = _request() + monkeypatch.setenv("AITER_DISABLE", "1") + assert not supports_kimi_k3_mxfp4_expert_handoff(request)