David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/cuda_kernels/attention.cu

2674 lines
121 KiB
Plaintext
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file attention.cu
* @brief On-die causal GQA flash attention for Foundry (fp32 and Q4_0 KV cache).
*
* Written from a behaviour specification (docs/cleanroom/f_cua_spec.md); no ggml or llama.cpp source was consulted.
* Scope of that rewrite: the Q4_0 KV-cache format defines, fc_kvq4_row_bytes,
* fc_kvq4_quant_block, fc_kvq4_dequant_elem, fc_kvq4_dequant_lane_chunk and
* fc_attn_flash_prefill_q4_kvs; the rest of the file is unchanged original work.
*
* Target: sm_37 (K80, Kepler)
* Compile: nvcc -ptx -arch=sm_37 attention.cu -o attention.ptx
*
* Replaces the HOST-side causal GQA softmax in the qwen35moe batched
* prefill (the O(N^2) 2-core bottleneck: 87.5% of an 8k prefill) with a
* flash-attention-style kernel: tiled QK^T -> ONLINE softmax (running
* max + running sum, the N x N score matrix is never materialized) ->
* PV accumulate, entirely on the die against a GPU-resident KV cache.
*
* Numerics: same math as the host reference (qwen35moe_forward_cuda.c
* host softmax) but reduced in tile order with running-max rescaling, so
* results drift ~1e-6 relative vs the host's global-max single pass —
* NOT bit-exact by design (gate: argmax-exact + logits within ~1e-3).
*
* Sliding-window attention (QWEN35_ATTN_WINDOW, StreamingLLM-style): each
* flash kernel takes attn_sink/attn_win. When attn_win > 0 the query at
* absolute position cur_pos attends ONLY the key set
* [0, attn_sink) U [cur_pos - attn_win + 1, cur_pos]
* (attention sinks + recent window); tiles that intersect neither are
* skipped BEFORE any staging/sync, turning prefill O(N^2) -> O(N*W). The
* online-softmax (m, l, acc) simply never sees the excluded keys, so the
* result is a correct softmax over the windowed set. When attn_win <= 0
* every tile is processed — byte-identical to full attention (the OFF /
* A-baseline path). A query always sees its own position (cur_pos is in
* the recent window for attn_win >= 1); an overlapping sink/window tile is
* processed once (the range predicate is an OR).
*
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
* FC_PTX_ATTENTION -> kernels fc_attn_qz_split, fc_attn_kv_append,
* fc_attn_flash_prefill (+ _q4 variants),
* fc_attn_flash_decode(_q4) + _combine,
* the *_dev device-position variants for
* CUDA-graph capture (C1), and the *_slots
* multi-user decode variants (L7)
*
* Kepler notes:
* - Pure fp32 FMA (no tensor cores); expf matches the host libm family.
* - 48 KB shared per block: the flash kernel's working set (q group +
* one padded K tile + softmax state) is 26 720 B for kv_groups=8/
* head_dim=256, so 4 blocks per SMX (114 688 B/SMX; shared -- not
* registers -- is what caps residency here). V is streamed through
* __ldg (read-only cache), never staged in shared.
* - __shfl_down_sync with width 16 for the dot-product lane reduction
* (sm_30+; the mask is ignored pre-Volta).
*
* @status None
* @issues None
* @todo Split into prefill and decode files (see the TODO in the
* NARROW-STAGING comment for the SASS-identity gate it needs).
*/
/* Flash-kernel tiling: keys staged per tile, lanes cooperating on one
* q.k dot, and the register-accumulator cap on GQA group size. */
#define FC_ATTN_TILE_KEYS 16
#define FC_ATTN_DOT_LANES 16
#define FC_ATTN_MAX_GROUPS 8
/* K-tile row stride, in floats. The tile is read as k_s[ky*stride + lane +
* i*16] with ky = tid/16 and lane = tid%16, so the two half-warps of a warp
* carry IDENTICAL lanes and ADJACENT ky. At the shipped head_dim = 256 an
* unpadded row stride is exactly 8 x 32 banks, which puts both half-warps on
* the same banks at different addresses: a 2-way conflict on every k-operand
* LDS. Padding the row by half a bank cycle breaks it for any head_dim that
* is a multiple of 16. Costs TILE_KEYS*16 floats (1 KB) of dynamic shared and
* does not change occupancy at the shipped shape (26 720 B x 4 blocks of
* 114 688 B/SMX). MUST match fc_attn_flash_shared_bytes() in cuda/fc_cuda.c. */
#define FC_ATTN_K_PAD 16
#define FC_ATTN_K_STRIDE(hd) ((hd) + FC_ATTN_K_PAD)
/* ============================================================
* 4-bit KV cache (Q4_0 blocks)
*
* A K or V row of head_dim floats is stored as head_dim/32 blocks of 18 bytes:
* bytes 0-1 scale d, binary16 bit pattern, little-endian
* bytes 2-17 16 code bytes; byte 2+j holds element j in its low nibble and
* element j+16 in its high nibble
* A code c in 0..15 decodes to (c - 8) * d. The host (Pantheon pn_kvq4_*)
* reads and writes the same bytes, so the layout is an interface.
* ============================================================ */
#include <cuda_fp16.h>
#define FC_KVQ4_BLK 32
#define FC_KVQ4_BLK_BYTES 18
/* fp16 sink + recent-window exemption: the attention sink (first `sink`
* positions) and the last `win` tokens are kept at fp16 in small side buffers
* instead of q4, so greedy argmax stops flipping on near-ties. The q4 body
* still holds every position; recent rows are read from the fp16 window ring
* while recent and from q4 once they age out (no rolling). Buffers are sized to
* these maxima (must match host PN_KVQ4_WMAX / PN_KVQ4_SINKMAX).
*
* THE RING LAPS AT FC_KVQ4_WMAX, NOT AT `win`. The reader's window is only
* `win` wide, but a BATCHED prefill call appends ALL n rows before ANY query
* row reads, so a ring lapped at `win` lets row t's own window keys be
* overwritten by rows t+1..n-1 of the same call — row t then attends to K/V
* from positions AFTER it, a causality violation (see fc_attn_flash_prefill_q4
* for the guard and Pantheon results/q4_win_ring_hazard.md for the history).
* Lapping at WMAX instead keeps live every key a batch of up to
* (WMAX - win) rows can still need, and costs nothing: the buffers were
* always WMAX rows per head — only `win` of them were ever being used. */
#define FC_KVQ4_WMAX 256
#define FC_KVQ4_SINKMAX 8
/* Predicate: key position s read from the fp16 shadow (sink or recent window)
* for a query at cur_pos. Matches host pn_kvq4_is_fp16 exactly. */
__device__ __forceinline__ int fc_kvq4_is_fp16(int s, int cur_pos, int sink,
int win) {
return (s < sink) || (win > 0 && s > cur_pos - win);
}
/* Bytes in one Q4_0 row of head_dim elements (host mirror: FC_ATTN_KVQ4_ROW_BYTES). */
__device__ __forceinline__ int fc_kvq4_row_bytes(int head_dim) {
return (head_dim / FC_KVQ4_BLK) * FC_KVQ4_BLK_BYTES;
}
/* Max head_dim the append kernels stage in shared memory (CUDA's block-size
* ceiling; every dispatcher launches block = head_dim threads). */
#define FC_KVQ4_MAX_HEAD_DIM 1024
/* Quantize one 32-element block x32 (shared memory, already synced) into the
* 18 bytes at blk_ptr. Called by all 32 lanes of a warp with the same block;
* lane 0 writes the scale, lanes 0..15 one code byte each, so every byte of the
* block is written exactly once.
*
* Each lane folds the block maximum itself instead of reducing across the warp:
* on sm_37 a shuffle from an inactive source lane returns 0, which corrupted the
* high nibbles in a cooperative version. The serial fold also matches the host
* quantizer's order bit for bit. The scale is stored negative (d = amax / -8).
* The order convert -> +8 -> clamp is part of the output for out-of-range
* products (NaN scale, reciprocal overflow); keep it. */
__device__ __forceinline__ void fc_kvq4_quant_block(const float *x32,
unsigned char *blk_ptr,
int lane) {
float amax = 0.0f;
for (int j = 0; j < FC_KVQ4_BLK; j++)
amax = fmaxf(amax, fabsf(x32[j]));
const float d = amax / -8.0f;
const float id = (d != 0.0f) ? 1.0f / d : 0.0f;
if (lane == 0) {
/* One 16-bit store. Every block starts at an even offset (18-byte
* stride), so the scale is 2-byte aligned. Do not split it into two
* byte stores: ptxas 11.4 (sm_37) compiles st.u8 of the f16 register
* as a value conversion (F2I.U32.F16.TRUNC) and writes 0 as the low
* byte of every (negative) scale. */
*(__half *)blk_ptr = __float2half(d);
}
if (lane < FC_KVQ4_BLK / 2) {
int q_lo = (int)rintf(x32[lane] * id) + 8;
int q_hi = (int)rintf(x32[lane + FC_KVQ4_BLK / 2] * id) + 8;
if (q_lo < 0)
q_lo = 0;
else if (q_lo > 15)
q_lo = 15;
if (q_hi < 0)
q_hi = 0;
else if (q_hi > 15)
q_hi = 15;
blk_ptr[2 + lane] = (unsigned char)(q_lo | (q_hi << 4));
}
}
/* Decode element e of a Q4_0 row. The scale is read as two byte loads. */
__device__ __forceinline__ float fc_kvq4_dequant_elem(const unsigned char *row, int e) {
const unsigned char *blk = row + (e / FC_KVQ4_BLK) * FC_KVQ4_BLK_BYTES;
const unsigned short d_bits = (unsigned short)(blk[0] | (blk[1] << 8));
const float d = __half2float(__ushort_as_half(d_bits));
const int local = e & (FC_KVQ4_BLK - 1);
const int half_blk = FC_KVQ4_BLK / 2;
const int code_off = local < half_blk ? local : local - half_blk;
const int shift = local < half_blk ? 0 : 4;
const int code = (blk[2 + code_off] >> shift) & 0xF;
return (float)(code - 8) * d;
}
/* Sliding-window tile predicate: is the key tile [j0, j0+tk) attended by the
* query at absolute position cur_pos? True when attn_win <= 0 (full attention),
* OR the tile intersects the sink region [0, attn_sink), OR (attn_win > 0) it
* intersects the recent window [cur_pos - attn_win + 1, cur_pos]. The test is
* uniform across the block (depends only on j0/tk/cur_pos), so branching on it
* before a __syncthreads() is safe. */
__device__ __forceinline__ int fc_attn_tile_in_window(int j0, int tk,
int cur_pos,
int attn_sink,
int attn_win) {
if (attn_win <= 0)
return 1;
if (j0 < attn_sink)
return 1;
int win_lo = cur_pos - attn_win + 1;
if (win_lo < 0)
win_lo = 0;
/* intersects [win_lo, cur_pos]: tile end > win_lo AND tile start <= cur_pos */
return (j0 + tk > win_lo) && (j0 <= cur_pos);
}
extern "C" {
/*
* Split the joint Q+gate projection into contiguous q and z tensors.
*
* qg: [n][n_head][2*head_dim] fp32 — per head: [q head_dim | z head_dim]
* q, z: [n][n_head][head_dim] fp32 (written)
* head_dim: per-head dimension
* total: n * n_head * head_dim (elements of each output)
*
* Replaces the host-side split loop (a_qg D2H + memcpy + a_q H2D).
* 1D launch, one thread per output element.
*/
__global__ void fc_attn_qz_split(const float *qg, float *q, float *z,
int head_dim, int total) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total) return;
int dd = i % head_dim;
int th = i / head_dim; /* flattened (position, head) */
size_t src = (size_t)th * 2 * head_dim + dd;
q[i] = qg[src];
z[i] = qg[src + head_dim];
}
/*
* Append n roped K/V rows to the GPU-resident ring KV cache.
*
* k_new, v_new: [n][n_head_kv][head_dim] fp32 (post QK-norm + RoPE)
* k_cache, v_cache: [n_head_kv][max_seq][head_dim] fp32 (written) — same
* layout as the host cache so host<->GPU syncs are
* contiguous per-head row ranges
* pos0: global position of row 0 (row t lands at
* (pos0 + t) % max_seq, the host ring rule)
*
* Grid (n, n_head_kv), block head_dim threads.
*/
__global__ void fc_attn_kv_append(const float *k_new, const float *v_new,
float *k_cache, float *v_cache,
int n_head_kv, int head_dim, int max_seq,
int pos0) {
int t = blockIdx.x;
int h = blockIdx.y;
int dd = threadIdx.x;
if (dd >= head_dim) return;
int cache_pos = (pos0 + t) % max_seq;
size_t src = ((size_t)t * n_head_kv + h) * head_dim + dd;
size_t dst = ((size_t)h * max_seq + cache_pos) * head_dim + dd;
k_cache[dst] = k_new[src];
v_cache[dst] = v_new[src];
}
/*
* Causal GQA flash attention over the GPU KV cache (prefill mode).
*
* q: [n][n_head][head_dim] fp32 (post QK-norm + RoPE);
* n_head = gridDim.y * kv_groups, group g of kv head h is
* global head h*kv_groups + g (the host GQA map qh/kv_groups)
* k_cache: [n_head_kv][max_seq][head_dim] fp32 (rows [0, pos0+n) live)
* v_cache: same layout
* out: [n][n_head][head_dim] fp32 (written; softmax(qK^T*scale)V)
* kv_groups: q heads per kv head (<= FC_ATTN_MAX_GROUPS)
* pos0: global position of query row 0 — query t attends cache rows
* [0, pos0+t] (causal); caller guarantees pos0+n <= max_seq
* (no ring wrap inside the window)
* scale: 1/sqrt(head_dim)
*
* Grid (n, n_head_kv), block head_dim threads (multiple of 16), dynamic
* shared = (kv_groups*head_dim + TILE*FC_ATTN_K_STRIDE(head_dim) +
* 2*kv_groups*TILE + 3*kv_groups) floats.
*
* Per K tile: the block's kv_groups q heads score the tile's keys with
* 16-lane cooperative dots (each KV row is read ONCE per tile and reused
* across the whole q group — the GQA lever), one thread per head folds
* the tile into the running (max, sum) state, then every thread
* accumulates its output dim across the tile's V rows with the standard
* online-softmax rescale. Fully-masked tiles are never visited (the key
* loop stops at the query's causal bound).
*/
__global__ void fc_attn_flash_prefill(const float *q, const float *k_cache,
const float *v_cache, float *out,
int kv_groups, int head_dim,
int max_seq, int pos0, int attn_sink,
int attn_win, float scale) {
const int t = blockIdx.x; /* query row in the batch */
const int kvh = blockIdx.y; /* kv head */
const int tid = threadIdx.x; /* == output dim owned */
const int G = kv_groups;
const int seq_len = pos0 + t + 1; /* causal window */
const int cur_pos = seq_len - 1; /* this query's absolute position */
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh; /* [G][head_dim] */
float *k_s = q_s + (size_t)G * head_dim; /* [TILE][kstride] */
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride; /* [G][TILE] */
float *p_s = score_s + G * FC_ATTN_TILE_KEYS; /* [TILE][G] */
float *m_s = p_s + G * FC_ATTN_TILE_KEYS; /* [G] running max */
float *l_s = m_s + G; /* [G] running sum */
float *corr_s = l_s + G; /* [G] tile rescale */
/* Stage this block's q group; init the online-softmax state. */
const size_t n_head = (size_t)gridDim.y * G;
const float *q_base = q + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000); /* -inf */
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const float *k_head = k_cache + (size_t)kvh * max_seq * head_dim;
const float *v_head = v_cache + (size_t)kvh * max_seq * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES; /* key-in-tile this thread scores */
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES; /* dims per lane */
for (int j0 = 0; j0 < seq_len; j0 += FC_ATTN_TILE_KEYS) {
int tk = seq_len - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window before any
* staging/sync (uniform across the block — no __syncthreads() split). */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
/* Stage the K tile (coalesced; read once, reused by all G heads).
* Key-outer so the padded row stride costs no address arithmetic. */
for (int key = 0; key < tk; key++)
for (int e = tid; e < head_dim; e += blockDim.x)
k_s[(size_t)key * kstride + e] =
k_head[(size_t)(j0 + key) * head_dim + e];
__syncthreads();
/* QK^T: 16 lanes cooperate on each (head, key) dot.
*
* i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
/* Online-softmax fold, one thread per head: new running max, the
* rescale factor for the accumulators, this tile's probabilities. */
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
/* PV accumulate: thread tid owns out dim tid; V rows stream through
* the read-only cache (coalesced across the block). */
for (int j = 0; j < tk; j++) {
float vv = __ldg(&v_head[(size_t)(j0 + j) * head_dim + tid]);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads(); /* next tile overwrites k_s/score_s/p_s */
}
float *out_base = out + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
out_base[(size_t)g * head_dim + tid] = acc[g] / l_s[g];
}
/*
* Append n roped K/V rows to the GPU-resident Q4 ring KV cache (quantizing
* each row to Q4_0 as it lands). Same ring rule as fc_attn_kv_append.
*
* k_new, v_new: [n][n_head_kv][head_dim] fp32 (post QK-norm + RoPE)
* k_cache, v_cache: [n_head_kv][max_seq][row_bytes] Q4_0 bytes (written)
* pos0: global position of row 0 (caller guarantees pos0+n<=max_seq
* so cache_pos == pos0+t, no ring wrap in the window)
*
* Grid (n, n_head_kv), block head_dim threads (a multiple of 32; each 32-lane
* warp owns one Q4_0 block of the row).
*/
__global__ void fc_attn_kv_append_q4(const float *k_new, const float *v_new,
unsigned char *k_cache,
unsigned char *v_cache,
__half *k_sink, __half *v_sink,
__half *k_win, __half *v_win,
int n_head_kv, int head_dim, int max_seq,
int pos0, int sink, int win) {
int t = blockIdx.x;
int h = blockIdx.y;
int tid = threadIdx.x;
if (tid >= head_dim) return;
int cache_pos = pos0 + t;
int row_bytes = fc_kvq4_row_bytes(head_dim);
size_t src = ((size_t)t * n_head_kv + h) * head_dim + tid;
size_t dst_row = ((size_t)h * max_seq + cache_pos) * row_bytes;
int blk = tid / FC_KVQ4_BLK;
int lane = tid & (FC_KVQ4_BLK - 1);
/* Stage the row in shared memory so the (shuffle-free, Kepler-safe)
* block quantizer can read all 32 elements of its block. */
__shared__ float s_k[FC_KVQ4_MAX_HEAD_DIM];
__shared__ float s_v[FC_KVQ4_MAX_HEAD_DIM];
float kf = k_new[src];
float vf = v_new[src];
s_k[tid] = kf;
s_v[tid] = vf;
__syncthreads();
fc_kvq4_quant_block(s_k + (size_t)blk * FC_KVQ4_BLK,
k_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
fc_kvq4_quant_block(s_v + (size_t)blk * FC_KVQ4_BLK,
v_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
/* Mirror the row into the fp16 sink / recent-window shadow (element tid). */
if (win > 0) {
size_t wi = ((size_t)h * FC_KVQ4_WMAX + (cache_pos % FC_KVQ4_WMAX)) * head_dim + tid;
k_win[wi] = __float2half(kf);
v_win[wi] = __float2half(vf);
}
if (cache_pos < sink) {
size_t si = ((size_t)h * FC_KVQ4_SINKMAX + cache_pos) * head_dim + tid;
k_sink[si] = __float2half(kf);
v_sink[si] = __float2half(vf);
}
}
/*
* Causal GQA flash attention over the Q4 GPU KV cache (prefill mode). Identical
* online-softmax math and tiling to fc_attn_flash_prefill; the only difference
* is K/V are read from the packed Q4_0 cache and dequantized in-kernel (the K
* tile into shared, V per element on read) so no fp32 KV is ever materialized.
*
* k_cache/v_cache: [n_head_kv][max_seq][row_bytes] Q4_0 bytes.
* Grid/shared/other args as fc_attn_flash_prefill.
*/
__global__ void fc_attn_flash_prefill_q4(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win,
float *out, int kv_groups,
int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink,
int attn_win, float scale) {
const int t = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int seq_len = pos0 + t + 1;
const int cur_pos = seq_len - 1; /* this query row's absolute position */
const int row_bytes = fc_kvq4_row_bytes(head_dim);
/* Liveness floor for the fp16 window ring. Every row of this call was
* appended BEFORE any row reads, so slot (p % WMAX) holds position p only
* while p + WMAX > last_pos, where last_pos is the highest position this
* call appended (row gridDim.x-1). Keys below that floor were lapped by a
* LATER row of the same batch; reading them would attend to a future
* token. They are still exact in the q4 body, so fall back to it. With the
* host capping a batch at (WMAX - win) rows this never fires; it is the
* backstop that keeps an oversized batch merely lossier, never acausal. */
const int win_live_lo = (cur_pos - t + (int)gridDim.x - 1) - FC_KVQ4_WMAX;
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh;
float *k_s = q_s + (size_t)G * head_dim;
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const size_t n_head = (size_t)gridDim.y * G;
const float *q_base = q + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache + (size_t)kvh * max_seq * row_bytes;
const unsigned char *v_head = v_cache + (size_t)kvh * max_seq * row_bytes;
const __half *k_win_head = k_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = 0; j0 < seq_len; j0 += FC_ATTN_TILE_KEYS) {
int tk = seq_len - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window (attn_sink/
* attn_win are the ATTENDED set, orthogonal to sink/win which pick the
* fp16-shadow vs q4-dequant read path). Uniform across the block. */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
/* Stage the K tile (read once, reused by all G heads): fp16 shadow for
* sink/recent keys, in-kernel q4 dequant for the aged-out middle. */
for (int key = 0; key < tk; key++) {
int kp = j0 + key;
float *krow = k_s + (size_t)key * kstride;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else if (win > 0 && kp > cur_pos - win && kp > win_live_lo) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else {
const unsigned char *s = k_head + (size_t)kp * row_bytes;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = fc_kvq4_dequant_elem(s, e);
}
}
__syncthreads();
/* QK^T, i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
/* PV accumulate: thread tid owns out dim tid; V rows dequantized on
* read (element tid of each V row). */
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink)
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
else if (win > 0 && vp > cur_pos - win && vp > win_live_lo)
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
else
vv = fc_kvq4_dequant_elem(v_head + (size_t)vp * row_bytes, tid);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
float *out_base = out + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
out_base[(size_t)g * head_dim + tid] = acc[g] / l_s[g];
}
/* ============================================================
* REGISTER-K prefill variant (FC_ATTN_KTILE_REG) — BIT-EXACT, DEFAULT OFF
*
* Same tiling, same online softmax, same q4 KV cache and the same accumulation
* ORDER as fc_attn_flash_prefill_q4. One thing changes: the staged K tile is
* DELETED and K lives in per-thread registers instead.
*
* Why that is free. In the QK dot, thread (ky, lane) reads exactly the
* head_dim/16 elements k[ky][lane + 16*i] of key ky and nothing else — the
* shared tile is written by one thread and read back by that same thread. It
* exists only because the STAGING loop assigns elements element-outer
* (`for e = tid; e < head_dim; e += blockDim.x`) while the CONSUMING loop reads
* them lane-strided. Re-assigning the staging to match the consumer removes the
* STS/LDS round trip and the tile with it:
*
* shared 26 720 -> 9 312 B residency 4 -> 7 blocks/SMX
* registers 72 (unchanged) stack frame 0, no LDL/STL
*
* Measured on the standalone microbenchmark (cuda_kernels/tests/exp_attn_qtile.cu
* arm `xq_q1_kr`, driven by cuda/tests/bench_attn_qtile.c): 1.51x on the kernel
* at the shipped shape, depth-flat 1 k -> 32 k, at both 666 and 875 MHz, and
* BIT-EXACT — 0 of 524 288 output floats differ from fc_attn_flash_prefill_q4.
* The factor decomposes as 1.194x structural (the round trip) x 1.269x from the
* residency gain. Q-tiling the same body was REFUTED there (Q_TILE=2 costs
* shared and comes out at 0.886x); this variant is Q_TILE=1 only.
*
* TWO THINGS ARE LOAD-BEARING AND MUST NOT BE "SIMPLIFIED":
*
* 1. The HOISTED dequant (fc_kvq4_dequant_lane_chunk) is the enabler, not a
* detail. Calling fc_kvq4_dequant_elem once per element instead costs
* 105 REGISTERS, which caps residency back at 4 blocks/SMX and collapses
* the measured gain from 1.51x to 1.068x (microbench arm `xq_q1_krn`).
* 2. The online-softmax rescale stays FUSED in the PV loop at j == 0, exactly
* where fc_attn_flash_prefill_q4 puts it, because ptxas contracts
* `acc *= corr; acc += p0*v0` into a single `acc = fma(acc, corr, p0*v0)` —
* one rounding instead of two. Hoisting it out of the j loop is
* algebraically identical and moved 148 099 of 524 288 output floats by up
* to 4.0e-7 relative (1-ulp class), measured. Keep it at j == 0.
*
* CONTRACT (enforced host-side in fc_attn_flash_prefill_q4_ctx):
* - head_dim <= FC_ATTN_KREG_MAX_CHUNK * FC_ATTN_DOT_LANES (= 256), so the
* kreg[] array is bounded by a COMPILE-TIME constant. An array bounded by
* the runtime `chunk` is never promoted to registers: it lands in local
* memory and ptxas prints "0 bytes spill" beside a nonzero stack frame.
* - FC_ATTN_DOT_LANES == 16 (the hoisted dequant's nibble mapping).
* - attn_win <= 0 (full attention). The tile predicate here is the shipped
* one and a windowed prefill should behave identically, but that is
* UNMEASURED, so the host declines attn_win > 0 and uses the shipped
* kernel.
* ============================================================ */
#if FC_ATTN_DOT_LANES != 16
#error "fc_kvq4_dequant_lane_chunk assumes FC_ATTN_DOT_LANES == 16"
#endif
/* Compile-time bound on head_dim / FC_ATTN_DOT_LANES; 16 covers head_dim 256. */
#define FC_ATTN_KREG_MAX_CHUNK 16
/* Decode the elements a QK-dot lane owns, lane + 16*i for i < chunk, into
* kreg[i]. Element lane + 16*i sits in block i/2 at code byte 2 + lane: the low
* nibble for even i, the high nibble for odd i. So one scale and one code byte
* per block serve two kreg entries. kreg[chunk..] is left untouched.
*
* Unrolled over the compile-time bound so every kreg index is a constant; a
* runtime index would demote kreg to local memory. Requires lane < 16 and
* chunk <= FC_ATTN_KREG_MAX_CHUNK. Bit-equal to fc_kvq4_dequant_elem. */
__device__ __forceinline__ void fc_kvq4_dequant_lane_chunk(
const unsigned char *row, int lane, int chunk, float *kreg) {
#pragma unroll
for (int b = 0; b < FC_ATTN_KREG_MAX_CHUNK / 2; b++) {
if (2 * b < chunk) {
const unsigned char *bp = row + (size_t)b * FC_KVQ4_BLK_BYTES;
const unsigned short d_bits = (unsigned short)(bp[0] | (bp[1] << 8));
const float d = __half2float(__ushort_as_half(d_bits));
const unsigned int codes = bp[2 + lane];
kreg[2 * b] = (float)((int)(codes & 0xF) - 8) * d;
if (2 * b + 1 < chunk)
kreg[2 * b + 1] = (float)((int)((codes >> 4) & 0xF) - 8) * d;
}
}
}
/*
* Causal GQA flash attention over the Q4 GPU KV cache (prefill mode), with the
* K tile held in REGISTERS instead of shared memory. Bit-exact against
* fc_attn_flash_prefill_q4; see the block comment above for the contract, the
* measured factor and the two things that must not be changed.
*
* Args, grid and block are identical to fc_attn_flash_prefill_q4 — only the
* dynamic-shared request differs (fc_attn_flash_kreg_shared_bytes in
* cuda/fc_cuda.c).
*/
__global__ void fc_attn_flash_prefill_q4_kr(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win,
float *out, int kv_groups,
int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink,
int attn_win, float scale) {
const int t = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int seq_len = pos0 + t + 1;
const int cur_pos = seq_len - 1;
const int row_bytes = fc_kvq4_row_bytes(head_dim);
const int win_live_lo = (cur_pos - t + (int)gridDim.x - 1) - FC_KVQ4_WMAX;
/* No K tile: q group + scores + probs + m/l/corr only. MUST match
* fc_attn_flash_kreg_shared_bytes() in cuda/fc_cuda.c. */
extern __shared__ float shk[];
float *q_s = shk;
float *score_s = q_s + (size_t)G * head_dim;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const size_t n_head = (size_t)gridDim.y * G;
const float *q_base = q + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache + (size_t)kvh * max_seq * row_bytes;
const unsigned char *v_head = v_cache + (size_t)kvh * max_seq * row_bytes;
const __half *k_win_head = k_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES; /* key-in-tile this thread scores */
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = 0; j0 < seq_len; j0 += FC_ATTN_TILE_KEYS) {
int tk = seq_len - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
/* Load key ky's own 16 elements straight into registers — the same read
* path fc_attn_flash_prefill_q4's staging loop uses (fp16 shadow for
* sink/recent keys, q4 dequant for the aged-out middle), just assigned
* to the thread that consumes it. No shared write, so no sync. */
float kreg[FC_ATTN_KREG_MAX_CHUNK];
if (ky < tk) {
const int kp = j0 + ky;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++)
if (i < chunk)
kreg[i] = __half2float(s[lane + i * FC_ATTN_DOT_LANES]);
} else if (win > 0 && kp > cur_pos - win && kp > win_live_lo) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++)
if (i < chunk)
kreg[i] = __half2float(s[lane + i * FC_ATTN_DOT_LANES]);
} else {
fc_kvq4_dequant_lane_chunk(k_head + (size_t)kp * row_bytes,
lane, chunk, kreg);
}
}
/* QK^T, i-outer / g-inner. Same accumulation order over e as
* fc_attn_flash_prefill_q4 — the K operand simply comes from a register
* instead of from k_s. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++) {
if (i < chunk) {
const int e = lane + i * FC_ATTN_DOT_LANES;
const float kve = kreg[i];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj;
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
/* PV accumulate — byte-for-byte the shipped loop, INCLUDING the rescale
* fused at j == 0 (see point 2 of the block comment). */
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink)
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
else if (win > 0 && vp > cur_pos - win && vp > win_live_lo)
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
else
vv = fc_kvq4_dequant_elem(v_head + (size_t)vp * row_bytes, tid);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
float *out_base = out + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
out_base[(size_t)g * head_dim + tid] = acc[g] / l_s[g];
}
/* ============================================================
* SHARED-V prefill variant (FC_ATTN_VSTAGE) — BIT-EXACT, DEFAULT OFF
*
* fc_attn_flash_prefill_q4_kr with one thing changed: the key tile's V rows are
* staged into shared memory AS RAW Q4 BYTES once per tile, and the per-key
* dequant reads shared instead of global.
*
* WHY, and it is the mirror image of the register-K finding. Register-K deleted
* the shared K tile because its reuse was 1:1 — one writer, and the only reader
* was the same thread. V's reuse is 32:1 and it had NO tile: a V row is
* row_bytes (144 at the shipped shape), and the block issued 256 threads x 3
* byte loads = 768 load instructions to read it, because all 32 threads of a
* Q4_0 block load the same 2-byte fp16 scale and each nibble byte serves 2
* threads. The staging pattern (none) did not match the consuming pattern
* (block-wide). So the same question that said "delete the tile" for K says
* "add one" for V.
*
* MEASURED, die 6 pinned 666 MHz, noise floor 0.301 % over 10 processes, three
* interleaved reps per arm, all bit-EXACT on 524 288 outputs:
*
* VIND (the address-arithmetic half alone, V still in global) 1.0418x
* THIS (VIND + the shared V tile) 1.1497x
* the shared-V half on its own 1.0959x
*
* DEPTH-FLAT: 1.1453 / 1.1455 / 1.1479 at 8 192 / 14 080 / 32 000 tokens.
*
* It is FREE in residency. Shared goes 9 312 -> 11 616 B and 7 blocks/SMX is
* unchanged because 7 is the REGISTER cap at 72 registers x 256 threads
* (floor(131072/(72*256)) = 7; 80 registers would give 6). Registers stay at
* 72 — COUNTED, and it was this lever's whole risk. Staging is placed in the
* SOFTMAX SHADOW: the online softmax is `if (tid < G)`, so 248 of 256 threads
* are idle there, and the copy is covered by the __syncthreads() that already
* separates the softmax from the PV loop. No barrier is added.
*
* WHY IT PAYS MORE THAN ITS INSTRUCTION COUNT: the per-key 64-bit global address
* chain (29 of the PV loop's 97 executed instructions, 15 of them loop-invariant
* and recomputed all 16 times) collapses to 3 instructions of 32-bit shared
* addressing, and the three LD.E.U8 become ONE LDS.U16 plus one LDS.U8 — in
* shared ptxas CAN prove the 2-byte alignment of the fp16 scale that it cannot
* prove for v_head + vp*row_bytes with a runtime row_bytes, so the PRMT that
* reassembled it disappears. But the counted instruction cut is only 15.7 % and
* the measured gain is 15.0 %, i.e. instruction reduction converts at ~27 % on
* this kernel: the rest is LATENCY. The old form exposed 16 serial global round
* trips per warp per tile at 3 loads in flight; this one exposes 3 at tile entry
* and reads shared thereafter.
*
* THREE THINGS ARE LOAD-BEARING AND MUST NOT BE "SIMPLIFIED":
*
* 1. Everything fc_attn_flash_prefill_q4_kr's own block comment says is still
* load-bearing here, because this kernel is that kernel: the HOISTED K
* dequant keeps registers at 72, and the online-softmax rescale stays FUSED
* in the PV loop at j == 0 where ptxas contracts it to
* acc = fma(acc, corr, p0*v0). Hoisting it is algebraically identical and
* moves 148 099 of 524 288 outputs by up to 4.0e-7 relative.
* 2. The three staging rounds are written OUT rather than looped under a
* #pragma unroll. A rolled runtime-bounded copy is the exact pathology the
* matvec staging gate priced at 1.426x for 1 -> 8 loads in flight, and a
* dropped unroll would compile it silently — nvcc emitted "this pragma must
* immediately precede a statement" for the looped version.
* 3. The dequant is INLINE, not a call to fc_kvq4_dequant_elem. Passing a
* shared pointer into a function taking a generic pointer lets ptxas fall
* back to generic LD instead of LDS. Inline, the SASS is LDS.U16 + LDS.U8 —
* verified in cuobjdump, which is the only reason the latency win exists.
*
* CONTRACT (enforced host-side in fc_attn_flash_prefill_q4_ctx):
* - everything fc_attn_flash_prefill_q4_kr requires, plus
* - row_bytes % 4 == 0, so the tile copies as 32-bit words. True for every
* head_dim that is a multiple of 64; head_dim is pinned to exactly 256 by
* the kr contract, giving row_bytes 144.
* - the staged tile must fit: shared <= 48 KB, checked by the host.
* ============================================================ */
/* Words in the staged V tile at the contracted shape, and the number of
* blockDim-strided rounds needed to copy it. head_dim is pinned to
* FC_ATTN_TILE_KEYS * FC_ATTN_DOT_LANES by the kr contract, and a block has
* head_dim threads, so both are compile-time constants. */
#define FC_ATTN_VS_ROW_BYTES ((FC_ATTN_TILE_KEYS * FC_ATTN_DOT_LANES / \
FC_KVQ4_BLK) * FC_KVQ4_BLK_BYTES)
#define FC_ATTN_VS_TILE_WORDS ((FC_ATTN_TILE_KEYS * FC_ATTN_VS_ROW_BYTES) / 4)
#define FC_ATTN_VS_THREADS (FC_ATTN_TILE_KEYS * FC_ATTN_DOT_LANES)
#define FC_ATTN_VS_ROUNDS ((FC_ATTN_VS_TILE_WORDS + FC_ATTN_VS_THREADS - 1) \
/ FC_ATTN_VS_THREADS)
#if FC_ATTN_VS_ROUNDS != 3
#error "fc_attn_flash_prefill_q4_kvs stages the V tile in exactly 3 rounds"
#endif
#if (FC_ATTN_TILE_KEYS * FC_ATTN_VS_ROW_BYTES) % 4 != 0
#error "the staged V tile must be a whole number of 32-bit words"
#endif
/* Falsification hook for the in-model digest, default OFF. See the use site. */
#ifndef FC_ATTN_VSTAGE_POISON
#define FC_ATTN_VSTAGE_POISON 0
#endif
/*
* Causal GQA flash attention over the Q4 GPU KV cache (prefill mode):
* fc_attn_flash_prefill_q4_kr with each key tile's V rows staged in shared
* memory as raw Q4 bytes and decoded from there. Bit-exact against
* fc_attn_flash_prefill_q4 and _kr; see the block comment above for why and for
* the contract.
*
* Args, grid and block are identical to fc_attn_flash_prefill_q4; dynamic
* shared is fc_attn_flash_kvs_shared_bytes in cuda/fc_cuda.c.
*/
__global__ void fc_attn_flash_prefill_q4_kvs(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win,
float *out, int kv_groups,
int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink,
int attn_win, float scale) {
const int t = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int seq_len = pos0 + t + 1;
const int cur_pos = seq_len - 1;
const int row_bytes = fc_kvq4_row_bytes(head_dim);
/* Ring slots below the batch's highest appended position - WMAX are lapped. */
const int win_live_lo = (cur_pos - t + (int)gridDim.x - 1) - FC_KVQ4_WMAX;
/* kr's layout plus the staged V tile after corr_s. MUST match
* fc_attn_flash_kvs_shared_bytes() in cuda/fc_cuda.c. */
extern __shared__ float shv[];
float *q_s = shv;
float *score_s = q_s + (size_t)G * head_dim;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
unsigned int *v_sh = (unsigned int *)(corr_s + G);
const size_t n_head = (size_t)gridDim.y * G;
const float *q_base = q + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache + (size_t)kvh * max_seq * row_bytes;
const unsigned char *v_head = v_cache + (size_t)kvh * max_seq * row_bytes;
const __half *k_win_head = k_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES; /* key-in-tile this thread scores */
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
/* Where element tid sits inside a staged V row; fixed per thread. */
const int vlocal = tid & (FC_KVQ4_BLK - 1);
const int vhalf = FC_KVQ4_BLK / 2;
const int vbyte = 2 + (vlocal < vhalf ? vlocal : vlocal - vhalf);
const int vshift = vlocal < vhalf ? 0 : 4;
const int vqoff = (tid / FC_KVQ4_BLK) * FC_KVQ4_BLK_BYTES;
for (int j0 = 0; j0 < seq_len; j0 += FC_ATTN_TILE_KEYS) {
int tk = seq_len - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
/* K for key ky straight into registers, as in _kr. */
float kreg[FC_ATTN_KREG_MAX_CHUNK];
if (ky < tk) {
const int kp = j0 + ky;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++)
if (i < chunk)
kreg[i] = __half2float(s[lane + i * FC_ATTN_DOT_LANES]);
} else if (win > 0 && kp > cur_pos - win && kp > win_live_lo) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++)
if (i < chunk)
kreg[i] = __half2float(s[lane + i * FC_ATTN_DOT_LANES]);
} else {
fc_kvq4_dequant_lane_chunk(k_head + (size_t)kp * row_bytes,
lane, chunk, kreg);
}
}
/* QK^T, i-outer / g-inner; every thread joins the shuffles. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
#pragma unroll
for (int i = 0; i < FC_ATTN_KREG_MAX_CHUNK; i++) {
if (i < chunk) {
const int e = lane + i * FC_ATTN_DOT_LANES;
const float kve = kreg[i];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj;
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
/* Stage the tile's V rows as words while the softmax threads work; the
* barrier below covers both. Three rounds written out so all three
* loads are in flight before the first store. Every load is guarded: a
* partial last tile is shorter than one round and must not read past
* the end of v_cache. */
{
const unsigned int *vsrc =
(const unsigned int *)(v_head + (size_t)j0 * row_bytes);
const int vwords = (tk * row_bytes) >> 2;
const int w0_idx = tid;
const int w1_idx = tid + FC_ATTN_VS_THREADS;
const int w2_idx = tid + 2 * FC_ATTN_VS_THREADS;
const unsigned int w0 = w0_idx < vwords ? vsrc[w0_idx] : 0u;
const unsigned int w1 = w1_idx < vwords ? vsrc[w1_idx] : 0u;
const unsigned int w2 = w2_idx < vwords ? vsrc[w2_idx] : 0u;
if (w0_idx < vwords)
v_sh[w0_idx] = w0;
if (w1_idx < vwords)
v_sh[w1_idx] = w1;
if (w2_idx < vwords)
v_sh[w2_idx] = w2;
}
__syncthreads();
/* PV accumulate with the rescale fused at j == 0, as in _kr. */
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink) {
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
} else if (win > 0 && vp > cur_pos - win && vp > win_live_lo) {
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
} else {
/* Inline on purpose: through a generic-pointer helper ptxas
* loses the shared address space and emits generic loads. */
const unsigned char *bp =
(const unsigned char *)v_sh + j * row_bytes + vqoff;
const unsigned short d_bits = (unsigned short)(bp[0] | (bp[1] << 8));
const float d = __half2float(__ushort_as_half(d_bits));
const int code = (bp[vbyte] >> vshift) & 0xF;
vv = (float)(code - 8) * d;
#if FC_ATTN_VSTAGE_POISON
vv = __fmul_rn(vv, 1.00000011920928955078125f);
#endif
}
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
float *out_base = out + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
out_base[(size_t)g * head_dim + tid] = acc[g] / l_s[g];
}
/* ============================================================
* NARROW-STAGING prefill variants (FC_ATTN_KTILE_FP16) — LOSSY, DEFAULT OFF
*
* Same tiling, same online softmax and the same q4 KV cache as
* fc_attn_flash_prefill_q4. ONE thing changes now:
*
* 1. The i-outer QK dot with eight REGISTER partials — which used to be a
* difference here — is on the DEFAULT path since the lossless landing, so
* all six shipped flash kernels carry it. It is what makes (2) affordable
* (a narrow tile costs one widening per K LOAD, and this cuts loads 8x),
* but it is no longer part of what this flag turns on.
* 2. The staged K tile (and, at level 2, the staged q group) is __half
* instead of float. That is the whole point: the K tile is 17 408 B of
* the 26 720 B footprint, and residency on this kernel is capped by
* SHARED MEMORY, never by registers (72 regs would allow 7 blocks/SMX).
* 26 720 -> 18 016 B lifts residency 4 -> 6 blocks/SMX; also narrowing q
* gives 13 920 B and 7 blocks/SMX.
*
* NUMERICS — this is NOT bit-exact, which is why it is behind a flag:
* - sink and recent-window keys are ALREADY fp16 in memory (k_sink/k_win);
* staging them as __half is bit-exact for those keys and today's widening
* to float is pure waste.
* - aged-out keys are q4-sourced (fp16 scale x a 4-bit level), so rounding
* them to fp16 costs <= 2^-11 relative — orders below the q4 error the
* cache already carries.
* - level 2 additionally rounds the q group, a live activation. Strictly
* lossier and gated separately.
* Measured max relative output error at the shipped shape: 1.3e-5 (level 1),
* 1.7e-5 (level 2). See Pantheon results/prefill_attention_ktile_fp16.md.
*
* TODO: attention.cu is past the 1 000-line guidance and now carries THREE
* prefill families (exact, narrow-staged, register-K) plus the whole decode
* family. Splitting it into a prefill file and a decode file is still the right
* move, and it needs a SASS-identity check as its gate: extracting code out of a
* flash body is NOT free here even behind __forceinline__. Pulling the
* register-K path's three-way K load into one helper was tried and measured to
* CHANGE that kernel's SASS (same 72 registers, same 0-byte stack frame,
* different instruction stream), so it was reverted rather than shipped on top
* of a measurement it would have invalidated. Same reason the register-K body is
* a copy of fc_attn_flash_prefill_q4's rather than a shared template.
* See Pantheon results/prefill_attention_register_ktile.md §8 gap 14.
* ============================================================ */
/* K-tile row stride in ELEMENTS is FC_ATTN_K_STRIDE(head_dim) for both the
* float and the __half tile: two __halfs inside one 32-bit word are
* conflict-free on sm_37 (the sub-word rule), and the +16-element pad still
* puts the two half-warps of a warp on disjoint bank ranges.
*
* The shared body is a TEMPLATE, which cannot carry C linkage, so the file's
* enclosing `extern "C"` block closes here and reopens for the two entry
* points — whose PTX names must stay unmangled, since cuModuleGetFunction
* looks them up by string. */
} /* extern "C" */
template <bool Q_FP16>
__device__ __forceinline__ void fc_attn_flash_prefill_q4_narrow(
const float *q, const unsigned char *k_cache,
const unsigned char *v_cache, const __half *k_sink,
const __half *v_sink, const __half *k_win, const __half *v_win,
float *out, int kv_groups, int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink, int attn_win, float scale)
{
const int t = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int seq_len = pos0 + t + 1;
const int cur_pos = seq_len - 1;
const int row_bytes = fc_kvq4_row_bytes(head_dim);
const int win_live_lo = (cur_pos - t + (int)gridDim.x - 1) - FC_KVQ4_WMAX;
/* Byte-addressed shared layout so the fp32 softmax state stays 4-byte
* aligned whichever tiles are narrowed. MUST match
* fc_attn_flash_narrow_shared_bytes() in cuda/fc_cuda.c. */
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ char shb[];
const int q_bytes = (int)(G * head_dim *
(Q_FP16 ? sizeof(__half) : sizeof(float)));
const int k_bytes = (int)(FC_ATTN_TILE_KEYS * kstride * sizeof(__half));
float *q_f = (float *)shb;
__half *q_h = (__half *)shb;
__half *k_h = (__half *)(shb + q_bytes);
float *score_s = (float *)(shb + q_bytes + k_bytes);
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const size_t n_head = (size_t)gridDim.y * G;
const float *q_base = q + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
for (int g = 0; g < G; g++) {
float qv = q_base[(size_t)g * head_dim + tid];
if (Q_FP16) q_h[(size_t)g * head_dim + tid] = __float2half(qv);
else q_f[(size_t)g * head_dim + tid] = qv;
}
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache + (size_t)kvh * max_seq * row_bytes;
const unsigned char *v_head = v_cache + (size_t)kvh * max_seq * row_bytes;
const __half *k_win_head = k_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = 0; j0 < seq_len; j0 += FC_ATTN_TILE_KEYS) {
int tk = seq_len - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
/* Stage the K tile as __half. The fp16 shadow rows copy across with no
* conversion at all (they were fp16 to begin with); only the q4 body
* rounds. */
for (int key = 0; key < tk; key++) {
int kp = j0 + key;
const size_t base = (size_t)key * kstride;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
k_h[base + e] = s[e];
} else if (win > 0 && kp > cur_pos - win && kp > win_live_lo) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
k_h[base + e] = s[e];
} else {
const unsigned char *s = k_head + (size_t)kp * row_bytes;
for (int e = tid; e < head_dim; e += blockDim.x)
k_h[base + e] = __float2half(fc_kvq4_dequant_elem(s, e));
}
}
__syncthreads();
/* QK^T, i-outer: one K load (and one widening) per element per tile,
* eight register partials across the GQA group. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const size_t kb = (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kv = __half2float(k_h[kb + e]);
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G) {
size_t qi = (size_t)g * head_dim + e;
float qv = Q_FP16 ? __half2float(q_h[qi]) : q_f[qi];
part[g] += qv * kv;
}
}
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj;
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
/* PV accumulate — unchanged from fc_attn_flash_prefill_q4: V streams
* through __ldg and is never staged, so it is untouched by all this. */
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink)
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
else if (win > 0 && vp > cur_pos - win && vp > win_live_lo)
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
else
vv = fc_kvq4_dequant_elem(v_head + (size_t)vp * row_bytes, tid);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
float *out_base = out + ((size_t)t * n_head + (size_t)kvh * G) * head_dim;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
out_base[(size_t)g * head_dim + tid] = acc[g] / l_s[g];
}
extern "C" {
/* FC_ATTN_KTILE_FP16=1: __half K tile, fp32 q group. 18 016 B at the shipped
* shape -> 6 blocks/SMX. */
__global__ void fc_attn_flash_prefill_q4_kh(
const float *q, const unsigned char *k_cache,
const unsigned char *v_cache, const __half *k_sink,
const __half *v_sink, const __half *k_win, const __half *v_win,
float *out, int kv_groups, int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink, int attn_win, float scale)
{
fc_attn_flash_prefill_q4_narrow<false>(
q, k_cache, v_cache, k_sink, v_sink, k_win, v_win, out, kv_groups,
head_dim, max_seq, pos0, sink, win, attn_sink, attn_win, scale);
}
/* FC_ATTN_KTILE_FP16=2: __half K tile AND __half q group. 13 920 B ->
* 7 blocks/SMX. Strictly lossier than level 1 — q is an activation. */
__global__ void fc_attn_flash_prefill_q4_kqh(
const float *q, const unsigned char *k_cache,
const unsigned char *v_cache, const __half *k_sink,
const __half *v_sink, const __half *k_win, const __half *v_win,
float *out, int kv_groups, int head_dim, int max_seq, int pos0,
int sink, int win, int attn_sink, int attn_win, float scale)
{
fc_attn_flash_prefill_q4_narrow<true>(
q, k_cache, v_cache, k_sink, v_sink, k_win, v_win, out, kv_groups,
head_dim, max_seq, pos0, sink, win, attn_sink, attn_win, scale);
}
/* ============================================================
* Decode-mode flash attention (QWEN35_FLASH_DECODE)
*
* A single query (the new token at absolute position seq_len-1) attends the
* full cached context [0, seq_len). Reusing fc_attn_flash_prefill with npos=1
* launches only n_head_kv blocks (=2 for this model) — the K80's ~13 SMX sit
* idle. So the KV sequence is SPLIT across blocks: grid (n_splits, n_head_kv),
* each block runs the same tiled online-softmax over its key slice for the
* whole GQA group and writes a PARTIAL {running max m, running sum l,
* UNnormalized accumulator acc} to scratch. fc_attn_flash_decode_combine then
* merges the n_splits partials per head with the standard max-rescale, which
* is algebraically identical to a single-pass softmax (only ~1e-6 tile drift,
* same class as prefill; gate: argmax-exact vs the host decode path).
*
* Partial scratch layout, indexed by GLOBAL head qh = kvh*G + g:
* m_out[qh*n_splits + split], l_out[qh*n_splits + split]
* acc_out[(qh*n_splits + split)*head_dim + dim]
* ============================================================ */
/* Split plan from the context length: ~kps keys per split, capped at `cap`
* splits, split_len rounded up to the tile width. MUST stay formula-identical
* to the host plan in qwen35moe_forward_cuda.c gpu_attn_block_decode_flash —
* the *_dev kernels below recompute the plan on-device (from a device-resident
* position) so a captured CUDA graph replays correctly at every position, and
* the recomputed plan must equal the host plan for the same seq_len or the
* graphed and direct paths stop being bit-identical. */
__device__ __forceinline__ int fc_attn_dec_plan(int seq_len, int kps, int cap,
int *split_len_out) {
int n_splits = (seq_len + kps - 1) / kps;
if (n_splits < 1)
n_splits = 1;
if (n_splits > cap)
n_splits = cap;
int split_len = (seq_len + n_splits - 1) / n_splits;
split_len = (split_len + 15) & ~15;
*split_len_out = split_len;
return (seq_len + split_len - 1) / split_len;
}
/* Merge the n_splits per-head partials into the final attention output.
* One block per head; thread tid owns output dim tid. Shared verbatim by the
* host-arg and device-pos combine kernels. */
__device__ __forceinline__ void fc_attn_flash_decode_combine_body(
const float *m_in, const float *l_in, const float *acc_in, float *out,
int head_dim, int n_splits) {
const int head = blockIdx.x;
const int tid = threadIdx.x;
if (tid >= head_dim) return;
const float *m_h = m_in + (size_t)head * n_splits;
const float *l_h = l_in + (size_t)head * n_splits;
const float *acc_h = acc_in + (size_t)head * n_splits * head_dim;
float M = -__int_as_float(0x7f800000);
for (int s = 0; s < n_splits; s++)
M = fmaxf(M, m_h[s]);
float o = 0.0f, L = 0.0f;
for (int s = 0; s < n_splits; s++) {
float w = expf(m_h[s] - M); /* empty split: m=-inf -> w=0 */
L += l_h[s] * w;
o += acc_h[(size_t)s * head_dim + tid] * w;
}
out[(size_t)head * head_dim + tid] = o / L;
}
/* Merge the n_splits per-head partials into the final attention output.
* Grid (n_head), block head_dim threads (thread tid owns output dim tid). */
__global__ void fc_attn_flash_decode_combine(const float *m_in,
const float *l_in,
const float *acc_in, float *out,
int head_dim, int n_splits) {
fc_attn_flash_decode_combine_body(m_in, l_in, acc_in, out, head_dim,
n_splits);
}
/* Device-position combine (CUDA-graph capture): seq_len = pos[0] + 1 is read
* from device memory and the split plan recomputed with fc_attn_dec_plan, so
* the launch itself is position-invariant. Bit-exact to
* fc_attn_flash_decode_combine at the same position. */
__global__ void fc_attn_flash_decode_combine_dev(const float *m_in,
const float *l_in,
const float *acc_in,
float *out, int head_dim,
const int *pos, int kps,
int splits_cap) {
int seq_len = pos[0] + 1;
int split_len;
int n_splits = fc_attn_dec_plan(seq_len, kps, splits_cap, &split_len);
fc_attn_flash_decode_combine_body(m_in, l_in, acc_in, out, head_dim,
n_splits);
}
/*
* Decode split-attention over the fp32 GPU KV cache. Same online-softmax math
* and tiling as fc_attn_flash_prefill, but the query is the single decode
* position and the key loop is restricted to this block's split slice
* [split*split_len, min(seq_len, (split+1)*split_len)). Writes the split's
* partial {m, l, acc} for every head in the group (see the section header).
*
* q: [n_head][head_dim] fp32 (post QK-norm + RoPE), single position.
* Grid (n_splits, n_head_kv), block head_dim threads (multiple of 16),
* dynamic shared identical to fc_attn_flash_prefill.
*/
__device__ __forceinline__ void fc_attn_flash_decode_body(
const float *q, const float *k_cache, const float *v_cache,
float *m_out, float *l_out, float *acc_out, int kv_groups,
int head_dim, int max_seq, int seq_len, int n_splits, int split_len,
int attn_sink, int attn_win, float scale) {
const int split = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int cur_pos = seq_len - 1; /* decode query's absolute position */
int k_lo = split * split_len;
int k_hi = k_lo + split_len;
if (k_hi > seq_len) k_hi = seq_len;
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh;
float *k_s = q_s + (size_t)G * head_dim;
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const float *q_base = q + (size_t)kvh * G * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const float *k_head = k_cache + (size_t)kvh * max_seq * head_dim;
const float *v_head = v_cache + (size_t)kvh * max_seq * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = k_lo; j0 < k_hi; j0 += FC_ATTN_TILE_KEYS) {
int tk = k_hi - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window. */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
for (int key = 0; key < tk; key++)
for (int e = tid; e < head_dim; e += blockDim.x)
k_s[(size_t)key * kstride + e] =
k_head[(size_t)(j0 + key) * head_dim + e];
__syncthreads();
/* QK^T, i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
for (int j = 0; j < tk; j++) {
float vv = __ldg(&v_head[(size_t)(j0 + j) * head_dim + tid]);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
size_t po = ((size_t)kvh * G + g) * n_splits + split;
acc_out[po * head_dim + tid] = acc[g];
if (tid == 0) {
m_out[po] = m_s[g];
l_out[po] = l_s[g];
}
}
}
}
__global__ void fc_attn_flash_decode(const float *q, const float *k_cache,
const float *v_cache, float *m_out,
float *l_out, float *acc_out,
int kv_groups, int head_dim, int max_seq,
int seq_len, int n_splits, int split_len,
int attn_sink, int attn_win, float scale) {
fc_attn_flash_decode_body(q, k_cache, v_cache, m_out, l_out, acc_out,
kv_groups, head_dim, max_seq, seq_len, n_splits,
split_len, attn_sink, attn_win, scale);
}
/*
* Device-position decode split-attention (CUDA-graph capture): the query's
* position is read from device memory (pos[0], seq_len = pos[0] + 1) and the
* split plan recomputed on-device, so every launch argument and the grid are
* position-invariant — capture once, replay every token. Launched with a
* FIXED grid (splits_cap, n_head_kv); blocks past the live plan exit before
* touching scratch (the partial layout is strided by the LIVE n_splits, same
* as the host-arg kernel, so combine sees an identical layout). Bit-exact to
* fc_attn_flash_decode at the same position. Caller guarantees no ring wrap
* (pos + 1 <= max_seq).
*/
__global__ void fc_attn_flash_decode_dev(const float *q, const float *k_cache,
const float *v_cache, float *m_out,
float *l_out, float *acc_out,
int kv_groups, int head_dim,
int max_seq, const int *pos, int kps,
int splits_cap, int attn_sink,
int attn_win, float scale) {
int seq_len = pos[0] + 1;
int split_len;
int n_splits = fc_attn_dec_plan(seq_len, kps, splits_cap, &split_len);
if ((int)blockIdx.x >= n_splits)
return;
fc_attn_flash_decode_body(q, k_cache, v_cache, m_out, l_out, acc_out,
kv_groups, head_dim, max_seq, seq_len, n_splits,
split_len, attn_sink, attn_win, scale);
}
/*
* Decode split-attention over the Q4 GPU KV cache. Identical to
* fc_attn_flash_decode but K/V are read from the packed Q4_0 cache (with the
* fp16 sink / recent-window shadow) and dequantized in-kernel, matching
* fc_attn_flash_prefill_q4. Writes the same partial layout; the combine
* kernel is shared.
*/
__device__ __forceinline__ void fc_attn_flash_decode_q4_body(
const float *q, const unsigned char *k_cache,
const unsigned char *v_cache, const __half *k_sink,
const __half *v_sink, const __half *k_win, const __half *v_win,
float *m_out, float *l_out, float *acc_out, int kv_groups,
int head_dim, int max_seq, int seq_len, int n_splits, int split_len,
int sink, int win, int attn_sink, int attn_win, float scale) {
const int split = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int G = kv_groups;
const int cur_pos = seq_len - 1;
const int row_bytes = fc_kvq4_row_bytes(head_dim);
int k_lo = split * split_len;
int k_hi = k_lo + split_len;
if (k_hi > seq_len) k_hi = seq_len;
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh;
float *k_s = q_s + (size_t)G * head_dim;
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const float *q_base = q + (size_t)kvh * G * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache + (size_t)kvh * max_seq * row_bytes;
const unsigned char *v_head = v_cache + (size_t)kvh * max_seq * row_bytes;
const __half *k_win_head = k_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win + (size_t)kvh * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink + (size_t)kvh * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = k_lo; j0 < k_hi; j0 += FC_ATTN_TILE_KEYS) {
int tk = k_hi - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window. */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
for (int key = 0; key < tk; key++) {
int kp = j0 + key;
float *krow = k_s + (size_t)key * kstride;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else if (win > 0 && kp > cur_pos - win) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else {
const unsigned char *s = k_head + (size_t)kp * row_bytes;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = fc_kvq4_dequant_elem(s, e);
}
}
__syncthreads();
/* QK^T, i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink)
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
else if (win > 0 && vp > cur_pos - win)
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
else
vv = fc_kvq4_dequant_elem(v_head + (size_t)vp * row_bytes, tid);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
size_t po = ((size_t)kvh * G + g) * n_splits + split;
acc_out[po * head_dim + tid] = acc[g];
if (tid == 0) {
m_out[po] = m_s[g];
l_out[po] = l_s[g];
}
}
}
}
__global__ void fc_attn_flash_decode_q4(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win, float *m_out,
float *l_out, float *acc_out,
int kv_groups, int head_dim,
int max_seq, int seq_len, int n_splits,
int split_len, int sink, int win,
int attn_sink, int attn_win,
float scale) {
fc_attn_flash_decode_q4_body(q, k_cache, v_cache, k_sink, v_sink, k_win,
v_win, m_out, l_out, acc_out, kv_groups,
head_dim, max_seq, seq_len, n_splits,
split_len, sink, win, attn_sink, attn_win,
scale);
}
/* Device-position Q4 decode split-attention (CUDA-graph capture). Same
* contract as fc_attn_flash_decode_dev; bit-exact to fc_attn_flash_decode_q4
* at the same position. */
__global__ void fc_attn_flash_decode_q4_dev(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win, float *m_out,
float *l_out, float *acc_out,
int kv_groups, int head_dim,
int max_seq, const int *pos,
int kps, int splits_cap, int sink,
int win, int attn_sink,
int attn_win, float scale) {
int seq_len = pos[0] + 1;
int split_len;
int n_splits = fc_attn_dec_plan(seq_len, kps, splits_cap, &split_len);
if ((int)blockIdx.x >= n_splits)
return;
fc_attn_flash_decode_q4_body(q, k_cache, v_cache, k_sink, v_sink, k_win,
v_win, m_out, l_out, acc_out, kv_groups,
head_dim, max_seq, seq_len, n_splits,
split_len, sink, win, attn_sink, attn_win,
scale);
}
/*
* Device-position single-row KV append (CUDA-graph capture): the decode
* token's ring slot is computed from device-resident pos[0] instead of a host
* scalar, so the recorded launch replays correctly at every position. Grid
* (1, n_head_kv) — decode appends exactly one row. Bit-exact to
* fc_attn_kv_append with pos0 = pos[0], npos = 1.
*/
__global__ void fc_attn_kv_append_dev(const float *k_new, const float *v_new,
float *k_cache, float *v_cache,
int n_head_kv, int head_dim, int max_seq,
const int *pos) {
int t = blockIdx.x;
int h = blockIdx.y;
int dd = threadIdx.x;
if (dd >= head_dim) return;
int cache_pos = (pos[0] + t) % max_seq;
size_t src = ((size_t)t * n_head_kv + h) * head_dim + dd;
size_t dst = ((size_t)h * max_seq + cache_pos) * head_dim + dd;
k_cache[dst] = k_new[src];
v_cache[dst] = v_new[src];
}
/* ============================================================
* Multi-user decode attention (QWEN35_FLASH_SLOTS)
*
* The batched-decode "slots" path: K independent users, each decoding ONE
* token at its OWN position against its OWN KV ring. The slot kernels below
* move the whole per-step attention on-die — replacing the host per-user
* softmax loop and its 2 syncs + PCIe round-trips per layer (the K=8
* steady-state ceiling, multi-tenant-speed-opt.md §3.2).
*
* KV layout: one ring PER SLOT, each slice byte-identical to the single-user
* ring — [slot][n_head_kv][max_seq][row] (row = head_dim fp32 or Q4_0
* bytes + per-slot fp16 sink/window shadows).
*
* Split plan: every user recomputes its OWN split plan from its position
* (fc_attn_dec_plan(seq_len_u), the exact single-user formula), so each
* user's per-split (m, l, acc) partials are bit-identical to the
* single-token decode kernel at the same position. The launch is uniform
* (grid.x = the max plan across users; partial stride = gridDim.x); blocks
* past a user's own plan write EMPTY partials (m=-inf, l=0, acc=0), which
* the combine's max-rescale weights to exactly zero — merging them adds
* only +0.0f terms, so the merged output is bit-identical to combining the
* live splits alone. Combine reuses fc_attn_flash_decode_combine with
* gridDim.x = nslots * n_head (partials/out are (slot, head)-major).
* ============================================================ */
/*
* Append each slot's single roped K/V row to its own fp32 ring at its own
* position (pos[slot], device int32 array — the RoPE position buffer).
*
* k_new, v_new: [nslots][n_head_kv][head_dim] fp32
* k_cache, v_cache: [nslots][n_head_kv][max_seq][head_dim] fp32 (written)
*
* Grid (nslots, n_head_kv), block head_dim threads. Row lands at
* pos[slot] % max_seq (the host ring rule).
*/
__global__ void fc_attn_kv_append_slots(const float *k_new, const float *v_new,
float *k_cache, float *v_cache,
int n_head_kv, int head_dim,
int max_seq, const int *pos) {
int u = blockIdx.x; /* slot */
int h = blockIdx.y;
int dd = threadIdx.x;
if (dd >= head_dim) return;
int cache_pos = pos[u] % max_seq;
size_t src = ((size_t)u * n_head_kv + h) * head_dim + dd;
size_t dst = (((size_t)u * n_head_kv + h) * max_seq + cache_pos) *
head_dim + dd;
k_cache[dst] = k_new[src];
v_cache[dst] = v_new[src];
}
/* Per-slot Q4 ring append (quantizing in-kernel) + per-slot fp16 sink /
* recent-window shadow mirror. Same contract as fc_attn_kv_append_q4 per
* slot; caller guarantees no ring wrap (pos[u] + 1 <= max_seq).
*
* k_cache/v_cache: [nslots][n_head_kv][max_seq][row_bytes] Q4_0 bytes
* k_sink/v_sink: [nslots][n_head_kv][FC_KVQ4_SINKMAX][head_dim] fp16
* k_win/v_win: [nslots][n_head_kv][FC_KVQ4_WMAX][head_dim] fp16
*
* Grid (nslots, n_head_kv), block head_dim threads (multiple of 32). */
__global__ void fc_attn_kv_append_slots_q4(const float *k_new,
const float *v_new,
unsigned char *k_cache,
unsigned char *v_cache,
__half *k_sink, __half *v_sink,
__half *k_win, __half *v_win,
int n_head_kv, int head_dim,
int max_seq, const int *pos,
int sink, int win) {
int u = blockIdx.x; /* slot */
int h = blockIdx.y;
int tid = threadIdx.x;
if (tid >= head_dim) return;
int cache_pos = pos[u];
int row_bytes = fc_kvq4_row_bytes(head_dim);
size_t src = ((size_t)u * n_head_kv + h) * head_dim + tid;
size_t dst_row = (((size_t)u * n_head_kv + h) * max_seq + cache_pos) *
row_bytes;
int blk = tid / FC_KVQ4_BLK;
int lane = tid & (FC_KVQ4_BLK - 1);
/* Shared staging for the shuffle-free block quantizer (Kepler-safe). */
__shared__ float s_k[FC_KVQ4_MAX_HEAD_DIM];
__shared__ float s_v[FC_KVQ4_MAX_HEAD_DIM];
float kf = k_new[src];
float vf = v_new[src];
s_k[tid] = kf;
s_v[tid] = vf;
__syncthreads();
fc_kvq4_quant_block(s_k + (size_t)blk * FC_KVQ4_BLK,
k_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
fc_kvq4_quant_block(s_v + (size_t)blk * FC_KVQ4_BLK,
v_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
/* Mirror the row into this slot's fp16 sink / recent-window shadow. */
if (win > 0) {
size_t wi = (((size_t)u * n_head_kv + h) * FC_KVQ4_WMAX +
(cache_pos % FC_KVQ4_WMAX)) * head_dim + tid;
k_win[wi] = __float2half(kf);
v_win[wi] = __float2half(vf);
}
if (cache_pos < sink) {
size_t si = (((size_t)u * n_head_kv + h) * FC_KVQ4_SINKMAX +
cache_pos) * head_dim + tid;
k_sink[si] = __float2half(kf);
v_sink[si] = __float2half(vf);
}
}
/*
* Multi-user decode split-attention over the fp32 per-slot KV rings. Each
* block runs one (split, kv-head, slot) slice of one user's decode attention
* with the user's OWN split plan recomputed on-device from pos[slot] — the
* per-split math is bit-identical to fc_attn_flash_decode at that position.
*
* q: [nslots][n_head][head_dim] fp32 (post QK-norm + RoPE), one row/slot.
* Partials indexed by GLOBAL head gh = slot*n_head + kvh*G + g, strided by
* the LAUNCH width gridDim.x:
* m_out[gh*gridDim.x + split], l_out[...], acc_out[(...)*head_dim + dim]
* Off-plan blocks (split >= this user's n_splits) write EMPTY partials.
*
* Grid (launch_splits, n_head_kv, nslots), block head_dim threads, dynamic
* shared identical to fc_attn_flash_prefill. Caller guarantees no ring wrap
* (pos[u] + 1 <= max_seq for every slot).
*/
__global__ void fc_attn_flash_decode_slots(const float *q,
const float *k_cache,
const float *v_cache, float *m_out,
float *l_out, float *acc_out,
int kv_groups, int head_dim,
int max_seq, const int *pos,
int kps, int splits_cap,
int attn_sink, int attn_win,
float scale) {
const int split = blockIdx.x;
const int kvh = blockIdx.y;
const int slot = blockIdx.z;
const int tid = threadIdx.x;
const int G = kv_groups;
const int n_head_kv = gridDim.y;
const int stride = gridDim.x; /* uniform partial stride */
const int seq_len = pos[slot] + 1;
const int cur_pos = seq_len - 1;
/* This user's own split plan (the single-user formula). */
int split_len;
int n_splits = fc_attn_dec_plan(seq_len, kps, splits_cap, &split_len);
/* Partial base: global head (slot, kvh, g=0), strided by the launch. */
const size_t head0 = ((size_t)slot * n_head_kv + kvh) * G;
if (split >= n_splits) {
/* Off-plan for this (shorter-context) user: EMPTY partial. The
* combine weights m=-inf splits to exactly 0 (see section header). */
for (int g = 0; g < G; g++) {
size_t po = (head0 + g) * stride + split;
acc_out[po * head_dim + tid] = 0.0f;
if (tid == 0) {
m_out[po] = -__int_as_float(0x7f800000);
l_out[po] = 0.0f;
}
}
return;
}
int k_lo = split * split_len;
int k_hi = k_lo + split_len;
if (k_hi > seq_len) k_hi = seq_len;
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh;
float *k_s = q_s + (size_t)G * head_dim;
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const float *q_base = q + head0 * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const float *k_head = k_cache +
((size_t)slot * n_head_kv + kvh) * max_seq * head_dim;
const float *v_head = v_cache +
((size_t)slot * n_head_kv + kvh) * max_seq * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = k_lo; j0 < k_hi; j0 += FC_ATTN_TILE_KEYS) {
int tk = k_hi - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window. */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
for (int key = 0; key < tk; key++)
for (int e = tid; e < head_dim; e += blockDim.x)
k_s[(size_t)key * kstride + e] =
k_head[(size_t)(j0 + key) * head_dim + e];
__syncthreads();
/* QK^T, i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
for (int j = 0; j < tk; j++) {
float vv = __ldg(&v_head[(size_t)(j0 + j) * head_dim + tid]);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
size_t po = (head0 + g) * stride + split;
acc_out[po * head_dim + tid] = acc[g];
if (tid == 0) {
m_out[po] = m_s[g];
l_out[po] = l_s[g];
}
}
}
}
/*
* Multi-user decode split-attention over the per-slot Q4 KV rings (with the
* per-slot fp16 sink / recent-window shadows). Identical plan/partial scheme
* to fc_attn_flash_decode_slots; K/V dequantized in-kernel like
* fc_attn_flash_decode_q4.
*/
__global__ void fc_attn_flash_decode_slots_q4(const float *q,
const unsigned char *k_cache,
const unsigned char *v_cache,
const __half *k_sink,
const __half *v_sink,
const __half *k_win,
const __half *v_win,
float *m_out, float *l_out,
float *acc_out, int kv_groups,
int head_dim, int max_seq,
const int *pos, int kps,
int splits_cap, int sink,
int win, int attn_sink,
int attn_win, float scale) {
const int split = blockIdx.x;
const int kvh = blockIdx.y;
const int slot = blockIdx.z;
const int tid = threadIdx.x;
const int G = kv_groups;
const int n_head_kv = gridDim.y;
const int stride = gridDim.x;
const int seq_len = pos[slot] + 1;
const int cur_pos = seq_len - 1;
const int row_bytes = fc_kvq4_row_bytes(head_dim);
int split_len;
int n_splits = fc_attn_dec_plan(seq_len, kps, splits_cap, &split_len);
const size_t head0 = ((size_t)slot * n_head_kv + kvh) * G;
if (split >= n_splits) {
for (int g = 0; g < G; g++) {
size_t po = (head0 + g) * stride + split;
acc_out[po * head_dim + tid] = 0.0f;
if (tid == 0) {
m_out[po] = -__int_as_float(0x7f800000);
l_out[po] = 0.0f;
}
}
return;
}
int k_lo = split * split_len;
int k_hi = k_lo + split_len;
if (k_hi > seq_len) k_hi = seq_len;
const int kstride = FC_ATTN_K_STRIDE(head_dim);
extern __shared__ float sh[];
float *q_s = sh;
float *k_s = q_s + (size_t)G * head_dim;
float *score_s = k_s + (size_t)FC_ATTN_TILE_KEYS * kstride;
float *p_s = score_s + G * FC_ATTN_TILE_KEYS;
float *m_s = p_s + G * FC_ATTN_TILE_KEYS;
float *l_s = m_s + G;
float *corr_s = l_s + G;
const float *q_base = q + head0 * head_dim;
for (int g = 0; g < G; g++)
q_s[(size_t)g * head_dim + tid] = q_base[(size_t)g * head_dim + tid];
if (tid < G) {
m_s[tid] = -__int_as_float(0x7f800000);
l_s[tid] = 0.0f;
}
__syncthreads();
/* Compile-time width, so ptxas keeps acc[] in REGISTERS. Every loop that
* indexes acc[] is unrolled to FC_ATTN_MAX_GROUPS with an `if (g < G)`
* guard; a loop bounded by the RUNTIME kv_groups cannot be unrolled, the
* array lands in LOCAL memory, and the PV inner loop then pays an
* LDL.128/STL.128 round trip per key. ptxas reports that as a 32-byte
* "stack frame" and prints "0 bytes spill stores" right beside it --
* a never-promoted local array is not counted as a spill. */
float acc[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
acc[g] = 0.0f;
const unsigned char *k_head = k_cache +
((size_t)slot * n_head_kv + kvh) * max_seq * row_bytes;
const unsigned char *v_head = v_cache +
((size_t)slot * n_head_kv + kvh) * max_seq * row_bytes;
const __half *k_win_head = k_win +
((size_t)slot * n_head_kv + kvh) * FC_KVQ4_WMAX * head_dim;
const __half *v_win_head = v_win +
((size_t)slot * n_head_kv + kvh) * FC_KVQ4_WMAX * head_dim;
const __half *k_sink_head = k_sink +
((size_t)slot * n_head_kv + kvh) * FC_KVQ4_SINKMAX * head_dim;
const __half *v_sink_head = v_sink +
((size_t)slot * n_head_kv + kvh) * FC_KVQ4_SINKMAX * head_dim;
const int ky = tid / FC_ATTN_DOT_LANES;
const int lane = tid % FC_ATTN_DOT_LANES;
const int chunk = head_dim / FC_ATTN_DOT_LANES;
for (int j0 = k_lo; j0 < k_hi; j0 += FC_ATTN_TILE_KEYS) {
int tk = k_hi - j0;
if (tk > FC_ATTN_TILE_KEYS) tk = FC_ATTN_TILE_KEYS;
/* Sliding window: skip tiles outside sink U recent-window. */
if (!fc_attn_tile_in_window(j0, tk, cur_pos, attn_sink, attn_win))
continue;
for (int key = 0; key < tk; key++) {
int kp = j0 + key;
float *krow = k_s + (size_t)key * kstride;
if (kp < sink) {
const __half *s = k_sink_head + (size_t)kp * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else if (win > 0 && kp > cur_pos - win) {
const __half *s = k_win_head +
(size_t)(kp % FC_KVQ4_WMAX) * head_dim;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = __half2float(s[e]);
} else {
const unsigned char *s = k_head + (size_t)kp * row_bytes;
for (int e = tid; e < head_dim; e += blockDim.x)
krow[e] = fc_kvq4_dequant_elem(s, e);
}
}
__syncthreads();
/* QK^T, i-outer / g-inner: each thread reads its chunk of the staged K
* row ONCE per tile into a register and fans it across the whole GQA
* group, instead of re-reading k_s once per group (G=8 -> 1 LDS per
* element). BIT-EXACT: every part[g] still accumulates over e in the
* same ascending order the g-outer form used, so this is pure operand
* motion. part[] is indexed only by the unrolled compile-time g, for
* the same reason acc[] is -- a runtime-bounded loop would push it to
* LOCAL memory and cost an LDL/STL round trip per element. */
{
float part[FC_ATTN_MAX_GROUPS];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
part[g] = 0.0f;
if (ky < tk) {
const float *kv = k_s + (size_t)ky * kstride;
for (int i = 0; i < chunk; i++) {
int e = lane + i * FC_ATTN_DOT_LANES;
float kve = kv[e];
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++)
if (g < G)
part[g] += q_s[(size_t)g * head_dim + e] * kve;
}
}
/* G is uniform across the block, so the guarded unroll keeps the
* shuffle warp-uniform exactly as the runtime-bounded loop did. */
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
float p = part[g];
for (int off = FC_ATTN_DOT_LANES / 2; off > 0; off >>= 1)
p += __shfl_down_sync(0xffffffffu, p, (unsigned)off,
FC_ATTN_DOT_LANES);
if (lane == 0 && ky < tk)
score_s[g * FC_ATTN_TILE_KEYS + ky] = p * scale;
}
}
}
__syncthreads();
if (tid < G) {
const float *sc = score_s + tid * FC_ATTN_TILE_KEYS;
float m_new = m_s[tid];
for (int j = 0; j < tk; j++)
m_new = fmaxf(m_new, sc[j]);
float corr = expf(m_s[tid] - m_new);
float sum = 0.0f;
for (int j = 0; j < tk; j++) {
float pj = expf(sc[j] - m_new);
p_s[(size_t)j * G + tid] = pj; /* [TILE][G]: PV reads a run */
sum += pj;
}
l_s[tid] = l_s[tid] * corr + sum;
m_s[tid] = m_new;
corr_s[tid] = corr;
}
__syncthreads();
for (int j = 0; j < tk; j++) {
int vp = j0 + j;
float vv;
if (vp < sink)
vv = __half2float(v_sink_head[(size_t)vp * head_dim + tid]);
else if (win > 0 && vp > cur_pos - win)
vv = __half2float(v_win_head[(size_t)(vp % FC_KVQ4_WMAX) * head_dim + tid]);
else
vv = fc_kvq4_dequant_elem(v_head + (size_t)vp * row_bytes, tid);
const float *pj_s = p_s + (size_t)j * G;
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
if (j == 0)
acc[g] *= corr_s[g];
acc[g] += pj_s[g] * vv;
}
}
}
__syncthreads();
}
#pragma unroll
for (int g = 0; g < FC_ATTN_MAX_GROUPS; g++) {
if (g < G) {
size_t po = (head0 + g) * stride + split;
acc_out[po * head_dim + tid] = acc[g];
if (tid == 0) {
m_out[po] = m_s[g];
l_out[po] = l_s[g];
}
}
}
}
/* Device-position single-row Q4 KV append (CUDA-graph capture). Bit-exact to
* fc_attn_kv_append_q4 with pos0 = pos[0], npos = 1 (caller guarantees no
* ring wrap, matching the host-arg kernel's contract). */
__global__ void fc_attn_kv_append_q4_dev(const float *k_new,
const float *v_new,
unsigned char *k_cache,
unsigned char *v_cache,
__half *k_sink, __half *v_sink,
__half *k_win, __half *v_win,
int n_head_kv, int head_dim,
int max_seq, const int *pos, int sink,
int win) {
int t = blockIdx.x;
int h = blockIdx.y;
int tid = threadIdx.x;
if (tid >= head_dim) return;
int cache_pos = pos[0] + t;
int row_bytes = fc_kvq4_row_bytes(head_dim);
size_t src = ((size_t)t * n_head_kv + h) * head_dim + tid;
size_t dst_row = ((size_t)h * max_seq + cache_pos) * row_bytes;
int blk = tid / FC_KVQ4_BLK;
int lane = tid & (FC_KVQ4_BLK - 1);
/* Shared staging for the shuffle-free block quantizer (Kepler-safe). */
__shared__ float s_k[FC_KVQ4_MAX_HEAD_DIM];
__shared__ float s_v[FC_KVQ4_MAX_HEAD_DIM];
float kf = k_new[src];
float vf = v_new[src];
s_k[tid] = kf;
s_v[tid] = vf;
__syncthreads();
fc_kvq4_quant_block(s_k + (size_t)blk * FC_KVQ4_BLK,
k_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
fc_kvq4_quant_block(s_v + (size_t)blk * FC_KVQ4_BLK,
v_cache + dst_row + (size_t)blk * FC_KVQ4_BLK_BYTES,
lane);
/* Mirror the row into the fp16 sink / recent-window shadow (element tid). */
if (win > 0) {
size_t wi = ((size_t)h * FC_KVQ4_WMAX + (cache_pos % FC_KVQ4_WMAX)) * head_dim + tid;
k_win[wi] = __float2half(kf);
v_win[wi] = __float2half(vf);
}
if (cache_pos < sink) {
size_t si = ((size_t)h * FC_KVQ4_SINKMAX + cache_pos) * head_dim + tid;
k_sink[si] = __float2half(kf);
v_sink[si] = __float2half(vf);
}
}
} /* extern "C" */
Reference in New Issue View Git Blame Copy Permalink