`vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops` ¶

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.

Functions:

ballot –

The wave's lanes where pred holds, a bit a lane.
batched_rounds –

A CTA's strided loop over elements v = tid + THREADS i < n (n a Python
bf16_pair –

Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).
bf2_f32 –

Packed bf16 pair word -> (f32 low, f32 high).
bf_hi –

The f32 of a packed bf16 pair word's high half.
bf_lo –

The f32 of a packed bf16 pair word's low half.
block_excl_scan –

(exclusive prefix of v over the CTA in thread order, CTA total); buf
block_max –

Max of v over the CTA's WAVES waves; red needs WAVES words.
butterfly –

Butterfly reduction over the given xor offsets, in the given order.
div_rn –

x / d rounded to nearest, bit-identical to the IEEE division, given
fp8_pack4 –

Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.
fp8x8_bf16_pk –

Eight fp8 of two dwords -> bf16 MFMA operand (exact) by
fresh –

v (a per-lane Int32) through an opaque move, here: the lane math built
global_load –

words i32 (1, 2 or 4) at the 64-bit global address addr. A pool
hw_rsq –

v_rsq_f32 (~1 ulp, no libm range fixups).
kernel_symbol –

Kernel symbol that profilers demangle to atom::<stem>_<params>. The
lane_gather –

Lane src_lane's value of v, a source lane a lane (ds_bpermute).
lanes_below –

How many of the lanes in ballot bal are below lane.
ld_dev –

words i32 (1, 2 or 4) of a buffer from word i, at device scope
ld_i32 –

Word i of a buffer, a plain Int32 load.
lds_bytes –

The LDS an fx.struct class takes, alignment included.
load_ptr64 –

Entry i of an i64 table in memory, wave-uniform (each half through
memrealtime –

s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).
mfma_fp8 –

16x16x32 fp8 x fp8 MFMA, unscaled: a / b an i64 of 8 E4M3 each.
permlane_swap –

v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half
readlane –

Lane src_lane's value of v (wave-uniform).
row_live –

Whether row k of the step is a request's: its batch_id_per_q_token
row_shr –

DPP row_shr:k of an f32: lane - k of the same 16-lane row, else 0.
row_sum –

Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at red
rows_to_lds –

rows rows of row_words words (a multiple of 4) read at device scope
traced –

Apply flydsl's AST rewriting to a helper, so its while / if over
unlikely –

cond marked cold (llvm.expect): LLVM lays the branch out of the hot
wave_incl_scan –

Inclusive prefix sum of an Int32 over the wave, in lane order.
xred –

op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.
xshfl –

Value of lane lane ^ off: offsets 32 / 16 via v_permlane*_swap, in-row

`_xpartner(v, off, lane)` ¶

Lane lane ^ off's v: DPP within a row, v_permlane*_swap across rows (xshfl takes shuffle_xor there).