diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/__init__.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/__init__.py new file mode 100644 index 000000000000..c5d91f411add --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 persistent FlyDSL decode kernels for MI355X (gfx950), TP8. + +Ported from vLLM (vllm/models/kimi_k3/amd/mono, Apache-2.0). ``common/`` is in +turn adapted from ROCm/ATOM ``atom/mono`` at a526f0d (MIT); AMD's notice is +kept in those files. + +Nothing is imported eagerly: pulling in a submodule pulls in ``flydsl`` and +``aiter``, which a caller only probing availability should not pay for. +""" diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/__init__.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/__init__.py new file mode 100644 index 000000000000..f81492660596 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 mono decode: attention.""" diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/back.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/back.py new file mode 100644 index 000000000000..b6cfb8d3a441 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/back.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2a ``attn_post``: one Kimi-K3 KDA layer from the gated core output to the +MoE's input row -- o_proj, its TP all-reduce inside the kernel, and the MLP +AttnRes seam. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``: + + gemv CTAs 56 .. o_proj (448 x 16 rows, K = 1536 local): bf16 GEMV of the + core rows (plain loads: the previous launch wrote them) -> + bf16 partial -> pushed to every rank's AR region (system + scope, this rank's slot) + attn CTAs 0 .. 55 per 128 columns: the 8 ranks' partials summed in rank + order (fp32 -> bf16: the custom all-reduce's result) -> + LDS delta; then K1's AttnRes stages with the MLP's weights: + prefix + delta -> prefix out, the sources' softmax mix, + post_attention_layernorm -> the MoE input x + +The AR region is double-buffered by the step epoch's parity: a rank a step +ahead writes the other half, so no peer's unread partials are overwritten. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import gpu, range_constexpr +from flydsl.expr.typing import Int32, Int64, T + +from sglang.kernels.ops.moe.k3_mono_flydsl.attention.kda import ( + _SOURCES, + COLS, + HIDDEN, + LDS_PAD, + NA, + PROJ, + ROWS, + stage_gate, + stage_mix, + stage_slice, + step_tag, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.debug import ( + stamp, + stamp_begin, + stamp_flush, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + CM_NT, + bf16_round, + ld_i32, + mfma_bf16, + row_sum, + rsrc, + traced, + uniform, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import ( + BLOCKS, + THREADS, + WAVES, + KernelAbi, + key_tuple, + pair_layout, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ranks import peer_bases, sum_partials +from sglang.kernels.ops.moe.k3_mono_flydsl.common.sync import Mailbox, preg, shift, sreg + +TP = 8 +OROWS = HIDDEN # o_proj rows (full hidden, a rank's K slice) +OTASKS = OROWS // ROWS # 448 +OK = PROJ # K = 1536 a rank +OKCH = OK // 128 # 12 chunks +GEMV0 = NA # the first gemv CTA +GEMV_N = BLOCKS - NA # 200 +TL_POINTS = 6 # start, gemv, pushed, reduced, slice/gate, mix + +_STREAM = fx.Stream(None) + + +@dataclass(frozen=True) +class AttnPostBuild: + tokens: int + nblocks: int # the MLP AttnRes block sources + eps: float = 1e-5 + out_eps: float = 1e-5 + timeline: bool = False + + +ABI = KernelAbi( + ( + "core", + "w_o", + "prefix", + "blocks", + "blk_sm", + "blk_sr", + "ares_nw", + "ares_qk", + "in_nw", + "x_out", + "scratch", + "peers", + "rank", + "epoch", + "layer", + "tl", + ) +) + + +def ar_half_pairs(s): + """One parity's AR region: [src rank][token][column pair].""" + return TP * s * HIDDEN // 2 + + +def peer_bytes(max_tokens): + return 2 * ar_half_pairs(max_tokens) * 8 + + +def scratch_layout(key: AttnPostBuild) -> dict: + s, ns = key.tokens, key.nblocks + 1 + pairs = { + "part": NA * s * ns * 2, + "wgt": s * (ns + 1), + "msq": NA * s, + "xrdy": NA, + } + return pair_layout(pairs.items()) + + +def scratch_bytes(key: AttnPostBuild) -> int: + return max(o + n for o, n in scratch_layout(key).values()) + + +def o_loads(c, rg): + """Rows 16 rg .. of o_proj, this wave's K chunks (wave, wave + 8 < 12): a lane + 64 B of its row a chunk.""" + lane, wave = c["lane"], c["wave"] + r_w = rsrc(c["args"]["w_o"]) + row = rg * ROWS + lane % ROWS + g = lane // ROWS + out = [] + for i in range(2): + ch = fx.min(wave + WAVES * i, OKCH - 1) + base = (row * OK + ch * 128 + g * 8) // 2 + out.append( + [ + fx.Vector( + bo.buffer_load( + r_w, + base + 16 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for q in range(4) + ] + ) + return out + + +@traced +def load_core(c): + """The core rows (S x 1536 bf16) -> LDS, rows padded.""" + s, tid = c["S"], c["tid"] + words = s * OK // 8 + row = OK + LDS_PAD + for i in range_constexpr((words + THREADS - 1) // THREADS): + e = fx.min(tid + THREADS * i, words - 1) + v = fx.Vector( + bo.buffer_load(rsrc(c["args"]["core"]), e * 4, vec_width=4, dtype=T.i32) + ) + t = e // (OK // 8) + k = e % (OK // 8) * 8 + if tid + THREADS * i < words: + fx.ptr_store(v, c["xl"] + (t * row + k) // 2) + gpu.barrier() + + +@traced +def stage_o(c, rg, ops): + """Rows 16 rg ..: GEMV -> bf16 partials -> every rank's AR region, this + rank's slot (thread: token, row pair; a put a peer).""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + g = lane // ROWS + t = fx.min(lane % ROWS, s - 1) + row = OK + LDS_PAD + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(2): + ch = wave + WAVES * i + live = ch < OKCH + k = fx.min(ch, OKCH - 1) * 128 + g * 8 + for q in range_constexpr(4): + xb = fx.Vector( + fx.ptr_load( + c["xl"] + (t * row + k + 32 * q) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xb = fx.Vector.from_elements( + [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 + ) + acc = mfma_bf16( + ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + if tid < s * (ROWS // 2): + tt = tid // (ROWS // 2) + rp = tid % (ROWS // 2) + v0 = bf16_round(row_sum(red, 2 * rp, tt)) + v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) + at = (c["rank"] * s + tt) * HIDDEN + rg * ROWS + 2 * rp + for p in range_constexpr(TP): + c["put_bf"](c["ar_peer"][p], at, [v0, v1]) + gpu.barrier() + + +@traced +def stage_reduce(c, task): + """Columns 128 task ..: the 8 ranks' partials, rank order, -> bf16 delta in + LDS (thread: token, column pair).""" + s, tid = c["S"], c["tid"] + if tid < s * (COLS // 2): + t = tid // (COLS // 2) + cp = tid % (COLS // 2) + col = task * COLS + 2 * cp + lo, hi = sum_partials( + c["poll"], c["ar_own"], lambda src: ((src * s + t) * HIDDEN + col) // 2, TP + ) + fx.ptr_store(bf16_round(lo), c["delta_lds"] + (t * COLS + 2 * cp)) + fx.ptr_store(bf16_round(hi), c["delta_lds"] + (t * COLS + 2 * cp + 1)) + gpu.barrier() + + +_BUILDS: dict = {} + + +def build(key: AttnPostBuild): + if key in _BUILDS: + return _BUILDS[key] + s, ns = key.tokens, key.nblocks + 1 + assert s * (ROWS // 2) <= THREADS and s * (COLS // 2) <= THREADS + lay = scratch_layout(key) + half = ar_half_pairs(s) * 8 + + @fx.struct + class Smem: + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + xl: fx.Array[fx.Int32, s * (OK + LDS_PAD) // 2, 16] + vals: fx.Array[fx.Float32, s * ns * COLS, 16] + dl: fx.Array[fx.Float32, s * COLS, 16] + tls: fx.Array[fx.Int64, TL_POINTS, 16] + + keyed = key_tuple(key, _SOURCES) + name = f"k3_mono_attn_post_s{s}_nb{key.nblocks}_t{int(key.timeline)}" + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def attn_post( + core: Int64, + w_o: Int64, + prefix: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + x_out: Int64, + scratch: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + tl: Int64, + ): + _ = keyed + tid = fx.thread_idx.x + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(Smem).peek() + mb = Mailbox(step_tag(epoch, layer) - 1) + par = uniform(ld_i32(epoch, 0)) % 2 + bases = peer_bases(peers, TP) + c = { + "S": s, + "NS": ns, + "NB": key.nblocks, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rank": rank, + "delta": True, + "debug": False, + "eps": key.eps, + "out_eps": key.out_eps, + "put": mb.put, + "put_bf": mb.put_bf, + "poll": mb.poll, + "red": lds.red.ptr, + "xl": lds.xl.ptr, + "vals": lds.vals.ptr, + "delta_lds": lds.dl.ptr, + "xbuf": x_out, + "ar_peer": [shift(preg(b, 0, "ar"), fx.Int64(par) * half) for b in bases], + "ar_own": shift(preg(load_own(peers, rank), 0, "ar"), fx.Int64(par) * half), + "args": { + "core": core, + "w_o": w_o, + "prefix": prefix, + "blocks": blocks, + "blk_sm": blk_sm, + "blk_sr": blk_sr, + "ares_nw": ares_nw, + "ares_qk": ares_qk, + "in_nw": in_nw, + "x_dbg": 0, + }, + } + for region in ("part", "wgt", "msq", "xrdy"): + c[region] = sreg(scratch, lay[region][0], region) + on, tls = key.timeline, lds.tls.ptr + stamp_begin(on, tls, tid, TL_POINTS) + if bid >= GEMV0: + gi = bid - GEMV0 + # every task's weights in flight first: nothing here polls + ops = [o_loads(c, fx.min(gi + GEMV_N * j, OTASKS - 1)) for j in range(3)] + load_core(c) + stamp(on, tls, tid, 1) + for j in range_constexpr(3): + if gi + GEMV_N * j < OTASKS: + stage_o(c, gi + GEMV_N * j, ops[j]) + stamp(on, tls, tid, 2) + if bid < NA: + stage_reduce(c, bid) + stamp(on, tls, tid, 3) + stage_slice(c, bid) + if bid < s: + stage_gate(c, bid) + stamp(on, tls, tid, 4) + stage_mix(c, bid) + stamp(on, tls, tid, 5) + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + + @flyc.jit + def launch( + core: Int64, + w_o: Int64, + prefix: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + x_out: Int64, + scratch: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + attn_post( + core, + w_o, + prefix, + blocks, + blk_sm, + blk_sr, + ares_nw, + ares_qk, + in_nw, + x_out, + scratch, + peers, + rank, + epoch, + layer, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + ABI.check(attn_post, launch) + _BUILDS[key] = launch + return launch + + +def load_own(peers, rank): + """This rank's own peer buffer base (its entry of the table).""" + from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import load_ptr64 + + return load_ptr64(peers, rank) + + +def attn_post( + key: AttnPostBuild, + *, + core, + w_o, + prefix, + blocks, + ares_nw, + ares_qk, + in_nw, + x_out, + scratch, + peers, + rank, + epoch, + layer, + tl=None, +): + """Host launcher. ``peers``: the int64 table of every rank's peer buffer + (at least ``peer_bytes(S)`` each); ``prefix`` updated in place.""" + assert w_o.shape == (OROWS, OK) and w_o.is_contiguous() + assert prefix.is_contiguous() and core.is_contiguous() and x_out.is_contiguous() + f = build(key) + f( + core.data_ptr(), + w_o.data_ptr(), + prefix.data_ptr(), + blocks.data_ptr(), + blocks.stride(0), + blocks.stride(1), + ares_nw.data_ptr(), + ares_qk.data_ptr(), + in_nw.data_ptr(), + x_out.data_ptr(), + scratch.data_ptr(), + peers.data_ptr(), + rank, + epoch.data_ptr(), + layer, + 0 if tl is None else tl.data_ptr(), + stream=torch.cuda.current_stream(), + ) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/kda.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/kda.py new file mode 100644 index 000000000000..1e2e21dffc8e --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/attention/kda.py @@ -0,0 +1,1142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K1 ``kda_pre``: one Kimi-K3 KDA layer from the AttnRes seam to the gated core +output, for a spec-verify step (R requests x L tokens, S = R L rows). + +One launch a layer, ``BLOCKS`` CTAs x ``THREADS`` threads, in three roles so that +no CTA that polls a mailbox has weight loads in flight (a poll's load retires +behind every earlier load): + + attn CTAs 0 .. 55 slice (128 hidden columns each): prefix (+ delta) -> + bf16 prefix out, per-source partial sum of squares and + qk dot (PART); gate (token t on CTA t): softmax weights + (WGT); mix: the softmax mix, its partial sum of squares + (MSQ), every slice's -> output RMSNorm -> bf16 x rows + (plain device-scope stores, an XRDY flag a slice) + gemv CTAs 56 .. in_proj, two 16-row groups each: the first group's + weights issued at kernel start, the second's as x lands + -> q|k|v|g|f_a|beta (PROJ) + kda tasks (R x 12 heads x 8 v-chunks) on CTAs 0 ..: inputs that + need no hand-off loaded first (metadata, conv window, + state rows, f_b rows), then conv update of the owned + channels (q / k out through QK), f_b MFMA + lower-bound + gate, l2norm, the recurrence on 16 state rows (a state a + token written), partial norms (ONRM), gated RMSNorm -> + core_out + +Rounding follows the original path where it is elementwise (the AttnRes prefix +add, every bf16 hand-off between the original kernels); split-K reductions +differ only in their order. +""" + +import math +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import Int32, Int64, T + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.debug import ( + stamp, + stamp_begin, + stamp_flush, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + CM_DEV, + CM_NT, + bf_hi, + bf_lo, + butterfly, + hw_exp2, + hw_rsq, + ld_i32, + mfma_bf16, + row_sum, + rsrc, + traced, + uniform, + wave_max, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import ( + BLOCKS, + THREADS, + WAVES, + KernelAbi, + first_task, + key_tuple, + pair_layout, + source_digest, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.sync import Mailbox, publish, sreg + +HIDDEN = 7168 +NA = 56 # attn CTAs +COLS = HIDDEN // NA # 128 columns a slice +QC = COLS // 4 # threads a (token, source) row: 4 columns each +GSL = 7 # slices a gate / mix-norm thread sums (NA = 8 x 7) +HD = 128 +NH = 12 # local heads at TP8 +PROJ = NH * HD # 1536 +QKV = 3 * PROJ +G2_0 = QKV +FA_0 = G2_0 + PROJ # 6144 +BETA_0 = FA_0 + HD # 6272 +NPROJ = 6288 # in_proj rows a rank holds (6284 + 4 pad) +ROWS = 16 +GEMV_TASKS = NPROJ // ROWS # 393 +GEMV_CTAS = (GEMV_TASKS + 1) // 2 # 197 +KCH = 128 # K a chunk: 4 MFMAs of 32 +NKCH = HIDDEN // KCH # 56 +PER_WAVE = NKCH // WAVES # 7 +VCH = 16 # state rows a KDA task +NVCH = HD // VCH # 8 +CONV_W = 4 +LOG2E = 1.0 / math.log(2.0) +LDS_PAD = 8 # bf16 a padded x row +# timeline points: start, slice, gate, mix, load_x, gemv0, gemv1, k.stage, +# k.conv, k.gates, k.recur, k.norm +TL_POINTS = 15 + +assert NA + GEMV_CTAS <= BLOCKS and NA == 8 * GSL + +_STREAM = fx.Stream(None) +# every mono source file: in every build's JIT cache key +_SOURCES = source_digest(".") + + +@dataclass(frozen=True) +class KdaPreBuild: + tokens: int # S + qlen: int # L, tokens a request + nblocks: int # AttnRes block sources (the prefix is one more) + delta: bool + state_len: int # conv state entries: CONV_W - 1 + num_spec + eps: float = 1e-5 # AttnRes norm + out_eps: float = 1e-5 # input_layernorm + onorm_eps: float = 1e-5 + lower_bound: float = -5.0 + debug: bool = False + timeline: bool = False + write_idx: int = -1 # a block-write layer's block: the updated prefix stored there + kda0: int = 0 # the first KDA task's CTA + xp: int = 0 # experiment switch (timing only) + + +ABI = KernelAbi( + ( + "prefix", + "delta", + "blocks", + "blk_sm", + "blk_sr", + "ares_nw", + "ares_qk", + "in_nw", + "w_in", + "w_fb", + "conv_w", + "conv_st", + "cs_seq", + "cs_dim", + "cs_tok", + "a_log", + "dt_bias", + "on_w", + "rstate", + "rs_seq", + "st_idx", + "st_stride", + "num_acc", + "core_out", + "x_dbg", + "proj_dbg", + "scratch", + "layer", + "epoch", + "tl", + ) +) + + +def scratch_layout(key: KdaPreBuild) -> dict: + s, ns = key.tokens, key.nblocks + 1 + assert s % key.qlen == 0 + pairs = { + "part": NA * s * ns * 2, + "wgt": s * (ns + 1), + "msq": NA * s, + "xrdy": NA, + "proj": s * NPROJ, + "qk": s * NH * 2 * HD, + "onrm": s * NH * NVCH, + } + lay = pair_layout(pairs.items()) + end = max(o + n for o, n in lay.values()) + # the plain bf16 x rows (not pairs), 16 B aligned + lay["x"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) + return lay + + +def scratch_bytes(key: KdaPreBuild) -> int: + return max(o + n for o, n in scratch_layout(key).values()) + + +TAG_SLOTS = 256 # a step's mailbox tags: epoch x 256 + a launch's slot + + +def step_tag(epoch, slot): + """This launch's mailbox tag: the step's epoch (a device counter the runner + bumps once a step, in the graph) and the launch's slot, so no step's or + other launch's pair ever passes for this one's -- nothing is zeroed.""" + e = uniform(ld_i32(epoch, 0)) + return (e % (1 << 22)) * TAG_SLOTS + slot + 1 + + +def ld_bf(ptr, i, cm=0): + return fx.Float32( + fx.BFloat16( + bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16, cache_modifier=cm) + ) + ) + + +def ld_bf4(ptr, i): + """bf16 elements i .. i + 4 (i a multiple of 4) as four f32.""" + w = fx.Vector(bo.buffer_load(rsrc(ptr), i // 2, vec_width=2, dtype=T.i32)) + return [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])] + + +def bf4_words(v): + """Four f32 -> two words of packed bf16.""" + return fx.Vector.from_elements(v, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def bf16r(v): + return fx.Float32(fx.Float32(v).to(fx.BFloat16)) + + +def sigmoid(z): + return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) + + +def f32_of(w): + return w.bitcast(fx.Float32) + + +def sum_polled(vals): + acc = f32_of(vals[0][0]) + for k in range(1, len(vals)): + acc = acc + f32_of(vals[k][0]) + return acc + + +# ---------------------------------------------------------------- attn CTAs +@traced +def stage_slice(c, task): + """Columns 128 task ..: every token's sources into LDS ``vals`` (4 columns a + thread), the prefix (+ delta) rounded to bf16 and stored back, per-(token, + source) partial sum of squares and qk dot -> PART.""" + s, ns, nb, tid, lane = c["S"], c["NS"], c["NB"], c["tid"], c["lane"] + a = c["args"] + c0 = task * COLS + n = s * ns * QC # threads' work items + for i in range_constexpr((n + THREADS - 1) // THREADS): + e = fx.min(tid + THREADS * i, n - 1) + live = (tid + THREADS * i) < n + row = e // QC # (token, source) + t = row // ns + src = row % ns + col = c0 + (e % QC) * 4 + is_pre = src == nb + blk = ld_bf4( + a["blocks"], t * a["blk_sm"] + fx.min(src, nb - 1) * a["blk_sr"] + col + ) + pre = ld_bf4(a["prefix"], t * HIDDEN + col) + if const_expr(c["delta"]): + if const_expr(c.get("delta_lds") is not None): + # the delta reduced into LDS first (K2a: the all-reduced o_proj) + dl = [ + fx.ptr_load(c["delta_lds"] + (t * COLS + (e % QC) * 4 + j)) + for j in range(4) + ] + else: + dl = ld_bf4(a["delta"], t * HIDDEN + col) + if const_expr(c.get("reset", False)): + # a block-write layer's MLP seam: the prefix restarts at the + # attention's output (the input prefix is not read) + pre = dl + else: + pre = [bf16r(pre[j] + dl[j]) for j in range(4)] + if live & is_pre: + bo.buffer_store( + bf4_words(pre), rsrc(a["prefix"]), (t * HIDDEN + col) // 2 + ) + if const_expr(c.get("write_idx", -1) >= 0): # noqa: SIM102 (traced: `and` does not combine traced values) + # the block-write layer's attention seam: the updated prefix is + # block write_idx + if live & is_pre: + bo.buffer_store( + bf4_words(pre), + rsrc(a["blocks"]), + (t * a["blk_sm"] + c["write_idx"] * a["blk_sr"] + col) // 2, + ) + v = [ + live.select(is_pre.select(pre[j], blk[j]), fx.Float32(0.0)) + for j in range(4) + ] + w = ld_bf4(a["ares_nw"], col) + q = ld_bf4(a["ares_qk"], col) + # a thread past the end holds a clamped copy of the last element: it + # must not store (it raced the owner with zeros) + if live: + for j in range_constexpr(4): + fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j)) + sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3] + dt = ( + v[0] * (w[0] * q[0]) + + v[1] * (w[1] * q[1]) + + v[2] * (w[2] * q[2]) + + v[3] * (w[3] * q[3]) + ) + sq = butterfly(sq, (16, 8, 4, 2, 1)) + dt = butterfly(dt, (16, 8, 4, 2, 1)) + if live & (lane % 32 == 0): + at = ((task * s + t) * ns + src) * 2 + c["put"](c["part"], at, sq) + c["put"](c["part"], at + 1, dt) + gpu.barrier() + + +@traced +def stage_gate(c, t): + """Token t: the 56 slices' partials -> source weights e_s and their sum.""" + s, ns, tid, lane, wave, red = ( + c["S"], + c["NS"], + c["tid"], + c["lane"], + c["wave"], + c["red"], + ) + nq = ns * 2 + if tid < 8 * nq: + g = tid // nq + q = tid % nq + vals = c["poll"]( + [(c["part"], (((g * GSL + k) * s + t) * ns) * 2 + q, 1) for k in range(GSL)] + ) + fx.ptr_store(sum_polled(vals), red + (g * nq + q)) + gpu.barrier() + if wave == 0: + src = fx.min(lane, ns - 1) + sq = fx.ptr_load(red + src * 2) + dt = fx.ptr_load(red + src * 2 + 1) + for g in range_constexpr(1, 8): + sq = sq + fx.ptr_load(red + (g * nq + src * 2)) + dt = dt + fx.ptr_load(red + (g * nq + src * 2 + 1)) + rstd = hw_rsq(sq * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["eps"])) + logit = dt * rstd * fx.Float32(LOG2E) + logit = (lane < ns).select(logit, fx.Float32(-3.0e38)) + m = wave_max(logit) + ex = (lane < ns).select(hw_exp2(logit - m), fx.Float32(0.0)) + den = butterfly(ex, (32, 16, 8, 4, 2, 1)) + if lane < ns: + c["put"](c["wgt"], t * (ns + 1) + lane, ex) + if lane == 0: + c["put"](c["wgt"], t * (ns + 1) + ns, den) + gpu.barrier() + + +GATE_OFF = 512 # red words of stage_gate_local's group sums + + +@traced +def stage_gate_local(c): + """Every token's source weights computed in this CTA (``stage_gate``'s sums + in its order, so every CTA's are bit-identical): the 56 slices' partials + polled in one round trip (no gate CTA, no WGT hand-off) -> red[t (ns + 1) + + src], the sum at + ns.""" + s, ns, tid, lane, wave, red = ( + c["S"], + c["NS"], + c["tid"], + c["lane"], + c["wave"], + c["red"], + ) + nq = ns * 2 + n = 8 * s * nq # (group, token, partial) items + per = (n + THREADS - 1) // THREADS + specs = [] + for r in range_constexpr(per): + item = fx.min(tid + THREADS * r, n - 1) + g = item // (s * nq) + t = (item // nq) % s + q = item % nq + specs += [ + (c["part"], (((g * GSL + k) * s + t) * ns) * 2 + q, 1) for k in range(GSL) + ] + got = c["poll"](specs, per * GSL) + for r in range_constexpr(per): + acc = f32_of(got[r * GSL][0]) + for k in range_constexpr(1, GSL): + acc = acc + f32_of(got[r * GSL + k][0]) + if tid + THREADS * r < n: + fx.ptr_store(acc, red + (GATE_OFF + tid + THREADS * r)) + gpu.barrier() + # wave t: token t's weights (stage_gate's math) + if wave < s: + t = wave + src = fx.min(lane, ns - 1) + base = GATE_OFF + t * nq + sq = fx.ptr_load(red + (base + src * 2)) + dt = fx.ptr_load(red + (base + src * 2 + 1)) + for g in range_constexpr(1, 8): + sq = sq + fx.ptr_load(red + (base + g * s * nq + src * 2)) + dt = dt + fx.ptr_load(red + (base + g * s * nq + src * 2 + 1)) + rstd = hw_rsq(sq * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["eps"])) + logit = dt * rstd * fx.Float32(LOG2E) + logit = (lane < ns).select(logit, fx.Float32(-3.0e38)) + m = wave_max(logit) + ex = (lane < ns).select(hw_exp2(logit - m), fx.Float32(0.0)) + den = butterfly(ex, (32, 16, 8, 4, 2, 1)) + if lane < ns: + fx.ptr_store(ex, red + (t * (ns + 1) + lane)) + if lane == 0: + fx.ptr_store(den, red + (t * (ns + 1) + ns)) + gpu.barrier() + + +@traced +def stage_mix(c, task): + """Columns 128 task ..: the softmax mix of ``vals`` (4 columns a thread), its + sum of squares (MSQ), every slice's -> rstd, then bf16 x = mixed rstd w -> + the plain x rows, and the slice's XRDY flag.""" + s, ns, tid, lane, red = c["S"], c["NS"], c["tid"], c["lane"], c["red"] + a = c["args"] + c0 = task * COLS + nw = s * (ns + 1) + if const_expr(c.get("gate_local", False)): + stage_gate_local(c) + else: + if tid < nw: + v = c["poll"]([(c["wgt"], tid, 1)])[0][0] + fx.ptr_store(f32_of(v), red + tid) + gpu.barrier() + mine = tid < s * QC + t = fx.min(tid // QC, s - 1) + cc = (tid % QC) * 4 + mixed = [] + for j in range_constexpr(4): + acc = fx.Float32(0.0) + for src in range_constexpr(ns): + acc = acc + fx.ptr_load(red + (t * (ns + 1) + src)) * fx.ptr_load( + c["vals"] + ((t * ns + src) * COLS + cc + j) + ) + mixed.append(acc / fx.ptr_load(red + (t * (ns + 1) + ns))) + sq = ( + mixed[0] * mixed[0] + + mixed[1] * mixed[1] + + mixed[2] * mixed[2] + + mixed[3] * mixed[3] + ) + sq = butterfly(sq, (16, 8, 4, 2, 1)) + if mine & (lane % 32 == 0): + c["put"](c["msq"], task * s + t, sq) + gpu.barrier() + if tid < 8 * s: + g = tid // s + tt = tid % s + vals = c["poll"]([(c["msq"], (g * GSL + k) * s + tt, 1) for k in range(GSL)]) + fx.ptr_store(sum_polled(vals), red + (128 + g * s + tt)) + gpu.barrier() + tot = fx.ptr_load(red + (128 + t)) + for g in range_constexpr(1, 8): + tot = tot + fx.ptr_load(red + (128 + g * s + t)) + r = hw_rsq(tot * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["out_eps"])) + wn = ld_bf4(a["in_nw"], c0 + cc) + xw = bf4_words([mixed[j] * r * wn[j] for j in range(4)]) + if mine: + at = (t * HIDDEN + c0 + cc) // 2 + bo.buffer_store(xw, rsrc(c["xbuf"]), at, cache_modifier=CM_DEV) + if const_expr(c["debug"]): + bo.buffer_store(xw, rsrc(a["x_dbg"]), at) + publish(c["put"], c["xrdy"], task, fx.Int32(1), tid == 0) + gpu.barrier() + + +# ---------------------------------------------------------------- gemv CTAs +def gemv_loads(c, rg, chunks): + """This wave's weight operands of rows 16 rg .. for its K chunks ``chunks`` + (indices into wave, wave + 8, ..): a lane 64 contiguous bytes of its row a + chunk (4 dwordx4). ``rg`` past the last group loads the last group.""" + lane, wave = c["lane"], c["wave"] + r_w = rsrc(c["args"]["w_in"]) + row = fx.min(rg, GEMV_TASKS - 1) * ROWS + lane % ROWS + g = lane // ROWS + out = [] + for i in chunks: + k = (wave + WAVES * i) * KCH + g * 8 + base = (row * HIDDEN + k) // 2 # dwords + out.append( + [ + fx.Vector( + bo.buffer_load( + r_w, + base + 16 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for q in range(4) + ] + ) + return out + + +@traced +def load_x(c): + """Every slice's x ready -> the bf16 x rows into LDS (rows padded).""" + s, tid = c["S"], c["tid"] + if tid < NA: + c["poll"]([(c["xrdy"], tid, 1)]) + gpu.barrier() + words = s * HIDDEN // 8 # 16 B chunks + row = HIDDEN + LDS_PAD + for i in range_constexpr((words + THREADS - 1) // THREADS): + e = fx.min(tid + THREADS * i, words - 1) + v = fx.Vector( + bo.buffer_load( + rsrc(c["xbuf"]), e * 4, vec_width=4, dtype=T.i32, cache_modifier=CM_DEV + ) + ) + t = e // (HIDDEN // 8) + k = e % (HIDDEN // 8) * 8 + if tid + THREADS * i < words: + fx.ptr_store(v, c["xl"] + (t * row + k) // 2) + gpu.barrier() + + +@traced +def stage_gemv(c, rg, ops): + """Rows 16 rg .. of in_proj against the LDS x -> bf16 PROJ (and proj_dbg).""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + g = lane // ROWS + t = fx.min(lane % ROWS, s - 1) + row = HIDDEN + LDS_PAD + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(PER_WAVE): + k = (wave + WAVES * i) * KCH + g * 8 + for q in range_constexpr(4): + xb = fx.Vector( + fx.ptr_load( + c["xl"] + (t * row + k + 32 * q) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + acc = mfma_bf16( + ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + if (tid < ROWS * s) & (rg < GEMV_TASKS): + r = tid % ROWS + tt = tid // ROWS + v = bf16r(row_sum(red, r, tt)) + c["put"](c["proj"], tt * NPROJ + rg * ROWS + r, v) + if const_expr(c["debug"]): + bo.buffer_store( + v.to(fx.BFloat16), + rsrc(c["args"]["proj_dbg"]), + tt * NPROJ + rg * ROWS + r, + ) + gpu.barrier() + + +# ---------------------------------------------------------------- kda +def kda_stage(c, req, h, ch): + """The task's inputs that need no hand-off, issued before its first poll: + metadata, the owned channels' conv window / weights, the initial state rows, + the f_b rows of this wave and the per-head constants.""" + L, sl, tid, lane, wave = c["L"], c["SL"], c["tid"], c["lane"], c["wave"] + a = c["args"] + k = {} + # every value a buffer resource is built on wave-uniform (else a waterfall) + k["acc_idx"] = uniform(ld_i32(a["num_acc"], req)) - 1 + k["slot_c"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"])) + k["slot0"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"] + k["acc_idx"])) + k["slots"] = [ + uniform(ld_i32(a["st_idx"], req * a["st_stride"] + j)) for j in range(L) + ] + # conv: thread < 48 owns channel (part, d) + ct = fx.min(tid, 3 * VCH - 1) + part = ct // VCH + chan = part * PROJ + h * HD + ch * VCH + ct % VCH + cs = a["conv_st"] + fx.Int64(fx.max(k["slot_c"], 0)) * fx.Int64(a["cs_seq"]) * 2 + co = chan * a["cs_dim"] # the channel's first entry (bf16 elements) + k["cs"], k["co"], k["chan"], k["part"] = cs, co, chan, part + k["cw"] = fx.Vector( + bo.buffer_load(rsrc(a["conv_w"]), chan * CONV_W, vec_width=4, dtype=T.f32) + ) + k["win"] = [ + ld_bf(cs, co + (k["acc_idx"] + j) * a["cs_tok"]) for j in range(CONV_W - 1) + ] + k["old"] = [ + ld_bf(cs, co + (k["acc_idx"] + 1 + j) * a["cs_tok"]) for j in range(sl - L) + ] + # recurrence: thread -> state row v = tid / 32, k columns (tid % 32) 4 .. + vr = ch * VCH + tid // 32 + k0 = (tid % 32) * 4 + k["head_off"] = (h * HD + vr) * HD + k0 + base0 = a["rstate"] + fx.Int64(fx.max(k["slot0"], 0)) * fx.Int64(a["rs_seq"]) * 4 + k["st"] = fx.Vector( + bo.buffer_load(rsrc(base0), k["head_off"], vec_width=4, dtype=T.f32) + ) + # f_b: wave w's rows 16 w .. of head h, a lane 64 B of its row + fr = h * HD + wave * 16 + lane % 16 + fb = (fr * HD + (lane // 16) * 32) // 2 + k["fb"] = [ + fx.Vector(bo.buffer_load(rsrc(a["w_fb"]), fb + 4 * q, vec_width=4, dtype=T.i32)) + for q in range(4) + ] + k["aa"] = hw_exp2(ld_f32(a["a_log"], h) * fx.Float32(LOG2E)) + k["dtb"] = [ + ld_f32(a["dt_bias"], h * HD + wave * 16 + 4 * (lane // 16) + i) + for i in range(4) + ] + return k + + +@traced +def kda_conv(c, k, req, h, ch): + """Owned channels: the spec conv update of request ``req``'s L tokens -> new + conv state; q / k out through QK, v to LDS.""" + L, sl, tid = c["L"], c["SL"], c["tid"] + a = c["args"] + stamp(c["on"], c["tls"], tid, 12) + if tid < 3 * VCH: + xs = c["poll"]( + [(c["proj"], (req * L + j) * NPROJ + k["chan"], 1) for j in range(L)] + ) + xv = [f32_of(xs[j][0]) for j in range(L)] + stamp(c["on"], c["tls"], tid, 13) + rs = rsrc(k["cs"]) + if const_expr(c["xp"] == 2): + # force the staged conv window before timing the stores + fx.ptr_store(k["old"][0] + k["win"][0], c["red"] + tid) + stamp(c["on"], c["tls"], tid, 13) + # A negative slot is the "no state" marker, matching SGLang's + # null_block_id = -1 convention; row 0 is an ordinary slot. The address + # above already clamps with max(slot, 0), so the two agree. + if (k["slot_c"] >= 0) & (c["xp"] != 1): + for j in range_constexpr(sl): + nv = k["old"][j] if j < sl - L else xv[j - (sl - L)] + bo.buffer_store( + fx.Float32(nv).to(fx.BFloat16), rs, k["co"] + j * a["cs_tok"] + ) + stamp(c["on"], c["tls"], tid, 14) + seq = k["win"] + xv + for j in range_constexpr(L): + acc = fx.Float32(0.0) + for q in range_constexpr(CONV_W): + acc = acc + k["cw"][q] * seq[j + q] + o = bf16r(acc * sigmoid(acc)) + if k["part"] < 2: + c["put"]( + c["qk"], + ((req * NH + h) * L + j) * 2 * HD + + k["part"] * HD + + ch * VCH + + tid % VCH, + o, + ) + if k["part"] == 2: + fx.ptr_store(o, c["vloc"] + (j * VCH + tid % VCH)) + gpu.barrier() + + +@traced +def kda_gates(c, k, req, h): + """f_a, beta and the head's q / k (one poll batch a thread); g1 = f_b f_a by + MFMA (a wave 16 rows) -> lower-bound gate -> decay in LDS; q / k l2-normed + (q scaled).""" + L, tid, lane, wave = c["L"], c["tid"], c["lane"], c["wave"] + n_fa = L * HD + n = n_fa + L * 2 * HD + per = (n + THREADS - 1) // THREADS + specs = [] + for i in range_constexpr(per): + e = fx.min(tid + THREADS * i, n - 1) + fa_at = (req * L + e // HD) * NPROJ + FA_0 + e % HD + qk_at = (req * NH + h) * L * 2 * HD + (e - n_fa) + specs.append((e, fa_at, qk_at)) + # two regions, two batches: the f_a pairs, then the q / k pairs + got_fa = c["poll"]( + [(c["proj"], fx.min(fa, (req * L + L) * NPROJ - 1), 1) for _, fa, _ in specs] + ) + got_qk = c["poll"]( + [(c["qk"], fx.max(qk, (req * NH + h) * L * 2 * HD), 1) for _, _, qk in specs] + ) + for i in range_constexpr(per): + e = specs[i][0] + if tid + THREADS * i < n_fa: + fx.ptr_store(f32_of(got_fa[i][0]), c["fa"] + e) + if (tid + THREADS * i >= n_fa) & (tid + THREADS * i < n): + fx.ptr_store(f32_of(got_qk[i][0]), c["qkl"] + (e - n_fa)) + if tid < L: + b = c["poll"]([(c["proj"], (req * L + tid) * NPROJ + BETA_0 + h, 1)])[0][0] + fx.ptr_store(sigmoid(f32_of(b)), c["beta"] + tid) + gpu.barrier() + # g1 rows 16 wave .. x tokens: A = f_b rows (staged), B = f_a [token][K] + t = fx.min(lane % 16, L - 1) + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for q in range_constexpr(4): + kb = (lane // 16) * 32 + 8 * q + xs = [fx.ptr_load(c["fa"] + (t * HD + kb + e)) for e in range(8)] + xb = fx.Vector.from_elements(xs, fx.Float32).to(fx.BFloat16) + acc = mfma_bf16(k["fb"][q].bitcast(fx.BFloat16), xb, acc) + for i in range_constexpr(4): + kk = wave * 16 + 4 * (lane // 16) + i + g = bf16r(acc[i]) + k["dtb"][i] + gate = fx.Float32(c["lb"]) * sigmoid(k["aa"] * g) + if lane % 16 < L: + fx.ptr_store(hw_exp2(gate * fx.Float32(LOG2E)), c["decay"] + (t * HD + kk)) + # l2norm: q / k row r (2 L rows of 128) on wave r % 8 + for rr in range_constexpr((2 * L + WAVES - 1) // WAVES): + row = fx.min(wave + WAVES * rr, 2 * L - 1) + x0 = fx.ptr_load(c["qkl"] + (row * HD + lane)) + x1 = fx.ptr_load(c["qkl"] + (row * HD + 64 + lane)) + ss = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) + sc = hw_rsq(ss + fx.Float32(1e-6)) + sc = (row % 2 == 0).select(sc * fx.Float32(HD**-0.5), sc) + if wave + WAVES * rr < 2 * L: + fx.ptr_store(x0 * sc, c["qkl"] + (row * HD + lane)) + fx.ptr_store(x1 * sc, c["qkl"] + (row * HD + 64 + lane)) + gpu.barrier() + + +@traced +def kda_recur(c, k, req, h, ch): + """State rows 16 ch .. of head h: the recurrence over the L tokens, a state + written a token, bf16 outputs to LDS ``ol`` and their partial norms (ONRM).""" + L, tid, lane = c["L"], c["tid"], c["lane"] + a = c["args"] + v = tid // 32 + k0 = (tid % 32) * 4 + # See kda_conv: negative is the "no state" marker, row 0 is a real slot. + live0 = k["slot0"] >= 0 + stv = [live0.select(k["st"][e], fx.Float32(0.0)) for e in range(4)] + for j in range_constexpr(L): + qv = [fx.ptr_load(c["qkl"] + ((2 * j) * HD + k0 + e)) for e in range(4)] + kv = [fx.ptr_load(c["qkl"] + ((2 * j + 1) * HD + k0 + e)) for e in range(4)] + dk = [fx.ptr_load(c["decay"] + (j * HD + k0 + e)) for e in range(4)] + stv = [stv[e] * dk[e] for e in range(4)] + p = stv[0] * kv[0] + stv[1] * kv[1] + stv[2] * kv[2] + stv[3] * kv[3] + p = butterfly(p, (16, 8, 4, 2, 1)) + vn = (fx.ptr_load(c["vloc"] + (j * VCH + v)) - p) * fx.ptr_load(c["beta"] + j) + stv = [stv[e] + vn * kv[e] for e in range(4)] + o = stv[0] * qv[0] + stv[1] * qv[1] + stv[2] * qv[2] + stv[3] * qv[3] + o = bf16r(butterfly(o, (16, 8, 4, 2, 1))) + o = live0.select(o, fx.Float32(0.0)) + slot = k["slots"][j] + if (slot >= 0) & live0: + dst = rsrc(a["rstate"] + fx.Int64(slot) * fx.Int64(a["rs_seq"]) * 4) + bo.buffer_store( + fx.Vector.from_elements(stv, fx.Float32), dst, k["head_off"] + ) + if lane % 32 == 0: + fx.ptr_store(o, c["ol"] + (j * VCH + v)) + gpu.barrier() + if tid < L: + ss = fx.Float32(0.0) + for e in range_constexpr(VCH): + x = fx.ptr_load(c["ol"] + (tid * VCH + e)) + ss = ss + x * x + c["put"](c["onrm"], ((req * NH + h) * L + tid) * NVCH + ch, ss) + gpu.barrier() + + +@traced +def kda_norm(c, req, h, ch): + """Rows 16 ch ..: gated RMSNorm over the head's 128 (every chunk's partial) + -> core_out.""" + L, tid = c["L"], c["tid"] + a = c["args"] + if tid < L * VCH: + j = tid // VCH + e = tid % VCH + vr = ch * VCH + e + got = c["poll"]( + [(c["onrm"], ((req * NH + h) * L + j) * NVCH + q, 1) for q in range(NVCH)] + + [(c["proj"], (req * L + j) * NPROJ + G2_0 + h * HD + vr, 1)] + ) + ss = sum_polled(got[:NVCH]) + r = hw_rsq(ss * fx.Float32(1.0 / HD) + fx.Float32(c["onorm_eps"])) + y = ( + fx.ptr_load(c["ol"] + (j * VCH + e)) + * r + * ld_bf(a["on_w"], vr) + * sigmoid(f32_of(got[NVCH][0])) + ) + bo.buffer_store( + y.to(fx.BFloat16), rsrc(a["core_out"]), (req * L + j) * PROJ + h * HD + vr + ) + gpu.barrier() + + +@traced +def kda_tasks(c, on, tls): + s, L, tid, bid = c["S"], c["L"], c["tid"], c["bid"] + for task in range(first_task(bid, c["kda0"]), (s // L) * NH * NVCH, BLOCKS): + req = task // (NH * NVCH) + h = (task // NVCH) % NH + ch = task % NVCH + k = kda_stage(c, req, h, ch) + stamp(on, tls, tid, 7) + kda_conv(c, k, req, h, ch) + stamp(on, tls, tid, 8) + kda_gates(c, k, req, h) + stamp(on, tls, tid, 9) + kda_recur(c, k, req, h, ch) + stamp(on, tls, tid, 10) + kda_norm(c, req, h, ch) + stamp(on, tls, tid, 11) + + +# ---------------------------------------------------------------- kernel +_BUILDS: dict = {} + + +def build(key: KdaPreBuild): + if key in _BUILDS: + return _BUILDS[key] + s, L = key.tokens, key.qlen + ns = key.nblocks + 1 + lay = scratch_layout(key) + xrow = HIDDEN + LDS_PAD + assert s * xrow * 2 <= 136 * 1024, "x rows exceed LDS" + assert L <= 16 and key.nblocks >= 1 + + @fx.struct + class AttnLds: + vals: fx.Array[fx.Float32, s * ns * COLS, 16] + + @fx.struct + class GemvLds: + xl: fx.Array[fx.Int32, s * xrow // 2, 16] + + @fx.struct + class KdaLds: + fa: fx.Array[fx.Float32, L * HD, 16] + qkl: fx.Array[fx.Float32, L * 2 * HD, 16] + decay: fx.Array[fx.Float32, L * HD, 16] + vloc: fx.Array[fx.Float32, L * VCH, 16] + ol: fx.Array[fx.Float32, L * VCH, 16] + beta: fx.Array[fx.Float32, 16, 16] + + @fx.union + class RoleLds: + attn: AttnLds + gemv: GemvLds + kda: KdaLds + + @fx.struct + class Smem: + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + tls: fx.Array[fx.Int64, TL_POINTS, 16] + + keyed = key_tuple(key, _SOURCES) + name = ( + f"k3_mono_kda_pre_s{s}_l{L}_nb{key.nblocks}_d{int(key.delta)}" + f"_g{int(key.debug)}_t{int(key.timeline)}" + ) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def kda_pre( + prefix: Int64, + delta: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + w_in: Int64, + w_fb: Int64, + conv_w: Int64, + conv_st: Int64, + cs_seq: Int32, + cs_dim: Int32, + cs_tok: Int32, + a_log: Int64, + dt_bias: Int64, + on_w: Int64, + rstate: Int64, + rs_seq: Int32, + st_idx: Int64, + st_stride: Int32, + num_acc: Int64, + core_out: Int64, + x_dbg: Int64, + proj_dbg: Int64, + scratch: Int64, + layer: Int32, + epoch: Int64, + tl: Int64, + ): + _ = keyed + tid = fx.thread_idx.x + bid = fx.block_idx.x + alloc = fx.SharedAllocator() + lds = alloc.allocate(Smem).peek() + role = alloc.allocate(RoleLds) + al, gl, kl = role.attn.peek(), role.gemv.peek(), role.kda.peek() + mb = Mailbox(step_tag(epoch, layer) - 1) + c = { + "S": s, + "L": L, + "NS": ns, + "NB": key.nblocks, + "SL": key.state_len, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "delta": key.delta, + "debug": key.debug, + "eps": key.eps, + "out_eps": key.out_eps, + "onorm_eps": key.onorm_eps, + "lb": key.lower_bound, + "kda0": key.kda0, + "xp": key.xp, + "write_idx": key.write_idx, + "gate_local": False, + "put": mb.put, + "poll": mb.poll, + "red": lds.red.ptr, + "vals": al.vals.ptr, + "xl": gl.xl.ptr, + "fa": kl.fa.ptr, + "qkl": kl.qkl.ptr, + "decay": kl.decay.ptr, + "vloc": kl.vloc.ptr, + "ol": kl.ol.ptr, + "beta": kl.beta.ptr, + "xbuf": scratch + fx.Int64(lay["x"][0]), + "args": { + "prefix": prefix, + "delta": delta, + "blocks": blocks, + "blk_sm": blk_sm, + "blk_sr": blk_sr, + "ares_nw": ares_nw, + "ares_qk": ares_qk, + "in_nw": in_nw, + "w_in": w_in, + "w_fb": w_fb, + "conv_w": conv_w, + "conv_st": conv_st, + "cs_seq": cs_seq, + "cs_dim": cs_dim, + "cs_tok": cs_tok, + "a_log": a_log, + "dt_bias": dt_bias, + "on_w": on_w, + "rstate": rstate, + "rs_seq": rs_seq, + "st_idx": st_idx, + "st_stride": st_stride, + "num_acc": num_acc, + "core_out": core_out, + "x_dbg": x_dbg, + "proj_dbg": proj_dbg, + }, + } + for region in ("part", "wgt", "msq", "xrdy", "proj", "qk", "onrm"): + c[region] = sreg(scratch, lay[region][0], region) + on, tls = key.timeline, lds.tls.ptr + c["on"], c["tls"] = on, tls + stamp_begin(on, tls, tid, TL_POINTS) + if bid < NA: + stage_slice(c, bid) + stamp(on, tls, tid, 1) + if bid < s: + stage_gate(c, bid) + stamp(on, tls, tid, 2) + stage_mix(c, bid) + stamp(on, tls, tid, 3) + if bid >= NA: + gi = bid - NA + rg0 = 2 * gi + rg1 = 2 * gi + 1 + ops0 = gemv_loads(c, rg0, range(PER_WAVE)) + load_x(c) + stamp(on, tls, tid, 4) + # the second group's first half in flight under the first's MFMAs + ops1 = gemv_loads(c, rg1, range(4)) + stage_gemv(c, rg0, ops0) + stamp(on, tls, tid, 5) + ops1 = ops1 + gemv_loads(c, rg1, range(4, PER_WAVE)) + stage_gemv(c, rg1, ops1) + stamp(on, tls, tid, 6) + kda_tasks(c, on, tls) + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + + @flyc.jit + def launch( + prefix: Int64, + delta: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + w_in: Int64, + w_fb: Int64, + conv_w: Int64, + conv_st: Int64, + cs_seq: Int32, + cs_dim: Int32, + cs_tok: Int32, + a_log: Int64, + dt_bias: Int64, + on_w: Int64, + rstate: Int64, + rs_seq: Int32, + st_idx: Int64, + st_stride: Int32, + num_acc: Int64, + core_out: Int64, + x_dbg: Int64, + proj_dbg: Int64, + scratch: Int64, + layer: Int32, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + kda_pre( + prefix, + delta, + blocks, + blk_sm, + blk_sr, + ares_nw, + ares_qk, + in_nw, + w_in, + w_fb, + conv_w, + conv_st, + cs_seq, + cs_dim, + cs_tok, + a_log, + dt_bias, + on_w, + rstate, + rs_seq, + st_idx, + st_stride, + num_acc, + core_out, + x_dbg, + proj_dbg, + scratch, + layer, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + ABI.check(kda_pre, launch) + _BUILDS[key] = launch + return launch + + +def ptr(t): + return 0 if t is None else t.data_ptr() + + +def kda_pre( + key: KdaPreBuild, + *, + prefix, + delta, + blocks, + ares_nw, + ares_qk, + in_nw, + w_in, + w_fb, + conv_w, + conv_state, + a_log, + dt_bias, + on_w, + rstate, + st_idx, + num_acc, + core_out, + scratch, + layer, + epoch, + x_dbg=None, + proj_dbg=None, + tl=None, +): + """Host launcher. ``conv_state`` [slots, dim, state_len] or a transposed view + of [slots, state_len, dim] (strides read); ``rstate`` [slots, 12, 128, 128] + fp32; ``st_idx`` [R, >= L] int32; ``scratch`` zeroed before the step's first + layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and + ``layer`` this launch's slot below ``TAG_SLOTS``.""" + assert w_in.shape == (NPROJ, HIDDEN) and w_in.is_contiguous() + assert rstate.dtype == torch.float32 and rstate[0].numel() == NH * HD * HD + assert rstate[0].is_contiguous(), ( + "a slot's state must be dense; slots may be strided" + ) + assert conv_state.dtype == torch.bfloat16 + f = build(key) + f( + prefix.data_ptr(), + ptr(delta), + blocks.data_ptr(), + blocks.stride(0), + blocks.stride(1), + ares_nw.data_ptr(), + ares_qk.data_ptr(), + in_nw.data_ptr(), + w_in.data_ptr(), + w_fb.data_ptr(), + conv_w.data_ptr(), + conv_state.data_ptr(), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + a_log.data_ptr(), + dt_bias.data_ptr(), + on_w.data_ptr(), + rstate.data_ptr(), + rstate.stride(0), + st_idx.data_ptr(), + st_idx.stride(0), + num_acc.data_ptr(), + core_out.data_ptr(), + ptr(x_dbg), + ptr(proj_dbg), + scratch.data_ptr(), + layer, + epoch.data_ptr(), + ptr(tl), + stream=torch.cuda.current_stream(), + ) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/__init__.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/__init__.py new file mode 100644 index 000000000000..0ce41a0e09d7 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 mono decode: common.""" diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/debug.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/debug.py new file mode 100644 index 000000000000..4c5571be4dc1 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/debug.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/stamps.py +# atom/mono/runtime/debug.py +"""A timeline build's per-CTA phase stamps (debug only; ``ATOM_MONO_TIMELINE``, +saved by ``atom.mono.runtime.timeline``): the s_memrealtime (100 MHz) at which +the CTA passes point k, kept in LDS until the kernel ends -- a global store per +stamp would sit on every later ``s_waitcnt vmcnt(0)``. + +``on`` is the build's compile-time switch: off, every call traces to nothing. +""" + +import zlib + +import flydsl.expr as fx +import torch +from flydsl.expr import const_expr + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import memrealtime, traced +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import BLOCKS, DIAG_WORDS + + +@traced +def stamp_begin(on, tls, tid, points): + """Zero the LDS record, then point 0.""" + if const_expr(on): # noqa: SIM102 + if tid < points: # the stamping thread's wave: ordered + fx.ptr_store(fx.Int64(0), tls + tid) + stamp(on, tls, tid, 0) + + +@traced +def stamp(on, tls, tid, k): + # compile-time gate outside, traced condition inside: they cannot be one `and` + if const_expr(on): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) + + +@traced +def stamp_flush(on, tls, tl, tid, bid, points): + """The CTA's record -> ``tl`` [CTA][points] int64.""" + if const_expr(on): # noqa: SIM102 + if tid < points: + fx.generic_store( + fx.inttoptr( + fx.PointerType.get(fx.Int64.ir_type, fx.AddressSpace.Global, 8), + tl + fx.Int64((bid * points + tid) * 8), + ), + fx.ptr_load(tls + tid), + ) + + +# An ``ATOM_MONO_DEBUG`` step's report: the mailbox waits that gave up +# (``sync._spin_bounded`` records, one a CTA) and the ranks the +# step fence gave up on (``StepMailboxes.missing_ranks``). Synchronizes: eager only. + + +DIAG_BYTES = BLOCKS * DIAG_WORDS * 4 + + +def region_id(name: str) -> int: + """A mailbox region's id in a wait record: its name's CRC-32 (31 bits), the + same in every kernel and table that names it.""" + return zlib.crc32(name.encode()) & 0x7FFFFFFF + + +def region_ids(names) -> dict[str, int]: + """Name -> ``region_id``: a debug build's ``Mailbox`` table.""" + return {n: region_id(n) for n in names} + + +def region_namer(names): + """``region_id`` -> name for ``given_up_waits`` (an unknown id as itself).""" + by_id = {region_id(n): n for n in names} + return lambda rid: by_id.get(rid, rid) + + +def given_up_waits(diag: torch.Tensor, name_of) -> list[str]: + """Each CTA's first wait of the step that gave up, the earliest first: + ``diag`` the records' ``DIAG_BYTES`` bytes, ``name_of(region_id)`` the name a + record's region goes by.""" + torch.accelerator.synchronize() + records = diag.view(torch.int32).view(BLOCKS, DIAG_WORDS) + records = records[records[:, 0] != 0].tolist() + t0 = min((r[6] for r in records), default=0) + return [ + f"+{(at - t0) / 100:.0f} us CTA {cta}: {name_of(rid)} pair {idx} waited for" + f" tag {tag}, saw {seen}" + for _, rid, idx, tag, seen, cta, at, _ in sorted(records, key=lambda r: r[6]) + ] + + +def raise_if_given_up(what: str, rank: int, rows: int, waits, mailboxes) -> None: + """Raise with ``waits`` and the fence's missing ranks, if any.""" + lines = list(waits) + [ + f"step fence: rank {r} never arrived" for r in mailboxes.missing_ranks() + ] + if lines: + raise RuntimeError( + f"{what} rank {rank}: {len(lines)} waits gave up ({rows}-row step):\n" + + "\n".join(lines) + ) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/mx.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/mx.py new file mode 100644 index 000000000000..8c12ca9aea94 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/mx.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/mx.py +"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale +codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight +tile's operand loads and the scaled MFMA that takes them. + +An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import rocdl +from flydsl.expr.typing import T + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import CM_NT + +FP8_MAX = 448.0 +# the f8f6f4 MFMA's operand formats (cbsz / blgp) +FP8, FP4 = 0, 4 +UNIT_SCALE = 127 # the E8M0 code of 1.0 + + +def clamp_fp8(v): + """``v`` clamped to the E4M3 range.""" + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def code_ceil(v): + """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), + clamped to [1, 254]: the codes a scale may take.""" + bits = v.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + + +def pow2(code): + """The fp32 power of two a biased exponent stands for.""" + return (code << 23).bitcast(fx.Float32) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale + bit for bit, a multiply where the division is a dozen instructions.""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def fp4_tile_load(r_w, base, rg, k, st, at): + """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` + (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a + step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). + Nontemporal: a weight a step reads once.""" + return fx.Vector( + bo.buffer_load( + r_w, + base + rg * (2 * k) + st * 256 + at, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): + """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 + ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its + byte ``sb_byte``).""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] + ) + ) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ops.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ops.py new file mode 100644 index 000000000000..787bba5651b9 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ops.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ops.py +"""Device-side primitives of the mono kernels: wave reductions and shuffles, the +approximate hardware math, bf16 / fp8 packing, buffer resources, the +execution-model constants and the kernel symbol naming. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs +and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``. +""" + +from __future__ import annotations + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import gpu +from flydsl.expr import math as fmath +from flydsl.expr import range_constexpr, rocdl +from flydsl.expr.typing import T, as_ir_value + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import THREADS, WAVES + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 + + +CM_SYS = 17 + + +# NT = 2: nontemporal. A weight a layer reads once streams past the caches it +# would only evict the step's reused data from (K2b at 48 rows: 143 -> 133 us) +CM_NT = 2 + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``atom::_``. The + amdhsa assembler rejects ``::`` in a symbol, so it is Itanium-mangled, as + aiter's ``_ZN5aiter...E`` kernels are.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN4atom{len(name)}{name}E" + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def readlane(v, src_lane): + """Lane ``src_lane``'s value of v (wave-uniform).""" + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def unlikely(cond): + """``cond`` marked cold (llvm.expect): LLVM lays the branch out of the hot + path. A rare path's code inline among the hot stages costs them even when it + never runs (the long-context selection: S = 16, 3k context +1.4 us).""" + return fx.Boolean( + llvm.call_intrinsic( + ir.IntegerType.get_signless(1), + "llvm.expect.i1", + [as_ir_value(cond), as_ir_value(fx.Boolean(False))], + [], + [], + ) + ) + + +def fresh(v): + """``v`` (a per-lane Int32) through an opaque move, here: the lane math built + on it is computed at its use, where LLVM would otherwise hoist it out of a task + loop or CSE it across stages into registers held for the whole kernel (spills, + and a kernel with scratch is refused: ``runtime.widths``).""" + return fx.Int32( + llvm.inline_asm( + T.i32, [fx.Int32(v).ir_value()], "v_mov_b32 $0, $1", "=v,v", + has_side_effects=True, + ) + ) # fmt: skip + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS). Powers of two only: any other + offset would take the lane ^ 1 path.""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +@traced +def block_max(v, lane, wave, red): + """Max of v over the CTA's WAVES waves; ``red`` needs WAVES words.""" + wv = wave_max(v) + gpu.barrier() + if lane == 0: + fx.ptr_store(wv, red + wave) + gpu.barrier() + t = fx.ptr_load(red + 0) + for w in range_constexpr(1, WAVES): + t = fx.max(t, fx.ptr_load(red + w)) + gpu.barrier() + return t + + +def row_shr(v, lane, k): + """DPP row_shr:k of an f32: lane - k of the same 16-lane row, else 0.""" + moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) + return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) + + +def _xpartner(v, off, lane): + """Lane ``lane ^ off``'s v: DPP within a row, v_permlane*_swap across rows + (``xshfl`` takes ``shuffle_xor`` there).""" + if off < 16: + return xshfl(v, off) + a, b = permlane_swap(off, v, v) + return ((lane & off) != 0).select(a, b) + + +def wave_incl_scan(v, lane): + """Inclusive prefix sum of an Int32 over the wave, in lane order.""" + blk = v + for off in (1, 2, 4, 8, 16, 32): + o = _xpartner(blk, off, lane) + v = v + ((lane & off) != 0).select(o, 0) + blk = blk + o + return v + + +@traced +def block_excl_scan(v, lane, wave, buf): + """(exclusive prefix of v over the CTA in thread order, CTA total); ``buf`` + needs WAVES words.""" + incl = wave_incl_scan(v, lane) + gpu.barrier() + if lane == 63: + fx.ptr_store(incl, buf + wave) + gpu.barrier() + off = fx.Int32(0) + tot = fx.Int32(0) + for w in range_constexpr(WAVES): + cnt = fx.ptr_load(buf + w) + off = off + (w < wave).select(cnt, 0) + tot = tot + cnt + return incl - v + off, tot + + +def global_load(addr, words=1, nt=False): + """``words`` i32 (1, 2 or 4) at the 64-bit global address ``addr``. A pool + past 4 GiB is reached only this way: a buffer load's offset (index x stride + included) is 32 bits. ``nt``: nontemporal.""" + i32 = ir.IntegerType.get_signless(32) + ty = i32 if words == 1 else ir.VectorType.get([words], i32) + ptr = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr))) + return llvm.LoadOp( + ty, ptr.result, alignment=4 * words, nontemporal=nt or None + ).result + + +def ld_i32(ptr, i): + """Word ``i`` of a buffer, a plain Int32 load.""" + return fx.Int32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32)) + + +def row_live(batch_ids, k): + """Whether row k of the step is a request's: its ``batch_id_per_q_token`` + entry, -1 on a CUDA-graph pad row (wave-uniform for a uniform k). A pad row + runs like any other but joins no decision another row's work depends on.""" + return uniform(ld_i32(batch_ids, k)) >= 0 + + +def ld_dev(ptr, i, words=1): + """``words`` i32 (1, 2 or 4) of a buffer from word ``i``, at device scope + (past the per-XCD L2).""" + return bo.buffer_load( + rsrc(ptr), i, vec_width=words, dtype=T.i32, cache_modifier=CM_DEV + ) + + +LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python + int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. + A round whose ``use`` stores to global memory keeps the next round's loads + behind it (they may alias), so the rounds' loads would go out in series. A + round past n loads element n - 1 (in range) and uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def rows_to_lds(tid, src, rows, row_words, dst, dst_row): + """``rows`` rows of ``row_words`` words (a multiple of 4) read at device scope + from address ``src`` -> LDS ``dst``, a row every ``dst_row`` words: 16 B a + lane, the loads batched (``batched_rounds``).""" + assert row_words % 4 == 0, row_words + batched_rounds( + tid, rows * row_words // 4, + lambda v: fx.Vector(ld_dev(src, v * 4, 4)), + lambda v, w: fx.ptr_store( + w, dst + (v * 4 // row_words * dst_row + v * 4 % row_words) + ), + ) # fmt: skip + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` + + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in + order; all of them by default).""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def mfma_fp8(a, b, c): + """16x16x32 fp8 x fp8 MFMA, unscaled: ``a`` / ``b`` an i64 of 8 E4M3 each.""" + return fx.Vector( + rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a, b, c, 0, 0, 0]) + ) + + +def ballot(pred): + """The wave's lanes where ``pred`` holds, a bit a lane.""" + return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) + + +def popcount(x): + return fx.Int32(fx.ctpop(x)) + + +def lanes_below(bal, lane): + """How many of the lanes in ballot ``bal`` are below ``lane``.""" + return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + + +def load_ptr64(table, i): + """Entry ``i`` of an i64 table in memory, wave-uniform (each half through + ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" + pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) + + +def fp8x8_bf16_pk(words): + """Eight fp8 of two dwords -> bf16 MFMA operand (exact) by + v_cvt_scalef32_pk_bf16_fp8 (scale 1), two values an instruction.""" + out = [] + for w in words: + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(1.0)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32).bitcast(fx.BFloat16) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def lds_bytes(struct) -> int: + """The LDS an ``fx.struct`` class takes, alignment included.""" + return struct.__dsl_size_of__() + + +def bf_lo(w): + """The f32 of a packed bf16 pair word's low half.""" + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + """The f32 of a packed bf16 pair word's high half.""" + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return bf_lo(w), bf_hi(w) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/peer_memory.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/peer_memory.py new file mode 100644 index 000000000000..40fed56d7367 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/peer_memory.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/peer_memory.py +# atom/mono/runtime/consensus.py +"""Symmetric peer memory for in-kernel TP reductions.""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a + non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no rank + overwrites a peer's unread data, is the model's layout's to say. + + ``bytes`` is the whole buffer: the kernels tag every pair with the step's + epoch, so it is never zeroed between steps. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so a + runner reaches it only after ``tp_agree``. A failed mapping on any rank raises + ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + # a view, not an owner: close() frees the memory under it + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: # every rank must still reach the agreement + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def kernel_args(self) -> dict: + """A kernel's view of the group, by ABI name: this rank's buffer + (``sym``), the table of every rank's (``peers``) and the TP ``rank``.""" + return { + "sym": self.local, + "peers": self.addresses.data_ptr(), + "rank": self.rank, + } + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent. Only once no + kernel of any rank can still touch it (the runner is being dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + self.bytes = None + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 + + +# Decisions every TP rank must take alike. +# +# A mono kernel waits in-kernel on every peer rank, so a rank that turned mono off +# alone -- or failed to build what the others launch -- leaves the others spinning +# on the GPU for ever. Each such decision is a local verdict agreed on over the TP +# group's CPU group before anything collective (peer handshake, launch) follows. +# + + +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" + + +def tp_agree(ok: bool, group) -> bool: + """True when every rank of ``group`` passed ``ok`` true. Every rank must call + it at the same point; ``group`` is a CPU (gloo) process group, or None for a + single rank.""" + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +def bind_agreed(bind, group) -> None: + """``bind()`` on this rank (raising ``MonoUnsupported`` to refuse), then every + rank's verdict agreed: a rank that raised re-raises, the others raise too.""" + try: + bind() + except Exception: + tp_agree(False, group) + raise + if not tp_agree(True, group): + raise MonoUnsupported("another TP rank refused mono") diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/plan.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/plan.py new file mode 100644 index 000000000000..53ca755ebf22 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/plan.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/plan/build_key.py +# atom/mono/plan/execution.py +# atom/mono/plan/layout.py +# atom/mono/runtime/abi.py +# atom/mono/plan/trace.py +"""A kernel build's parameters, all of them in FlyDSL's JIT cache key. + +FlyDSL keys a compiled kernel by the scalar values (int, float, bool, str, enum, +tuple) its launcher's closures capture; an object in a closure is not keyed, and +a parameter used only by build-time Python (a placement table, a layout) is in no +closure at all. Two builds differing only there would share one binary -- every +TP rank ran rank 3's kernel once that way. So a build's parameters are one frozen +dataclass of scalars, ``key_tuple`` turns it into one (name, value) tuple, and the +launcher references that tuple: every field is keyed, whatever the body uses. + +The tuple also carries the model's kernel sources (``source_digest``): FlyDSL +keys a helper's source only when its walk reaches it, and it does not reach a +helper called inside a traced function's dynamic loop or branch (the rewriter +moves those bodies into nested functions it does not scan) nor one outside the +kernel's directory -- an edit there kept the old binary. +""" + +import enum +import functools +import hashlib +import inspect +from contextlib import contextmanager +from dataclasses import dataclass, field, fields +from enum import Enum +from pathlib import Path + +_SCALARS = (int, float, bool, str, enum.Enum) +_ROOT = Path(__file__).resolve().parents[1] + + +@functools.cache +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the ``mono`` package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these.""" + h = hashlib.sha256() + for rel in paths: + root = _ROOT / rel + files = [root] if root.is_file() else sorted(root.rglob("*.py")) + if not files: + raise FileNotFoundError(f"no kernel sources at {root}") + for f in files: + h.update(str(f.relative_to(_ROOT)).encode()) + h.update(f.read_bytes()) + return h.hexdigest()[:16] + + +def key_tuple(key, sources: str) -> tuple: + """``key``'s (field name, value) pairs, every value a scalar, then + ("sources", ``sources``): the model's ``source_digest``.""" + out = [] + for f in fields(key): + value = getattr(key, f.name) + if not isinstance(value, _SCALARS): + raise TypeError( + f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" + " a scalar to reach the JIT cache key" + ) + out.append((f.name, value)) + out.append(("sources", sources)) + return tuple(out) + + +def symbol_params(key) -> dict: + """The fields that name the kernel symbol (``metadata={"sym": short name}``), + short name -> value, in field order.""" + return { + f.metadata["sym"]: getattr(key, f.name) + for f in fields(key) + if "sym" in f.metadata + } + + +# The execution model every mono kernel is written for. +# +# One CTA per CU of an MI355X, all of them co-resident -- a spin wait on another +# CTA can only make progress if that CTA is running, so the runner refuses a GPU +# with another CU count -- and THREADS threads each. +# + +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 +LDS_BYTES = 160 * 1024 # a CU's LDS: one CTA's whole + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of + CTAs past 0: placements that add task counts run past the grid), then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# A kernel's mailbox regions laid out in its scratch or peer buffer. + +PAIR_BYTES = 8 # a (value, tag) pair +# a debug build's wait record, one a CTA: flag, region id, pair index, tag, +# seen tag, CTA, clock, pad (``sglang.kernels.ops.moe.k3_mono_flydsl.common.sync._spin_bounded``) +DIAG_WORDS = 8 + + +def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: + """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), + each region from an ``align``-byte boundary at or past ``start``.""" + out, off = {}, start + for name, pairs in items: + off = (off + align - 1) // align * align + out[name] = (off, pairs * PAIR_BYTES) + off += pairs * PAIR_BYTES + return out + + +# A kernel's argument list, named once. +# +# A FlyDSL kernel and its launcher each spell out their positional parameters, and +# the host passes them by position: an argument out of order is a pointer read as +# another, silently. ``KernelAbi`` is the one list of names; the build checks the +# kernel's and the launcher's parameters against it, and the host packs by name. +# + + +def _python_params(fn) -> tuple[str, ...]: + """Parameter names of a ``@flyc.kernel`` / ``@flyc.jit`` object's Python + function (FlyDSL keeps it as ``_func`` / ``func``).""" + py = getattr(fn, "func", None) or getattr(fn, "_func", None) + if py is None: + raise TypeError(f"no Python function behind {fn!r}") + return tuple(inspect.signature(py).parameters) + + +@dataclass(frozen=True) +class KernelAbi: + names: tuple[str, ...] + + def check(self, kernel, launcher) -> None: + """The kernel takes exactly ``names``; the launcher too, then ``stream``.""" + got = _python_params(kernel) + if got != self.names: + raise TypeError(f"kernel parameters {got} != ABI {self.names}") + got = _python_params(launcher) + if got != (*self.names, "stream"): + raise TypeError(f"launcher parameters {got} != ABI {self.names} + stream") + + def pack(self, values: dict) -> list: + """``values`` (every name, nothing else) in ABI order.""" + try: + if len(values) == len(self.names): + return [values[n] for n in self.names] + except KeyError: + pass + extra = sorted(set(values) - set(self.names)) + missing = sorted(set(self.names) - set(values)) + raise TypeError(f"ABI mismatch: missing {missing}, unexpected {extra}") + + def zeros(self) -> list: + """Every argument 0: what a compile-only call passes.""" + return [0] * len(self.names) + + +# A kernel's mailbox accesses, recorded while FlyDSL traces it. +# +# Tracing runs the kernel body's Python once, so every mailbox put / poll executes +# once as Python: noting it here costs the device nothing (no IR is emitted). The +# stage currently being traced is set by ``enter_stage``, called where the program +# starts each stage, so the record is ``(stage, region, put | poll)`` in program +# order. +# +# Everything here is plain functions on a module-level recorder: FlyDSL's AST +# rewriter carries every local whose method is called inside a dynamic if / for as +# loop state, so a kernel body must not call methods on these objects. +# + + +class Space(Enum): + SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs + PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs + + +@dataclass(frozen=True) +class Access: + stage: str + region: str + space: Space + kind: str # "put" | "poll" + + +@dataclass +class Record: + stages: list[str] = field(default_factory=list) # in the order entered + accesses: list[Access] = field(default_factory=list) + unnamed: list[tuple[str, str]] = field(default_factory=list) # (stage, kind) + + +_record: Record | None = None + + +@contextmanager +def recording(): + """Record the accesses of kernels traced inside this block.""" + global _record + outer, _record = _record, Record() + try: + yield _record + finally: + _record = outer + + +def enter_stage(name: str) -> None: + if _record is not None: + _record.stages.append(name) + + +def _stage() -> str: + assert _record is not None + return _record.stages[-1] if _record.stages else "" + + +def note(region: str, space: Space, kind: str) -> None: + if _record is not None: + _record.accesses.append(Access(_stage(), region, space, kind)) + + +def note_unnamed(kind: str) -> None: + """A put / poll on a raw address: the check reports it (every mailbox access of + a checked kernel names its region).""" + if _record is not None: + _record.unnamed.append((_stage(), kind)) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ranks.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ranks.py new file mode 100644 index 000000000000..6c578cc01ccd --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/ranks.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ranks.py +"""A TP group seen from inside a kernel: every rank's peer buffer, and an +all-reduce's sum of the ranks' partials. ``tp`` is a build-time constant.""" + +from flydsl.expr import range_constexpr + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import bf2_f32, load_ptr64 + + +def peer_bases(peers, tp): + """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), + loaded once: a put's address must not wait on a load of its own.""" + return [load_ptr64(peers, p) for p in range(tp)] + + +def sum_partials(poll, own, pair_of, tp): + """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox + region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> + (lo sum, hi sum). Every rank sums in the same order, so they agree.""" + ws = poll([(own, pair_of(src), 1) for src in range(tp)]) + acc0, acc1 = bf2_f32(ws[0][0]) + for src in range_constexpr(1, tp): + lo, hi = bf2_f32(ws[src][0]) + acc0 = acc0 + lo + acc1 = acc1 + hi + return acc0, acc1 diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/sync.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/sync.py new file mode 100644 index 000000000000..94aa75505f97 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/common/sync.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/sync.py +"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and +between GPUs (system scope), every address a named region (``Addr``) recorded +when the kernel is traced (``trace``).""" + +from __future__ import annotations + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + CM_DEV, + CM_SYS, + bf16_pair, + memrealtime, + rsrc, + traced, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import DIAG_WORDS, Space, note + +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +# a debug build's mailbox wait gives up after this long (100 MHz ticks: 10 s) +DIAG_TICKS = 10 * 100_000_000 + + +@traced +def _spin_bounded(load_all, pending, diag, region_id, index, tag): + """``_spin`` giving up after DIAG_TICKS: the CTA's record in ``diag`` names the + wait (the runner raises with it after the step), and the kernel runs on.""" + t0 = memrealtime() + v = load_all() + late = fx.Int32(0) + while pending(v) & (late == 0): + rocdl.s_sleep(1) + v = load_all() + late = (memrealtime() - t0 > DIAG_TICKS).select(fx.Int32(1), fx.Int32(0)) + if late != 0: + bid = fx.block_idx.x + # the CTA's first give-up of the step: the later ones wait on its fallout + seen = bo.buffer_load( + rsrc(diag), + bid * DIAG_WORDS, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_SYS, + ) + if fx.Int32(seen) == 0: + at = fx.Int32(memrealtime() & 0x7FFFFFFF) + record = [1, region_id, index, tag, v[1], bid, at, 0] + for h in range_constexpr(2): # 16 B a store + bo.buffer_store( + fx.Vector.from_elements( + [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32 + ), + rsrc(diag), + bid * DIAG_WORDS + 4 * h, + cache_modifier=CM_SYS, + ) + return v + + +@dataclass(frozen=True) +class Addr: + """A mailbox region's address: the region's name and space, recorded when the + kernel is traced (``trace``), and the traced address. The space also + sets the hand-off's scope: device for SCRATCH, system for PEER.""" + + region: str + space: Space + value: object + + +def sreg(base, offset, region): + """Scratch region ``region`` at byte ``offset`` of ``base``.""" + return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) + + +def preg(base, offset, region): + """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer + ``base``.""" + return Addr(region, Space.PEER, base + fx.Int64(offset)) + + +def shift(a, delta): + """``a`` moved by ``delta`` bytes (a traced or Python integer), same region.""" + return Addr(a.region, a.space, a.value + delta) + + +# the store / load cache policy of a hand-off in each space +_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} + + +def _noted(a, kind): + note(a.region, a.space, kind) + return a.value + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The tag is ``layer + 1``: it tells the + layers of one step apart, and the runner zeroes every mailbox at a step's + start (``StepMailboxes.begin_step``), so a pair from an earlier step never + carries a live tag. Every address is an ``Addr``: its space picks the scope. + """ + + def __init__(self, layer, diag=None, region_ids=None): + """``diag`` (a debug build): the address of the per-CTA wait records, every + wait bounded (``_spin_bounded``); ``region_ids``: region name -> the id a + record names it by.""" + self.tag = layer + 1 + self.diag = diag + self.region_ids = region_ids or {} + + def put(self, a, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(_noted(a, "put")), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_bf(self, a, i, vs): + """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one + store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(_noted(a, "put")), + i, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_words(self, a, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(_noted(a, "put")), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], + all of one space. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + spaces = {a.space for a, _, _ in specs} + assert len(spaces) == 1, f"one poll batch spans {spaces}" + cm = _SCOPE_CM[spaces.pop()] + bases = [_noted(a, "poll") for a, _, _ in specs] + tag = self.tag + + def load_all(): + words = [] + for b, (_, i, n) in zip(bases, specs): + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + if const_expr(self.diag is None): + v = _spin(load_all, pending) + else: + rid = self.region_ids.get(specs[0][0].region, -1) + v = _spin_bounded(load_all, pending, self.diag, rid, specs[0][1], tag) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +@traced +def publish(put, a, i, v, who): + """This CTA's plain stores made readable by every CTA, then (thread ``who``) + the mailbox pair (``a``, i) := v. The stores are drained first: a flag may + not overtake the data it announces, and a workgroup-scope release fence does + not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit + mode: the workgroup's waves share one CU, the readers here do not).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if who: + put(a, i, v) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/layer.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/layer.py new file mode 100644 index 000000000000..4cae28ee524c --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/layer.py @@ -0,0 +1,605 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2 ``layer_post``: one Kimi-K3 layer from the attention's core output (the +o_proj input) to the layer's MoE output, in one launch -- K2a's o_proj, its +all-reduce and the MLP AttnRes seam, then K2b's latent MoE (``moe``). + +As ATOM's ``layer_post``, the two halves share the launch so that the MoE's +static weights stream while the AttnRes chain (partials -> reduce -> slice -> +gate -> mix) is still settling: the CTAs that will run the router, latent and +shared gate/up GEMVs issue their weights before they first wait. + + CTAs 0 .. 55 K2a's reduce + AttnRes (slice, gate, mix -> x rows + XRDY); + then route (0 .. S-1) and the shared SiTU (8 .. 8+S-1) + CTAs 56 .. 159 o_proj (448 x 16 rows) -> partials to every rank; then + the shared gate/up (96 tasks, weights right after o_proj) + CTAs 160 .. 243 router parts (28) and latent parts (56): weights first + every CTA ug; down (0 .. 223) or the shared down (224 .. 255); + lnorm (56 .. 83), up (0 .. 55), reduce (200 .. 255) +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import Int32, Int64 + +from sglang.kernels.ops.moe.k3_mono_flydsl.attention import back as k2a +from sglang.kernels.ops.moe.k3_mono_flydsl.attention.kda import ( + _SOURCES, + COLS, + HIDDEN, + NA, + stage_gate, + stage_mix, + stage_slice, + step_tag, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.debug import ( + region_ids, + stamp, + stamp_begin, + stamp_flush, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + ld_i32, + load_ptr64, + traced, + uniform, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import ( + BLOCKS, + THREADS, + WAVES, + KernelAbi, + key_tuple, + pair_layout, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ranks import peer_bases +from sglang.kernels.ops.moe.k3_mono_flydsl.common.sync import Mailbox, preg, shift, sreg +from sglang.kernels.ops.moe.k3_mono_flydsl.stages import gemv as gemv +from sglang.kernels.ops.moe.k3_mono_flydsl.stages import moe as k2b + +O0, O1 = NA, BLOCKS # o_proj CTAs: every CTA past AttnRes's +O_N = O1 - O0 # 200 +OTASKS = k2a.OTASKS # 448 +O_ROUNDS = (OTASKS + O_N - 1) // O_N # 5 +E0 = 160 # router / latent CTAs (their weights in flight under their o_proj) +H0 = 8 # the shared SiTU's CTAs (after their AttnRes work) +SD0 = k2b.DOWN_TASKS # 224: the shared down's CTAs +# timeline points: 0 start, 1 o_proj done / AR reduced, 2 early done / slice, +# 3 ug, 4 down/sdown, 5 lnorm, 6 up, 7 reduce, 8 gate, 9 mix (x published), +# 10 route, 11 route table, 12 latent quant, 13 route starts +TL_POINTS = 17 # + 14 (unused), 15 INTER loaded, 16 down done + +assert E0 + k2b.R_TASKS + k2b.L_TASKS <= BLOCKS and k2b.G_TASKS <= E0 - O0 + +REGIONS = ("part", "wgt", "msq", "xrdy", "ar") + k2b.REGIONS + +_STREAM = fx.Stream(None) + + +@dataclass(frozen=True) +class K2Build: + tokens: int + nblocks: int # the MLP AttnRes block sources + eps: float = 1e-5 # mlp_res_norm + out_eps: float = 1e-5 # post_attention_layernorm + ln_eps: float = 1e-5 # the latent RMSNorm + reduce: bool = True # the final reduce here (else the next layer's K1) + reset: bool = False # a block-write layer: the prefix restarts at the o_proj output + diag: bool = False + timeline: bool = False + xp: int = 0 # timing experiments (7: route without its selection) + guint: bool = False # w13's AITER layout; see stages.moe.MoeBuild + + +ABI = KernelAbi( + ( + "core", + "w_o", + "prefix", + "blocks", + "blk_sm", + "blk_sr", + "ares_nw", + "ares_qk", + "in_nw", + "w_gate", + "bias", + "w_ld", + "w_sgu", + "w_sd", + "w13", + "w13s", + "w2", + "w2s", + "ln_w", + "w_up", + "out", + "scratch", + "peers", + "rank", + "epoch", + "layer", + "diag", + "tl", + ) +) + + +def peer_layout(s): + """K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes)).""" + pairs = {"ar": k2a.ar_half_pairs(s), **k2b.peer_pairs(s)} + return pair_layout(pairs.items()) + + +def peer_bytes(s): + lay = peer_layout(s) + return 2 * max(o + n for o, n in lay.values()) + + +def scratch_layout(key: K2Build) -> dict: + s = key.tokens + a = k2a.scratch_layout(k2a.AttnPostBuild(tokens=s, nblocks=key.nblocks)) + end = max(o + n for o, n in a.values()) + b = k2b.scratch_layout(k2b.MoeBuild(tokens=s)) + lay = dict(a) + for name, (o, n) in b.items(): + lay[name] = ((end + 15) // 16 * 16 + o, n) + end = max(o + n for o, n in lay.values()) + lay["xrow"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) + return lay + + +def scratch_bytes(key: K2Build) -> int: + return max(o + n for o, n in scratch_layout(key).values()) + + +@traced +def stage_oproj(c, gi): + """o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in + flight at once (one memory latency, not three), then the MFMAs.""" + ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)] + k2a.load_core(c) + for j in range_constexpr(O_ROUNDS): + t = gi + O_N * j + if t < OTASKS: + k2a.stage_o(c, t, ops[j]) + + +_BUILDS: dict = {} + + +def build(key: K2Build): + if key in _BUILDS: + return _BUILDS[key] + s, ns = key.tokens, key.nblocks + 1 + assert 1 <= s <= 8 + lay = scratch_layout(key) + play = peer_layout(s) + half = peer_bytes(s) // 2 + + @fx.struct + class RouteLds: + route: fx.Array[fx.Int32, s * 2 * k2b.TOPK, 16] + flag: fx.Array[fx.Int32, k2b.E, 16] + expert: fx.Array[fx.Int32, k2b.MAX_U, 16] + kof: fx.Array[fx.Int32, k2b.MAX_U * 8, 16] + scan: fx.Array[fx.Int32, WAVES, 16] + + @fx.struct + class EarlyLds: + xl: fx.Array[fx.Int32, s * (HIDDEN + gemv.LDS_PAD) // 2, 16] + + @fx.struct + class AttnLds: + vals: fx.Array[fx.Float32, s * ns * COLS, 16] + dl: fx.Array[fx.Float32, s * COLS, 16] + + @fx.struct + class UgLds: + latq: fx.Array[fx.Int32, s * k2b.LAT // 8, 16] + latsc: fx.Array[fx.Int32, s * k2b.LAT // 32, 16] + av: fx.Array[fx.Float32, THREADS, 16] + + @fx.struct + class DownLds: + interq: fx.Array[fx.Int32, s * k2b.TOPK * 48, 16] + intsc: fx.Array[fx.Int32, s * k2b.TOPK * k2b.UG_GROUPS, 16] + contrib: fx.Array[fx.Float32, s * k2b.TOPK * k2b.ROWS, 16] + + @fx.union + class StageLds: + early: EarlyLds + attn: AttnLds + ug: UgLds + down: DownLds + + @fx.struct + class Smem: + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + tls: fx.Array[fx.Int64, TL_POINTS, 16] + + keyed = key_tuple(key, _SOURCES) + name = ( + f"k3_mono_k2_s{s}_nb{key.nblocks}_x{int(key.reset)}_r{int(key.reduce)}" + f"_g{int(key.diag)}_t{int(key.timeline)}" + ) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k2( + core: Int64, + w_o: Int64, + prefix: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + w_gate: Int64, + bias: Int64, + w_ld: Int64, + w_sgu: Int64, + w_sd: Int64, + w13: Int64, + w13s: Int64, + w2: Int64, + w2s: Int64, + ln_w: Int64, + w_up: Int64, + out: Int64, + scratch: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + diag: Int64, + tl: Int64, + ): + _ = keyed + tid = fx.thread_idx.x + bid = fx.block_idx.x + alloc = fx.SharedAllocator() + lds = alloc.allocate(Smem).peek() + rl = alloc.allocate(RouteLds).peek() + st = alloc.allocate(StageLds) + el, al, ul, dl = st.early.peek(), st.attn.peek(), st.ug.peek(), st.down.peek() + if const_expr(key.diag): + mb = Mailbox(step_tag(epoch, layer) - 1, diag, region_ids(REGIONS)) + else: + mb = Mailbox(step_tag(epoch, layer) - 1) + par = uniform(ld_i32(epoch, 0)) % 2 + bases = peer_bases(peers, TP_) + own = load_ptr64(peers, rank) + peer = {} + for region in ("rlog", "lat", "latr", "final", "ar"): + off = play[region][0] + peer[region] = [ + shift(preg(b, off, region), fx.Int64(par) * half) for b in bases + ] + peer[region + "_own"] = shift(preg(own, off, region), fx.Int64(par) * half) + c = { + "S": s, + "NS": ns, + "NB": key.nblocks, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rank": rank, + "delta": True, + "reset": key.reset, + "gate_local": False, + "debug": False, + "eps": key.eps, + "out_eps": key.out_eps, + "ln_eps": key.ln_eps, + "xp": key.xp, + "guint": key.guint, + "lazy_inter": True, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "red": lds.red.ptr, + "peer": peer, + "ar_peer": peer["ar"], + "ar_own": peer["ar_own"], + "xl": el.xl.ptr, + "vals": al.vals.ptr, + "delta_lds": al.dl.ptr, + "latq": ul.latq.ptr, + "latsc": ul.latsc.ptr, + "av": ul.av.ptr, + "interq": dl.interq.ptr, + "intsc": dl.intsc.ptr, + "contrib": dl.contrib.ptr, + "scan": rl.scan.ptr, + "rt": { + "route": rl.route.ptr, + "flag": rl.flag.ptr, + "expert": rl.expert.ptr, + "kof": rl.kof.ptr, + }, + "xbuf": scratch + fx.Int64(lay["xrow"][0]), + "xflags": NA, + "tl_on": key.timeline, + "tls": lds.tls.ptr, + "ln": scratch + fx.Int64(lay["ln"][0]), + "shp": scratch + fx.Int64(lay["shp"][0]), + "args": { + "core": core, + "w_o": w_o, + "prefix": prefix, + "blocks": blocks, + "blk_sm": blk_sm, + "blk_sr": blk_sr, + "ares_nw": ares_nw, + "ares_qk": ares_qk, + "in_nw": in_nw, + "x_dbg": 0, + "w_gate": w_gate, + "bias": bias, + "w_ld": w_ld, + "w_sgu": w_sgu, + "w_sd": w_sd, + "w13": w13, + "w13s": w13s, + "w2": w2, + "w2s": w2s, + "ln_w": ln_w, + "w_up": w_up, + "out": out, + }, + } + for region in ( + "part", + "wgt", + "msq", + "xrdy", + "sgu", + "h", + "rpart", + "lpart", + "route", + "inter", + "isc", + "lmsq", + "lnrdy", + "shrdy", + ): + c[region] = sreg(scratch, lay[region][0], region) + on, tls = key.timeline, lds.tls.ptr + stamp_begin(on, tls, tid, TL_POINTS) + # ---- o_proj on every CTA past AttnRes's + eb = bid - E0 + lb = eb - k2b.R_TASKS + is_r = (bid >= E0) & (eb < k2b.R_TASKS) + is_l = (bid >= E0 + k2b.R_TASKS) & (lb < k2b.L_TASKS) + # o_proj first (every AttnRes CTA waits on the slowest partial; loads + # retire in order), the early weights after: they land before x does + if is_r: + stage_oproj(c, bid - O0) + stamp(on, tls, tid, 1) + rops = k2b.early_loads(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS) + k2b.early_compute(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS, rops) + if is_l: + stage_oproj(c, bid - O0) + stamp(on, tls, tid, 1) + lops = k2b.early_loads(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS) + k2b.early_compute(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS, lops) + if (bid >= O0) & ~is_r & ~is_l: + stage_oproj(c, bid - O0) + stamp(on, tls, tid, 1) + gi = bid - O0 + if gi < k2b.G_TASKS: + k2b.early_compute(c, 2, gi, 0, k2b.early_loads(c, 2, gi, 0)) + # ---- AttnRes: the reduced o_proj -> x; then route / SiTU + if bid < NA: + k2a.stage_reduce(c, bid) + stamp(on, tls, tid, 1) + stage_slice(c, bid) + stamp(on, tls, tid, 2) + if bid < s: + stage_gate(c, bid) + stamp(on, tls, tid, 8) + stage_mix(c, bid) + stamp(on, tls, tid, 9) + stamp(on, tls, tid, 13) + if bid < s: + k2b.stage_route(c, bid) + stamp(on, tls, tid, 10) + if (bid >= H0) & (bid < H0 + s): + k2b.stage_h(c, bid - H0) + if bid >= NA: + stamp(on, tls, tid, 2) + # ---- the routed experts + # the latent (ready ~10 us before the route) quantized first + k2b.load_latq(c) + stamp(on, tls, tid, 12) + nu = k2b.load_route(c) + stamp(on, tls, tid, 11) + total = nu * k2b.UG_GROUPS + last = total - 1 + # reversed placement: the round's leftover tasks go to the high CTAs, + # whose early work ends first (the AttnRes / route CTAs start ug last) + task = BLOCKS - 1 - bid + ops = k2b.ug_loads(c, fx.min(task, last)) + while task < total: + nxt = task + BLOCKS + ops_n = k2b.ug_loads(c, fx.min(nxt, last)) + k2b.stage_ug(c, task, ops) + task = nxt + ops = ops_n + stamp(on, tls, tid, 3) + if bid < k2b.DOWN_TASKS: + dops = k2b.down_loads(c, bid, nu, fx.min(c["wave"], nu - 1)) + k2b.load_inter(c) + stamp(on, tls, tid, 15) + k2b.stage_down(c, bid, nu, dops) + stamp(on, tls, tid, 16) + k2b.run_sdown_rest(c, bid, w_sd) + k2b.run_sdown(c, bid, w_sd) + stamp(on, tls, tid, 4) + if (bid >= k2b.LN0) & (bid < k2b.LN0 + k2b.LN_TASKS): + k2b.stage_lnorm(c, bid - k2b.LN0) + stamp(on, tls, tid, 5) + if bid < k2b.UP_TASKS: + k2b.stage_up(c, bid) + stamp(on, tls, tid, 6) + if const_expr(key.reduce): # noqa: SIM102 (traced: `and` does not combine traced values) + if (bid >= k2b.RED0) & (bid < k2b.RED0 + k2b.RED_TASKS): + k2b.stage_reduce(c, bid - k2b.RED0) + stamp(on, tls, tid, 7) + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + + @flyc.jit + def launch( + core: Int64, + w_o: Int64, + prefix: Int64, + blocks: Int64, + blk_sm: Int32, + blk_sr: Int32, + ares_nw: Int64, + ares_qk: Int64, + in_nw: Int64, + w_gate: Int64, + bias: Int64, + w_ld: Int64, + w_sgu: Int64, + w_sd: Int64, + w13: Int64, + w13s: Int64, + w2: Int64, + w2s: Int64, + ln_w: Int64, + w_up: Int64, + out: Int64, + scratch: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + diag: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k2( + core, + w_o, + prefix, + blocks, + blk_sm, + blk_sr, + ares_nw, + ares_qk, + in_nw, + w_gate, + bias, + w_ld, + w_sgu, + w_sd, + w13, + w13s, + w2, + w2s, + ln_w, + w_up, + out, + scratch, + peers, + rank, + epoch, + layer, + diag, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + ABI.check(k2, launch) + _BUILDS[key] = launch + return launch + + +TP_ = k2b.TP + + +def ptr(t): + return 0 if t is None else t.data_ptr() + + +def k2_launch( + key: K2Build, + *, + core, + w_o, + prefix, + blocks, + ares_nw, + ares_qk, + in_nw, + w_gate, + bias, + w_ld, + w_sgu, + w_sd, + w13, + w13s, + w2, + w2s, + ln_w, + w_up, + out, + scratch, + peers, + rank, + epoch, + layer, + diag=None, + tl=None, +): + """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in + place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``.""" + assert ( + w_o.shape == (HIDDEN, k2a.OK) + and core.is_contiguous() + and prefix.is_contiguous() + ) + assert w_up.shape == (k2b.UP_N, k2b.LAT) and w_up.is_contiguous() + f = build(key) + f( + core.data_ptr(), + w_o.data_ptr(), + prefix.data_ptr(), + blocks.data_ptr(), + blocks.stride(0), + blocks.stride(1), + ares_nw.data_ptr(), + ares_qk.data_ptr(), + in_nw.data_ptr(), + w_gate.data_ptr(), + bias.data_ptr(), + w_ld.data_ptr(), + w_sgu.data_ptr(), + w_sd.data_ptr(), + w13.data_ptr(), + w13s.data_ptr(), + w2.data_ptr(), + w2s.data_ptr(), + ln_w.data_ptr(), + w_up.data_ptr(), + out.data_ptr(), + scratch.data_ptr(), + peers.data_ptr(), + rank, + epoch.data_ptr(), + layer, + ptr(diag), + ptr(tl), + stream=torch.cuda.current_stream(), + ) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/runner.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/runner.py new file mode 100644 index 000000000000..8deb318a5072 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/runner.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono decode's process-wide state: one scratch buffer every launch +shares (their mailbox tags differ), the device step epoch and the TP group's +peer memory. + +Mailbox tags carry the epoch the model bumps once a forward (``step_begin``, +which a CUDA graph captures), so no buffer is zeroed between steps. +""" + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.runtime_context import get_parallel + +if TYPE_CHECKING: + from sglang.kernels.ops.moe.k3_mono_flydsl.common.peer_memory import PeerBuffer + +MAX_TOKENS = 8 # the x rows fit one CTA's LDS + + +class _State: + scratch: torch.Tensor | None = None + zero_bytes = 0 + epoch: torch.Tensor | None = None + queue: torch.Tensor | None = None + peers: "PeerBuffer | None" = None + + +def alloc_scratch(device) -> None: + """At model init, outside any graph capture: the largest build's scratch + and the epoch.""" + if _State.scratch is not None: + return + from sglang.kernels.ops.moe.k3_mono_flydsl import layer + from sglang.kernels.ops.moe.k3_mono_flydsl.attention import kda + from sglang.kernels.ops.moe.k3_mono_flydsl.stages import moe + + n = max( + kda.scratch_bytes( + kda.KdaPreBuild( + tokens=MAX_TOKENS, + qlen=MAX_TOKENS, + nblocks=8, + delta=True, + state_len=3 + MAX_TOKENS, + ) + ), + moe.scratch_bytes(moe.MoeBuild(tokens=MAX_TOKENS)), + layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)), + ) + _State.zero_bytes = (n + 255) // 256 * 256 + _State.scratch = torch.zeros(_State.zero_bytes, dtype=torch.uint8, device=device) + _State.epoch = torch.zeros(1, dtype=torch.int32, device=device) + _State.queue = torch.zeros(moe.QUEUE_BYTES // 4, dtype=torch.int32, device=device) + + +def step_begin() -> None: + """Once a forward, before its first mono launch (a graph captures it).""" + if _State.epoch is not None: + _State.epoch.add_(1) + + +def peers(device) -> "PeerBuffer": + """The TP group's peer buffer. Collective: every rank builds the same layers + in the same order, so every rank reaches it at the first MoE layer.""" + if _State.peers is None: + from sglang.kernels.ops.moe.k3_mono_flydsl import layer + from sglang.kernels.ops.moe.k3_mono_flydsl.common.peer_memory import PeerBuffer + from sglang.kernels.ops.moe.k3_mono_flydsl.stages import moe + + tp = get_parallel().tp_group + nbytes = max(moe.peer_bytes(MAX_TOKENS), layer.peer_bytes(MAX_TOKENS)) + buf = PeerBuffer(nbytes, tp.cpu_group, tp.rank_in_group, tp.world_size, device) + buf.bytes.zero_() + _State.peers = buf + return _State.peers + + +def launch_args() -> dict: + """The state every launch takes.""" + assert _State.peers is not None and _State.epoch is not None + return { + "scratch": _State.scratch, + "peers": _State.peers.addresses, + "rank": get_parallel().tp_group.rank_in_group, + "epoch": _State.epoch, + } + + +def scratch_bytes() -> int: + return _State.zero_bytes + + +def scratch() -> torch.Tensor: + assert _State.scratch is not None + return _State.scratch + + +def epoch() -> torch.Tensor: + assert _State.epoch is not None + return _State.epoch + + +def queue() -> torch.Tensor: + assert _State.queue is not None + return _State.queue diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/__init__.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/__init__.py new file mode 100644 index 000000000000..41a2962dfba0 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 mono decode: stages.""" diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/gemv.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/gemv.py new file mode 100644 index 000000000000..f0323c9c52b3 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/gemv.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""bf16 GEMV pieces the K3 mono kernels share: 16 weight rows a task against up +to 16 token rows in LDS, one ``mfma_f32_16x16x32_bf16`` per 32 K. + +K is walked in 128-wide chunks, wave w taking chunks w, w + 8, ...; a lane reads +64 contiguous bytes of its row a chunk (4 dwordx4 at q 32 + g 8 elements: lane +group g = lane / 16), and the token rows are read at the same K. The K order +inside a chunk is a permutation both operands share, so the dot is unchanged. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import gpu, range_constexpr +from flydsl.expr.typing import T + +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + CM_DEV, + CM_NT, + mfma_bf16, + rsrc, + traced, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import THREADS, WAVES + +ROWS = 16 +KCH = 128 +LDS_PAD = 8 # bf16 a padded token row + + +def chunks(k): + """A wave's K chunks for K = ``k``: (count a wave, total).""" + n = k // KCH + return (n + WAVES - 1) // WAVES, n + + +def loads(lane, wave, w, row0, k, nt=True, ldk=None, k0=0): + """This wave's operands of weight rows ``row0`` .. + 16 (row-major, ``ldk`` + bf16 a row at ``w``, K columns ``k0`` .. ``k0 + k``): a list a chunk + (clamped into K: ``mfmas`` zeroes the extra).""" + per, n = chunks(k) + ldk = k if ldk is None else ldk + r_w = rsrc(w) + row = row0 + lane % ROWS + g = lane // ROWS + out = [] + for i in range(per): + ch = fx.min(wave + WAVES * i, n - 1) + base = (row * ldk + k0 + ch * KCH + g * 8) // 2 + out.append( + [ + fx.Vector( + bo.buffer_load( + r_w, + base + 16 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT if nt else 0, + ) + ) + for q in range(4) + ] + ) + return out + + +def mfmas(lane, wave, xl, xrow, k, s, ops): + """``loads``' operands against token rows 0 .. s - 1 of the LDS rows ``xl`` + (``xrow`` bf16 a row) -> this wave's 16x16 accumulator.""" + per, n = chunks(k) + g = lane // ROWS + t = fx.min(lane % ROWS, s - 1) + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(per): + ch = wave + WAVES * i + live = ch < n + kk = fx.min(ch, n - 1) * KCH + g * 8 + for q in range_constexpr(4): + xb = fx.Vector( + fx.ptr_load( + xl + (t * xrow + kk + 32 * q) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + if const_live(per, n): + xb = fx.Vector.from_elements( + [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 + ) + acc = mfma_bf16( + ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc + ) + return acc + + +def const_live(per, n): + """Whether some wave's last chunk is past K (its MFMA must add zero).""" + return per * WAVES != n + + +@traced +def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0): + """Columns k0 .. k0 + k of token rows 0 .. s - 1 of the bf16 [s, ldk] at + ``src`` -> LDS (``xrow`` a row), 16 B a lane.""" + ldk = k if ldk is None else ldk + words = s * k // 8 + for i in range_constexpr((words + THREADS - 1) // THREADS): + e = fx.min(tid + THREADS * i, words - 1) + t = e // (k // 8) + kk = e % (k // 8) * 8 + v = fx.Vector( + bo.buffer_load( + rsrc(src), + (t * ldk + k0 + kk) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=cm, + ) + ) + if tid + THREADS * i < words: + fx.ptr_store(v, xl + (t * xrow + kk) // 2) + gpu.barrier() + + +def dev_rows_to_lds(tid, src, k, s, xl, xrow): + """``rows_to_lds`` at device scope: rows another CTA of this launch wrote.""" + rows_to_lds(tid, src, k, s, xl, xrow, CM_DEV) diff --git a/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/moe.py b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/moe.py new file mode 100644 index 000000000000..f1e570edb769 --- /dev/null +++ b/python/sglang/kernels/ops/moe/k3_mono_flydsl/stages/moe.py @@ -0,0 +1,1515 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2b ``moe``: one Kimi-K3 latent-MoE layer, from the MoE input row x to the +layer's output, every TP reduction inside the kernel. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S <= 8 rows (a +spec-verify step). The original path (vLLM ROCm, see MOE_SPEC.md) replicates the +router and the latent down projection on every rank; here each rank computes +its eighth and the rows are all-gathered over the peer buffer. + + early (CTAs 0 .. 130, one task each, x from the input row; weights first): + router (7 x 16 rows of this rank's 112 experts): fp32 logits -> every + rank's RLOG + latent (28 x 16 rows of this rank's 448): bf16 -> every rank's LAT + sgu (96 x 16 rows of the shared gate | up): bf16 -> SGU + route (S, CTAs 131 ..): sigmoid (v_rcp), + bias, top-16, the pre-bias + sigmoids renormalized (IEEE 1 / sum) -> ROUTE (every rank alike) + sdown (448 x 16 rows, K 768, CTAs 131 ..): SiTU of SGU (vLLM's kernel) -> + bf16 -> GEMV -> bf16 partial; this rank's up columns -> SHP (plain, + SHRDY flags), the rest -> every rank's FINAL + ug (|U| x 12, every CTA): the union U of the step's picks; an expert's + 32 gate + 32 up rows against the MXFP4 latent (aiter's inline quant), + SiTUv2 (aiter's fp32 epilogue), MXFP4 -> INTER / ISC a pick + down (224 x 16 latent rows, every CTA): w2 of every expert of U against + its picks' INTER, bf16(w acc) folded in each token's top-k order with + a bf16 rounding an add -> every rank's LATR + lnorm (28 x 128 columns): the ranks' LATR in rank order, bf16, RMSNorm + (the fused AR + RMSNorm form) -> LN (plain, LNRDY flags) + up (56 x 16 rows of this rank's 896): LN GEMV + SHP -> bf16 -> every + rank's FINAL + reduce (56 x 128 columns, ``reduce``): FINAL in rank order -> bf16 out + +Peer regions are double-buffered by the epoch's parity. +""" + +import math +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Int32, Int64, T, as_ir_value + +from sglang.kernels.ops.moe.k3_mono_flydsl.attention.kda import _SOURCES, step_tag +from sglang.kernels.ops.moe.k3_mono_flydsl.common.debug import ( + region_ids, + stamp, + stamp_begin, + stamp_flush, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.mx import ( + FP4, + UNIT_SCALE, + fp4_tile_load, + mfma_scaled, + scale_index, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ops import ( + CM_DEV, + ballot, + bf16_round, + bf_hi, + bf_lo, + block_excl_scan, + butterfly, + hw_exp2, + hw_rcp, + hw_rsq, + lanes_below, + ld_i32, + load_ptr64, + popcount, + row_sum, + rsrc, + traced, + uniform, + wave_sum, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.plan import ( + BLOCKS, + THREADS, + WAVES, + KernelAbi, + key_tuple, + pair_layout, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.common.ranks import peer_bases, sum_partials +from sglang.kernels.ops.moe.k3_mono_flydsl.common.sync import ( + Mailbox, + preg, + publish, + shift, + sreg, +) +from sglang.kernels.ops.moe.k3_mono_flydsl.stages import gemv as gemv + +REGIONS = ( + "sgu", + "h", + "rpart", + "lpart", + "route", + "inter", + "isc", + "lmsq", + "lnrdy", + "shrdy", + "rlog", + "lat", + "latr", + "final", +) + + +TP = 8 +HIDDEN = 7168 +LAT = 3584 +E = 896 +TOPK = 16 +RI = 384 # routed intermediate a rank +SI = 768 # shared intermediate a rank +ROWS = 16 +LOG2E = 1.0 / math.log(2.0) +BETA, LBETA = 4.0, 25.0 + +R_GROUPS = E // TP // ROWS # 7 router row groups a rank +R_PARTS = 4 # K quarters: the router is on the path to the route +L_GROUPS = LAT // TP // ROWS # 28 latent row groups a rank +L_PARTS = 2 # K halves: the latent is on the path to ug +R_TASKS = R_GROUPS * R_PARTS # 28 +L_TASKS = L_GROUPS * L_PARTS # 56 +G_TASKS = 2 * SI // ROWS # 96 +EARLY = R_TASKS + L_TASKS + G_TASKS # 180 +LATE0 = EARLY # route from here +H_CTA = LATE0 + 8 # the shared experts' SiTU +SPARE0 = H_CTA + 8 # 196: CTAs with no early task run the shared down first +SD_TASKS = HIDDEN // ROWS # 448 +UP_N = HIDDEN // TP # 896 columns a rank up-projects +UP_TASKS = UP_N // ROWS # 56 +LN_TASKS = LAT // 128 # 28 +LN0 = 56 # lnorm's CTAs +RED_TASKS = HIDDEN // 128 # 56 +RED0 = 200 +UG_GROUPS = RI // 32 # 12 +W13_K = LAT # 3584 +W13_STEPS = W13_K // 128 # 28 +W2_STEPS = RI // 128 # 3 +DOWN_TASKS = LAT // ROWS # 224 +DOWN_BATCH = 5 # a wave's experts whose w2 is in flight together +MAX_U = 8 * TOPK +TL_POINTS = 8 # start, early, sdown, ug, down, lnorm, up, reduce + +_STREAM = fx.Stream(None) + + +@dataclass(frozen=True) +class MoeBuild: + tokens: int + eps: float = 1e-5 # the latent RMSNorm + reduce: bool = True # the final reduce in this launch (standalone) + timeline: bool = False + diag: bool = False # bounded waits recording into ``diag`` (debug) + xp: int = 0 # timing experiments: 1 ug without its epilogue, 2 ug loads only + # debug: run phases < stop only (1 early, 2 route/sdown, 3 ug, 4 down, + # 5 lnorm, 6 up) + stop: int = 99 + # w13's AITER layout: False is gate/up separated (``gate_up=False``, what + # the kernel was written for), True is gate/up interleaved, which is what + # SGLang's Mxfp4MoEMethod produces for Kimi-K3 under AITER_SITUV2_A8W4. + # Only the 16-row block index differs, so it costs one remap in ug_loads. + guint: bool = False + + +ABI = KernelAbi( + ( + "x", + "w_gate", + "bias", + "w_ld", + "w_sgu", + "w_sd", + "w13", + "w13s", + "w2", + "w2s", + "ln_w", + "w_up", + "out", + "scratch", + "queue", + "peers", + "rank", + "epoch", + "layer", + "diag", + "tl", + ) +) + + +def peer_pairs(s): + """One parity's peer regions, pairs: name -> count.""" + return { + "rlog": TP * s * (E // TP), + "lat": s * LAT // 2, + "latr": TP * s * LAT // 2, + "final": TP * s * HIDDEN // 2, + } + + +def peer_layout(s): + return pair_layout(peer_pairs(s).items()) + + +def peer_bytes(s): + lay = peer_layout(s) + return 2 * max(o + n for o, n in lay.values()) + + +def scratch_layout(key: MoeBuild) -> dict: + s = key.tokens + pairs = { + "sgu": s * SI, + "h": s * SI // 2, + "rpart": R_GROUPS * R_PARTS * s * ROWS, + "lpart": L_GROUPS * L_PARTS * s * ROWS, + "route": s * 2 * TOPK, + "inter": s * TOPK * 48, + "isc": s * TOPK * UG_GROUPS, + "lmsq": LN_TASKS * s, + "lnrdy": LN_TASKS, + "shrdy": UP_TASKS, + } + lay = pair_layout(pairs.items()) + end = max(o + n for o, n in lay.values()) + lay["ln"] = ((end + 15) // 16 * 16, s * LAT * 2) + end = lay["ln"][0] + lay["ln"][1] + lay["shp"] = ((end + 15) // 16 * 16, s * UP_N * 2) + return lay + + +# the task queues, their own buffer: [slot 256][epoch parity][ug, down] counters +# (persistent words: a shared scratch's other launches would overwrite them) +QUEUE_BYTES = 256 * 2 * 2 * 4 + + +def scratch_bytes(key: MoeBuild) -> int: + return max(o + n for o, n in scratch_layout(key).values()) + + +def f32_of(w): + return w.bitcast(fx.Float32) + + +def tanh_libm(x): + """The ``tanhf`` of x to a few ulp: (1 - e) / (1 + e), e = exp(-2|x|), x's sign.""" + e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) + t = (fx.Float32(1.0) - e) / (fx.Float32(1.0) + e) + return (x < fx.Float32(0.0)).select(-t, t) + + +def situ_vllm(g, u): + """The ``situ_and_mul`` of vLLM on bf16 g, u: + ((2 tanh(g/4)) (1 + tanh(g/2))) (25 tanh(u/25)).""" + a = fx.Float32(0.5 * BETA) * tanh_libm(g * fx.Float32(1.0 / BETA)) + a = a * (fx.Float32(1.0) + tanh_libm(g * fx.Float32(0.5))) + return a * (fx.Float32(LBETA) * tanh_libm(u * fx.Float32(1.0 / LBETA))) + + +def tanh_aiter(x): + """The ``act.tanh_f32`` of aiter: (1 - e) rcp(1 + e), e = exp2(|x| (-2 log2 e)).""" + e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) + t = (fx.Float32(1.0) - e) * hw_rcp(fx.Float32(1.0) + e) + return (x > fx.Float32(0.0)).select(t, -t) + + +def situ_aiter(g, u): + """The situv2 epilogue of aiter on fp32 accumulators.""" + sg = hw_rcp(fx.Float32(1.0) + hw_exp2(g * fx.Float32(-LOG2E))) + a = (fx.Float32(BETA) * tanh_aiter(g * fx.Float32(1.0 / BETA))) * sg + return (a * fx.Float32(LBETA)) * tanh_aiter(u * fx.Float32(1.0 / LBETA)) + + +def e8m0_roundup(amax): + """The ``_e8m0_roundup(amax, 6)`` of aiter: ceil code of amax * f32(1/6), + at most 254.""" + wi = (amax * fx.Float32(1.0 / 6.0)).bitcast(fx.Int32) + b = fx.Int32(fx.Uint32(wi + fx.Int32(0x7FFFFF)) >> fx.Uint32(23)) & 0xFF + return fx.min(b, fx.Int32(254)) + + +def pk_fp4_f32(old, a, b, scale, j): + """Two f32 / scale -> E2M1 nibbles into byte j of ``old`` (RNE, saturating).""" + return fx.Int32( + rocdl.cvt_scalef32_pk_fp4_f32( + T.i32, + as_ir_value(fx.Int32(old)), + as_ir_value(fx.Float32(a)), + as_ir_value(fx.Float32(b)), + as_ir_value(fx.Float32(scale)), + j, + ) + ) + + +def pk_fp4_bf16(old, word, scale, j): + """A packed bf16 pair word / scale -> E2M1 nibbles into byte j of ``old``.""" + v = fx.Vector.from_elements([fx.Int32(word)], fx.Int32).bitcast(fx.BFloat16) + return fx.Int32( + rocdl.cvt_scalef32_pk_fp4_bf16( + T.i32, + as_ir_value(fx.Int32(old)), + as_ir_value(v), + as_ir_value(fx.Float32(scale)), + j, + ) + ) + + +# ---------------------------------------------------------------- early +@traced +def early_out(c, kind, i, q, red): + """Task (i, q) of ``kind`` (0 router, 1 latent, 2 sgu): its 16 rows. A K-split + kind's parts q > 0 publish fp32 partials; part 0 adds them in part order and + publishes the rows (router: fp32 to every rank's RLOG, latent: bf16 to every + rank's LAT); sgu bf16 to SGU.""" + s, tid = c["S"], c["tid"] + if const_expr(kind == 2): + if tid < (ROWS // 2) * s: + rp = tid % (ROWS // 2) + t = tid // (ROWS // 2) + v0 = bf16_round(row_sum(red, 2 * rp, t)) + v1 = bf16_round(row_sum(red, 2 * rp + 1, t)) + c["put_bf"](c["sgu"], t * 2 * SI + i * ROWS + 2 * rp, [v0, v1]) + else: + parts = R_PARTS if kind == 0 else L_PARTS + region = c["rpart"] if kind == 0 else c["lpart"] + if tid < ROWS * s: + r = tid % ROWS + t = tid // ROWS + v = row_sum(red, r, t) + at = lambda qq: ((i * parts + qq) * s + t) * ROWS + r # noqa: E731 + if q > 0: + c["put"](region, at(q), v) + if q == 0: + got = c["poll"]([(region, at(qq), 1) for qq in range(1, parts)]) + for qq in range_constexpr(parts - 1): + v = v + f32_of(got[qq][0]) + fx.ptr_store(v, c["av"] + tid) + gpu.barrier() + if const_expr(kind == 0): + if (q == 0) & (tid < ROWS * s): + r = tid % ROWS + t = tid // ROWS + v = fx.ptr_load(c["av"] + tid) + for p in range_constexpr(TP): + c["put"]( + c["peer"]["rlog"][p], + (c["rank"] * s + t) * (E // TP) + i * ROWS + r, + v, + ) + else: + if (q == 0) & (tid < (ROWS // 2) * s): + rp = tid % (ROWS // 2) + t = tid // (ROWS // 2) + v0 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp))) + v1 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp + 1))) + col = c["rank"] * (LAT // TP) + i * ROWS + 2 * rp + for p in range_constexpr(TP): + c["put_bf"](c["peer"]["lat"][p], t * LAT + col, [v0, v1]) + gpu.barrier() + + +def early_shape(kind): + """(K a part, parts) of an early kind: router a quarter, latent a half, sgu all.""" + return { + 0: (HIDDEN // R_PARTS, R_PARTS), + 1: (HIDDEN // L_PARTS, L_PARTS), + 2: (HIDDEN, 1), + }[kind] + + +def early_loads(c, kind, i, q): + """Early task (i, q) of ``kind``'s weights for this wave: issued as early as + the CTA can, then ``early_compute`` takes them.""" + a = c["args"] + klen, _ = early_shape(kind) + if kind == 0: + w, row0 = a["w_gate"], c["rank"] * (E // TP) + i * ROWS + elif kind == 1: + w, row0 = a["w_ld"], c["rank"] * (LAT // TP) + i * ROWS + else: + w, row0 = a["w_sgu"], i * ROWS + return gemv.loads(c["lane"], c["wave"], w, row0, klen, ldk=HIDDEN, k0=q * klen) + + +@traced +def early_compute(c, kind, i, q, ops): + """Early task (i, q): x's K slice (the input row, or -- ``xflags`` -- the + row this launch's AttnRes stage publishes) against ``ops`` -> ``early_out``.""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + klen, _ = early_shape(kind) + xrow = klen + gemv.LDS_PAD + if const_expr(c.get("xflags") is not None): + if tid < c["xflags"]: + c["poll"]([(c["xrdy"], tid, 1)]) + gpu.barrier() + gemv.rows_to_lds( + tid, c["xbuf"], klen, s, c["xl"], xrow, CM_DEV, ldk=HIDDEN, k0=q * klen + ) + else: + gemv.rows_to_lds( + tid, c["args"]["x"], klen, s, c["xl"], xrow, ldk=HIDDEN, k0=q * klen + ) + acc = gemv.mfmas(lane, wave, c["xl"], xrow, klen, s, ops) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + early_out(c, kind, i, q, red) + + +@traced +def stage_early(c, bid): + """K2b alone: CTA bid's early task (router parts, latent parts, sgu).""" + lb = bid - R_TASKS + gb = bid - R_TASKS - L_TASKS + if bid < R_TASKS: + early_compute( + c, + 0, + bid // R_PARTS, + bid % R_PARTS, + early_loads(c, 0, bid // R_PARTS, bid % R_PARTS), + ) + if (bid >= R_TASKS) & (bid < R_TASKS + L_TASKS): + early_compute( + c, + 1, + lb // L_PARTS, + lb % L_PARTS, + early_loads(c, 1, lb // L_PARTS, lb % L_PARTS), + ) + if bid >= R_TASKS + L_TASKS: + early_compute(c, 2, gb, 0, early_loads(c, 2, gb, 0)) + + +# ---------------------------------------------------------------- route +def _okey(v): + """An order-preserving uint32 key of an f32 (larger float, larger key).""" + b = v.bitcast(fx.Int32) + neg = b < fx.Int32(0) + return fx.Uint32(neg.select(~b, b ^ fx.Int32(-2147483648))) + + +@traced +def stage_route(c, t): + """Token t (wave 0): 896 logits, 14 a lane (expert lane + 64 i) -> sigmoid + (v_rcp) + bias -> the top-16 by a radix select (the 16th largest key by a + 32-step bit search of ballot counts; equal keys: the lower experts) -> the + 16 ranked by score, descending (aiter's slot order) -> the pre-bias + sigmoids renormalized in that order (IEEE 1 / sum).""" + lane, wave, red = c["lane"], c["wave"], c["red"] + a = c["args"] + if wave == 0: + s = c["S"] + per = E // 64 + specs = [] + for i in range_constexpr(per): + e = lane + 64 * i + src = e // (E // TP) + specs.append( + (c["peer"]["rlog_own"], (src * s + t) * (E // TP) + e % (E // TP), 1) + ) + got = c["poll"](specs) + keys, sigs = [], [] + for i in range_constexpr(per): + e = lane + 64 * i + lg = f32_of(got[i][0]) + sg = hw_rcp(fx.Float32(1.0) + hw_exp2(lg * fx.Float32(-LOG2E))) + sc = sg + bo_ld_f32(a["bias"], e) + # a NaN logit (a warmup's garbage row) ranks last + keys.append(_okey((sc == sc).select(sc, fx.Float32(-3.0e38)))) + sigs.append((sg == sg).select(sg, fx.Float32(0.0))) + + def count_ge(th): + n = fx.Int32(0) + for i in range_constexpr(per): + n = n + popcount(ballot(keys[i] >= th)) + return n + + th = fx.Uint32(0) + for bit in range_constexpr(31, -1, -1): + cand = th | fx.Uint32(1 << bit) + th = (count_ge(cand) >= TOPK).select(cand, th) + # selected: key > th, then the lowest experts with key == th + n_gt = fx.Int32(0) + for i in range_constexpr(per): + n_gt = n_gt + popcount(ballot(keys[i] > th)) + need = fx.Int32(TOPK) - n_gt + eq_before = fx.Int32(0) + sel_before = fx.Int32(0) + for i in range_constexpr(per): + eqb = ballot(keys[i] == th) + eq_rank = eq_before + lanes_below(eqb, lane) + sel = (keys[i] > th) | ((keys[i] == th) & (eq_rank < need)) + selb = ballot(sel) + pos = sel_before + lanes_below(selb, lane) + if sel: + fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos * 3) + fx.ptr_store( + fx.Int32(lane + 64 * i).bitcast(fx.Float32), red + (pos * 3 + 1) + ) + fx.ptr_store(sigs[i], red + (pos * 3 + 2)) + eq_before = eq_before + popcount(eqb) + sel_before = sel_before + popcount(selb) + # rank the 16 by key descending, the lower expert first on a tie + q = fx.min(lane, TOPK - 1) + mk = fx.Uint32(fx.ptr_load(red + q * 3).bitcast(fx.Int32)) + me = fx.ptr_load(red + (q * 3 + 1)).bitcast(fx.Int32) + ms = fx.ptr_load(red + (q * 3 + 2)) + rank = fx.Int32(0) + for j in range_constexpr(TOPK): + ok = fx.Uint32(fx.ptr_load(red + j * 3).bitcast(fx.Int32)) + oe = fx.ptr_load(red + (j * 3 + 1)).bitcast(fx.Int32) + before = (ok > mk) | ((ok == mk) & (oe < me)) + rank = rank + before.select(fx.Int32(1), fx.Int32(0)) + # the sigmoids in rank order, summed serially (aiter's renorm) + if lane < TOPK: + fx.ptr_store(ms, red + (64 + rank)) + total = fx.Float32(0.0) + for j in range_constexpr(TOPK): + total = total + fx.ptr_load(red + (64 + j)) + inv = fx.Float32(1.0) / total + if lane < TOPK: + c["put"]( + c["route"], t * 2 * TOPK + rank, fx.min(fx.max(me, 0), fx.Int32(E - 1)) + ) + c["put"](c["route"], t * 2 * TOPK + TOPK + rank, ms * inv) + gpu.barrier() + + +def bo_ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +@traced +def load_route(c): + """ROUTE -> LDS (ids, weights); the union U: slot_of[e], expert_of[slot], + kof[slot][t] (the token's top-k index, -1 if not picked); returns |U|.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rt = c["rt"] + n = s * 2 * TOPK + if tid < n: + v = c["poll"]([(c["route"], tid, 1)])[0][0] + fx.ptr_store(v, rt["route"] + tid) + for i in range_constexpr((E + THREADS - 1) // THREADS): + e = tid + THREADS * i + if e < E: + fx.ptr_store(fx.Int32(0), rt["flag"] + e) + for i in range_constexpr((MAX_U * 8 + THREADS - 1) // THREADS): + j = tid + THREADS * i + if j < MAX_U * 8: + fx.ptr_store(fx.Int32(-1), rt["kof"] + j) + gpu.barrier() + if tid < s * TOPK: + t = tid // TOPK + k = tid % TOPK + e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) + fx.ptr_store(fx.Int32(1), rt["flag"] + e) + gpu.barrier() + e0 = fx.min(2 * tid, E - 1) + f0 = (2 * tid < E).select(fx.ptr_load(rt["flag"] + e0), fx.Int32(0)) + f1 = (2 * tid + 1 < E).select( + fx.ptr_load(rt["flag"] + fx.min(2 * tid + 1, E - 1)), fx.Int32(0) + ) + off, tot = block_excl_scan(f0 + f1, lane, wave, c["scan"]) + gpu.barrier() + if f0 == 1: + fx.ptr_store(off, rt["flag"] + 2 * tid) + fx.ptr_store(fx.Int32(2 * tid), rt["expert"] + off) + if f1 == 1: + fx.ptr_store(off + f0, rt["flag"] + 2 * tid + 1) + fx.ptr_store(fx.Int32(2 * tid + 1), rt["expert"] + off + f0) + gpu.barrier() + if tid < s * TOPK: + t = tid // TOPK + k = tid % TOPK + e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) + slot = fx.ptr_load(rt["flag"] + e) + fx.ptr_store(fx.Int32(k), rt["kof"] + (slot * 8 + t)) + gpu.barrier() + return uniform(tot) + + +@traced +def load_latq(c): + """LAT (every rank's 448 columns) -> MXFP4 as aiter's inline quant: a + thread a (token, 32-block): amax of |bf16| bits, the ceil code of amax / 6, + v_cvt_scalef32_pk_fp4_bf16 -> latq words, latsc codes.""" + s, tid = c["S"], c["tid"] + nblk = s * (LAT // 32) + for i in range_constexpr((nblk + THREADS - 1) // THREADS): + b = fx.min(tid + THREADS * i, nblk - 1) + t = b // (LAT // 32) + kb = b % (LAT // 32) + base = (t * LAT + kb * 32) // 2 + got = c["poll"]([(c["peer"]["lat_own"], base + j, 1) for j in range(16)]) + ws = [got[j][0] for j in range(16)] + am = fx.Int32(0) + for w in ws: + am = fx.max(am, w & 0x7FFF) + am = fx.max(am, (w >> 16) & 0x7FFF) + amax = (am << 16).bitcast(fx.Float32) + code = e8m0_roundup(amax) + scale = (code << 23).bitcast(fx.Float32) + if tid + THREADS * i < nblk: + for q in range_constexpr(4): + pk = fx.Int32(0) + for j in range_constexpr(4): + pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j) + fx.ptr_store(pk, c["latq"] + (t * (LAT // 8) + kb * 4 + q)) + fx.ptr_store(code, c["latsc"] + (t * (LAT // 32) + kb)) + gpu.barrier() + + +# ---------------------------------------------------------------- ug +def w13_scale_word(e, rg, st, lane): + """The dword of aiter's gate_up=False ``shuffle_scale`` holding (row rg 16 + + lane % 16, K block 4 st + lane / 16) for st's pair and rg's pair: byte + (st % 2) 2 + rg % 2.""" + n1 = rg // 2 + k1 = st // 2 + return ( + e * (2 * RI * (W13_K // 32) // 4) + + ((n1 * (W13_STEPS // 2) + k1) * 4 + lane // 16) * 16 + + lane % 16 + ) + + +def ug_row_group(c, g, wave): + """This wave's 16-row block of w13 for 32-column group ``g``. + + ``rgl`` picks gate row group 0/1 (waves 0, 1) or up row group 0/1 (2, 3). + AITER's gate_up=False layout keeps the 384 gate rows before the 384 up + rows; gate_up=True alternates the two 16-row blocks. The scale word packs + an adjacent pair of blocks and selects the byte by the low bit, so both the + load and the byte select have to use the same number -- hence one helper. + """ + rgl = wave % 4 + j, up = 2 * g + rgl % 2, rgl // 2 + return 2 * j + up if c["guint"] else up * (RI // ROWS) + j + + +def ug_loads(c, task): + """Task ``task``'s w13 operands for this wave (its row group, its K half) + and their scale words: issued a task ahead (``stage_ug`` takes them).""" + lane, wave = c["lane"], c["wave"] + a, rt = c["args"], c["rt"] + slot = task // UG_GROUPS + g = task % UG_GROUPS + e = uniform(fx.ptr_load(rt["expert"] + slot)) + half = wave // 4 + rg = ug_row_group(c, g, wave) + r_w = rsrc(a["w13"]) + base = e * (2 * RI * W13_K // 8) + steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] + av = [fp4_tile_load(r_w, base, rg, W13_K, st, lane * 4) for st in steps] + r_s = rsrc(a["w13s"]) + sw = [ + fx.Int32( + bo.buffer_load( + r_s, w13_scale_word(e, rg, st, lane), vec_width=1, dtype=T.i32 + ) + ) + for st in steps[::2] + ] + return av + sw + + +@traced +def stage_ug(c, task, ops): + """Slot task / 12, 32-column group task % 12: gate rows 32 g .. and up rows + 384 + 32 g .. of the slot's expert (``ops``: ``ug_loads``) against the MXFP4 + latent; SiTUv2; MXFP4 a 32-column group -> INTER / ISC of each picking + token.""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + rt = c["rt"] + slot = task // UG_GROUPS + g = task % UG_GROUPS + half = wave // 4 + rg = ug_row_group(c, g, wave) + steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] + av = ops[: len(steps)] + sw = ops[len(steps) :] + t = fx.min(lane % 16, s - 1) + kb = lane // 16 + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(len(steps) if c["xp"] != 2 else 0): + st = steps[i] + bw = fx.Vector( + fx.ptr_load( + c["latq"] + (t * (LAT // 8) + st * 16 + kb * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + sb = fx.ptr_load(c["latsc"] + (t * (LAT // 32) + st * 4 + kb)) + # the code's byte (st % 2) 2 + rg % 2 of the word; rg % 2 is the wave's + sa = (sw[i // 2] >> ((st % 2) * 16 + (rg % 2) * 8)) & 0xFF + acc = mfma_scaled(av[i], bw, acc, sa, sb, FP4, FP4, 0, 0) + if const_expr(c["xp"] == 2): + # loads only: keep them live + acc = fx.Vector.from_elements( + [ + acc[d] + av[-1][d].bitcast(fx.Float32) * fx.Float32(0.0) + for d in range(4) + ], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + if const_expr(c["xp"] != 0): + return + # thread (token, column): g / u summed over the K halves, SiTUv2 + tt = tid // 32 + col = tid % 32 + live = tt < s + tt = fx.min(tt, s - 1) + first = col < 16 # gate / up row group 0 (waves 0, 2) or 1 (1, 3); +4: K's 2nd half + r = col % 16 + gv = first.select( + row_sum(red, r, tt, waves=(0, 4)), row_sum(red, r, tt, waves=(1, 5)) + ) + uv = first.select( + row_sum(red, r, tt, waves=(2, 6)), row_sum(red, r, tt, waves=(3, 7)) + ) + v = situ_aiter(gv, uv) + amax = butterfly(abs(v), (16, 8, 4, 2, 1), fx.max) + code = e8m0_roundup(amax) + gpu.barrier() + fx.ptr_store(v, c["av"] + tid) + gpu.barrier() + k = fx.ptr_load(rt["kof"] + (slot * 8 + tt)) + if live & (col < 4) & (k >= 0): + scale = (code << 23).bitcast(fx.Float32) + pk = fx.Int32(0) + for j in range_constexpr(4): + x0 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j)) + x1 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j + 1)) + pk = pk_fp4_f32(pk, x0, x1, scale, j) + pick = tt * TOPK + k + c["put_words"](c["inter"], pick * 48 + g * 4 + col, [pk]) + if col == 0: + c["put"](c["isc"], pick * UG_GROUPS + g, code) + gpu.barrier() + + +# ---------------------------------------------------------------- queues +@traced +def take(c, which): + """The launch's next task of queue ``which`` (0 ug, 1 down) past the first + static round: an agent-scope atomic on this slot's counter of this epoch's + parity (the other parity's is reset by this launch for the next step).""" + if c["tid"] == 0: + at = c["queue"] + fx.Int64(((c["slot"] * 2 + c["par"]) * 2 + which) * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + task = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(task), c["qslot"]) + gpu.barrier() + v = fx.ptr_load(c["qslot"]) + gpu.barrier() + return BLOCKS + v + + +@traced +def reset_queues(c): + """The next step's counters of this slot (the other parity) to zero.""" + if (c["bid"] == 0) & (c["tid"] < 2): + bo.buffer_store( + fx.Int32(0), + rsrc(c["queue"]), + (c["slot"] * 2 + 1 - c["par"]) * 2 + c["tid"], + cache_modifier=CM_DEV, + ) + + +# ---------------------------------------------------------------- down +@traced +def load_inter(c): + """Every pick's INTER words and ISC codes -> LDS.""" + s, tid = c["S"], c["tid"] + nw = s * TOPK * 48 + per = (nw + THREADS - 1) // THREADS + got = c["poll"]( + [(c["inter"], fx.min(tid + THREADS * i, nw - 1), 1) for i in range(per)] + ) + for i in range_constexpr(per): + if tid + THREADS * i < nw: + fx.ptr_store(got[i][0], c["interq"] + (tid + THREADS * i)) + nc = s * TOPK * UG_GROUPS + perc = (nc + THREADS - 1) // THREADS + gc = c["poll"]( + [(c["isc"], fx.min(tid + THREADS * i, nc - 1), 1) for i in range(perc)] + ) + for i in range_constexpr(perc): + if tid + THREADS * i < nc: + fx.ptr_store(gc[i][0], c["intsc"] + (tid + THREADS * i)) + gpu.barrier() + + +def w2_scale(c, e, rgd, st, lane): + """The e8m0 code of w2 (row e 3584 + rgd 16 + lane % 16, K block 4 st + + lane / 16) in ``e8m0_shuffle``'s layout (16 codes a row).""" + row = e * LAT + rgd * ROWS + lane % 16 + idx = scale_index(row, st * 4 + lane // 16, 16) + w = fx.Int32( + bo.buffer_load(rsrc(c["args"]["w2s"]), idx // 4, vec_width=1, dtype=T.i32) + ) + return (w >> ((idx % 4) * 8)) & 0xFF + + +def poll_inter(c, pick, kb): + """A lane's INTER words (3 steps x 4) and ISC codes of ``pick``, polled at + its use: a slot whose ug tasks are done is multiplied while later slots' + still run (no wait for all of U).""" + return c["poll"]( + [ + (c["inter"], pick * 48 + st * 16 + kb * 4 + d, 1) + for st in range(W2_STEPS) + for d in range(4) + ] + + [(c["isc"], pick * UG_GROUPS + st * 4 + kb, 1) for st in range(W2_STEPS)] + ) + + +def down_loads(c, rgd, nu, b0): + """The w2 tiles and scale codes of a wave's batch from slot b0 (slots b0, + b0 + 8, ...; past |U| the last slot's): flat, ``stage_down`` unpacks.""" + lane = c["lane"] + r_w = rsrc(c["args"]["w2"]) + rt = c["rt"] + slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] + es = [uniform(fx.ptr_load(rt["expert"] + sl)) for sl in slots] + av = [ + fp4_tile_load(r_w, es[j] * (LAT * RI // 8), rgd, RI, st, lane * 4) + for j in range(DOWN_BATCH) + for st in range(W2_STEPS) + ] + sa = [ + w2_scale(c, es[j], rgd, st, lane) + for j in range(DOWN_BATCH) + for st in range(W2_STEPS) + ] + return av + sa + + +@traced +def stage_down(c, rgd, nu, ops0): + """Latent rows 16 rgd ..: every slot's w2 against its picks (wave w the slots + w, w + 8, ...; a batch of DOWN_BATCH in flight, the next issued under this + one's MFMAs; ``ops0`` the first, issued before INTER was waited for) -> + bf16(w acc) a (token, k) in LDS -> folded in top-k order with a bf16 + rounding an add -> every rank's LATR.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rt = c["rt"] + t = fx.min(lane % 16, s - 1) + kb = lane // 16 + nbs = DOWN_BATCH * W2_STEPS + b0 = wave + ops = ops0 + while b0 < nu: + nb = b0 + WAVES * DOWN_BATCH + ops_n = down_loads(c, rgd, nu, fx.min(nb, nu - 1)) + slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] + for j in range_constexpr(DOWN_BATCH): + live = (b0 + WAVES * j) < nu + k = fx.ptr_load(rt["kof"] + (slots[j] * 8 + t)) + picked = (k >= 0) & (lane % 16 < s) + pick = t * TOPK + fx.max(k, 0) + acc = fx.Vector.filled(4, 0.0, fx.Float32) + # this slot's picks' INTER, polled here (not all of U first): a slot + # whose ug tasks are done is multiplied while later ones still run + for st in range_constexpr(W2_STEPS): + bw = fx.Vector( + fx.ptr_load( + c["interq"] + (pick * 48 + st * 16 + kb * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + scode = fx.ptr_load(c["intsc"] + (pick * UG_GROUPS + st * 4 + kb)) + bw = fx.Vector.from_elements( + [picked.select(bw[d], fx.Int32(0)) for d in range(4)], fx.Int32 + ) + sb = picked.select(scode, fx.Int32(UNIT_SCALE)) + acc = mfma_scaled( + ops[j * W2_STEPS + st], + bw, + acc, + ops[nbs + j * W2_STEPS + st], + sb, + FP4, + FP4, + 0, + 0, + ) + # lane: rows 4 (lane / 16) + i of token lane % 16 + if live & picked: + wt = fx.ptr_load(rt["route"] + (t * 2 * TOPK + TOPK + k)).bitcast( + fx.Float32 + ) + for i in range_constexpr(4): + row = 4 * (lane // 16) + i + fx.ptr_store( + bf16_round(acc[i] * wt), + c["contrib"] + ((t * TOPK + k) * ROWS + row), + ) + b0 = nb + ops = ops_n + gpu.barrier() + if tid < (ROWS // 2) * s: + rp = tid % (ROWS // 2) + tt = tid // (ROWS // 2) + y0 = fx.Float32(0.0) + y1 = fx.Float32(0.0) + for k in range_constexpr(TOPK): + y0 = bf16_round( + y0 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp)) + ) + y1 = bf16_round( + y1 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp + 1)) + ) + for p in range_constexpr(TP): + c["put_bf"]( + c["peer"]["latr"][p], + (c["rank"] * s + tt) * LAT + rgd * ROWS + 2 * rp, + [y0, y1], + ) + gpu.barrier() + + +# ---------------------------------------------------------------- shared down +@traced +def stage_h(c, t): + """Token t's shared h = SiTU(g, u) (vLLM's situ_and_mul, bf16) of SGU, once + -> H; a thread a column pair, its polls in one batch.""" + tid = c["tid"] + n = SI // 2 + if tid < n: + got = c["poll"]( + [ + (c["sgu"], (t * 2 * SI) // 2 + tid, 1), + (c["sgu"], (t * 2 * SI + SI) // 2 + tid, 1), + ] + ) + g0, g1 = bf_lo(got[0][0]), bf_hi(got[0][0]) + u0, u1 = bf_lo(got[1][0]), bf_hi(got[1][0]) + c["put_bf"](c["h"], t * SI + 2 * tid, [situ_vllm(g0, u0), situ_vllm(g1, u1)]) + gpu.barrier() + + +@traced +def load_h(c): + """H -> LDS rows (for the shared down GEMV), once a CTA.""" + s, tid = c["S"], c["tid"] + hrow = SI + gemv.LDS_PAD + n = s * SI // 2 + per = (n + THREADS - 1) // THREADS + got = c["poll"]([(c["h"], fx.min(tid + THREADS * i, n - 1), 1) for i in range(per)]) + for i in range_constexpr(per): + pidx = tid + THREADS * i + t = fx.min(pidx, n - 1) // (SI // 2) + cp = fx.min(pidx, n - 1) % (SI // 2) + if pidx < n: + fx.ptr_store(got[i][0], c["xl"] + (t * hrow + 2 * cp) // 2) + gpu.barrier() + + +@traced +def stage_sdown(c, j, ops): + """Shared down rows 16 j .. (K 768) against h in LDS (``load_h``): bf16 + partial -> SHP (this rank's up columns) or every rank's FINAL.""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + hrow = SI + gemv.LDS_PAD + acc = gemv.mfmas(lane, wave, c["xl"], hrow, SI, s, ops) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + lo = c["rank"] * UP_TASKS + mine = (j >= lo) & (j < lo + UP_TASKS) + if tid < (ROWS // 2) * s: + rp = tid % (ROWS // 2) + tt = tid // (ROWS // 2) + v0 = bf16_round(row_sum(red, 2 * rp, tt)) + v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) + if mine: + w = ( + fx.Vector.from_elements([v0, v1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store( + w, + rsrc(c["shp"]), + (tt * UP_N + (j - lo) * ROWS + 2 * rp) // 2, + cache_modifier=CM_DEV, + ) + else: + for p in range_constexpr(TP): + c["put_bf"]( + c["peer"]["final"][p], + (c["rank"] * s + tt) * HIDDEN + j * ROWS + 2 * rp, + [v0, v1], + ) + publish( + c["put"], + c["shrdy"], + fx.max(fx.min(j - lo, UP_TASKS - 1), 0), + fx.Int32(1), + (tid == 0) & mine, + ) + gpu.barrier() + + +@traced +def run_sdown(c, bid, w_sd): + """This rank's 56 shared-down tasks whose columns the up stage adds to (k -> + row group (k + 56 rank) % 448), on the CTAs past the down tasks: up waits + on them.""" + load_h(c) + lo = c["rank"] * UP_TASKS + if bid >= DOWN_TASKS: + for k in range(bid - DOWN_TASKS, UP_TASKS, BLOCKS - DOWN_TASKS): + j = (k + lo) % SD_TASKS + stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) + + +@traced +def run_sdown_rest(c, bid, w_sd): + """The other 392 shared-down tasks (only the final reduce waits on them), on + the CTAs past lnorm / up's (84 ..), after their down: off the latent-norm -> + up critical path.""" + r0 = LN0 + LN_TASKS # 84: past the lnorm / up CTAs + if bid >= r0: + load_h(c) + lo = c["rank"] * UP_TASKS + for k in range(UP_TASKS + bid - r0, SD_TASKS, BLOCKS - r0): + j = (k + lo) % SD_TASKS + stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) + + +# ---------------------------------------------------------------- lnorm / up / reduce +@traced +def stage_lnorm(c, j): + """Columns 128 j ..: the ranks' LATR (rank order, fp32 -> bf16), every + task's sum of squares -> rstd -> bf16(acc rstd w) -> LN, LNRDY.""" + s, tid, lane, red = c["S"], c["tid"], c["lane"], c["red"] + a = c["args"] + mine = tid < s * 64 + t = fx.min(tid // 64, s - 1) + cp = tid % 64 + col = j * 128 + 2 * cp + lo, hi = sum_partials( + c["poll"], + c["peer"]["latr_own"], + lambda src: ((src * s + t) * LAT + col) // 2, + TP, + ) + lo, hi = bf16_round(lo), bf16_round(hi) + sq = wave_sum(lo * lo + hi * hi) + if mine & (lane == 0): + c["put"](c["lmsq"], j * s + t, sq) + gpu.barrier() + if mine & (lane < LN_TASKS): + v = c["poll"]([(c["lmsq"], lane * s + t, 1)])[0][0] + fx.ptr_store(f32_of(v), red + (t * 32 + lane)) + gpu.barrier() + tot = fx.ptr_load(red + t * 32) + for q in range_constexpr(1, LN_TASKS): + tot = tot + fx.ptr_load(red + (t * 32 + q)) + r = hw_rsq(tot * fx.Float32(1.0 / LAT) + fx.Float32(c["ln_eps"])) + w0 = fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col, vec_width=1, dtype=T.bf16)) + ) + w1 = fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col + 1, vec_width=1, dtype=T.bf16)) + ) + y = ( + fx.Vector.from_elements([lo * r * w0, hi * r * w1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + if mine: + bo.buffer_store(y, rsrc(c["ln"]), (t * LAT + col) // 2, cache_modifier=CM_DEV) + publish(c["put"], c["lnrdy"], j, fx.Int32(1), tid == 0) + gpu.barrier() + + +@traced +def stage_up(c, j): + """This rank's up rows 16 j .. of its shard (hidden rows 896 rank ..): LN GEMV + + SHP -> bf16 -> FINAL.""" + s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] + a = c["args"] + row0 = c["rank"] * UP_N + j * ROWS # the hidden column; w_up is the shard + ops = gemv.loads(lane, wave, a["w_up"], j * ROWS, LAT) + if tid < LN_TASKS: + c["poll"]([(c["lnrdy"], tid, 1)]) + if tid == 0: + c["poll"]([(c["shrdy"], j, 1)]) + gpu.barrier() + lrow = LAT + gemv.LDS_PAD + gemv.dev_rows_to_lds(tid, c["ln"], LAT, s, c["xl"], lrow) + acc = gemv.mfmas(lane, wave, c["xl"], lrow, LAT, s, ops) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + gpu.barrier() + if tid < (ROWS // 2) * s: + rp = tid % (ROWS // 2) + tt = tid // (ROWS // 2) + w = fx.Int32( + bo.buffer_load( + rsrc(c["shp"]), + (tt * UP_N + j * ROWS + 2 * rp) // 2, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + v0 = bf16_round(bf_lo(w) + row_sum(red, 2 * rp, tt)) + v1 = bf16_round(bf_hi(w) + row_sum(red, 2 * rp + 1, tt)) + for p in range_constexpr(TP): + c["put_bf"]( + c["peer"]["final"][p], + (c["rank"] * s + tt) * HIDDEN + row0 + 2 * rp, + [v0, v1], + ) + gpu.barrier() + + +@traced +def stage_reduce(c, j): + """Columns 128 j ..: FINAL in rank order -> bf16 out.""" + s, tid = c["S"], c["tid"] + if tid < s * 64: + t = tid // 64 + cp = tid % 64 + col = j * 128 + 2 * cp + lo, hi = sum_partials( + c["poll"], + c["peer"]["final_own"], + lambda src: ((src * s + t) * HIDDEN + col) // 2, + TP, + ) + w = ( + fx.Vector.from_elements([lo, hi], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(w, rsrc(c["args"]["out"]), (t * HIDDEN + col) // 2) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel +_BUILDS: dict = {} + + +def build(key: MoeBuild): + if key in _BUILDS: + return _BUILDS[key] + s = key.tokens + assert 1 <= s <= 8 + lay = scratch_layout(key) + play = peer_layout(s) + half = peer_bytes(s) // 2 + + @fx.struct + class RouteLds: + route: fx.Array[fx.Int32, s * 2 * TOPK, 16] + flag: fx.Array[fx.Int32, E, 16] + expert: fx.Array[fx.Int32, MAX_U, 16] + kof: fx.Array[fx.Int32, MAX_U * 8, 16] + scan: fx.Array[fx.Int32, WAVES, 16] + + @fx.struct + class EarlyLds: + xl: fx.Array[fx.Int32, s * (HIDDEN + gemv.LDS_PAD) // 2, 16] + + @fx.struct + class UgLds: + latq: fx.Array[fx.Int32, s * LAT // 8, 16] + latsc: fx.Array[fx.Int32, s * LAT // 32, 16] + av: fx.Array[fx.Float32, THREADS, 16] + + @fx.struct + class DownLds: + interq: fx.Array[fx.Int32, s * TOPK * 48, 16] + intsc: fx.Array[fx.Int32, s * TOPK * UG_GROUPS, 16] + contrib: fx.Array[fx.Float32, s * TOPK * ROWS, 16] + + @fx.union + class StageLds: + early: EarlyLds + ug: UgLds + down: DownLds + + @fx.struct + class Smem: + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + tls: fx.Array[fx.Int64, TL_POINTS, 16] + qslot: fx.Array[fx.Int32, 4, 16] + + keyed = key_tuple(key, _SOURCES) + name = f"k3_mono_moe_s{s}_r{int(key.reduce)}_t{int(key.timeline)}" + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def moe( + x: Int64, + w_gate: Int64, + bias: Int64, + w_ld: Int64, + w_sgu: Int64, + w_sd: Int64, + w13: Int64, + w13s: Int64, + w2: Int64, + w2s: Int64, + ln_w: Int64, + w_up: Int64, + out: Int64, + scratch: Int64, + queue: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + diag: Int64, + tl: Int64, + ): + _ = keyed + tid = fx.thread_idx.x + bid = fx.block_idx.x + alloc = fx.SharedAllocator() + lds = alloc.allocate(Smem).peek() + rl = alloc.allocate(RouteLds).peek() + st_lds = alloc.allocate(StageLds) + el, ul, dl = st_lds.early.peek(), st_lds.ug.peek(), st_lds.down.peek() + if const_expr(key.diag): + mb = Mailbox(step_tag(epoch, layer) - 1, diag, region_ids(REGIONS)) + else: + mb = Mailbox(step_tag(epoch, layer) - 1) + par = uniform(ld_i32(epoch, 0)) % 2 + bases = peer_bases(peers, TP) + own = load_ptr64(peers, rank) + peer = {} + for region in ("rlog", "lat", "latr", "final"): + off = play[region][0] + peer[region] = [ + shift(preg(b, off, region), fx.Int64(par) * half) for b in bases + ] + peer[region + "_own"] = shift(preg(own, off, region), fx.Int64(par) * half) + c = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rank": rank, + "ln_eps": key.eps, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "red": lds.red.ptr, + "peer": peer, + "xl": el.xl.ptr, + "latq": ul.latq.ptr, + "latsc": ul.latsc.ptr, + "av": ul.av.ptr, + "interq": dl.interq.ptr, + "intsc": dl.intsc.ptr, + "contrib": dl.contrib.ptr, + "scan": rl.scan.ptr, + "rt": { + "route": rl.route.ptr, + "flag": rl.flag.ptr, + "expert": rl.expert.ptr, + "kof": rl.kof.ptr, + }, + "ln": scratch + fx.Int64(lay["ln"][0]), + "shp": scratch + fx.Int64(lay["shp"][0]), + "queue": queue, + "qslot": lds.qslot.ptr, + "slot": layer, + "par": par, + "xp": key.xp, + "guint": key.guint, + "lazy_inter": True, + "args": { + "x": x, + "w_gate": w_gate, + "bias": bias, + "w_ld": w_ld, + "w_sgu": w_sgu, + "w_sd": w_sd, + "w13": w13, + "w13s": w13s, + "w2": w2, + "w2s": w2s, + "ln_w": ln_w, + "w_up": w_up, + "out": out, + }, + } + for region in ( + "sgu", + "h", + "rpart", + "lpart", + "route", + "inter", + "isc", + "lmsq", + "lnrdy", + "shrdy", + ): + c[region] = sreg(scratch, lay[region][0], region) + on, tls = key.timeline, lds.tls.ptr + stamp_begin(on, tls, tid, TL_POINTS) + if const_expr(key.stop > 1): # noqa: SIM102 (traced: `and` does not combine traced values) + if bid < EARLY: + stage_early(c, bid) + stamp(on, tls, tid, 1) + if const_expr(key.stop <= 2): + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + return + if (bid >= LATE0) & (bid < LATE0 + s): + stage_route(c, bid - LATE0) + if (bid >= H_CTA) & (bid < H_CTA + s): + stage_h(c, bid - H_CTA) + stamp(on, tls, tid, 2) + if const_expr(key.stop <= 3): + return + nu = load_route(c) + if const_expr(key.stop == 35): + # debug: the route table, |U| and U into ``out`` (as int32), CTA 0 + if bid == 0: + n = s * 2 * TOPK + if tid < n: + bo.buffer_store(fx.ptr_load(c["rt"]["route"] + tid), rsrc(out), tid) + if tid < MAX_U: + bo.buffer_store( + fx.ptr_load(c["rt"]["expert"] + tid), rsrc(out), n + 1 + tid + ) + if tid == 0: + bo.buffer_store(nu, rsrc(out), n) + return + load_latq(c) + if const_expr(key.stop == 36): + return + # static round-robin: deterministic, no state across launches (a + # counter carried between steps went stale when a step skipped K2b) + total = nu * UG_GROUPS + last = total - 1 + # reversed placement: the round's leftover tasks go to the high CTAs, + # whose early work ends first (the AttnRes / route CTAs start ug last) + task = BLOCKS - 1 - bid + ops = ug_loads(c, fx.min(task, last)) + while task < total: + nxt = task + BLOCKS + # the next task's weights in flight under this one's MFMAs and + # epilogue (past the end: the last task's, cache hits) + ops_n = ug_loads(c, fx.min(nxt, last)) + stage_ug(c, task, ops) + task = nxt + ops = ops_n + stamp(on, tls, tid, 3) + if const_expr(key.stop <= 4): + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + return + if bid < DOWN_TASKS: + # the first w2 batch in flight before INTER (every ug) is waited for + dops = down_loads(c, bid, nu, fx.min(c["wave"], nu - 1)) + load_inter(c) + stage_down(c, bid, nu, dops) + run_sdown_rest(c, bid, w_sd) + # the shared down (needed only by up) on the CTAs with no down task: + # after their ug, so every CTA's ug starts as the route lands + run_sdown(c, bid, w_sd) + stamp(on, tls, tid, 4) + if const_expr(key.stop <= 5): + return + if (bid >= LN0) & (bid < LN0 + LN_TASKS): + stage_lnorm(c, bid - LN0) + stamp(on, tls, tid, 5) + if const_expr(key.stop <= 6): + return + if bid < UP_TASKS: + stage_up(c, bid) + stamp(on, tls, tid, 6) + if const_expr(key.reduce): # noqa: SIM102 (traced: `and` does not combine traced values) + if (bid >= RED0) & (bid < RED0 + RED_TASKS): + stage_reduce(c, bid - RED0) + stamp(on, tls, tid, 7) + stamp_flush(on, tls, tl, tid, bid, TL_POINTS) + + @flyc.jit + def launch( + x: Int64, + w_gate: Int64, + bias: Int64, + w_ld: Int64, + w_sgu: Int64, + w_sd: Int64, + w13: Int64, + w13s: Int64, + w2: Int64, + w2s: Int64, + ln_w: Int64, + w_up: Int64, + out: Int64, + scratch: Int64, + queue: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + layer: Int32, + diag: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + moe( + x, + w_gate, + bias, + w_ld, + w_sgu, + w_sd, + w13, + w13s, + w2, + w2s, + ln_w, + w_up, + out, + scratch, + queue, + peers, + rank, + epoch, + layer, + diag, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + ABI.check(moe, launch) + _BUILDS[key] = launch + return launch + + +def moe( + key: MoeBuild, + *, + x, + w_gate, + bias, + w_ld, + w_sgu, + w_sd, + w13, + w13s, + w2, + w2s, + ln_w, + w_up, + out, + scratch, + queue, + peers, + rank, + epoch, + layer, + tl=None, + diag=None, +): + """Host launcher; every weight as vLLM holds it after loading (MOE_SPEC.md + §5), the experts' shuffled. ``queue``: QUEUE_BYTES of zeros, this launch's + alone (no other kernel writes it).""" + assert queue.numel() * queue.element_size() >= QUEUE_BYTES + assert w_gate.shape == (E, HIDDEN) and w_ld.shape == (LAT, HIDDEN) + assert w_sgu.shape == (2 * SI, HIDDEN) and w_sd.shape == (HIDDEN, SI) + assert w_up.shape == (UP_N, LAT) and w_up.is_contiguous() + f = build(key) + f( + x.data_ptr(), + w_gate.data_ptr(), + bias.data_ptr(), + w_ld.data_ptr(), + w_sgu.data_ptr(), + w_sd.data_ptr(), + w13.data_ptr(), + w13s.data_ptr(), + w2.data_ptr(), + w2s.data_ptr(), + ln_w.data_ptr(), + w_up.data_ptr(), + out.data_ptr(), + scratch.data_ptr(), + queue.data_ptr(), + peers.data_ptr(), + rank, + epoch.data_ptr(), + layer, + 0 if diag is None else diag.data_ptr(), + 0 if tl is None else tl.data_ptr(), + stream=torch.cuda.current_stream(), + ) diff --git a/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py b/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py index 6a0157314e6c..4e1a29bdb2c4 100644 --- a/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py +++ b/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py @@ -184,6 +184,56 @@ def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict: return overrides +@_register_for("KimiK3ForConditionalGeneration", "KimiK3LinearForCausalLM") +def _kimi_k3_mono_decode_checks(server_args: Any, hf_config: Any) -> dict: + """Reject a serving config the ROCm mono decode launches can never serve. + + The launches themselves fall back per layer, so this changes nothing about + correctness. What it buys is the timing: every condition below is an + explicit CLI choice that contradicts an explicit env opt-in, and without + this the operator only learns from an info log after the weights have + loaded, minutes in. + + Off-platform stays a no-op on purpose: the env may be set globally, and a + CUDA host is not a contradiction, just somewhere the feature does not run. + """ + from sglang.srt.environ import envs + from sglang.srt.layers.moe.k3_mono_decode import MONO_TP_SIZE + from sglang.srt.utils import is_hip + + if envs.SGLANG_K3_MONO_LAYER.get() and not envs.SGLANG_K3_MONO_DECODE.get(): + raise ValueError( + "SGLANG_K3_MONO_LAYER=1 needs SGLANG_K3_MONO_DECODE=1: the layer " + "launch is built on the MoE launch's per-layer binding." + ) + if not (envs.SGLANG_K3_MONO_DECODE.get() and is_hip()): + return {} + + cfg = resolving_view(server_args) + # Mirrors mono_supported() in srt/layers/moe/k3_mono_decode.py, limited to + # what a server arg decides. + reasons = [] + if cfg.tp_size != MONO_TP_SIZE: + reasons.append(f"--tp-size must be {MONO_TP_SIZE} (got {cfg.tp_size})") + if cfg.moe_a2a_backend not in (None, "none"): + reasons.append( + "--moe-a2a-backend must be none " + f"(got {cfg.moe_a2a_backend!r}); the launch holds the experts " + "TP-sharded, not expert-parallel" + ) + # --enable-dp-attention is a deprecated alias that resolution turns into + # attn_dp_size, so on the pristine config the size is the knob to read. + if (cfg.attn_dp_size or 1) > 1 or cfg.enable_dp_attention: + reasons.append("--attn-dp-size must be 1 (DP attention is not supported)") + if reasons: + raise ValueError( + "SGLANG_K3_MONO_DECODE=1 cannot serve this configuration: " + + "; ".join(reasons) + + ". Unset the env to keep the default path." + ) + return {} + + @_register_for("KimiK3ForConditionalGeneration") def _kimi_k3_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict: # MoE runner default, independent of the attention-backend gate above. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c1bdf79f565f..eb88bfe3dd65 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1861,6 +1861,26 @@ class Envs: SGLANG_K3_FUSED_FRONT = EnvBool(True) # Use the ROCm radix-4 router for covered K3 top-k workloads. SGLANG_K3_RADIX4_TOPK = EnvBool(False) + # ROCm FlyDSL mono decode for K3: one persistent launch runs the whole MoE + # -- router, latent down, shared gate_up, MXFP4 experts, shared down, + # latent up and both TP all-reduces, the reduces in-kernel over a HIP-IPC + # peer buffer. The router and the latent down are TP sharded here, so the + # two widest decode weights are read 8x less. gfx950 + TP8 + decode steps + # of at most 8 rows; everything else keeps the existing path. See + # srt/layers/moe/k3_mono_decode.py. + SGLANG_K3_MONO_DECODE = EnvBool(False) + # Widen that launch to a whole layer: o_proj and its all-reduce, the MLP + # attention-residual seam and the MoE in one launch instead of three. + # Needs SGLANG_K3_MONO_DECODE; a layer it cannot serve keeps the MoE-only + # launch. + SGLANG_K3_MONO_LAYER = EnvBool(False) + # Widen it once more to the KDA half of the layer (the kernel's "K1"): + # the attention-residual seam, input_layernorm, the in-projection, the + # conv and the recurrence join K2's launch. Needs SGLANG_K3_MONO_LAYER. + # Not yet usable: MonoK1.forward still reads index tensors on the host, + # which a CUDA graph capture rejects. Kept off until that is fixed so the + # layer launch above stays usable on its own. + SGLANG_K3_MONO_K1 = EnvBool(False) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY = EnvInt(2) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS = EnvInt(2) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN = EnvInt(6144) diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index d155aa63b4cc..f60815112254 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -1873,3 +1873,39 @@ def _run_dspark_cutedsl_mtp( if apply_onorm: layer._k3_onorm_consumed = True return out + + def kimi_k3_mono_state( + self, layer: RadixLinearAttention, forward_batch: ForwardBatch + ) -> Optional[dict]: + """The speculative state of one KDA layer, for the K3 mono K1 launch. + + K1 replaces this layer's seam, in-projection, conv and recurrence in + one launch, so it needs what the verify path owns. The model side has + no route to these tensors, so expose them here -- the same derivation + ``_forward_target_verify`` opens with, and nothing more. Everything + about K1 itself stays in srt/layers/moe/k3_mono_decode.py. + + None when this batch has no speculative state to hand over. + """ + params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + inter_ssm = getattr(params, "intermediate_ssm", None) + windows = getattr(params, "intermediate_conv_window", None) + spec_info = getattr(forward_batch, "spec_info", None) + if inter_ssm is None or not windows or spec_info is None: + return None + num_reqs = forward_batch.req_pool_indices.shape[0] + slots = self.forward_metadata.mamba_cache_indices[:num_reqs] + return { + "committed_conv": params.conv[0], + "committed_ssm": params.temporal, + "inter_ssm": inter_ssm, + "inter_window": windows[0], + "slots": slots, + "rows": select_verify_intermediate_state_indices( + self.verify_intermediate_state_indices, + forward_batch.req_pool_indices, + slots >= 0, + self.req_to_token_pool.size, + ), + "num_draft": spec_info.draft_token_num, + } diff --git a/python/sglang/srt/layers/moe/k3_mono_decode.py b/python/sglang/srt/layers/moe/k3_mono_decode.py new file mode 100644 index 000000000000..5286915c5417 --- /dev/null +++ b/python/sglang/srt/layers/moe/k3_mono_decode.py @@ -0,0 +1,829 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0. +"""Kimi-K3 mono decode for ROCm TP8: the SGLang side of the FlyDSL kernels. + +One persistent launch runs the whole MoE at a small decode shape -- the router +and its top-16, the latent down and shared gate_up projections, the MXFP4 +experts, the shared down and latent up projections, and the two TP all-reduces, +which go over a HIP-IPC peer buffer rather than leaving the kernel. + +Two properties are worth knowing before reading the rest: + +* The router and the latent down projection are **TP sharded** here. Each rank + computes its eighth of the rows and the result is all-gathered in-kernel, so + the two widest weights on the decode critical path are read 8x less. The + unfused path replicates both. +* Nothing is repacked. The dense weights are read as BF16 in their checkpoint + layout, and the experts are read in whichever AITER A16W4 layout + ``Mxfp4MoEMethod`` left behind. So there is no prepared-weight copy and no + cost to the KV cache budget. + +``SGLANG_K3_MONO_LAYER`` widens that launch to the whole layer (``MonoK2``): +it then starts from the attention's TP-local output, so o_proj, the +attention-TP all-reduce and the MLP attention-residual seam join the MoE in +the same launch. The attention itself and the final residual add stay outside. + +Every rank must agree on whether to take this path: the kernels spin in-kernel +on each peer's mailbox, so a rank that opted out alone would leave the other +seven waiting for ever. ``_tp_agree`` makes the decision collective. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +import torch +import torch.distributed as dist + +from sglang.srt.runtime_context import get_parallel + +if TYPE_CHECKING: + from sglang.srt.layers.attn_residual import AttnResidual + from sglang.srt.models.kimi_k3 import KimiK3DecoderLayer, KimiK3MoE + +logger = logging.getLogger(__name__) + +# The kernels are built for at most this many rows, because the rows have to +# fit one CTA's LDS. At concurrency 1 a DSPARK-7 verify step is exactly 8. +MONO_MAX_ROWS = 8 +MONO_TP_SIZE = 8 +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +# o_proj's input width per rank: 12 heads of 128. +CORE_SIZE = 1536 +# The attention-residual bank rows the layer launch is built for. +MONO_MAX_BLOCKS = 8 + +_logged: set[str] = set() + + +class MonoUnsupported(Exception): + """This rank cannot serve the mono path.""" + + +def _log_once(message: str) -> None: + """Report a per-layer decision once. Every layer answers the same way.""" + if message not in _logged: + _logged.add(message) + logger.info("%s", message) + + +def _tp_agree(ok: bool) -> bool: + """Return True only when every TP rank passed ``ok``. + + Every rank has to reach this at the same point. A MIN all-reduce over the + gloo group is enough and keeps the decision off the compute stream. + """ + tp_group = get_parallel().tp_group + if tp_group.world_size == 1: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=tp_group.cpu_group) + return bool(flag.item()) + + +def _bind_agreed(build) -> None: + """Run ``build`` on every rank, and fail on all of them if any one fails.""" + try: + build() + except Exception: + _tp_agree(False) + raise + if not _tp_agree(True): + raise MonoUnsupported("another TP rank refused the mono path") + + +def step_begin() -> None: + """Advance the mailbox epoch. Call once per forward, before the first layer. + + The epoch is a device counter, so a captured graph replays the bump. It is + what lets a tag identify one launch of one step without anything being + zeroed in between. + """ + from sglang.kernels.ops.moe.k3_mono_flydsl import runner + + runner.step_begin() + + +def _is_decode_step(forward_batch) -> bool: + """Whether this forward is a decode-side step the launches are meant for. + + Shape alone does not say: a prefill of eight tokens looks exactly like a + DSPARK-7 verify step. Two reasons to ask anyway. The launches are a + latency optimization for the decode path and buy nothing on a short + prefill, and prefill is the phase that can run under a breakable CUDA + graph, whose segmenting these launches have not been validated against. + + Target-verify counts: with DSPARK that *is* the decode step. + """ + if forward_batch is None: + return False + mode = forward_batch.forward_mode + return mode.is_decode_or_idle() or mode.is_target_verify() + + +def _experts_are_gu_interleaved(experts) -> bool: + """Whether the experts carry the gate/up interleaved MXFP4 layout. + + ``is_shuffled`` alone cannot answer this: Mxfp4MoEMethod sets it on both of + its branches, and only the SiTU A8W4 branch interleaves. Reading those + bytes with the wrong layout does not fail, it returns wrong numbers, so ask + the same question Mxfp4MoEMethod asked. + """ + if getattr(experts.moe_runner_config, "activation", None) != "situ": + return False + from sglang.srt.layers.quantization.mxfp4 import ( + _aiter_situ_uses_gu_interleaved_weights, + ) + + return _aiter_situ_uses_gu_interleaved_weights() + + +def mono_supported(moe: KimiK3MoE) -> bool: + """Static eligibility, evaluated once per layer. + + Only conditions that are the same on every TP rank belong here: this is + read without a collective, so a per-rank answer would deadlock the peers. + """ + from sglang.srt.environ import envs + from sglang.srt.layers.dp_attention import is_dp_attention_enabled + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + from sglang.srt.utils import is_hip + + if not envs.SGLANG_K3_MONO_DECODE.get() or not is_hip(): + return False + + experts = getattr(moe, "experts", None) + shared = getattr(moe, "shared_experts", None) + up_proj = getattr(moe, "routed_expert_up_proj", None) + checks = ( + (moe.use_latent_moe, "not a latent MoE layer"), + (moe.tp_size == MONO_TP_SIZE, f"needs TP{MONO_TP_SIZE}"), + (get_moe_a2a_backend().is_none(), "needs --moe-a2a-backend none"), + # Read the predicate, not a layer field: the field this used to cache + # went away with the layer-boundary refactor. + (not is_dp_attention_enabled(), "incompatible with DP attention"), + (shared is not None, "needs shared experts"), + (moe.routed_expert_norm is not None, "needs the latent RMSNorm"), + (up_proj is not None, "needs the latent up projection"), + (moe.moe_hidden_size == LATENT_SIZE, f"latent width must be {LATENT_SIZE}"), + (moe.routed_scaling_factor == 1.0, "routed scaling must be 1.0"), + (experts is not None and experts.top_k == 16, "top-k must be 16"), + ( + experts is not None and getattr(experts, "num_local_experts", None) == 896, + "expert count per rank must be 896", + ), + ( + experts is not None + and getattr(experts, "intermediate_size_per_partition", None) == 384, + "expert width per rank must be 384", + ), + ( + experts is not None + and getattr(experts.w13_weight, "is_shuffled", False) + and getattr(experts.w2_weight, "is_shuffled", False), + "experts must be in AITER's preshuffled MXFP4 layout", + ), + ( + shared is not None and shared.down_proj.weight.dtype == torch.bfloat16, + "shared down projection must be BF16", + ), + ( + up_proj is not None + and up_proj.weight.shape[0] == HIDDEN_SIZE + and up_proj.weight.is_contiguous(), + "latent up projection must be a contiguous [hidden, latent]", + ), + ) + for ok, why in checks: + if not ok: + # The operator asked for this explicitly, so say why it did not + # happen rather than falling back in silence. + _log_once(f"Kimi-K3 mono decode off: {why}") + return False + return True + + +class MonoMoe: + """Per-layer binding of one ``KimiK3MoE`` to the fused MoE launch. + + Everything is set up here, eagerly, before graph capture. That ordering is + required, not an optimization: the decode path replays a captured graph, so + whichever branch ``KimiK3MoE.forward`` takes at capture time is the branch + every replay runs. It also has to be before capture because the HIP-IPC + handshake allocates and runs a collective, neither of which is legal under + capture. + """ + + def __init__(self, moe: KimiK3MoE) -> None: + from sglang.kernels.ops.moe.k3_mono_flydsl import runner + + self._moe = moe + self._layer_idx = moe.layer_idx + self._rank = get_parallel().tp_rank + self._guint = _experts_are_gu_interleaved(moe.experts) + self._w_up: Optional[torch.Tensor] = None + self._ready = False + device = moe.gate.weight.device + + def build(): + runner.alloc_scratch(device) + runner.peers(device) + + try: + _bind_agreed(build) + self._ready = True + except Exception as error: + logger.warning( + "Kimi-K3 mono decode disabled on layer %d: %s", self._layer_idx, error + ) + + def w_up(self) -> torch.Tensor: + """This rank's rows of the replicated latent up projection. + + Slicing the first dimension of a contiguous weight keeps it contiguous, + so this is a view, not a copy. + """ + if self._w_up is None: + from sglang.kernels.ops.moe.k3_mono_flydsl.stages.moe import UP_N + + self._w_up = self._moe.routed_expert_up_proj.weight.narrow( + 0, self._rank * UP_N, UP_N + ) + return self._w_up + + def covered(self, x: torch.Tensor, forward_batch) -> bool: + """Whether this batch can take the fused launch. + + Fail closed. The kernel reads these tensors as raw addresses, so an + unexpected dtype or stride would be read as if it were the expected one + rather than raising. + """ + return ( + self._ready + and _is_decode_step(forward_batch) + and 0 < x.size(0) <= MONO_MAX_ROWS + and x.size(1) == HIDDEN_SIZE + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the whole MoE for one decode step. Returns the all-reduced sum + of the routed and shared experts; the caller still owns the residual.""" + from sglang.kernels.ops.moe.k3_mono_flydsl import runner + from sglang.kernels.ops.moe.k3_mono_flydsl.stages import moe as moe_stage + + _log_once("Kimi-K3 mono decode engaged on a decode step") + m = self._moe + experts = m.experts + out = torch.empty_like(x) + moe_stage.moe( + moe_stage.MoeBuild( + tokens=x.size(0), + eps=m.routed_expert_norm.variance_epsilon, + guint=self._guint, + ), + x=x, + w_gate=m.gate.weight, + bias=m.gate.e_score_correction_bias, + w_ld=m.routed_expert_down_proj.weight, + w_sgu=m.shared_experts.gate_up_proj.weight, + w_sd=m.shared_experts.down_proj.weight, + w13=experts.w13_weight, + w13s=experts.w13_weight_scale, + w2=experts.w2_weight, + w2s=experts.w2_weight_scale, + ln_w=m.routed_expert_norm.weight, + w_up=self.w_up(), + out=out, + queue=runner.queue(), + **runner.launch_args(), + # K1 owns tags below 128; the MoE launches take the upper half. + layer=128 + self._layer_idx, + ) + return out + + +def maybe_build_mono(moe: KimiK3MoE) -> Optional[MonoMoe]: + """Return this layer's mono binding, or None when it does not apply.""" + if not mono_supported(moe): + return None + mono = MonoMoe(moe) + if not mono._ready: + return None + _log_once("Kimi-K3 mono decode armed (TP8, decode steps of at most 8 rows)") + return mono + + +def k2_supported(layer: KimiK3DecoderLayer, mono: MonoMoe) -> bool: + """Static eligibility of one decoder layer for the whole-layer launch. + + The MoE binding is already a precondition, so only what the wider launch + adds is asked here: the attention seam it absorbs and the o_proj it runs + itself. Same rule as ``mono_supported`` -- every condition has to answer + the same on every TP rank. + """ + from sglang.srt.environ import envs + from sglang.srt.utils.common import require_mlp_sync + + if not envs.SGLANG_K3_MONO_LAYER.get(): + return False + if not layer.use_attn_residuals: + # prev_valid_blocks and the seam itself only exist with the stream. + _log_once("Kimi-K3 mono layer off: needs the attention-residual stream") + return False + + attn = layer.self_attn + o_proj = getattr(attn, "o_proj", None) + weight = getattr(o_proj, "weight", None) + checks = ( + # The seam the launch replaces is the plain one. SP-MoE reduce-scatters + # o_proj's output and the fused all-reduce completes it out of a + # symmetric buffer; the launch does neither. + (not layer._sp_moe, "incompatible with SP-MoE"), + (not layer.all_reduce_fusion, "incompatible with the fused all-reduce"), + # mlp-sync (DP attention or a MoE a2a backend) pads extend batches to a + # multiple of attn_tp_size; the launch has no padded-row handling. + # Read the predicate, not a layer field: the field this used to cache + # went away with the layer-boundary refactor. + (not require_mlp_sync(), "incompatible with mlp-sync padding"), + ( + hasattr(attn, "forward_gated_core"), + "attention cannot stop before o_proj", + ), + ( + weight is not None + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN_SIZE, CORE_SIZE) + and weight.is_contiguous(), + f"o_proj must be a contiguous BF16 [{HIDDEN_SIZE}, {CORE_SIZE}]", + ), + ( + o_proj is not None and getattr(o_proj, "bias", None) is None, + "o_proj must have no bias", + ), + ( + _k2_blocks(layer) <= MONO_MAX_BLOCKS, + f"at most {MONO_MAX_BLOCKS} attention-residual blocks", + ), + ) + for ok, why in checks: + if not ok: + _log_once(f"Kimi-K3 mono layer off: {why}") + return False + return mono is not None + + +def _k2_blocks(layer: KimiK3DecoderLayer) -> int: + """Bank rows readable at this layer's MLP seam. + + A block-write layer snapshots its own pre-attention prefix on the + attention side, so by the MLP seam it has one row more than it inherited. + """ + return layer.prev_valid_blocks + int(layer.is_block_write_layer) + + +class MonoK2: + """Per-layer binding of one ``KimiK3DecoderLayer`` to the layer launch. + + The launch starts from the attention's TP-local output (o_proj's input) + and ends at the layer's MoE output, so it owns o_proj, the attention-TP + all-reduce, the MLP attention-residual seam and the MoE. What stays + outside is the attention itself and the final residual add. + """ + + def __init__(self, layer: KimiK3DecoderLayer, mono: MonoMoe) -> None: + self._layer = layer + self._mono = mono + self._nblocks = _k2_blocks(layer) + # The kernel updates `prefix` in place. A block-write layer restarts it + # at the attention output (the kernel writes rather than accumulates); + # every other layer starts it from the stream's written residual, which + # the boundary lends out read-only, so this class owns the buffer in + # both modes and copies that residual in. Allocate before graph capture. + if True: + weight = layer.self_attn.o_proj.weight + self._prefix = torch.empty( + MONO_MAX_ROWS, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device=weight.device, + ) + + def covered( + self, + normed: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + attn_res: AttnResidual, + forward_batch, + nvb: Optional[int] = None, + ) -> bool: + """Whether this batch can take the layer launch. + + ``nvb`` overrides the bank count to check against, for a caller that + knows how many rows will be live by the time the launch runs. + + Fail closed, for the reason ``MonoMoe.covered`` fails closed: the + kernel reads every tensor as a raw address, so a wrong dtype or stride + is read as if it were the expected one. + """ + rows = normed.size(0) + if not ( + self._mono._ready + and _is_decode_step(forward_batch) + and 0 < rows <= MONO_MAX_ROWS + and normed.size(1) == HIDDEN_SIZE + and normed.dtype == torch.bfloat16 + and normed.is_contiguous() + ): + return False + # The kernel's two prefix modes are chosen at build time, so the + # caller's mode has to be the one this layer was built for. + if (prefix_sum is None) != self._layer.is_block_write_layer: + return False + if prefix_sum is not None and not ( + prefix_sum.shape == normed.shape + and prefix_sum.dtype == torch.bfloat16 + and prefix_sum.is_contiguous() + ): + return False + bank = attn_res.block_residual + if nvb is None: + nvb = attn_res.num_valid_blocks + return ( + nvb == self._nblocks + and bank.dtype == torch.bfloat16 + and bank.size(0) >= rows + and bank.size(2) == HIDDEN_SIZE + and bank.stride(2) == 1 + ) + + def forward( + self, + normed: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + attn_res: AttnResidual, + positions: torch.Tensor, + forward_batch, + zero_allocator, + core: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run attention, then the whole layer tail in one launch. + + Returns the layer's output: the updated prefix plus the MoE output. + ``core`` lets K1 hand over the attention output it already produced. + """ + from sglang.kernels.ops.moe.k3_mono_flydsl import layer as k2 + from sglang.kernels.ops.moe.k3_mono_flydsl import runner + + _log_once("Kimi-K3 mono layer engaged on a decode step") + layer = self._layer + moe = layer.mlp + experts = moe.experts + if core is None: + core = layer._run_self_attn( + normed, positions, forward_batch, zero_allocator, core=True + ) + rows = core.size(0) + assert core.shape == (rows, CORE_SIZE) and core.dtype == torch.bfloat16 + core = core.contiguous() + + prefix = self._prefix[:rows] + if prefix_sum is not None: + # Not `prefix = prefix_sum`: that tensor is the residual stream's, + # lent out read-only, and the kernel updates its prefix in place. + prefix.copy_(prefix_sum) + out = torch.empty_like(prefix) + k2.k2_launch( + k2.K2Build( + tokens=rows, + nblocks=self._nblocks, + eps=layer.mlp_res_norm.variance_epsilon, + out_eps=layer.post_attention_layernorm.variance_epsilon, + ln_eps=moe.routed_expert_norm.variance_epsilon, + reset=layer.is_block_write_layer, + guint=self._mono._guint, + ), + core=core, + w_o=layer.self_attn.o_proj.weight, + prefix=prefix, + blocks=attn_res.block_residual, + ares_nw=layer.mlp_res_norm.weight, + ares_qk=layer.mlp_res_proj.weight.view(HIDDEN_SIZE), + in_nw=layer.post_attention_layernorm.weight, + w_gate=moe.gate.weight, + bias=moe.gate.e_score_correction_bias, + w_ld=moe.routed_expert_down_proj.weight, + w_sgu=moe.shared_experts.gate_up_proj.weight, + w_sd=moe.shared_experts.down_proj.weight, + w13=experts.w13_weight, + w13s=experts.w13_weight_scale, + w2=experts.w2_weight, + w2s=experts.w2_weight_scale, + ln_w=moe.routed_expert_norm.weight, + w_up=self._mono.w_up(), + out=out, + **runner.launch_args(), + # Same tag as this layer's MoE-only launch: only one of them runs + # for a given layer in a given step. + layer=128 + layer.layer_idx, + ) + return out + prefix + + def maybe_forward( + self, + normed: torch.Tensor, + forward_batch, + positions: torch.Tensor, + zero_allocator, + ) -> Optional[torch.Tensor]: + """The layer's output from the whole-layer launch, or None to decline. + + The caller holds only the stage boundary's compute input; the prefix and + the snapshot bank come from the batch, the way the unfused seam reads + them. Declining is the fail-closed answer and must happen before the + kernel touches anything: a rank that opts out alone leaves the other + seven spinning on its mailbox. + """ + from sglang.srt.layers.layer_boundary.residual import batch as residual_batch + + bank = getattr(self._layer, "_attn_bank", None) + attn_res = bank.current if bank is not None else None + if attn_res is None: + return None + # A block-write layer restarts the prefix, so it takes none. + prefix_sum = ( + None + if self._layer.is_block_write_layer + else residual_batch.written_residual(forward_batch) + ) + if not self.covered(normed, prefix_sum, attn_res, forward_batch): + return None + return self.forward( + normed, prefix_sum, attn_res, positions, forward_batch, zero_allocator + ) + + +def maybe_build_k2(layer: KimiK3DecoderLayer) -> Optional[MonoK2]: + """Return this layer's whole-layer binding, or None when it does not apply.""" + mono = getattr(layer.mlp, "_mono", None) + if mono is None or not k2_supported(layer, mono): + return None + _log_once("Kimi-K3 mono layer armed (o_proj + AttnRes seam + MoE in one launch)") + return MonoK2(layer, mono) + + +# K1 shapes, from kernels/ops/moe/k3_mono_flydsl/attention/kda.py. +K1_CONV_W = 4 +K1_WIN = K1_CONV_W - 1 +K1_NPROJ = 6288 +K1_HEADS = 12 +K1_HEAD_DIM = 128 +K1_QKV_DIM = 3 * CORE_SIZE # the conv channels: q, k and v +# The launch rebases its conv buffer every call, writing qlen entries after the +# carried window. We own the buffer, so size it for the widest step. +K1_STATE_LEN = MONO_MAX_ROWS + K1_WIN + + +def k1_supported(layer: KimiK3DecoderLayer) -> bool: + """Static eligibility of one KDA layer for the K1 launch. + + Every weight below is a view of something SGLang already built -- in + particular ``_qkvgbfa_layer.weight`` is exactly K1's in-projection layout + ([q,k,v,g | f_a | b | pad], 6288 rows), so nothing is repacked. + """ + from sglang.srt.environ import envs + + if not (envs.SGLANG_K3_MONO_LAYER.get() and envs.SGLANG_K3_MONO_K1.get()): + return False + attn = layer.self_attn + # K1 is the KDA half of the model; MLA layers keep their own attention and + # only take K2. Check before anything else -- the conditions below read + # KDA-only attributes, and a tuple of them is built eagerly. + inner = getattr(attn, "attn", None) + if inner is None or not hasattr(attn, "qkv_conv1d"): + return False + inproj = getattr(attn, "_qkvgbfa_layer", None) + conv = attn.qkv_conv1d + checks = ( + (getattr(attn, "use_full_rank_gate", False), "needs the full-rank gate"), + ( + inproj is not None + and tuple(inproj.weight.shape) == (K1_NPROJ, HIDDEN_SIZE) + and inproj.weight.dtype == torch.bfloat16, + "needs the merged ROCm in-projection (SGLANG_ROCM_K3_FUSE_KDA_INPROJ=1)", + ), + ( + getattr(attn, "local_num_heads", None) == K1_HEADS + and getattr(attn, "head_dim", None) == K1_HEAD_DIM, + f"needs {K1_HEADS} local heads of {K1_HEAD_DIM}", + ), + ( + conv.bias is None and conv.weight.dtype == torch.float32, + "conv1d must be bias-free FP32 (K1 has no bias term)", + ), + ( + getattr(inner, "lower_bound", None) is not None, + "needs the KDA gate lower bound", + ), + (layer.prev_valid_blocks <= MONO_MAX_BLOCKS, "too many residual blocks"), + ) + for ok, why in checks: + if not ok: + _log_once(f"Kimi-K3 mono K1 off: {why}") + return False + return True + + +class MonoK1: + """Per-layer binding of one KDA layer to the K1 launch. + + K1 runs from the attention-residual seam to the gated core output, so it + replaces aggregation 1, input_layernorm and the whole attention. Its output + is what ``forward_gated_core`` returns, which is what the layer launch (K2) + consumes -- the two chain with no glue. + + This class owns the conv buffer. SGLang's speculative window is not used. + The buffer is reseeded from the committed conv state on every launch, which + keeps ``num_acc`` at 1. That drops the deepest read index from + ``qlen + CONV_W - 2`` to ``CONV_W - 1``. The launch then depends on no + buffer that SGLang leaves stale. See test_kimi_k3_mono_kda_pre.py. + """ + + def __init__(self, layer: KimiK3DecoderLayer) -> None: + attn = layer.self_attn + self._layer = layer + self._nblocks = layer.prev_valid_blocks + self._write_idx = layer.prev_valid_blocks if layer.is_block_write_layer else -1 + self._w_in = attn._qkvgbfa_layer.weight + self._w_fb = attn.f_b_proj.weight + self._conv_w = attn.qkv_conv1d.weight.squeeze(1) + self._a_log = attn.A_log.detach().reshape(-1) + self._dt_bias = attn.dt_bias + self._on_w = attn.o_norm.weight.data.to(torch.bfloat16).contiguous() + self._lower_bound = float(attn.attn.lower_bound) + self._eps = layer.self_attention_res_norm.variance_epsilon + self._out_eps = layer.input_layernorm.variance_epsilon + self._onorm_eps = float(attn.o_norm.eps) + assert self._conv_w.shape == (K1_QKV_DIM, K1_CONV_W), self._conv_w.shape + # Ours, so it never goes stale: reseeded from the committed state every + # launch. One request's worth -- see covered(). + self._conv_buf = torch.empty( + 1, + K1_QKV_DIM, + K1_STATE_LEN, + dtype=torch.bfloat16, + device=self._w_in.device, + ) + + def covered(self, prefix: torch.Tensor, attn_res: AttnResidual) -> bool: + """Fail closed, like the other two launches: the kernel reads raw + addresses, so a wrong dtype or stride is read, not rejected.""" + bank = attn_res.block_residual + return ( + 0 < prefix.size(0) <= MONO_MAX_ROWS + and prefix.size(1) == HIDDEN_SIZE + and prefix.dtype == torch.bfloat16 + and prefix.is_contiguous() + and attn_res.num_valid_blocks == self._nblocks + and bank.dtype == torch.bfloat16 + and bank.size(2) == HIDDEN_SIZE + and bank.stride(2) == 1 + ) + + def forward(self, prefix, attn_res, forward_batch) -> Optional[torch.Tensor]: + """Returns the gated core output, or None when the batch is not served. + + None is a clean fallback: nothing has been written yet when it is + returned, so the caller just takes the normal path. + """ + from sglang.kernels.ops.moe.k3_mono_flydsl import runner + from sglang.kernels.ops.moe.k3_mono_flydsl.attention import kda + from sglang.srt.model_executor.forward_context import get_attn_backend + + backend = getattr(get_attn_backend(), "linear_attn_backend", None) + state = ( + backend.kimi_k3_mono_state(self._layer.self_attn.attn, forward_batch) + if hasattr(backend, "kimi_k3_mono_state") + else None + ) + if state is None: + _log_once("Kimi-K3 mono K1 skipped: no speculative state") + return None + qlen = state["num_draft"] + rows = state["rows"] + slots = state["slots"] + # Take the request count from the token rows, not from rows.numel(). + # A captured graph pads req_pool_indices out to the captured batch + # size, so that table can be longer than the batch. + # + # One request per launch. K1 indexes the conv buffer and the recurrent + # state with the same slot number. A private conv buffer therefore + # serves one request only. covered() caps the rows at MONO_MAX_ROWS, + # so a second request cannot fit anyway. + if qlen <= 0 or prefix.size(0) != qlen or rows.numel() < 1: + _log_once( + f"Kimi-K3 mono K1 skipped: {prefix.size(0)} rows for a " + f"{qlen}-token step, {rows.numel()} state rows" + ) + return None + inter_ssm = state["inter_ssm"] + if inter_ssm.dtype != torch.float32 or inter_ssm.shape[1] != qlen: + _log_once( + f"Kimi-K3 mono K1 skipped: intermediate_ssm is " + f"{inter_ssm.dtype} {tuple(inter_ssm.shape)}" + ) + return None + + # Every index below stays on the device. Reading one out with int() + # is a host sync. A CUDA graph capture rejects a host sync, and this + # path only runs inside a captured decode graph. + row = rows[:1].long() + slot = slots[:1].long() + st_idx = (row * qlen).view(1, 1) + torch.arange( + qlen, device=prefix.device, dtype=torch.long + ).view(1, qlen) + # acc_idx = 0: the window sits at offset 0 and the initial recurrent + # state is the one at st_idx[0]. Seed both from the committed state. + # No buffer of this class then has to stay live across steps. + num_acc = torch.ones(1, dtype=torch.int32, device=prefix.device) + + window = state["committed_conv"].index_select(0, slot) + assert window.shape[-1] == K1_QKV_DIM and window.shape[-2] == K1_WIN, ( + f"unexpected committed conv layout {tuple(window.shape)}" + ) + self._conv_buf[:, :, :K1_WIN] = window.transpose(-1, -2) + flat_ssm = inter_ssm.reshape(-1, *inter_ssm.shape[2:]) + flat_ssm.index_copy_( + 0, + st_idx[:, 0], + state["committed_ssm"].index_select(0, slot).to(flat_ssm.dtype), + ) + + core = torch.empty( + prefix.size(0), CORE_SIZE, dtype=torch.bfloat16, device=prefix.device + ) + kda.kda_pre( + kda.KdaPreBuild( + tokens=prefix.size(0), + qlen=qlen, + nblocks=self._nblocks, + delta=False, + state_len=K1_STATE_LEN, + eps=self._eps, + out_eps=self._out_eps, + onorm_eps=self._onorm_eps, + lower_bound=self._lower_bound, + write_idx=self._write_idx, + ), + prefix=prefix, + delta=None, + blocks=attn_res.block_residual, + ares_nw=self._layer.self_attention_res_norm.weight, + ares_qk=self._layer.self_attention_res_proj.weight.view(HIDDEN_SIZE), + in_nw=self._layer.input_layernorm.weight, + w_in=self._w_in, + w_fb=self._w_fb, + conv_w=self._conv_w, + # Expanded so that its slot stride is zero. K1 indexes the conv + # buffer with the number it indexes the recurrent state with. + # That number is a row of SGLang's scratch, not of this buffer. + # A zero stride folds every such number onto the single row here, + # so the expanded size does not matter. + conv_state=self._conv_buf.expand(2, -1, -1), + a_log=self._a_log, + dt_bias=self._dt_bias, + on_w=self._on_w, + rstate=flat_ssm, + st_idx=st_idx.to(torch.int32), + num_acc=num_acc, + core_out=core, + scratch=runner.scratch(), + layer=self._layer.layer_idx, + epoch=runner.epoch(), + ) + # Hand the step's conv inputs to SGLang's speculative window so its own + # accept-time commit (fused_conv_window_scatter_with_mask) advances the + # committed state. view[s, j, k] aliases phys[s, j + k], so writing + # column CONV_W - 2 of steps 1.. lays down one entry each. + # Indexed with a tensor, not an int, for the reason above. The two + # writes are disjoint in the physical buffer -- the first lands on + # entries 0 .. CONV_W - 2, the second on the rest -- so the view's + # overlapping windows cannot alias between them. + view = state["inter_window"] + view[row, 0] = window + view[row, 1:, K1_WIN - 1] = self._conv_buf[ + :, :, K1_WIN : K1_STATE_LEN - 1 + ].transpose(-1, -2) + _log_once("Kimi-K3 mono K1 engaged on a decode step") + return core + + +def maybe_build_k1(layer: KimiK3DecoderLayer) -> Optional[MonoK1]: + """Return this layer's K1 binding, or None when it does not apply.""" + if not k1_supported(layer): + return None + _log_once("Kimi-K3 mono K1 armed (KDA layers, seam + attention in one launch)") + return MonoK1(layer) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 1979a2e76d30..3fe4b999af24 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -474,6 +474,8 @@ def __init__( # True when _front_w merges only [gate, routed_expert_down_proj] (the EP # a2a pair) rather than the three-way fused-front weight. self._front_is_ep_pair = False + # Set by _prepare_mono() once the weights are in their final layout. + self._mono = None self.moe_hidden_size = ( config.routed_expert_hidden_size if self.use_latent_moe else hidden_size ) @@ -751,6 +753,20 @@ def _merge_front_weights(self) -> None: ): self.__dict__.pop(prop, None) + def _prepare_mono(self) -> None: + """Bind the ROCm mono decode launch to this layer, if it applies. + + Called from precompile_kernels_after_loading. That hook runs after + every quant method's process_weights_after_loading, so the experts are + in their final preshuffled layout here. They are not yet in + post_load_weights. The hook is still before the KV cache is sized and + before any graph capture. + """ + from sglang.srt.layers.moe.k3_mono_decode import maybe_build_mono + + if self._mono is None: + self._mono = maybe_build_mono(self) + @cached_property def _routed_needs_reduce(self): return self.tp_size > 1 and get_moe_a2a_backend().is_none() @@ -1477,14 +1493,24 @@ def forward( hidden_states: torch.Tensor, *, prefix_sum: Optional[torch.Tensor] = None, + forward_batch: Optional[ForwardBatch] = None, ) -> torch.Tensor: """A latent MoE, which writes the next stream itself, takes the pending prefix_sum: folded into the 3-way JIT tail add when covered, plain adds otherwise (bit-identical either way). A non-latent MoE leaves the - residual add to its stage boundary and takes none.""" + residual add to its stage boundary and takes none. + + ``forward_batch`` only feeds the mono launch's decode-step predicate; a + caller that omits it declines the launch (the predicate fails closed).""" num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - if hidden_states.shape[0] > 0 and self._eligible_for_fused_front: + if self._mono is not None and self._mono.covered(hidden_states, forward_batch): + # The fused launch returns the all-reduced routed + shared sum. + # The caller still adds the residual. K2 absorbs that add. + out = self._mono.forward(hidden_states) + if prefix_sum is not None: + out = out + prefix_sum + elif hidden_states.shape[0] > 0 and self._eligible_for_fused_front: out = self._forward_fused(hidden_states, prefix_sum=prefix_sum) else: out = self._forward_unfused(hidden_states, prefix_sum=prefix_sum) @@ -1887,7 +1913,14 @@ def _may_fuse_kda_inproj(self) -> bool: that a raw row-cat would silently drop.""" if not (_is_hip and envs.SGLANG_ROCM_K3_FUSE_KDA_INPROJ.get()): return False - if not (self.do_fuse_qkvbfg and self.use_full_rank_gate): + # Deliberately not do_fuse_qkvbfg. That flag also means "quant_config + # is unset", which is stricter than this merge needs. The K3 + # checkpoint carries a quantization_config for the MoE experts. Its + # attention linears still resolve to UnquantizedLinearMethod. The real + # question is whether these three weights are mergeable, and + # _is_unquantized_mergeable below asks that directly. The part of + # do_fuse_qkvbfg still needed here is its attention-TP condition. + if not (self.use_full_rank_gate and self.attn_tp_size == self.tp_size): return False # Block-FP8 in-proj needs dequantized BF16 buffers; a raw row-cat # would drop the scales. Leave fusion to the split [f_a|b] path. @@ -2077,13 +2110,23 @@ def forward_qkvbfg_fused( ) return qkv, beta, forget_gate, g_proj_states - def forward( + def forward_gated_core( self, + *, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, + **kwargs, ) -> torch.Tensor: + """The TP-local attention output after the output gate, before o_proj. + + Shape is ``[num_tokens, num_local_heads * head_dim]``. o_proj is left + to the caller, which also owns the attention-TP reduction that o_proj + would do: ``o_proj(forward_gated_core(x)) == forward(x)``. The fused + decode path folds o_proj into a larger kernel, so it needs the + attention to stop here. + """ defer_f_b = ( self._kda_hip_fused_decode_ready and forward_batch.forward_mode.is_decode() ) @@ -2131,7 +2174,21 @@ def forward( if not fused_onorm: norm_gate = g_proj_states.unflatten(-1, (-1, self.head_dim)) core_attn_out = self.o_norm(core_attn_out, norm_gate) - core_attn_out = core_attn_out.squeeze(0).flatten(-2) + return core_attn_out.squeeze(0).flatten(-2) + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + core_attn_out = self.forward_gated_core( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) if self.all_reduce_fusion: out = _k3_symm_o_proj_out(self.o_proj, core_attn_out) partial, _ = self.o_proj(core_attn_out, output_tensor=out) @@ -2142,6 +2199,10 @@ def forward( class KimiK3MLAAttention(DeepseekV2AttentionMLA): """MLA with output gate for K3. Gate is applied in TP-local space before o_proj.""" + # Set by forward_gated_core() for the length of one forward; the gated + # o_proj wrapper below writes the pre-o_proj tensor into it. + _core_sink: Optional[List[torch.Tensor]] = None + def __init__( self, config, @@ -2297,6 +2358,14 @@ def _gated_o_proj_forward(x, *args, **kwargs): # Even a skipped gate must close its capture branch. torch.cuda.current_stream().wait_stream(self._gate_pending_stream) self._gate_pending_stream = None + sink = self._core_sink + if sink is not None and not isinstance(x, tuple): + # forward_gated_core(): hand the gated TP-local output back + # and skip o_proj, whose GEMM and attention-TP reduction + # belong to the caller. The MLA cores reach o_proj from + # several places and only unpack a 2-tuple from it. + sink.append(x) + return x, None return _orig_o_proj_forward(x, *args, **kwargs) self.o_proj.forward = _gated_o_proj_forward @@ -2360,6 +2429,40 @@ def forward( positions, hidden_states, forward_batch, zero_allocator, **kwargs ) + def forward_gated_core( + self, + *, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + **kwargs, + ) -> torch.Tensor: + """The TP-local attention output after the output gate, before o_proj. + + Shape is ``[num_tokens, num_local_heads * v_head_dim]``. o_proj is left + to the caller, which also owns the attention-TP reduction that o_proj + would do: ``o_proj(forward_gated_core(x)) == forward(x)``. The fused + decode path folds o_proj into a larger kernel, so it needs the + attention to stop here. + + o_proj is called from inside the DeepseekV2AttentionMLA cores, so the + split point is the same instance-level wrapper that applies the gate. + """ + assert self.use_output_gate, "forward_gated_core needs the gated o_proj wrapper" + sink: List[torch.Tensor] = [] + self._core_sink = sink + try: + self.forward( + positions, hidden_states, forward_batch, zero_allocator, **kwargs + ) + finally: + self._core_sink = None + if not sink: + # Empty batch, or a gate skipped because o_proj got a tuple input. + raise RuntimeError("MLA forward returned without reaching o_proj") + return sink[0] + def _is_moe_layer(config: KimiLinearConfig, layer_idx: int) -> bool: """Whether a layer runs a MoE rather than the dense MLP: past the dense @@ -2508,6 +2611,12 @@ def __init__( super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx + # Set by _prepare_mono_k2() / _prepare_mono_k1() once the weights are + # in their final layout. The bank is kept because the mono layer launch + # reads the snapshots itself, in place of the seam it replaces. + self._mono_k2 = None + self._mono_k1 = None + self._attn_bank = attn_bank self._is_moe_layer = _is_moe_layer(config, layer_idx) # SP-MoE: with an EP a2a backend and attention TP, a MoE layer runs on # this rank's attention-TP shard of the rows. @@ -2638,12 +2747,31 @@ def _ffn_input_fusions(self): ) return () + def _prepare_mono_k2(self) -> None: + """Bind the ROCm whole-layer mono launch to this layer, if it applies. + + Called from precompile_kernels_after_loading, right after the MoE + binding this one builds on. The experts are in their final preshuffled + layout by then. That point is still before the KV cache is sized and + before any graph capture. + """ + from sglang.srt.layers.moe.k3_mono_decode import maybe_build_k2 + + if self._mono_k2 is None: + self._mono_k2 = maybe_build_k2(self) + if self._mono_k2 is not None and self._mono_k1 is None: + # K1 only helps where K2 already runs: it feeds K2 its core. + from sglang.srt.layers.moe.k3_mono_decode import maybe_build_k1 + + self._mono_k1 = maybe_build_k1(self) + def _run_self_attn( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, + core: bool = False, ) -> torch.Tensor: # DP attention: idle ranks (padded to the global shape) have no # attention metadata; pass hidden_states through shape-preserving @@ -2663,7 +2791,12 @@ def _run_self_attn( attn_inputs = AttentionInputs(hidden_states, forward_batch, qkv_latent_func) get_attn_tp_context().set_attn_inputs(attn_inputs) - result = self.self_attn( + # core=True stops before o_proj, for a caller that runs o_proj itself + # (the mono layer launch folds it into the kernel). + attn_forward = ( + self.self_attn.forward_gated_core if core else self.self_attn.__call__ + ) + result = attn_forward( hidden_states=hidden_states, positions=positions, forward_batch=forward_batch, @@ -2686,18 +2819,32 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, capture_gathered=capture_gathered ) + # The whole layer tail in one launch: o_proj, the attention-TP + # all-reduce, the FFN input's seam and the MoE. It returns the layer's + # output with the residual already folded in, so the stream is written + # rather than finished through ffn_boundary. None means it declined. + if self._mono_k2 is not None: + mono_out = self._mono_k2.maybe_forward( + hidden_states, forward_batch, positions, zero_allocator + ) + if mono_out is not None: + return residual_batch.set_written(mono_out, forward_batch) hidden_states = self._run_self_attn( hidden_states, positions, forward_batch, zero_allocator ) hidden_states = self.attn_boundary.finish(hidden_states, forward_batch) hidden_states = self.ffn_boundary.prepare(hidden_states, forward_batch) + # Only a MoE takes forward_batch, and only to feed the mono launch's + # decode-step predicate; the dense MLP's forward does not accept it. + mlp_call = {"forward_batch": forward_batch} if self._is_moe_layer else {} if self._ffn_writes_stream: hidden_states = self.mlp( hidden_states, prefix_sum=residual_batch.written_residual(forward_batch), + **mlp_call, ) else: - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, **mlp_call) return self.ffn_boundary.finish(hidden_states, forward_batch) @@ -2714,6 +2861,9 @@ def __init__( self.config = config self.pp_group = get_parallel().pp_group self.dspark_layers_to_capture: Optional[list[int]] = None + # Set by precompile_kernels_after_loading once a layer binds the ROCm + # mono decode launch; gates the per-forward mailbox epoch bump. + self._mono_decode = False self._dp_attention = is_dp_attention_enabled() # Whether K3's SP collective is enabled on this rank, which the layers' # stage declarations depend on: settled once, before any is built. @@ -2811,6 +2961,12 @@ def forward( inputs_embeds: torch.Tensor | None = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: + if self._mono_decode: + # Advance the mono kernels' mailbox epoch once a forward. A device + # op, so a captured graph replays it. + from sglang.srt.layers.moe.k3_mono_decode import step_begin + + step_begin() if get_parallel().pp_group.is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -3422,6 +3578,21 @@ def _warm_cw(proj, norm): rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.") break + def precompile_kernels_after_loading(self) -> None: + # Runs after every quant method's process_weights_after_loading, so + # the expert weights are in their final preshuffled layout here. They + # are not in post_load_weights. This point is still before the KV + # cache is sized and before any graph capture. + armed = False + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, KimiK3MoE): + layer.mlp._prepare_mono() + layer._prepare_mono_k2() + armed = armed or layer.mlp._mono is not None + self.model._mono_decode = armed + class KimiK3ForConditionalGeneration(nn.Module): """K3 multimodal wrapper: MoonViT3d tower + KimiK3LinearForCausalLM.""" @@ -3516,6 +3687,10 @@ def post_load_weights(self): self.language_model.post_load_weights() def precompile_kernels_after_loading(self) -> None: + # The LM tower's share runs either way; only the vision kernels are + # skipped in language-only mode. + if self.language_model is not None: + self.language_model.precompile_kernels_after_loading() if self.config.language_only: return if self.vision_tower.precompile_fused_rope(): diff --git a/test/registered/amd/test_kimi_k3_mono_decode.py b/test/registered/amd/test_kimi_k3_mono_decode.py new file mode 100644 index 000000000000..9fa394031921 --- /dev/null +++ b/test/registered/amd/test_kimi_k3_mono_decode.py @@ -0,0 +1,578 @@ +"""TP8 correctness tests for the Kimi-K3 mono launches (ROCm / FlyDSL). + +For the MoE launch, two things are checked, and they answer different +questions: + +* against AITER's ``fused_moe`` stack, that the kernel computes the MoE; +* that the two AITER expert layouts -- gate/up separated, which the kernel was + written for, and gate/up interleaved, which SGLang's ``Mxfp4MoEMethod`` + produces -- give the same answer from the same logical weights. + +For the whole-layer launch, the same MoE check plus the part the wider launch +adds: o_proj, its all-reduce and the MLP attention-residual seam. Those two +halves are checked separately, against the kernel's own seam output, so that a +rounding difference in the seam cannot flip a top-16 choice and turn a correct +MoE into a failure that reads like a MoE bug. See ``_worker_k2``. + +Run directly with ``python test/registered/amd/test_kimi_k3_mono_decode.py``; +it needs 8 GPUs. +""" + +import socket +import unittest + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=300, suite="stage-c-test-large-8-gpu-amd-mi35x") + +TP_SIZE = 8 +ROWS = 8 +COS_MIN = 0.999 +BF = torch.bfloat16 +# Attention-residual bank: rows the launch is built for, and the live rows the +# layer test uses. Keeping them apart exercises the production bank stride. +MAX_BLOCKS = 8 +NVB = 3 + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _randn(gen, *shape, scale=1.0, dtype=BF): + return ( + torch.randn(*shape, generator=gen, device=gen.device, dtype=torch.float32) + * scale + ).to(dtype) + + +def _cos(a, b): + return float( + F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0, eps=1e-12) + ) + + +def _all_reduce_fp32(t): + """The ranks' BF16 partials summed in FP32, as the custom all-reduce does.""" + t = t.float() + dist.all_reduce(t) + return t.to(BF) + + +def _situ(g, u): + g, u = g.float(), u.float() + a = (2.0 * torch.tanh(g * 0.25)) * (1.0 + torch.tanh(g * 0.5)) + return (a * (25.0 * torch.tanh(u * 0.04))).to(BF) + + +def _expert_bytes(gen, device, *, gate_up_interleaved): + """The routed experts in one of AITER's two A16W4 layouts. + + The logical bytes are drawn from a fixed generator, so the two layouts hold + the same weights and must give the same MoE output. Scale bytes stay near + E8M0 1.0: the exponent is raw with bias 127, and uniform bytes would reach + 2^128 and make the layer non-finite. + """ + from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 + + from sglang.kernels.ops.moe.k3_mono_flydsl.stages.moe import LAT, RI, E + + def u8(*shape, lo=0, hi=256): + return torch.randint( + lo, hi, shape, generator=gen, device=device, dtype=torch.int32 + ).to(torch.uint8) + + gu = gate_up_interleaved + return { + "w13": shuffle_weight_a16w4(u8(E, 2 * RI, LAT // 2), 16, gu), + "w13s": shuffle_scale_a16w4( + u8(E, 2 * RI, LAT // 32, lo=118, hi=123).view(-1, LAT // 32), E, gu + ), + "w2": shuffle_weight_a16w4(u8(E, LAT, RI // 2), 16, False), + "w2s": shuffle_scale_a16w4( + u8(E, LAT, RI // 32, lo=118, hi=123).view(-1, RI // 32), E, False + ), + } + + +def _weights(rank, device, *, gate_up_interleaved): + from sglang.kernels.ops.moe.k3_mono_flydsl.stages.moe import ( + HIDDEN, + LAT, + SI, + UP_N, + E, + ) + + # The router, the latent down and the latent up are replicated across + # ranks; the shared and routed experts are TP shards. + gc = torch.Generator(device=device).manual_seed(7) + gr = torch.Generator(device=device).manual_seed(100 + rank) + ge = torch.Generator(device=device).manual_seed(100 + rank) + + w_up = _randn(gc, HIDDEN, LAT, scale=0.02) + out = { + "w_gate": _randn(gc, E, HIDDEN, scale=0.02), + "bias": _randn(gc, E, scale=0.02, dtype=torch.float32), + "w_ld": _randn(gc, LAT, HIDDEN, scale=0.02), + "ln_w": (1 + _randn(gc, LAT, scale=0.1, dtype=torch.float32)).to(BF), + "w_up_shard": w_up[rank * UP_N : (rank + 1) * UP_N].contiguous(), + "w_sgu": _randn(gr, 2 * SI, HIDDEN, scale=0.02), + "w_sd": _randn(gr, HIDDEN, SI, scale=0.03), + } + out.update(_expert_bytes(ge, device, gate_up_interleaved=gate_up_interleaved)) + return out + + +def _reference(x, w, rank): + """The MoE as SGLang's unfused decode path computes it.""" + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe + + from sglang.kernels.ops.moe.k3_mono_flydsl.stages.moe import SI, TOPK, UP_N + + sig = torch.sigmoid(torch.mm(x.float(), w["w_gate"].float().t())) + ids = torch.topk(sig + w["bias"], TOPK, dim=-1, sorted=True).indices + wts = sig.gather(-1, ids) + wts = wts / wts.sum(-1, keepdim=True) + gu = F.linear(x, w["w_sgu"]) + shared = F.linear(_situ(gu[:, :SI], gu[:, SI:]), w["w_sd"]) + # The kernel takes the experts as raw bytes, which is how SGLang holds + # them; fused_moe dispatches on the dtype and on an is_shuffled marker + # that a view drops, so tag the same storage for the reference only. + fp4 = torch.float4_e2m1fn_x2 + w13, w2 = w["w13"].view(fp4), w["w2"].view(fp4) + w13.is_shuffled = w2.is_shuffled = True + y = fused_moe( + F.linear(x, w["w_ld"]), + w13, + w2, + wts.float(), + ids.to(torch.int32), + None, + ActivationType.Situv2, + QuantType.per_1x32, + False, + w["w13s"], + w["w2s"], + None, + None, + dtype=BF, + gate_mode="separated", + beta=4.0, + linear_beta=25.0, + ) + lat = _all_reduce_fp32(y).float() + ln = lat * torch.rsqrt(lat.pow(2).mean(-1, keepdim=True) + 1e-5) * w["ln_w"] + lo = rank * UP_N + up = ln.to(BF).float() @ w["w_up_shard"].float().t() + shared[:, lo : lo + UP_N] = (shared[:, lo : lo + UP_N].float() + up).to(BF) + return _all_reduce_fp32(shared) + + +def _run_moe(mod, w, x, *, peers, epoch, rank, layer, guint): + out = torch.empty_like(x) + key = ( + mod.MoeBuild(tokens=x.size(0), guint=guint) + if guint + else mod.MoeBuild(tokens=x.size(0)) + ) + mod.moe( + key, + x=x, + w_gate=w["w_gate"], + bias=w["bias"], + w_ld=w["w_ld"], + w_sgu=w["w_sgu"], + w_sd=w["w_sd"], + w13=w["w13"], + w13s=w["w13s"], + w2=w["w2"], + w2s=w["w2s"], + ln_w=w["ln_w"], + w_up=w["w_up_shard"], + out=out, + scratch=torch.zeros(mod.scratch_bytes(key), dtype=torch.uint8, device=x.device), + queue=torch.zeros(mod.QUEUE_BYTES // 4, dtype=torch.int32, device=x.device), + peers=peers.addresses, + rank=rank, + epoch=epoch, + layer=layer, + ) + return out + + +def _worker(rank, port, results, check_guint): + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=TP_SIZE + ) + peers = None + try: + from sglang.kernels.ops.moe.k3_mono_flydsl.common.peer_memory import PeerBuffer + from sglang.kernels.ops.moe.k3_mono_flydsl.stages import moe + + peers = PeerBuffer( + moe.peer_bytes(ROWS), dist.group.WORLD, rank, TP_SIZE, device + ) + peers.bytes.zero_() + epoch = torch.zeros(1, dtype=torch.int32, device=device) + + gx = torch.Generator(device=device).manual_seed(11) + x = _randn(gx, ROWS, moe.HIDDEN) + + w = _weights(rank, device, gate_up_interleaved=False) + epoch.add_(1) + out = _run_moe( + moe, w, x, peers=peers, epoch=epoch, rank=rank, layer=1, guint=False + ).clone() + torch.cuda.synchronize() + + gathered = [torch.empty_like(out.cpu()) for _ in range(TP_SIZE)] + dist.all_gather(gathered, out.cpu().contiguous()) + result = { + "finite": bool(torch.isfinite(out).all()), + "nonzero": bool(torch.count_nonzero(out) > 0), + "ranks_agree": all(torch.equal(gathered[0], g) for g in gathered[1:]), + "cos_vs_reference": _cos(out, _reference(x, w, rank)), + "error": "", + } + + if check_guint: + # Control: the same launch again, same weights. Any difference here + # is the harness or the kernel's step state, not the layout. + epoch.add_(1) + out_again = _run_moe( + moe, w, x, peers=peers, epoch=epoch, rank=rank, layer=1, guint=False + ).clone() + torch.cuda.synchronize() + result["cos_rerun"] = _cos(out_again, out) + + # Same logical weights, SGLang's gate/up interleaved layout. + w_gu = _weights(rank, device, gate_up_interleaved=True) + epoch.add_(1) + out_gu = _run_moe( + moe, w_gu, x, peers=peers, epoch=epoch, rank=rank, layer=1, guint=True + ).clone() + torch.cuda.synchronize() + result["cos_layouts"] = _cos(out_gu, out) + result["layouts_equal"] = bool(torch.equal(out_gu, out)) + + results[rank] = result + except Exception as error: # surfaced by the parent as a test failure + import traceback + + results[rank] = { + "error": f"{type(error).__name__}: {error}\n{traceback.format_exc()}" + } + finally: + if peers is not None: + peers.close() + dist.destroy_process_group() + + +def _rms(x, weight, eps): + """RMSNorm over the last dimension, in FP32 like the kernel.""" + x = x.float() + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight.float() + + +def _seam_reference(core, w_o, prefix_in, blocks, nvb, w, *, reset): + """o_proj + its all-reduce + the MLP attention-residual seam, in torch. + + Mirrors AttnResidual's eager path: score every source (the ``nvb`` banked + rows and the updated prefix) through the score RMSNorm and the 1-row + projection, softmax over the sources, mix, then the output RMSNorm. + Returns (updated prefix, the MoE input). + """ + attn = _all_reduce_fp32(F.linear(core, w_o)).float() + # A block-write layer restarts the prefix at the attention output. + prefix = (attn if reset else prefix_in.float() + attn).to(BF) + rows = torch.cat([blocks[:, :nvb, :], prefix.unsqueeze(1)], dim=1) + scores = (_rms(rows, w["ares_nw"], w["eps"]) * w["ares_qk"].float()).sum(-1) + probs = torch.softmax(scores, dim=-1) + mixed = (probs.unsqueeze(-1) * rows.float()).sum(dim=1) + return prefix, _rms(mixed, w["in_nw"], w["out_eps"]).to(BF) + + +def _layer_weights(rank, device, nvb): + """The layer launch's own weights, on top of the MoE's. + + o_proj is a TP shard, so it is drawn per rank; the attention-residual + stream and its two norms are replicated and must match on every rank. + """ + from sglang.kernels.ops.moe.k3_mono_flydsl.attention.back import OK + from sglang.kernels.ops.moe.k3_mono_flydsl.attention.kda import HIDDEN + + gc = torch.Generator(device=device).manual_seed(23) + gr = torch.Generator(device=device).manual_seed(200 + rank) + return { + "w_o": _randn(gr, HIDDEN, OK, scale=0.02), + "core": _randn(gr, ROWS, OK, scale=0.3), + # The bank is allocated for every block of the model, so only its + # first nvb rows are live -- that is the production stride. + "blocks": _randn(gc, ROWS, MAX_BLOCKS, HIDDEN, scale=0.3), + "prefix": _randn(gc, ROWS, HIDDEN, scale=0.3), + "ares_nw": (1 + _randn(gc, HIDDEN, scale=0.1)).to(BF), + "ares_qk": _randn(gc, HIDDEN, scale=0.02), + "in_nw": (1 + _randn(gc, HIDDEN, scale=0.1)).to(BF), + "eps": 1e-5, + "out_eps": 1e-5, + "nvb": nvb, + } + + +def _run_k2(mod, w, scratch, *, nvb, reset, peers, epoch, rank, layer): + """One layer launch. Returns (out, the updated prefix, the kernel's own + MoE input read back out of its scratch).""" + from sglang.kernels.ops.moe.k3_mono_flydsl.attention.kda import HIDDEN + + key = mod.K2Build( + tokens=ROWS, + nblocks=nvb, + eps=w["eps"], + out_eps=w["out_eps"], + reset=reset, + ) + # The kernel updates the prefix in place, so hand it a copy and keep the + # input for the reference. + prefix = w["prefix"].clone() + out = torch.empty_like(prefix) + mod.k2_launch( + key, + core=w["core"], + w_o=w["w_o"], + prefix=prefix, + blocks=w["blocks"], + ares_nw=w["ares_nw"], + ares_qk=w["ares_qk"], + in_nw=w["in_nw"], + w_gate=w["w_gate"], + bias=w["bias"], + w_ld=w["w_ld"], + w_sgu=w["w_sgu"], + w_sd=w["w_sd"], + w13=w["w13"], + w13s=w["w13s"], + w2=w["w2"], + w2s=w["w2s"], + ln_w=w["ln_w"], + w_up=w["w_up_shard"], + out=out, + scratch=scratch, + peers=peers.addresses, + rank=rank, + epoch=epoch, + layer=layer, + ) + torch.cuda.synchronize() + off, nbytes = mod.scratch_layout(key)["xrow"] + xrow = scratch[off : off + nbytes].view(BF).view(ROWS, HIDDEN) + return out.clone(), prefix, xrow.clone() + + +def _worker_k2(rank, port, results): + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=TP_SIZE + ) + peers = None + try: + from sglang.kernels.ops.moe.k3_mono_flydsl import layer as k2 + from sglang.kernels.ops.moe.k3_mono_flydsl.common.peer_memory import PeerBuffer + + peers = PeerBuffer(k2.peer_bytes(ROWS), dist.group.WORLD, rank, TP_SIZE, device) + peers.bytes.zero_() + epoch = torch.zeros(1, dtype=torch.int32, device=device) + scratch = torch.zeros( + k2.scratch_bytes(k2.K2Build(tokens=ROWS, nblocks=MAX_BLOCKS)), + dtype=torch.uint8, + device=device, + ) + + w = _weights(rank, device, gate_up_interleaved=False) + w.update(_layer_weights(rank, device, NVB)) + + result = {"error": ""} + seen = {} + # Both prefix modes: a plain layer accumulates into the prefix, a + # block-write layer restarts it at the attention output. + for tag, reset in (("accumulate", False), ("reset", True)): + epoch.add_(1) + out, prefix, xrow = _run_k2( + k2, + w, + scratch, + nvb=NVB, + reset=reset, + peers=peers, + epoch=epoch, + rank=rank, + layer=1, + ) + ref_prefix, ref_x = _seam_reference( + w["core"], w["w_o"], w["prefix"], w["blocks"], NVB, w, reset=reset + ) + gathered = [torch.empty_like(out.cpu()) for _ in range(TP_SIZE)] + dist.all_gather(gathered, out.cpu().contiguous()) + result[tag] = { + "finite": bool(torch.isfinite(out).all()), + "nonzero": bool(torch.count_nonzero(out) > 0), + "ranks_agree": all(torch.equal(gathered[0], g) for g in gathered[1:]), + # The seam: o_proj, its reduce and the aggregation. + "cos_prefix": _cos(prefix, ref_prefix), + "cos_x": _cos(xrow, ref_x), + # The MoE, given the seam output the kernel itself produced. + "cos_moe": _cos(out, _reference(xrow, w, rank)), + } + seen[tag] = xrow + + # Control: the two modes read the input prefix differently, so they + # must disagree. Without this a kernel that ignored the prefix + # entirely -- or a harness reading a stale scratch -- would still + # match both references and the test above would mean nothing. + result["modes_differ"] = _cos(seen["accumulate"], seen["reset"]) < COS_MIN + results[rank] = result + except Exception as error: # surfaced by the parent as a test failure + import traceback + + results[rank] = { + "error": f"{type(error).__name__}: {error}\n{traceback.format_exc()}" + } + finally: + if peers is not None: + peers.close() + dist.destroy_process_group() + + +class TestKimiK3MonoDecode(CustomTestCase): + check_guint = False + + def _run(self): + if torch.cuda.device_count() < TP_SIZE: + self.skipTest(f"needs {TP_SIZE} GPUs") + if not torch.version.hip: + self.skipTest("ROCm only") + manager = mp.Manager() + results = manager.dict() + mp.spawn( + _worker, + args=(_free_port(), results, self.check_guint), + nprocs=TP_SIZE, + join=True, + ) + return results + + def test_moe_matches_aiter_reference(self): + results = self._run() + r0 = results.get(0) or {} + if not r0.get("error"): + print( + f"rank0: cos vs reference {r0['cos_vs_reference']:.6f}" + + ( + f", rerun {r0['cos_rerun']:.6f}, layouts {r0['cos_layouts']:.6f}" + f" (equal={r0['layouts_equal']})" + if self.check_guint + else "" + ) + ) + for rank in range(TP_SIZE): + r = results.get(rank) + self.assertIsNotNone(r, f"rank {rank} produced no result") + self.assertEqual(r["error"], "", f"rank {rank} failed") + self.assertTrue(r["finite"], f"rank {rank} output is not finite") + self.assertTrue(r["nonzero"], f"rank {rank} output is all zero") + self.assertTrue(r["ranks_agree"], "TP ranks disagree after the reduce") + self.assertGreater( + r["cos_vs_reference"], + COS_MIN, + f"rank {rank}: cos {r['cos_vs_reference']:.6f} vs the AITER stack", + ) + if self.check_guint: + self.assertGreater( + r["cos_rerun"], + COS_MIN, + f"rank {rank}: re-running the same launch differs " + f"(cos {r['cos_rerun']:.6f}) -- not a layout problem", + ) + self.assertGreater( + r["cos_layouts"], + COS_MIN, + f"rank {rank}: the two expert layouts disagree " + f"(cos {r['cos_layouts']:.6f})", + ) + + +class TestKimiK3MonoLayer(CustomTestCase): + """The whole-layer launch: o_proj + its all-reduce + the MLP seam + the MoE.""" + + def test_layer_matches_torch_seam_and_aiter_moe(self): + if torch.cuda.device_count() < TP_SIZE: + self.skipTest(f"needs {TP_SIZE} GPUs") + if not torch.version.hip: + self.skipTest("ROCm only") + manager = mp.Manager() + results = manager.dict() + mp.spawn(_worker_k2, args=(_free_port(), results), nprocs=TP_SIZE, join=True) + + r0 = results.get(0) or {} + if not r0.get("error"): + for tag in ("accumulate", "reset"): + r = r0[tag] + print( + f"rank0 {tag}: cos prefix {r['cos_prefix']:.6f}, " + f"x {r['cos_x']:.6f}, moe {r['cos_moe']:.6f}" + ) + for rank in range(TP_SIZE): + r = results.get(rank) + self.assertIsNotNone(r, f"rank {rank} produced no result") + self.assertEqual(r["error"], "", f"rank {rank} failed") + self.assertTrue( + r["modes_differ"], + f"rank {rank}: the two prefix modes agree, so neither checks the " + "prefix -- the comparisons below prove nothing", + ) + for tag in ("accumulate", "reset"): + m = r[tag] + self.assertTrue(m["finite"], f"rank {rank} {tag}: output not finite") + self.assertTrue(m["nonzero"], f"rank {rank} {tag}: output all zero") + self.assertTrue( + m["ranks_agree"], f"{tag}: TP ranks disagree after the reduce" + ) + self.assertGreater( + m["cos_prefix"], + COS_MIN, + f"rank {rank} {tag}: the updated prefix differs " + f"(cos {m['cos_prefix']:.6f}) -- o_proj or its all-reduce", + ) + self.assertGreater( + m["cos_x"], + COS_MIN, + f"rank {rank} {tag}: the MoE input differs " + f"(cos {m['cos_x']:.6f}) -- the attention-residual seam", + ) + self.assertGreater( + m["cos_moe"], + COS_MIN, + f"rank {rank} {tag}: the MoE differs from the AITER stack " + f"(cos {m['cos_moe']:.6f}) given the kernel's own seam output", + ) + + +if __name__ == "__main__": + import sys + + if "--guint" in sys.argv: + sys.argv.remove("--guint") + TestKimiK3MonoDecode.check_guint = True + unittest.main() diff --git a/test/registered/amd/test_kimi_k3_mono_kda_pre.py b/test/registered/amd/test_kimi_k3_mono_kda_pre.py new file mode 100644 index 000000000000..08f120810206 --- /dev/null +++ b/test/registered/amd/test_kimi_k3_mono_kda_pre.py @@ -0,0 +1,518 @@ +"""Single-GPU correctness and contract tests for the Kimi-K3 K1 launch +(``kda_pre``, ROCm / FlyDSL): one KDA layer from the attention-residual seam to +the gated core output. + +K1 is rank-local -- its ABI carries no peer buffer and no rank -- so unlike the +MoE and whole-layer launches this needs one GPU, not eight. + +What is checked, and why each: + +* the seam, against a torch reference: aggregation then ``input_layernorm``. + Same formula as the layer launch's seam, so a divergence localizes to K1. +* the in-projection, against the seam output the kernel itself produced. Same + reasoning as the layer test: comparing the projection against the kernel's + own x keeps a rounding difference in the seam from being reported as a GEMM + bug. +* the **conv-state footprint**, which is the contract K1 would have to share + with SGLang's speculative conv-window pool. Nothing about wiring K1 into the + model is safe to design until the indices are measured, because FlyDSL's + ``buffer_load`` is OOB-clamped: reading past the buffer returns a clamped + address rather than faulting, so a window that is one entry short loses + precision silently. + +Run directly with ``python test/registered/amd/test_kimi_k3_mono_kda_pre.py``. +""" + +import unittest + +import torch +import torch.nn.functional as F + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +# One GPU: K1 is rank-local. gfx950 only, so the MI35x suite. The run itself +# is about 10 s; the budget covers the first FlyDSL build on a cold JIT cache. +register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd-mi35x") + +BF = torch.bfloat16 +COS_MIN = 0.999 +# One request of L tokens: the shape a DSPARK-7 verify step has at concurrency 1. +REQS = 1 +QLEN = 8 +NVB = 3 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +def _randn(gen, *shape, scale=1.0, dtype=BF): + return ( + torch.randn(*shape, generator=gen, device=gen.device, dtype=torch.float32) + * scale + ).to(dtype) + + +def _cos(a, b): + return float( + F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0, eps=1e-12) + ) + + +def _rms(x, weight, eps): + x = x.float() + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight.float() + + +def _seam_reference(prefix, blocks, nvb, ares_nw, ares_qk, in_nw, eps, out_eps): + """Aggregate the banked sources and the prefix, then input_layernorm.""" + rows = torch.cat([blocks[:, :nvb, :], prefix.unsqueeze(1)], dim=1) + scores = (_rms(rows, ares_nw, eps) * ares_qk.float()).sum(-1) + mixed = (torch.softmax(scores, dim=-1).unsqueeze(-1) * rows.float()).sum(dim=1) + return _rms(mixed, in_nw, out_eps).to(BF) + + +def _kda_reference(case): + """core_out in torch, from the kernel's own in-projection. + + Derived from the kernel source rather than from SGLang's KDA path: the two + agree on the math but not on argument conventions, so reproducing + ``kda_conv`` / ``kda_gates`` / ``kda_recur`` / ``kda_norm`` directly is + what makes a mismatch mean "the kernel is wrong" instead of "the harness + guessed the convention wrong". + + Feeds on ``proj_dbg`` for the same reason the layer test feeds on the + kernel's own x: a rounding difference upstream must not be reported here. + """ + kda = case.kda + L, hd, nh, proj_w = case.L, kda.HD, kda.NH, kda.PROJ + win = kda.CONV_W - 1 + acc_idx = case.num_acc - 1 + slot_c = int(case.st_idx[0, 0]) + slot0 = int(case.st_idx[0, acc_idx]) + proj = case.proj_dbg.float() + + # --- conv: the carried window then this step's tokens, then silu + xv = proj[:, : 3 * proj_w].t() # [3 * PROJ, L] + window = case.conv_state_in[slot_c][:, acc_idx : acc_idx + win].float() + seq = torch.cat([window, xv], dim=1) + cw = case.conv_w.float() + conv = torch.stack( + [sum(cw[:, q] * seq[:, j + q] for q in range(kda.CONV_W)) for j in range(L)], + dim=1, + ) + conv = (conv * torch.sigmoid(conv)).t() # [L, 3 * PROJ] + q_all, k_all, v_all = ( + conv[:, i * proj_w : (i + 1) * proj_w].reshape(L, nh, hd) for i in range(3) + ) + + def l2(x, scale): + return x * torch.rsqrt(x.pow(2).sum(-1, keepdim=True) + 1e-6) * scale + + q_all = l2(q_all, hd**-0.5) + k_all = l2(k_all, 1.0) + + fa = proj[:, kda.FA_0 : kda.FA_0 + hd] # [L, HD] + out = torch.empty(case.S, proj_w, dtype=torch.float32, device=proj.device) + for h in range(nh): + beta = torch.sigmoid(proj[:, kda.BETA_0 + h]) # [L] + g1 = fa.to(BF).float() @ case.w_fb[h * hd : (h + 1) * hd].float().t() + g = g1 + case.dt_bias[h * hd : (h + 1) * hd].float() + gate = -5.0 * torch.sigmoid(torch.exp(case.a_log[h].float()) * g) + decay = torch.exp(gate) # [L, HD], per-K + + state = case.rstate_in[slot0, h].float().clone() # [V, K] + o = torch.empty(L, hd, dtype=torch.float32, device=proj.device) + for j in range(L): + state = state * decay[j].unsqueeze(0) + delta = (v_all[j, h] - state @ k_all[j, h]) * beta[j] + state = state + delta.unsqueeze(1) * k_all[j, h].unsqueeze(0) + o[j] = state @ q_all[j, h] + + r = torch.rsqrt(o.pow(2).mean(-1, keepdim=True) + EPS) + g2 = proj[:, kda.G2_0 + h * hd : kda.G2_0 + (h + 1) * hd] + out[:, h * hd : (h + 1) * hd] = o * r * case.on_w.float() * torch.sigmoid(g2) + return out.to(BF) + + +class K1Case: + """One synthetic K1 launch: its inputs, and the tensors it wrote.""" + + def __init__(self, device, *, num_acc, state_len=None, write_idx=-1): + from sglang.kernels.ops.moe.k3_mono_flydsl.attention import kda + + self.kda = kda + self.device = device + self.L = QLEN + self.S = REQS * QLEN + # K1 rebases the conv window on every launch, reading up to + # acc_idx + CONV_W - 1 past the start. See test_conv_state_footprint. + self.state_len = state_len or (self.L + kda.CONV_W - 1) + self.num_acc = num_acc + self.write_idx = write_idx + + g = torch.Generator(device=device).manual_seed(5) + s, h, dim = self.S, kda.HIDDEN, 3 * kda.PROJ + self.prefix = _randn(g, s, h, scale=0.3) + self.blocks = _randn(g, s, MAX_BLOCKS, h, scale=0.3) + self.ares_nw = (1 + _randn(g, h, scale=0.1)).to(BF) + self.ares_qk = _randn(g, h, scale=0.02) + self.in_nw = (1 + _randn(g, h, scale=0.1)).to(BF) + self.w_in = _randn(g, kda.NPROJ, h, scale=0.02) + self.w_fb = _randn(g, kda.NH * kda.HD, kda.HD, scale=0.02) + self.conv_w = _randn(g, dim, kda.CONV_W, scale=0.3, dtype=torch.float32) + self.a_log = _randn(g, kda.NH, scale=0.1, dtype=torch.float32) + self.dt_bias = _randn(g, kda.NH * kda.HD, scale=0.1, dtype=torch.float32) + self.on_w = (1 + _randn(g, kda.HD, scale=0.1)).to(BF) + + # Slot 0 is an ordinary slot: the kernel's "no state" marker is a + # negative index, matching SGLang's null_block_id = -1. Starting at 0 + # is what covers the request that owns scratch row 0, which the pool + # hands out like any other. + slots = torch.arange(self.L, dtype=torch.int32, device=device) + self.st_idx = slots.view(REQS, self.L).contiguous() + self.num_acc_t = torch.full((REQS,), num_acc, dtype=torch.int32, device=device) + nslots = int(slots.max()) + 1 + # Always drawn at full width and truncated, so a short-buffer case is a + # prefix of the full one and every later draw is unchanged. + full = _randn(g, nslots, dim, self.L + kda.CONV_W - 1, scale=0.5) + self.conv_state = full[:, :, : self.state_len].contiguous() + self.conv_state_in = self.conv_state.clone() + self.rstate = _randn( + g, nslots, kda.NH, kda.HD, kda.HD, scale=0.05, dtype=torch.float32 + ) + self.rstate_in = self.rstate.clone() + + self.core_out = torch.zeros(self.S, kda.PROJ, dtype=BF, device=device) + self.x_dbg = torch.zeros(self.S, h, dtype=BF, device=device) + self.proj_dbg = torch.zeros(self.S, kda.NPROJ, dtype=BF, device=device) + + def run(self): + kda = self.kda + key = kda.KdaPreBuild( + tokens=self.S, + qlen=self.L, + nblocks=NVB, + delta=False, + state_len=self.state_len, + eps=EPS, + out_eps=EPS, + onorm_eps=EPS, + debug=True, + write_idx=self.write_idx, + ) + scratch = torch.zeros( + kda.scratch_bytes(key), dtype=torch.uint8, device=self.device + ) + epoch = torch.ones(1, dtype=torch.int32, device=self.device) + kda.kda_pre( + key, + prefix=self.prefix, + delta=None, + blocks=self.blocks, + ares_nw=self.ares_nw, + ares_qk=self.ares_qk, + in_nw=self.in_nw, + w_in=self.w_in, + w_fb=self.w_fb, + conv_w=self.conv_w, + conv_state=self.conv_state, + a_log=self.a_log, + dt_bias=self.dt_bias, + on_w=self.on_w, + rstate=self.rstate, + st_idx=self.st_idx, + num_acc=self.num_acc_t, + core_out=self.core_out, + scratch=scratch, + layer=1, + epoch=epoch, + x_dbg=self.x_dbg, + proj_dbg=self.proj_dbg, + ) + torch.cuda.synchronize() + return self + + +def _skip_reasons(): + if not torch.cuda.is_available(): + return "needs a GPU" + if not torch.version.hip: + return "ROCm only" + return None + + +class TestKimiK3MonoKdaPre(CustomTestCase): + @classmethod + def setUpClass(cls): + reason = _skip_reasons() + if reason: + raise unittest.SkipTest(reason) + cls.device = torch.device("cuda", 0) + torch.cuda.set_device(cls.device) + # acc_idx = num_acc - 1 reaches its maximum when every drafted token of + # the previous step was accepted. That is the case the conv window has + # to be widest for, so make it the default. + cls.case = K1Case(cls.device, num_acc=QLEN).run() + + def test_seam_matches_torch(self): + case = self.case + want = _seam_reference( + case.prefix, + case.blocks, + NVB, + case.ares_nw, + case.ares_qk, + case.in_nw, + EPS, + EPS, + ) + cos = _cos(case.x_dbg, want) + print(f"seam: cos {cos:.6f}") + self.assertGreater(cos, COS_MIN, "the attention-residual seam differs") + + def test_in_proj_matches_its_own_x(self): + case = self.case + # Rows past 6284 are padding the kernel never defines. + real = 6284 + want = F.linear(case.x_dbg, case.w_in[:real]) + cos = _cos(case.proj_dbg[:, :real], want) + print(f"in_proj: cos {cos:.6f}") + self.assertGreater(cos, COS_MIN, "the in-projection differs") + + def test_core_out_matches_torch(self): + case = self.case + out = case.core_out + self.assertTrue(bool(torch.isfinite(out).all()), "core_out is not finite") + self.assertTrue(bool(torch.count_nonzero(out) > 0), "core_out is all zero") + cos = _cos(out, _kda_reference(case)) + print(f"core_out: cos {cos:.6f}") + self.assertGreater(cos, COS_MIN, "the KDA recurrence differs") + + def test_every_acceptance_count(self): + """num_acc picks both the conv window offset and the initial state slot. + + The default case uses the widest offset; these cover the rest, since a + caller has no control over how many drafted tokens get accepted. + """ + for num_acc in range(1, QLEN + 1): + case = K1Case(self.device, num_acc=num_acc).run() + cos = _cos(case.core_out, _kda_reference(case)) + with self.subTest(num_acc=num_acc): + self.assertGreater(cos, COS_MIN, f"num_acc={num_acc}: cos {cos:.6f}") + print(f"acceptance counts 1..{QLEN}: all match") + + def test_slot_zero_is_a_real_slot(self): + """Row 0 of the speculative scratch must compute like any other row. + + SGLang's ``build_verify_intermediate_state_indices`` gives request slot + i scratch row i, so row 0 goes to whichever request holds req_pool + index 0 -- there is no way to avoid it. The default case already starts + its slots at 0; this pins the behaviour so the kernel's sentinel cannot + drift back to treating 0 as "no state". + """ + case = self.case + self.assertEqual(int(case.st_idx[0, 0]), 0, "the case must exercise slot 0") + cos = _cos(case.core_out, _kda_reference(case)) + self.assertGreater(cos, COS_MIN, f"slot 0 computed wrong (cos {cos:.6f})") + touched = [ + j + for j in range(case.rstate.shape[0]) + if not torch.equal(case.rstate[j], case.rstate_in[j]) + ] + self.assertIn(0, touched, "slot 0 was skipped by the state write") + + def test_negative_slot_is_the_sentinel(self): + """A negative slot means "no state": zero output, nothing written.""" + case = K1Case(self.device, num_acc=QLEN) + case.st_idx = torch.full_like(case.st_idx, -1) + case.run() + self.assertEqual( + int(torch.count_nonzero(case.core_out)), + 0, + "a negative slot should produce a zero core_out", + ) + self.assertTrue( + torch.equal(case.rstate, case.rstate_in), + "a negative slot should write no recurrent state", + ) + self.assertTrue( + torch.equal(case.conv_state, case.conv_state_in), + "a negative slot should write no conv state", + ) + print("negative slot: no output, no state written") + + def test_seeded_window_equals_the_rolling_convention(self): + """A caller may own the conv buffer and reseed it every step. + + K1's own convention is a rolling buffer it rebases each launch, with + ``num_acc`` saying where the live window starts. A caller that owns the + buffer instead can hand it the window directly at offset 0 and pass + ``num_acc = 1``. This has to compute the same thing, because that is + what lets K1 take its three carried values from SGLang's committed conv + state rather than from a buffer SGLang does not maintain. + + It also drops the highest read index from ``acc_idx + CONV_W - 1`` to + ``CONV_W - 1``, which is why a privately owned buffer sidesteps the + one-short problem entirely. + """ + kda = self.case.kda + win = kda.CONV_W - 1 + acc = 5 # an interior acceptance count, so the offsets really differ + rolling = K1Case(self.device, num_acc=acc).run() + + slot_c = int(rolling.st_idx[0, 0]) + acc_idx = acc - 1 + seeded = K1Case(self.device, num_acc=1) + # The window the rolling case used, placed at offset 0. + seeded.conv_state[slot_c][:, :win] = rolling.conv_state_in[slot_c][ + :, acc_idx : acc_idx + win + ] + # ... and the initial state the rolling case selected. + seeded.rstate[int(seeded.st_idx[0, 0])] = rolling.rstate_in[ + int(rolling.st_idx[0, acc_idx]) + ] + seeded.conv_state_in = seeded.conv_state.clone() + seeded.rstate_in = seeded.rstate.clone() + seeded.run() + + cos = _cos(seeded.core_out, rolling.core_out) + equal = bool(torch.equal(seeded.core_out, rolling.core_out)) + print(f"seeded vs rolling: cos {cos:.6f} (equal={equal})") + self.assertGreater( + cos, COS_MIN, "seeding the window at offset 0 is not equivalent" + ) + + def test_core_out_depends_on_the_initial_state(self): + """Control: the recurrence must actually read the slot it selects. + + ``num_acc`` picks which of the previous step's states to start from. If + core_out were the same with that slot perturbed, the reference above + would be matching a kernel that ignores its initial state, and the + test would prove nothing about the state wiring 5b depends on. + """ + case = self.case + acc_idx = case.num_acc - 1 + other = K1Case(self.device, num_acc=case.num_acc) + slot0 = int(other.st_idx[0, acc_idx]) + other.rstate[slot0] += 1.0 + other.rstate_in = other.rstate.clone() + other.run() + cos = _cos(other.core_out, case.core_out) + print(f"perturbed initial state: cos vs unperturbed {cos:.6f}") + self.assertLess( + cos, COS_MIN, "core_out ignored the initial state slot num_acc selects" + ) + self.assertGreater( + _cos(other.core_out, _kda_reference(other)), + COS_MIN, + "the reference does not track the perturbed initial state", + ) + + def test_conv_state_footprint(self): + """The contract a caller has to satisfy, measured rather than assumed. + + K1 rebases the conv window every launch: the new buffer is the three + entries starting one past the accepted token, then this step's L token + inputs. So it reads up to index ``acc_idx + CONV_W - 1`` of the old + buffer, and ``acc_idx`` reaches ``L - 1`` whenever every drafted token + was accepted. + """ + case = self.case + kda = case.kda + win = kda.CONV_W - 1 + acc_idx = case.num_acc - 1 + slot_c = int(case.st_idx[0, 0]) + before = case.conv_state_in[slot_c] + after = case.conv_state[slot_c] + + # The carried window: new[0 : win] == old[acc_idx + 1 : acc_idx + 1 + win] + self.assertTrue( + torch.equal(after[:, :win], before[:, acc_idx + 1 : acc_idx + 1 + win]), + "the carried conv window is not the entries past the accepted token", + ) + # The appended tokens: new[win:] == this step's conv channel inputs, + # which are the first 3 * PROJ rows of the in-projection. + want = case.proj_dbg[:, : 3 * kda.PROJ].t().contiguous() + self.assertTrue( + torch.equal(after[:, win:], want), + "the appended conv entries are not this step's token inputs", + ) + highest_read = acc_idx + win + print( + f"conv footprint: reads up to index {highest_read} of " + f"state_len {case.state_len} (acc_idx {acc_idx})" + ) + # The requirement this test exists to pin down. + self.assertEqual( + highest_read, + case.state_len - 1, + "state_len must be qlen + CONV_W - 1; a shorter buffer is read " + "out of bounds, which buffer_load clamps instead of faulting", + ) + + def test_recurrent_state_footprint(self): + """Which rstate slots K1 reads and writes. + + It starts from the slot of the last accepted token and writes one slot + per token of this step -- so a caller owes it ``st_idx[req, 0 : L]`` + with the previous step's states still in place. + """ + case = self.case + touched = [ + j + for j in range(case.rstate.shape[0]) + if not torch.equal(case.rstate[j], case.rstate_in[j]) + ] + want = sorted(int(s) for s in case.st_idx[0].tolist()) + print(f"rstate: wrote slots {touched}, st_idx holds {want}") + self.assertEqual(touched, want, "K1 did not write exactly st_idx[req, :L]") + + def test_one_short_state_len_is_wrong_and_silent(self): + """A buffer one entry short is read out of bounds, and nothing says so. + + This is the case a caller would hit by handing K1 SGLang's speculative + conv window, which is ``num_draft_tokens + CONV_W - 2`` wide -- one + less than K1 reads. The two cases below share the first entries, so the + in-bounds part of the carry must agree and only the entry that came + from past the end may differ. It does, and the launch still succeeds: + the failure is silent. + """ + kda = self.case.kda + win = kda.CONV_W - 1 + acc_idx = QLEN - 1 + slot_c = int(self.case.st_idx[0, 0]) + long_after = self.case.conv_state[slot_c] + + short = K1Case(self.device, num_acc=QLEN, state_len=QLEN + kda.CONV_W - 2).run() + short_after = short.conv_state[slot_c] + + self.assertTrue( + torch.equal(short_after[:, : win - 1], long_after[:, : win - 1]), + "the in-bounds part of the carried window should be unaffected", + ) + self.assertFalse( + torch.equal(short_after[:, win - 1], long_after[:, win - 1]), + "a one-short buffer happened to read the right value; this test " + "no longer demonstrates the out-of-bounds read", + ) + print( + f"short state_len {short.state_len}: carry entry {win - 1} differs " + f"(read index {acc_idx + win} does not exist), launch succeeded" + ) + + def test_block_write_layer_snapshots_the_prefix(self): + """write_idx >= 0 stores the updated prefix into that bank row. + + This is the seam the whole-layer launch's ``reset`` mode assumes has + already happened, so the two have to agree on the row. + """ + row = NVB + case = K1Case(self.device, num_acc=QLEN, write_idx=row).run() + self.assertTrue( + torch.equal(case.blocks[:, row, :], case.prefix), + "the block-write row does not hold the updated prefix", + ) + + +if __name__ == "__main__": + unittest.main()