benchmarks.SaturnNPU.kernel_library.smolvla_fused_attention

Fused attention (SDPA) kernel — matches MLIR variant_0_12_1024_64_bf16.

MLIR source: SaturnNPU/kernels/iree_linalg_ext.attention/variant_0_12_1024_64_bf16.mlir

Q: [12, 1024, 64] bf16 indexing (d0, d1, d3) → [batch, q_seq, head_dim] K: [12, 1024, 64] bf16 indexing (d0, d4, d3) → [batch, k_seq, head_dim] V: [12, 64, 1024] bf16 indexing (d0, d2, d4) → [batch, head_dim, k_seq] scale: scalar bf16 output: [12, 1024, 64] bf16

This demo tiles to one Q-block (32 rows) and two K tiles (k_seq=64 total, 32 per tile), with head_dim=64. Production wraps this body in loops over all 1024 q-rows and 12 heads.

m = -inf, l = 0, O = 0
for each K tile:
    S  = Q @ K^T * scale
    m' = max(m, rowmax(S))
    α  = exp(m − m')
    O  = α * O + exp(S − m') @ V
    l  = α * l + rowsum(exp(S − m'))
    m  = m'
output = O / l

DRAM layouts (all bf16, column-blocked so each vload gets a contiguous [32,16] chunk):

Q_DRAM: [32, 64] bf16 stored as [128, 16] (4 × [32,16] blocks, cols 0:16/16:32/32:48/48:64) KT_DRAM: K tile [32,64] pre-transposed → K^T[64,32] stored as [128,16] (top/bot × left/right) VT_DRAM: V_mlir tile [64,32] pre-transposed → V_std[32,64] stored as [128,16] OUT_DRAM: [32, 64] bf16 stored as [128, 16]

Q @ K^T (head_dim=64, two MXU passes accumulating over head_dim=32 per pass): Q_lo [32,32] fp8 @ K^T_top [32,32] fp8 → acc[0] (fresh) Q_hi [32,32] fp8 @ K^T_bot [32,32] fp8 → acc[0] (accumulate) vmatpop.bf16 writes acc[0][32,32] into registers v16 (left [32,16]) and v17 (right [32,16])

exp_s @ V (output is [32,64], two independent MXU passes): exp_s [32,32] fp8 @ V_left [32,32] fp8 → acc[0] (O left half) exp_s [32,32] fp8 @ V_right [32,32] fp8 → acc[1] (O right half)

bf16 → fp8 quantization uses the acc roundtrip (vmatpush.acc.bf16 + vmatpop.fp8). vmatpush.acc.bf16(vd=slot, vs1=v) reads v and v+1 as a [32,32] BF16 tile → acc[slot] vmatpop.fp8(vd=dst, vs1=slot) converts acc[slot][32,32] → fp8 → dst register [32,32]

-inf initialization: m (running row-max) is initialized to -100.0 via vli.all imm=-100. This is encoded as an integer immediate (vli.all calls torch.full(shape, -100, dtype=bf16) = -100.0_bf16). True bf16 -inf (0xFF80) cannot be encoded directly as a vli.all integer immediate. -100.0_bf16 is sufficient: all real attention logits after scaling (SCALE≈0.125) will be much larger than -100, so the first tile always overwrites the initial m value.