class SwapABDecodeKernel(VllmCuTeDSLJitKernel["SwapABDecodeKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
heads: int # live heads: 16 or 32
swa_width: int # columns of the SWA slot list
extra_width: int # columns of the compressed slot list
extra_record: int # compressed record bytes: 528 (fp8) or 288 (NVFP4)
swa_page: int # slots per SWA cache page
extra_page: int # slots per compressed cache page
@staticmethod
def kernel(compile_key: CompileKey) -> Any:
H = compile_key.heads
assert H in (16, 32), "live heads must be 16 or 32"
NK, NE = compile_key.swa_width, compile_key.extra_width
assert compile_key.extra_record in (_FP8_RECORD, _FP4_RECORD), (
"compressed records must be V4.1 fp8 or NVFP4"
)
EXTRA_FP4 = compile_key.extra_record == _FP4_RECORD
KP, EP = compile_key.swa_page, compile_key.extra_page
for page in (KP, EP):
assert page > 0 and page & (page - 1) == 0, "page sizes must be powers of 2"
KT = cdiv(NK, _SWA_ROWS)
ET = cdiv(NE, _TILE)
C = KT + ET
assert 1 <= C <= _MAX_CLUSTER, f"{C} KV tiles per token > {_MAX_CLUSTER}"
NC = H // 4 # score / P columns per thread in the softmax
S_COL = 8 * H # TMEM: O hi|lo for 4 M-blocks first, then S
TMEM_COLS = 256 if S_COL + H <= 256 else 512
P_ROW = 4 * H # bytes per P row: H bf16 hi + H bf16 lo
Q_PIECES = H * _HEAD_DIM // 8 # 16-byte Q pieces
QP = Q_PIECES // _THREADS
assert QP * _THREADS == Q_PIECES
U = (_HEAD_DIM // 32) * H # (head, 32-dim chunk) output units
UPC = cdiv(U, C)
NPASS = cdiv(UPC, _THREADS // 8)
STATS = C * 2 * H
UNIT_BYTES = 32 * 4 # one (head, 32-dim chunk) unit of the fp32 partial
IDESC_S = _tcgen05.make_bf16_idesc(_TILE, H)
IDESC_O = _tcgen05.make_bf16_idesc(
_TILE, 2 * H, transpose_A=True, transpose_B=True
)
DESC_KMAJ = _tcgen05.make_sdesc_128B_swizzle(0)
DESC_KV_MN = _tcgen05.make_sdesc_128B_swizzle(16384)
DESC_P = (
_tcgen05.make_sdesc_128B_swizzle(0)
if H == 32
else _make_sdesc_64B_swizzle(8 * P_ROW)
)
KB_Q = H * 128 # bytes per 64-dim k-block of Q
def stage_inputs(
fmt_fp4: bool,
cache,
idx,
lens,
q,
positions,
cos_sin,
b,
tid,
row0,
width: int,
page: int,
a_kv,
a_q,
a_valid,
a_cs,
rows: int,
):
"""Issue every global load, then Q RoPE -> SMEM and KV dequant -> SMEM.
``rows`` (128 or 64) slots of the list start at ``row0``; with 64 the
tile's upper half is left unwritten and flagged invalid. Returns 1 if
the tile has any candidate slot (CTA-uniform).
"""
page_shift = page.bit_length() - 1
row_bytes = 256 if fmt_fp4 else 512
sf_bytes = 32 if fmt_fp4 else 16
record = row_bytes + sf_bytes
# Gather mapping: TPR threads per row; lane % RPW = row within the warp,
# lane // RPW = part j, which owns 16-byte chunks j, j + TPR, ... .
TPR = _THREADS // rows
RPW = 32 // TPR
warp = tid // 32
lane = tid % 32
t = warp * RPW + lane % RPW
j = lane // RPW
# Slot ids do not wait for the length (one dependent round trip less).
rowpos = row0 + t
id_addr = (
idx.iterator.toint()
+ (Int64(b) * Int64(idx.stride[0]) + Int64(rowpos)) * 4
)
sid = _ldg_s32(id_addr, _i32(rowpos < width), Int32(-1))
limit = cutlass.min(lens[b], Int32(width))
pos = positions[b]
live = row0 < limit
in_range = rowpos < limit
# Q: 16-byte pieces of the chunk-interleaved layout
# q[b, (d // 16) * HP * 16 + h * 16 + d % 16]; only the last piece of a
# thread can reach the RoPE dims (static, see QP).
hp = Int64(q.shape[1])
q_base = q.iterator.toint() + Int64(b) * Int64(q.stride[0]) * 2
qw = []
qmeta = []
for k in range(QP):
i = tid + _THREADS * k
c16 = i // (2 * H)
rem = i % (2 * H)
h = rem // 2
half = rem % 2
g = q_base + (Int64(c16) * hp * 16 + Int64(h * 16 + half * 8)) * 2
qw.append(_ldg_v4(g, Int32(1)))
qmeta.append((c16 * 16 + half * 8, h))
# KV gather (predicated: invalid rows read nothing and stay zero).
slots = Int64(cache.shape[0]) * page
valid = in_range & (sid >= 0) & (Int64(sid) < slots)
vp = _i32(valid)
blk = sid >> page_shift
slot = sid & (page - 1)
base = cache.iterator.toint() + Int64(blk) * Int64(cache.stride[0])
data = base + Int64(slot * record)
sfa = data + row_bytes
# 16-byte chunk c = TPR k + j of the row: a warp-wide load covers
# TPR * 16 contiguous bytes of each of its rows.
n_chunks = (row_bytes // 16) // TPR
kv = []
for k in range(n_chunks):
kv.append(_ldg_v4(data + (Int64(TPR * k) + Int64(j)) * 16, vp))
sfw = _ldg_v4(sfa, vp)
if fmt_fp4:
sfw = sfw + _ldg_v4(sfa + 16, vp)
# cos | sin at the query position (after the KV loads: pos may arrive later
# than the slot ids). The whole row also goes to SMEM for the epilogue.
d_last = qmeta[QP - 1][0]
rope = d_last >= _ROPE_START
f0 = (d_last - _ROPE_START) // 2
cs_row = (
cos_sin.iterator.toint() + Int64(pos) * Int64(cos_sin.stride[0]) * 4
)
cosv = _ldg_f32x4(cs_row + Int64(f0) * 4, _i32(rope), 1.0)
sinv = _ldg_f32x4(cs_row + Int64(f0 + _ROPE_DIM // 2) * 4, _i32(rope), 0.0)
cs_lane = tid % 16
cs_pred = _i32(tid < 16)
csw = _ldg_v4(cs_row + Int64(cs_lane) * 16, cs_pred)
# Q RoPE (fp32, GPT-J pairs) -> bf16 -> SMEM [H, 512] K-major SW128.
for k in range(QP):
d0, h = qmeta[k]
words = list(qw[k])
if k == QP - 1:
for w in range(4):
qe, qo = cvt.bf16x2_to_fp32x2(words[w])
c, s = cosv[w], sinv[w]
re = _fma(qe, c, -_mul(qo, s))
ro = _fma(qo, c, _mul(qe, s))
words[w] = _sel(
rope, cvt.fp32x2_to_bf16x2(re, ro), words[w], Uint32
)
kb = d0 // 64
phys = ((d0 % 64) // 8) ^ (h % 8)
addr = a_q + kb * KB_Q + (h // 8) * 1024 + (h % 8) * 128 + phys * 16
_sts_v4(addr, words[0], words[1], words[2], words[3])
# KV dequant -> SMEM [128, 512] bf16, K-major SW128 (8 k-blocks of
# 128 rows x 128 B; 16-byte chunk c of row r stored at c ^ (r % 8)).
row_off = a_kv + (t // 8) * 1024 + (t % 8) * 128
if fmt_fp4:
# NVFP4: 16-byte chunk c = TPR k + j holds dims [32 c, 32 c + 32); its
# e4m3 scales are bytes 2 c and 2 c + 1 of the record's 32 scale
# bytes, i.e. word TPR / 2 * k + j // 2.
for k in range(n_chunks):
c = TPR * k + j
kb_off = row_off + (c // 2) * 16384
lc0 = (c % 2) * 4
word = sfw[(TPR // 2) * k]
for jj in range(1, TPR // 2):
word = _sel(
j // 2 == jj, sfw[(TPR // 2) * k + jj], word, Uint32
)
sh = Uint32((j % 2) * 16)
s0 = _e4m3_to_bf16x2((word >> sh) & Uint32(0xFF))
s1 = _e4m3_to_bf16x2((word >> (sh + Uint32(8))) & Uint32(0xFF))
for w in range(4):
o = _fp4x8_to_bf16x8(kv[k][w], s0 if w < 2 else s1)
_sts_v4(
kb_off + ((lc0 + w) ^ (t % 8)) * 16, o[0], o[1], o[2], o[3]
)
else:
# fp8: 16-byte chunk c = TPR k + j holds dims [16 c, 16 c + 16) (k-block
# c // 4, 16-byte SMEM chunks 2 (j % 4) + {0, 1}), scaled by ue8m0 byte
# g = c // 2 = TPR / 2 * k + j // 2 of the record's 16 scale bytes.
for k in range(n_chunks):
c = TPR * k + j
kb_off = row_off + (c // 4) * 16384
g0 = (TPR // 2) * k # static part of g; j // 2 adds 0..TPR/2-1
e = (
sfw[g0 // 4] >> Uint32(8 * (g0 % 4)) >> Uint32(8 * (j // 2))
) & Uint32(0xFF)
for w in range(2):
o = _fp8x8_to_bf16x8_ue8m0(kv[k][2 * w], kv[k][2 * w + 1], e)
_sts_v4(
kb_off + ((2 * (j % 4) + w) ^ (t % 8)) * 16,
o[0],
o[1],
o[2],
o[3],
)
_sts_s32(a_valid + t * 4, vp)
if rows < _TILE: # the unwritten rows never take part
_sts_s32(a_valid + (rows + tid % (_TILE - rows)) * 4, Int32(0))
_sts_v4p(a_cs + cs_lane * 16, csw[0], csw[1], csw[2], csw[3], cs_pred)
return _i32(live)
@cute.kernel
def device_kernel(
q: cute.Tensor,
swa: cute.Tensor,
swa_idx: cute.Tensor,
swa_len: cute.Tensor,
ext: cute.Tensor,
ext_idx: cute.Tensor,
ext_len: cute.Tensor,
positions: cute.Tensor,
cos_sin: cute.Tensor,
sink: cute.Tensor,
out: cute.Tensor,
out_sf: cute.Tensor,
scale_log2: Float32,
):
tid, _, _ = cute.arch.thread_idx()
rank, b, _ = cute.arch.block_idx()
warp = cute.arch.make_warp_uniform(tid // 32)
lane = tid % 32
smem = utils.SmemAllocator()
p_kv = smem.allocate(_TILE * _HEAD_DIM * 2, byte_alignment=1024)
p_q = smem.allocate(H * _HEAD_DIM * 2, byte_alignment=1024)
p_p = smem.allocate(_TILE * P_ROW, byte_alignment=1024)
p_valid = smem.allocate(_TILE * 4, byte_alignment=16)
p_red = smem.allocate(2 * 4 * H * 4, byte_alignment=16)
p_stat_all = smem.allocate(STATS * 4, byte_alignment=16)
p_cs = smem.allocate(_ROPE_DIM * 4, byte_alignment=16)
mbar = smem.allocate_array(Int64, 2, byte_alignment=8)
p_tmem = smem.allocate_array(Int32, 4, byte_alignment=16)
a_kv = p_kv.toint()
a_q = p_q.toint()
a_p = p_p.toint()
a_valid = p_valid.toint()
a_rmax = p_red.toint()
a_rsum = a_rmax + 4 * H * 4
# Own stats (m, l per head) live in this CTA's slot of the all-stats array.
a_stat = p_stat_all.toint() + rank * (2 * H * 4)
a_stat_all = p_stat_all.toint()
a_cs = p_cs.toint()
if tid == 0:
cute.arch.mbarrier_init(mbar, 1)
cute.arch.mbarrier_init(mbar + 1, 1)
cute.arch.mbarrier_init_fence()
# Every global read below may depend on the PDL predecessor.
cute.arch.griddepcontrol_wait()
# TMEM after the PDL wait: a predecessor on this SM may hold it.
if warp == 0:
cute.arch.alloc_tmem(TMEM_COLS, p_tmem)
cute.arch.relinquish_tmem_alloc_permit()
_tcgen05.fence_before_thread_sync()
# ---- 1-2: loads, Q RoPE, KV dequant into SMEM
tile_live = Int32(0)
if rank < KT:
tile_live = stage_inputs(
False,
swa,
swa_idx,
swa_len,
q,
positions,
cos_sin,
b,
tid,
rank * _SWA_ROWS,
NK,
KP,
a_kv,
a_q,
a_valid,
a_cs,
_SWA_ROWS,
)
else:
tile_live = stage_inputs(
EXTRA_FP4,
ext,
ext_idx,
ext_len,
q,
positions,
cos_sin,
b,
tid,
(rank - KT) * _TILE,
NE,
EP,
a_kv,
a_q,
a_valid,
a_cs,
_TILE,
)
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.sync_threads()
_tcgen05.fence_after_thread_sync()
tmem = _lds_s32(p_tmem.toint())
if tile_live != 0:
# ---- 3: S[128, H] = KV[128, 512] . Q[H, 512]^T
if warp == 0:
_tcgen05.fence_after_thread_sync()
s_adesc = DESC_KMAJ | Uint64((a_kv >> 4) & 0x3FFF)
s_bdesc = DESC_KMAJ | Uint64((a_q >> 4) & 0x3FFF)
for s_kb in cutlass.range_constexpr(_HEAD_DIM // 64):
for s_ks in cutlass.range_constexpr(4):
_tcgen05.mma_f16(
tmem + S_COL,
s_adesc + (s_kb * 16384 + s_ks * 32) // 16,
s_bdesc + (s_kb * KB_Q + s_ks * 32) // 16,
IDESC_S,
s_kb > 0 or s_ks > 0,
)
_tcgen05.commit(mbar)
# ---- 4: softmax over the tile's 128 rows (column-wise per head)
quad = warp % 4
cg = warp // 4
t = quad * 32 + lane
cute.arch.mbarrier_wait(mbar, 0)
_tcgen05.fence_after_thread_sync()
sv = cute.make_rmem_tensor(NC, Float32)
sv.store(_tcgen05.ld(quad * 32, tmem + S_COL + cg * NC, "32x32b", NC))
_tcgen05.wait_ld()
row_ok = _lds_s32(a_valid + t * 4) != 0
neg_inf = Float32(-math.inf)
xs = []
for i in cutlass.range_constexpr(NC):
xs.append(_sel(row_ok, sv[i] * scale_log2, neg_inf, Float32))
# Column max over the 128 rows: butterfly over the 32 lanes (all of
# this thread's columns in lockstep), then across the 4 lane quadrants.
mxs = list(xs)
for sh in cutlass.range_constexpr(5):
for i in cutlass.range_constexpr(NC):
mxs[i] = cute.arch.fmax(
mxs[i], cute.arch.shuffle_sync_bfly(mxs[i], offset=16 >> sh)
)
for k4 in cutlass.range_constexpr(NC // 4):
_sts_v4f(
a_rmax + (quad * H + cg * NC + 4 * k4) * 4,
mxs[4 * k4],
mxs[4 * k4 + 1],
mxs[4 * k4 + 2],
mxs[4 * k4 + 3],
_i32(lane == 0),
)
cute.arch.sync_threads()
cmax = []
for k4 in cutlass.range_constexpr(NC // 4):
v4 = list(_lds_v4f(a_rmax + (cg * NC + 4 * k4) * 4))
for qd in cutlass.range_constexpr(1, 4):
w4 = _lds_v4f(a_rmax + (qd * H + cg * NC + 4 * k4) * 4)
for e in cutlass.range_constexpr(4):
v4[e] = cute.arch.fmax(v4[e], w4[e])
for e in cutlass.range_constexpr(4):
cmax.append(
_sel(v4[e] == neg_inf, Float32(0.0), v4[e], Float32)
)
probs = []
for i in cutlass.range_constexpr(NC):
probs.append(_ex2(xs[i] - cmax[i]))
# P = hi + lo (bf16 each): one [128, 2H] MN-major B tile.
hi_w = []
lo_w = []
for i in cutlass.range_constexpr(NC // 2):
hw = cvt.fp32x2_to_bf16x2(probs[2 * i], probs[2 * i + 1])
h0, h1 = cvt.bf16x2_to_fp32x2(hw)
hi_w.append(hw)
lo_w.append(
cvt.fp32x2_to_bf16x2(probs[2 * i] - h0, probs[2 * i + 1] - h1)
)
if cutlass.const_expr(H == 32):
prow = a_p + (t // 8) * 1024 + (t % 8) * 128
_sts_v4(prow + ((cg ^ (t % 8)) * 16), *hi_w)
_sts_v4(prow + (((4 + cg) ^ (t % 8)) * 16), *lo_w)
else:
prow = a_p + (t // 8) * 512 + (t % 8) * 64 + (cg % 2) * 8
sw = (t % 8) // 2
_sts_v2(prow + (((cg // 2) ^ sw) * 16), *hi_w)
_sts_v2(prow + (((2 + cg // 2) ^ sw) * 16), *lo_w)
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.sync_threads()
# ---- 5: O[512, 2H] = KV^T[512, 128] . [P_hi | P_lo][128, 2H]
if warp == 0:
_tcgen05.fence_after_thread_sync()
o_adesc = DESC_KV_MN | Uint64((a_kv >> 4) & 0x3FFF)
o_bdesc = DESC_P | Uint64((a_p >> 4) & 0x3FFF)
for o_mb in cutlass.range_constexpr(_HEAD_DIM // _TILE):
for o_ks in cutlass.range_constexpr(_SWA_ROWS // 16):
_tcgen05.mma_f16(
tmem + o_mb * 2 * H,
o_adesc + (o_mb * 2 * 16384 + o_ks * 2048) // 16,
o_bdesc + (o_ks * 16 * P_ROW) // 16,
IDESC_O,
o_ks > 0,
)
# Compressed tiles: the other 64 K rows (an SWA tile's upper
# half is stale SMEM and must not be read).
if rank >= KT:
for o_mb in cutlass.range_constexpr(_HEAD_DIM // _TILE):
for o_ks in cutlass.range_constexpr(
_SWA_ROWS // 16, _TILE // 16
):
_tcgen05.mma_f16(
tmem + o_mb * 2 * H,
o_adesc + (o_mb * 2 * 16384 + o_ks * 2048) // 16,
o_bdesc + (o_ks * 16 * P_ROW) // 16,
IDESC_O,
True,
)
_tcgen05.commit(mbar + 1)
# Column sums (only the merge needs them) while MMA2 runs.
sums = list(probs)
for sh in cutlass.range_constexpr(5):
for i in cutlass.range_constexpr(NC):
sums[i] = sums[i] + cute.arch.shuffle_sync_bfly(
sums[i], offset=16 >> sh
)
for k4 in cutlass.range_constexpr(NC // 4):
_sts_v4f(
a_rsum + (quad * H + cg * NC + 4 * k4) * 4,
sums[4 * k4],
sums[4 * k4 + 1],
sums[4 * k4 + 2],
sums[4 * k4 + 3],
_i32(lane == 0),
)
cute.arch.sync_threads()
if tid < H:
st_m = _lds_f32(a_rmax + tid * 4)
st_l = _lds_f32(a_rsum + tid * 4)
for st_q in cutlass.range_constexpr(1, 4):
st_m = cute.arch.fmax(
st_m, _lds_f32(a_rmax + (st_q * H + tid) * 4)
)
st_l = st_l + _lds_f32(a_rsum + (st_q * H + tid) * 4)
st_m = _sel(st_m == Float32(-math.inf), Float32(0.0), st_m, Float32)
_sts_f32(a_stat + tid * 4, st_m)
_sts_f32(a_stat + (H + tid) * 4, st_l)
# ---- 6: O = hi + lo -> fp32 [H, 512] SMEM over the KV tile
cute.arch.mbarrier_wait(mbar + 1, 0)
_tcgen05.fence_after_thread_sync()
mb = warp // 4
ov = cute.make_rmem_tensor(2 * H, Float32)
ov.store(_tcgen05.ld(quad * 32, tmem + mb * 2 * H, "32x32b", 2 * H))
_tcgen05.wait_ld()
d = mb * _TILE + quad * 32 + lane
for h in cutlass.range_constexpr(H):
o = ov[h] + ov[H + h]
_sts_f32(a_kv + (h * _HEAD_DIM + d) * 4, o)
else:
if tid < H:
_sts_f32(a_stat + tid * 4, Float32(0.0))
_sts_f32(a_stat + (H + tid) * 4, Float32(0.0))
# Peers pull every partial unconditionally: hand them zeros.
zero = Uint32(0)
for z in cutlass.range_constexpr(H * _HEAD_DIM // (4 * _THREADS)):
_sts_v4(a_kv + (tid + z * _THREADS) * 16, zero, zero, zero, zero)
# ---- 7: every CTA's partial and stats are final.
cute.arch.cluster_arrive()
cute.arch.cluster_wait()
cute.arch.griddepcontrol_launch_dependents()
# ---- 8: pull this CTA's (head, 32-dim chunk) units from every peer's
# partial (fixed peer order: deterministic) while gathering their stats.
u_begin = rank * UPC
u_count = cutlass.min(Int32(UPC), Int32(U) - u_begin)
sub = tid % 8
pulls = []
for ps in cutlass.range_constexpr(NPASS):
uo = ps * (_THREADS // 8) + tid // 8
u = u_begin + cutlass.min(uo, u_count - 1)
src = a_kv + u * UNIT_BYTES + sub * 16
row = []
for qp in cutlass.range_constexpr(C):
row.append(_ld_peer_f32x4(src, Int32(qp), Int32(1)))
pulls.append(row)
for k in cutlass.range_constexpr(cdiv(STATS, _THREADS)):
i = cutlass.min(tid + _THREADS * k, Int32(STATS - 1))
peer = i // (2 * H)
v = _ld_peer_f32(p_stat_all.toint() + i * 4, peer)
_sts_f32(p_stat_all.toint() + i * 4, v, _i32(peer != rank))
cute.arch.sync_threads()
# ---- 9: merge this CTA's (head, 32-dim chunk) units (fixed peer order:
# deterministic; sink folded once), then the epilogue: inverse RoPE,
# MXFP8 cast, o-fused store.
out_base = out.iterator.toint() + Int64(b) * Int64(out.stride[0])
sf_base = out_sf.iterator.toint()
inf = Float32(math.inf)
for ps in cutlass.range_constexpr(NPASS):
uo = ps * (_THREADS // 8) + tid // 8
active = uo < u_count
uoc = cutlass.min(uo, u_count - 1)
u = u_begin + uoc
hh = u // 16
cc = u % 16
sk = sink[hh]
sk2 = sk * _LOG2E
parts = pulls[ps]
ms = []
ls = []
big = sk2
for qp in cutlass.range_constexpr(C):
mq = _lds_f32(a_stat_all + (qp * 2 * H + hh) * 4)
lq = _lds_f32(a_stat_all + (qp * 2 * H + H + hh) * 4)
ms.append(mq)
ls.append(lq)
big = _sel(lq > Float32(0.0), cute.arch.fmax(big, mq), big, Float32)
big = _sel((big == inf) | (big == -inf), Float32(0.0), big, Float32)
den = _ex2(sk2 - big)
acc = [Float32(0.0)] * 4
for qp in cutlass.range_constexpr(C):
f = _sel(
ls[qp] > Float32(0.0), _ex2(ms[qp] - big), Float32(0.0), Float32
)
den = den + f * ls[qp]
for e in cutlass.range_constexpr(4):
acc[e] = acc[e] + parts[qp][e] * f
ok = (den > Float32(0.0)) & (sk != inf)
for e in cutlass.range_constexpr(4):
acc[e] = _sel(ok, acc[e] / den, Float32(0.0), Float32)
if cutlass.const_expr(ps == NPASS - 1):
# Every pulled value is consumed: peers may leave once all of
# them got here too (split barrier, waited on before exit).
cute.arch.cluster_arrive_relaxed()
# Inverse GPT-J RoPE on dims >= 448 (identity elsewhere).
is_rope = cc >= _ROPE_START // 32
f0 = cutlass.max((cc - _ROPE_START // 32) * 16 + sub * 2, Int32(0))
cs2 = _lds_v2f(a_cs + f0 * 4)
sn2 = _lds_v2f(a_cs + (f0 + _ROPE_DIM // 2) * 4)
x = []
for pr in cutlass.range_constexpr(2):
cv = _sel(is_rope, cs2[pr], Float32(1.0), Float32)
sv1 = _sel(is_rope, sn2[pr], Float32(0.0), Float32)
oe, oo = acc[2 * pr], acc[2 * pr + 1]
x.append(_fma(oe, cv, _mul(sv1, oo)))
x.append(_fma(oo, cv, -_mul(sv1, oe)))
# MXFP8: ue8m0 = ceil(log2(amax / 448)) over the 32-dim chunk.
amax = cute.arch.fmax(
cute.arch.fmax(
cute.arch.fmax(x[0], -x[0]), cute.arch.fmax(x[1], -x[1])
),
cute.arch.fmax(
cute.arch.fmax(x[2], -x[2]), cute.arch.fmax(x[3], -x[3])
),
)
for off in (4, 2, 1):
amax = cute.arch.fmax(
amax, cute.arch.shuffle_sync_bfly(amax, offset=off)
)
amax = cute.arch.fmax(amax, Float32(1e-10))
ex = _ue8m0_rp(amax * Float32(1.0 / 448.0))
inv = recast_val((Uint32(254) - ex) << Uint32(23), Float32)
qb = cvt.fp32x4_to_fp8x4(x[0] * inv, x[1] * inv, x[2] * inv, x[3] * inv)
g = hh // WV_GROUP_SIZE
hg = hh % WV_GROUP_SIZE
ka = cc * WV_GROUP_SIZE + hg
oa = (
out_base
+ Int64(g) * Int64(out.stride[1])
+ Int64(ka * 32 + sub * 4)
)
_stg_b32(oa, qb, _i32(active))
sw_idx = (
Int64(b) * Int64(out_sf.stride[0])
+ Int64(g) * Int64(out_sf.stride[1])
+ Int64(ka // 4) * Int64(out_sf.stride[2])
)
_stg_b8(
sf_base + sw_idx * 4 + Int64(ka % 4), ex, _i32(active & (sub == 0))
)
# ---- 10: nobody leaves while a peer may still read its SMEM (the
# matching arrive follows this CTA's last pull).
cute.arch.cluster_wait()
if warp == 0:
cute.arch.dealloc_tmem(
cute.make_ptr(Float32, tmem, cute.AddressSpace.tmem), TMEM_COLS
)
@cute.jit
def host_entrypoint(
q: cute.Tensor,
swa: cute.Tensor,
swa_idx: cute.Tensor,
swa_len: cute.Tensor,
ext: cute.Tensor,
ext_idx: cute.Tensor,
ext_len: cute.Tensor,
positions: cute.Tensor,
cos_sin: cute.Tensor,
sink: cute.Tensor,
out: cute.Tensor,
out_sf: cute.Tensor,
scale_log2: Float32,
stream: CUstream,
):
device_kernel(
q,
swa,
swa_idx,
swa_len,
ext,
ext_idx,
ext_len,
positions,
cos_sin,
sink,
out,
out_sf,
scale_log2,
).launch(
grid=(C, q.shape[0], 1),
block=(_THREADS, 1, 1),
cluster=(C, 1, 1),
stream=stream,
use_pdl=True,
# 150+ KiB of SMEM leaves room for one CTA per SM anyway; saying
# so lets ptxas use up to 128 registers per thread.
min_blocks_per_mp=1,
)
return host_entrypoint
def dispatch(self, **compile_key_fields: int) -> CompileKey:
return self.CompileKey(**compile_key_fields)
def get_warmup_keys(self, **layer_shape: int) -> list[CompileKey]:
# The token count is a grid dimension, so one executor per layer shape
# serves every decode step.
return self._trace_dispatch(self.dispatch)(**layer_shape)
def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
def cache(page: int, record: int):
return make_fake_tensor(
Uint8,
(cute.sym_int(), page, record),
(cute.sym_int(), record, 1),
assumed_align=16,
)
# Independent leading dims: metadata buffers may be longer than s_q.
def slots(width: int):
return make_fake_tensor(
Int32, (cute.sym_int(), width), (cute.sym_int(), 1), assumed_align=4
)
def vec(dtype):
return make_fake_tensor(dtype, (cute.sym_int(),), (1,), assumed_align=4)
out_dim = WV_GROUP_SIZE * _HEAD_DIM
return (
make_fake_tensor(
BFloat16,
(cute.sym_int(), cute.sym_int(), _HEAD_DIM),
(cute.sym_int(), _HEAD_DIM, 1),
assumed_align=16,
),
cache(compile_key.swa_page, _FP8_RECORD),
slots(compile_key.swa_width),
vec(Int32),
cache(compile_key.extra_page, compile_key.extra_record),
slots(compile_key.extra_width),
vec(Int32),
make_fake_tensor(Int64, (cute.sym_int(),), (1,), assumed_align=8),
make_fake_tensor(
Float32,
(cute.sym_int(), _ROPE_DIM),
(cute.sym_int(), 1),
assumed_align=16,
),
vec(Float32),
make_fake_tensor(
Float8E4M3FN,
(cute.sym_int(), cute.sym_int(), out_dim),
(cute.sym_int(), out_dim, 1),
assumed_align=16,
),
make_fake_tensor(
Int32,
(cute.sym_int(), cute.sym_int(), out_dim // 128),
(cute.sym_int(), cute.sym_int(), cute.sym_int()),
assumed_align=4,
),
Float32(0.0),
)
@kernel_launcher
def __call__(
self,
*,
q: torch.Tensor,
swa_cache: torch.Tensor,
swa_idx: torch.Tensor,
swa_len: torch.Tensor,
extra_cache: torch.Tensor,
extra_idx: torch.Tensor,
extra_len: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
sink: torch.Tensor,
heads: int,
softmax_scale: float,
out_data: torch.Tensor,
out_scale: torch.Tensor,
) -> CuTeDSLLaunchSpec["SwapABDecodeKernel.CompileKey"]:
"""Decode attention of ``heads`` live heads, in place into ``out_*``.
The arguments are those of the mega kernel's decode call: ``q`` is
``[s_q, padded_heads, 512]`` bf16 in its chunk-interleaved layout, not
yet RoPEd; the caches are ``[num_blocks, page, record]`` uint8 views
of V4.1 records (a token's data, then its scales) whose pages may be
further apart than ``page * record`` bytes; the slot lists are
``[s_q, width]`` int32 slot ids (``block * page + offset``, ``-1``
invalid) with ``[s_q]`` int32 lengths; ``positions`` is int64. Only
the live heads' groups of ``out_data`` / ``out_scale`` are written.
"""
assert q.shape[1] >= heads
compile_key = self.dispatch(
heads=heads,
swa_width=swa_idx.shape[-1],
extra_width=extra_idx.shape[-1],
extra_record=extra_cache.shape[-1],
swa_page=swa_cache.shape[1],
extra_page=extra_cache.shape[1],
)
launch_args = (
q,
swa_cache,
swa_idx,
swa_len,
extra_cache,
extra_idx,
extra_len,
positions,
cos_sin_cache,
sink,
out_data,
out_scale,
softmax_scale * _LOG2E,
)
return compile_key, launch_args