`vllm.models.minimax_m3.amd.mono.drafter.device` ¶

Device helpers of the DSpark drafter kernels: memory, bounded waits, the MXFP8 GEMV and the row norms.

Every function is traced inside a @flyc.kernel body; the ones with control flow over traced values are @traced (FlyDSL rewrites only the kernel body's own AST). Conventions:

a global base is a wave-uniform Int64 address (a buffer resource); offsets are per lane, in elements;
tagged pairs are 8 bytes: (value word, tag word);
data other CTAs read after a flag is stored device-coherent (CM_DEV) and published by publish.

Functions:

bf16_bits –

bf16 bits of x (RNE) in the low half of an Int32.
bf16r –

x rounded to bf16 (RNE), as f32.
bfpair –

bf16(a) | bf16(b) << 16 as an Int32.
exp –

The natural exponent as exp2 of x log2(e) (denormal-safe llvm.exp2).
fail –

A timed-out wait sets the health word to [1, rank, site].
gst32 –

A 32-bit store at a per-lane 64-bit address.
hw_sqrt –

v_sqrt_f32 (~1 ulp), not the correctly rounded libm expansion.
ld –

Load n elements (i32 unless dtype) at element idx of base.
ld8 –

Byte idx of base, zero-extended.
ldn –

ld through a resource of nbytes: elements past it read 0 without a
mma8 –

Acc += A B, 16x16x128 e4m3 with e8m0 scales (in the low byte).
mma_chain –

One tile's partials over chunks 0..nc-1, in order.
mx_scale –

VLLM's MXFP8 block scale of amax am: (e8m0 exponent as f32, 1 / scale).
norm_row –

Gemma RMSNorm of a row: s = v (+ the residual in LDS res) becomes the
norm_w –

A thread's 16 norm weights (8 bf16-pair words).
opaque –

v through a side-effecting asm that also reads loop_var.
poll_pairs –

Wait for the n pairs (1, 2, 4 or 8) from pair idx of base to carry
poll_rows –

poll_pairs of every base in bases in one batch (one round trip a
publish –

Drain the CTA's stores, then set flag i of base to tag.
quant_store –

MXFP8 of a thread's 16 values -> e4m3 x8[t][16 tid ..], e8m0
red_store –

Wave wave's partial of tile t -> red[8 t + wave][row][token].
red_sum –

Tile t's (row, token) summed over waves 0..7, in order.
row_sum4 –

A thread's 16 values of row t: bf16 of the 4 source ranks' bf16 partials
row_x –

A thread's 16 values of row t of bf16 x_in[rows][6144].
st –

Store v (a Numeric or a Vector) at element idx of base.
w_chunk –

MFMA operand (8 words) and scale of row row, chunk c of wave.
wait_flags –

Wait for flags 0..n-1 of base (n <= 512), then a CTA barrier.
xq_chunk –

bf16 row row of x[16][K], chunk c of wave, quantized to MXFP8 in

`_spin(load_all, pending, dl, health)` ¶

Re-issue load_all until pending clears.

Every CHECK_EVERY rounds the clock is checked against dl and the health word read. Returns (values, failed): failed is 1 when the wait timed out, 2 when it gave up on a failure reported elsewhere.