diff --git a/CMakeLists.txt b/CMakeLists.txt index 970ded8cbd85..7d2db7fa9e65 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1306,6 +1306,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms_w8a8/instantiate_n4.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n5.cu" "csrc/rocm/moe_gemm_w4a16_wmma.cu" + "csrc/rocm/rdna35_causal_mha_attn.cu" "csrc/rocm/attention.cu") set(VLLM_ROCM_FLAGS ${VLLM_GPU_FLAGS}) diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index bbd1c621119f..8cfe212a6327 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -142,6 +142,20 @@ void gdn_chunked(torch::Tensor& q, torch::Tensor& k, torch::Tensor& v, torch::Tensor& cu_seqlens, torch::Tensor& out, torch::Tensor& final_state, double scale); +// Causal MHA decode attention for RDNA3.5 (defined in +// rdna35_causal_mha_attn.cu; real body is gfx11-only, stub elsewhere). Mutates +// out and the three partial buffers in place. Callers gate the shape (Python +// rdna35_causal_mha_attn.can_run), so an unsupported one raises via TORCH_CHECK +// rather than leaving out untouched. +void rdna35_causal_mha_attn(torch::Tensor& out, torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& block_table, torch::Tensor& seq_lens, + torch::Tensor& partial_out, + torch::Tensor& partial_max, + torch::Tensor& partial_sum, double scale, + double softcap); + void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, diff --git a/csrc/rocm/rdna35_causal_mha_attn.cu b/csrc/rocm/rdna35_causal_mha_attn.cu new file mode 100644 index 000000000000..3dca31ac649d --- /dev/null +++ b/csrc/rocm/rdna35_causal_mha_attn.cu @@ -0,0 +1,980 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Causal decode attention for MHA (num_q_heads == num_kv_heads) on AMD RDNA3.5 +// (gfx1151). +// +// The work is split over the KV axis and merged afterwards, in two kernels: +// +// rdna35_causal_mha one wave per (sequence, KV head group, KV segment). +// Walks its slice of the context and writes an fp32 +// partial: the running max, the running sum, and the +// unnormalised P.V accumulator. +// rdna35_mha_reduce one thread per output element. Merges the partials of +// a row with a log-sum-exp rescale and writes the fp16 or +// bf16 output. +// +// PHASES of the split kernel, in order: +// +// 1. Task decomposition. blockIdx.x unpacks to (sequence, head group, +// segment); the wave's lanes then split into HEADS_PER_WAVE groups, one +// per KV head, so a lane owns VEC2_PER_LANE elements of one head. +// 2. Q load. The wave's query rows are read once into registers and stay +// there for the whole KV walk. +// 3. KV burst load. UNROLL tokens of K and V are loaded per iteration. +// Adjacent KV heads are contiguous in the paged cache, so a head group +// reads as one run. Bursts that lie inside one page take a fast address +// path that hoists the block-table lookup out of the burst. +// 4. Scoring. Q.K per token as a dot2 chain, reduced across the lanes that +// share a KV token, then scaled, optionally soft-capped, and causally +// masked. +// 5. Online softmax. The burst's maximum is taken first and the +// accumulators are rescaled once per burst rather than once per token. +// 6. P.V accumulation. Also a dot2 chain, against the softmax weight packed +// as a one-hot pair. +// 7. Partial write. Each wave owns its (head, segment) partial outright, so +// it writes straight to global memory with no LDS staging and no barrier. +// +// Restrictions the host side MUST enforce before calling (see +// vllm/v1/attention/ops/rdna35_causal_mha_attn.py, which mirrors this list): +// fp16 or bf16, no quantized KV, Hq == Hkv, head_dim in {64,128,256,512}, +// num_q_tokens in {1,4}, causal, no sliding window / ALiBi / sinks, and a +// [num_blocks, 2, block_size, num_kv_heads, head_size] KV cache (NHD). + +#include +#include + +#include +#include +#include +#include + +// The kernels use RDNA3 wave32 DPP row_ror encodings, v_dot2 and v_exp_f32, +// which only exist on gfx11 device passes. __GFX11__ is predefined by the +// compiler on a gfx11 device pass (the same macro skinny_gemms / attention and +// moe_gemm_w4a16_wmma use). Compile the real bodies for gfx11 and on the host +// pass, which needs the full definition; emit stubs for any other device arch +// so the TU links in non-gfx11 and multi-arch builds. The op is never called +// off gfx11 -- Python gates on on_gfx1151(). +#if !defined(__HIPCC__) || defined(__GFX11__) + #define VLLM_RDNA35_MHA_REAL_BODY 1 +#endif + +#define RDNA35_MHA_INLINE __device__ __forceinline__ + +namespace rdna35_mha_common { + +// v_exp_f32 is natively base-2, so the online softmax works in the log2 domain +// and scale/softcap arrive pre-multiplied by log2(e); __expf would emit a fixup +// multiply on every call. +RDNA35_MHA_INLINE float exp2_fast(float x) { return __builtin_amdgcn_exp2f(x); } + +// vec2 is the 32-bit pair the dot2 instruction consumes. Adding bf16 means +// adding one specialisation here: gfx1151 has v_dot2_f32_bf16 as well. +template +struct elem_traits; + +template <> +struct elem_traits { + using vec2 = half2; + static RDNA35_MHA_INLINE float dot2(const vec2 a, const vec2 b, float acc) { + return __builtin_amdgcn_fdot2(a, b, acc, false); + } + static RDNA35_MHA_INLINE float lo(const vec2 v) { return __half2float(v.x); } + static RDNA35_MHA_INLINE float hi(const vec2 v) { return __half2float(v.y); } + static RDNA35_MHA_INLINE vec2 zero2() { return __float2half2_rn(0.f); } + static RDNA35_MHA_INLINE half from_float(float v) { + return __float2half_rn(v); + } + + // The softmax weight as the one-hot pair (weight, 0), so the P.V product can + // be a dot2. Same scheme as Triton's `acc += tl.dot(P.to(V.dtype), V)`, + // which rounds the weight to the KV dtype so the matmul takes both operands + // in fp16. + static RDNA35_MHA_INLINE vec2 pack_weight(float weight) { + return __halves2half2(__float2half_rn(weight), (half)0.f); + } + static RDNA35_MHA_INLINE float unpack_weight(const vec2 weight) { + return __half2float(weight.x); + } + + // acc.lo += weight*v.lo, acc.hi += weight*v.hi, as two dot2 against the + // one-hot weight. pack_weight builds (weight, 0), so dot2 against it selects + // v's LOW half; the swapped (0, weight) selects the high one. + static RDNA35_MHA_INLINE void accum_pv(float& lo_acc, float& hi_acc, + const vec2 weight, const vec2 v) { + lo_acc = __builtin_amdgcn_fdot2(weight, v, lo_acc, false); + hi_acc = + __builtin_amdgcn_fdot2(__lowhigh2highlow(weight), v, hi_acc, false); + } +}; + +// bf16. gfx1151 has v_dot2_f32_bf16 next to the fp16 one, so this mirrors the +// specialisation above. +// +// Accuracy is NOT the same as fp16 and the caller should not expect it to be: +// bf16 carries 8 mantissa bits against fp16's 11, so Q.K accumulates about 8x +// the rounding error per element. The accumulators and the entire softmax stay +// fp32 -- only the Q/K/V operands and the output are narrowed -- so the error +// is input quantisation, not a change of algorithm. The tests use a separate +// tolerance for it. +template <> +struct elem_traits<__hip_bfloat16> { + using vec2 = __hip_bfloat162; + // The builtin takes a native 2-wide short vector. HIP's short2 is a + // HIP_vector_type class, which does not convert to it, and __hip_bfloat162 is + // a struct so __builtin_bit_cast rejects it as not trivially copyable -- so + // read the two raw halves and build the vector explicitly. + using short2_native = short __attribute__((ext_vector_type(2))); + static RDNA35_MHA_INLINE short2_native as_short2(const vec2 v) { + short2_native out; + out.x = static_cast(__bfloat16_as_ushort(v.x)); + out.y = static_cast(__bfloat16_as_ushort(v.y)); + return out; + } + static RDNA35_MHA_INLINE float dot2(const vec2 a, const vec2 b, float acc) { + return __builtin_amdgcn_fdot2_f32_bf16(as_short2(a), as_short2(b), acc, + false); + } + static RDNA35_MHA_INLINE float lo(const vec2 v) { + return __bfloat162float(v.x); + } + static RDNA35_MHA_INLINE float hi(const vec2 v) { + return __bfloat162float(v.y); + } + static RDNA35_MHA_INLINE vec2 zero2() { + return __hip_bfloat162(__float2bfloat16(0.f), __float2bfloat16(0.f)); + } + static RDNA35_MHA_INLINE __hip_bfloat16 from_float(float v) { + return __float2bfloat16(v); + } + + // The P.V path, mirroring the fp16 traits: the softmax weight as the one-hot + // pair (weight, 0) so the product can be a dot2 rather than an fma_mix, which + // VOPD cannot pair. dot2 against (w, 0) selects v's LOW half; the swapped + // (0, w) selects the high one. + static RDNA35_MHA_INLINE vec2 pack_weight(float weight) { + return __hip_bfloat162(__float2bfloat16(weight), __float2bfloat16(0.f)); + } + static RDNA35_MHA_INLINE float unpack_weight(const vec2 weight) { + return __bfloat162float(weight.x); + } + static RDNA35_MHA_INLINE void accum_pv(float& lo_acc, float& hi_acc, + const vec2 weight, const vec2 v) { + const vec2 swapped(weight.y, weight.x); + lo_acc = __builtin_amdgcn_fdot2_f32_bf16(as_short2(weight), as_short2(v), + lo_acc, false); + hi_acc = __builtin_amdgcn_fdot2_f32_bf16(as_short2(swapped), as_short2(v), + hi_acc, false); + } +}; + +// Merges the per-segment partials into the final output, one thread per output +// element. Walking HEAD_DIM from a single wave per (head, m) row instead +// leaves only Hq waves at NUM_Q_TOKENS=1 on a 40-CU part, which measures pure +// latency; a thread per element raises the wave count by +// HEAD_DIM/THREADS_PER_BLOCK for the same bytes. +// +// The per-segment weight is hoisted into a register array computed once rather +// than recomputing exp2(seg_max - global_max) for every (segment, d) pair. An +// empty segment carries a -inf max, which becomes a zero weight rather than a +// skipped iteration, so the accumulation stays branch-free -- the split kernels +// always write a finite acc, so the zero never multiplies garbage. +// +// THREADS_PER_BLOCK must tile HEAD_DIM. Sets both the d-parallelism and +// how many blocks a row of HEAD_DIM splits into. +// +// head and seq get their own grid dimensions so the kernel never divides by +// the runtime Hq; blockIdx.x unpacks with compile-time constants only. +// +// The segment count is a runtime argument, not a template parameter. The rule +// that picks it emits values that are not powers of two (Hq=40 gives 6, Hq=48 +// gives 5), so templating it would force the dispatch to switch over an +// instantiated set and reject the rest. The per-segment maxes and weights are +// therefore not held in register arrays: pass one takes the maximum, pass two +// accumulates the sum and the output together, re-reading partial_max. That +// re-read is a few KB against the hundreds of MB of KV the split kernel moves. +// +// Both loops are unrolled by a fixed factor with independent accumulators, so +// the loads issue back to back rather than serialising on one accumulator. +// Only the unroll factor is compile-time; the trip count stays runtime, and a +// scalar tail handles counts that are not a multiple of it. +template +__global__ __launch_bounds__(THREADS_PER_BLOCK) void rdna35_mha_reduce( + const float* __restrict__ partial_out, + const float* __restrict__ partial_max, + const float* __restrict__ partial_sum, T* __restrict__ out, const int Hq, + const int num_kv_segments) { +#if !defined(VLLM_RDNA35_MHA_REAL_BODY) + // Non-gfx11 device pass: a stub so the TU links in multi-arch builds. The + // op is never called off gfx11 (Python gates on on_gfx1151()). + return; +#else + static_assert( + HEAD_DIM % THREADS_PER_BLOCK == 0 || THREADS_PER_BLOCK % HEAD_DIM == 0, + "THREADS_PER_BLOCK must tile HEAD_DIM"); + constexpr int DIM_CHUNKS = + (HEAD_DIM + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + const int chunk = blockIdx.x % DIM_CHUNKS; + const int m = (blockIdx.x / DIM_CHUNKS) % NUM_Q_TOKENS; + const int q_head = blockIdx.y; + const int seq = blockIdx.z; + const int d = chunk * THREADS_PER_BLOCK + (int)threadIdx.x; + + const size_t stat_base = + (size_t)(seq * Hq + q_head) * num_kv_segments * NUM_Q_TOKENS + m; + + // Unrolled by a fixed factor, with independent partial maxes so the loads + // issue back to back instead of serializing on one accumulator. UNROLL is a + // compile-time constant while the trip count stays runtime, which is the + // point: templating the trip count is what forced the dispatch to switch over + // an instantiated set of segment counts in the first place. + // + // This matters more than the traffic share suggests. The reduce moves + // 0.03% of the split kernel's bytes at ctx=65536 -- but 13% at ctx=128, where + // the KV walk is short and this pass is a real fraction of the kernel. A + // first version without the unroll cost 3.8% at ctx=128, decaying to 0.3% by + // ctx=8192, exactly tracking that curve. + constexpr int UNROLL = 4; + float gmax[UNROLL]; + #pragma unroll + for (int u = 0; u < UNROLL; ++u) gmax[u] = -INFINITY; + int seg = 0; + for (; seg + UNROLL <= num_kv_segments; seg += UNROLL) { + #pragma unroll + for (int u = 0; u < UNROLL; ++u) + gmax[u] = fmaxf( + gmax[u], partial_max[stat_base + (size_t)(seg + u) * NUM_Q_TOKENS]); + } + for (; seg < num_kv_segments; ++seg) + gmax[0] = + fmaxf(gmax[0], partial_max[stat_base + (size_t)seg * NUM_Q_TOKENS]); + float global_max = gmax[0]; + #pragma unroll + for (int u = 1; u < UNROLL; ++u) global_max = fmaxf(global_max, gmax[u]); + if (global_max == -INFINITY) global_max = 0.f; + + const size_t part_base = + ((size_t)(seq * Hq + q_head) * num_kv_segments * NUM_Q_TOKENS + m) * + HEAD_DIM + + d; + + // Sum and accumulate in one pass: the weight is consumed by both, so holding + // it across the two loops was the only thing the register array bought. + // Unrolled the same way, and for the same reason -- the fp32 adds here are a + // dependent chain, so independent accumulators let the loads overlap it. + float gsum[UNROLL], gacc[UNROLL]; + #pragma unroll + for (int u = 0; u < UNROLL; ++u) gsum[u] = gacc[u] = 0.f; + + seg = 0; + for (; seg + UNROLL <= num_kv_segments; seg += UNROLL) { + #pragma unroll + for (int u = 0; u < UNROLL; ++u) { + const size_t s = (size_t)(seg + u) * NUM_Q_TOKENS; + const float seg_max = partial_max[stat_base + s]; + // An empty segment carries a -inf max, which becomes a zero weight rather + // than a skipped iteration, so the accumulation stays branch-free. + const float weight = + (seg_max == -INFINITY) ? 0.f : exp2_fast(seg_max - global_max); + gsum[u] += partial_sum[stat_base + s] * weight; + gacc[u] += partial_out[part_base + s * HEAD_DIM] * weight; + } + } + for (; seg < num_kv_segments; ++seg) { + const size_t s = (size_t)seg * NUM_Q_TOKENS; + const float seg_max = partial_max[stat_base + s]; + const float weight = + (seg_max == -INFINITY) ? 0.f : exp2_fast(seg_max - global_max); + gsum[0] += partial_sum[stat_base + s] * weight; + gacc[0] += partial_out[part_base + s * HEAD_DIM] * weight; + } + float global_sum = gsum[0], acc = gacc[0]; + #pragma unroll + for (int u = 1; u < UNROLL; ++u) { + global_sum += gsum[u]; + acc += gacc[u]; + } + + // Raw v_rcp_f32: global_sum is a sum of exp2 values in [0,1] with at least + // one term equal to 1, so it lies in [1, num_kv_segments] -- 1 ulp is far + // inside the fp16 output this feeds. IEEE division would emit the full + // Newton-Raphson. + const float inv_sum = + global_sum > 0.f ? __builtin_amdgcn_rcpf(global_sum) : 0.f; + + // `seq` indexes the output as well as the partials. Without it every + // sequence in a batch wrote the same [NUM_Q_TOKENS, Hq, HEAD_DIM] rows and + // only the last block scheduled survived -- silently, since the launch is + // valid and every sequence's partials are correct. Q is read with the same + // stride in the split kernel, so both ends now agree on + // [num_seqs, NUM_Q_TOKENS, Hq, HEAD_DIM]. + // Via the traits rather than a cast: __hip_bfloat16 has no implicit + // conversion from float, so `(T)x` compiles for half and not for bf16. + out[(size_t)((seq * NUM_Q_TOKENS + m) * Hq + q_head) * HEAD_DIM + d] = + elem_traits::from_float(acc * inv_sum); +#endif // VLLM_RDNA35_MHA_REAL_BODY +} + +} // namespace rdna35_mha_common + +#define RDNA35_MHA_INLINE __device__ __forceinline__ + +namespace rdna35_mha { + +// exp2_fast, elem_traits and the reduce kernel are shared with the GQA and MQA +// kernels; see attn_wide_common.hip. +using rdna35_mha_common::elem_traits; +using rdna35_mha_common::exp2_fast; + +// Sums across the LANES lanes that share one KV token. The XOR butterfly +// leaves the total in every participating lane, and only the steps the group +// needs are emitted -- at LANES=8 that is three instructions, against the five +// a whole-wave reduction costs. Same shape as attn_gqa_wide's +// lane_group_sum; the two kernels agree on the encoding. +template +RDNA35_MHA_INLINE float lane_group_sum(float v) { + static_assert(LANES == 4 || LANES == 8 || LANES == 16 || LANES == 32, + "reduction supports 4, 8, 16 or 32 lanes per KV token"); + if constexpr (LANES >= 2) + v += __builtin_amdgcn_mov_dpp(v, 0x161, 0xf, 0xf, 1); // row_xmask:1 + if constexpr (LANES >= 4) + v += __builtin_amdgcn_mov_dpp(v, 0x162, 0xf, 0xf, 1); // row_xmask:2 + if constexpr (LANES >= 8) + v += __builtin_amdgcn_mov_dpp(v, 0x164, 0xf, 0xf, 1); // row_xmask:4 + if constexpr (LANES >= 16) + v += __builtin_amdgcn_mov_dpp(v, 0x168, 0xf, 0xf, 1); // row_xmask:8 + if constexpr (LANES == 32) v += __shfl_xor(v, 16); + return v; +} + +RDNA35_MHA_INLINE float softcap_score(float score, float cap, float inv_cap) { + const float x = score * inv_cap; + const float e = __expf(-2.f * fabsf(x)); + const float tanh_x = (1.f - e) / (1.f + e); + return cap * copysignf(tanh_x, x); +} + +// The segment count is a runtime argument rather than a template parameter: it +// appears only in index arithmetic, never in an array bound, so it costs no +// registers, while templating it would multiply the instantiation count by +// every value the dispatch offers. The cost is that a runtime divisor needs a +// magic-number reciprocal sequence where a compile-time one is a shift -- paid +// once per task, against a KV walk of thousands of tokens. +// +// The parameter order below is shared with attn_gqa_wide and attn_mqa_wide: +// element type, query tokens, head size, then whatever decomposition that +// kernel has, then UNROLL, USE_SOFTCAP, and the mode switches. +// +// HEADS_PER_WAVE exists because head_dim 64 needs it. A wave has 32 lanes and +// one KV token's head is HEAD_DIM halves, so at HEADS_PER_WAVE=1 each lane gets +// HEAD_DIM/32 halves -- at HEAD_DIM=512 that is 32 B (a b128 pair), but at 64 +// it is 4 B and the compiler emits global_load_b32, a quarter of the width the +// memory path wants. Worse, the whole 32-lane wave then reduces one dot2 per +// token, so the 5-step DPP chain is ~83% of the inner loop. +// +// The paged KV layout is [num_blocks, 2, block_size, num_kv_heads, head_size], +// so for a fixed token ADJACENT KV HEADS ARE ADJACENT IN MEMORY. MHA has one +// KV head per query head and therefore heads to spare: a wave takes +// HEADS_PER_WAVE of them as a single contiguous run, which both widens the load +// and shrinks the reduction to log2(32/HEADS_PER_WAVE) steps. +// +// HEADS_PER_WAVE 1 2 4 +// lanes per head 32 16 8 +// bytes per lane D/16 D/8 D/4 (D=64: 4, 8, 16) +// reduction steps 5 4 3 +// +// Only head_dim 64 groups: above it a single head already fills the lane, so +// there is no load width left to win and the tuned rule picks 1. +// +// T element type of Q/K/V +// NUM_Q_TOKENS query tokens (1 plain decode, 2-4 speculative) +// HEAD_DIM head size (64, 128, 256 or 512) +// HEADS_PER_WAVE adjacent KV heads one wave loads together +// UNROLL KV tokens loaded and scored per burst +// USE_SOFTCAP whether to apply the tanh cap +template +__global__ __launch_bounds__(32) void rdna35_causal_mha( + const T* __restrict__ q, // [num_seqs, NUM_Q_TOKENS, Hq, HEAD_DIM] + const T* __restrict__ k_cache, // [nblocks, 2, block_size, Hkv, HEAD_DIM] + const T* __restrict__ v_cache, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, + // [num_seqs, Hq, num_kv_segments, NUM_Q_TOKENS, HEAD_DIM] + float* __restrict__ partial_out, float* __restrict__ partial_max, + float* __restrict__ partial_sum, const int Hq, const int Hkv, + const int block_size, const int max_blocks_per_seq, + const int num_kv_segments, + // scale and softcap arrive pre-multiplied by log2(e) so the online softmax + // can use the natively base-2 v_exp_f32; the reduce kernel must match. + const float scale, const float softcap, const float inv_softcap, + const int block_stride, const int num_seqs, const int num_tasks) { + using Traits = elem_traits; + using vec2 = typename Traits::vec2; + // The 32 lanes split into HEADS_PER_WAVE groups, one per KV head, so a group + // is LANES_PER_HEAD lanes covering HEAD_DIM halves. + constexpr int LANES_PER_HEAD = 32 / HEADS_PER_WAVE; + constexpr int VEC2_PER_LANE = HEAD_DIM / (LANES_PER_HEAD * 2); + static_assert(LANES_PER_HEAD * HEADS_PER_WAVE == 32, + "HEADS_PER_WAVE must divide 32"); + static_assert(LANES_PER_HEAD >= 4, + "the reduction supports at most 8 heads per wave"); + static_assert(VEC2_PER_LANE * LANES_PER_HEAD * 2 == HEAD_DIM, + "HEAD_DIM must divide by LANES_PER_HEAD*2"); + + const int lane = threadIdx.x; + // Position within this lane's head, and which head of the group it serves. + // At HEADS_PER_WAVE=1 head_slot is 0 and elem is the whole lane index, i.e. + // exactly the old mapping. + const int elem = lane % LANES_PER_HEAD; + const int head_slot = lane / LANES_PER_HEAD; + + { + if (blockIdx.x >= (unsigned)num_tasks) return; + int task = blockIdx.x; + const int seg = task % num_kv_segments; + task /= num_kv_segments; + // The head axis counts GROUPS of HEADS_PER_WAVE, not heads. + const int head_group = task % (Hkv / HEADS_PER_WAVE); + task /= (Hkv / HEADS_PER_WAVE); + const int seq = task; + + // Adjacent KV heads are adjacent in memory, so the group is one contiguous + // run of HEADS_PER_WAVE*HEAD_DIM halves and each lane reads its slice of + // it. + const int kv_head = head_group * HEADS_PER_WAVE + head_slot; + const int q_head = kv_head; // MHA: one query head per KV head + + const int seq_len = seq_lens[seq]; + const int ctx_len = seq_len - NUM_Q_TOKENS; + const int tokens_per_seg = + (seq_len + num_kv_segments - 1) / num_kv_segments; + const int kv_beg = seg * tokens_per_seg; + const int kv_end = min(kv_beg + tokens_per_seg, seq_len); + // An empty segment still walks the epilogue with a (-inf, 0) contribution + // rather than skipping ahead: the log-sum-exp already handles that, and + // there is no LDS merge here to leave a peer waiting. + const bool seg_empty = kv_beg >= kv_end; + + vec2 q_regs[NUM_Q_TOKENS][VEC2_PER_LANE]; +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) { + // `seq` indexes Q, matching the output write in rdna35_mha_reduce. Both + // omitted it before, so a batch read one sequence's Q and wrote one + // sequence's rows while the partials -- which always carried seq -- were + // per-sequence and correct. Only num_seqs == 1 was ever exercised. + const vec2* q_row = (const vec2*)__builtin_assume_aligned( + q + (size_t)((seq * NUM_Q_TOKENS + m) * Hq + q_head) * HEAD_DIM, 16); +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + q_regs[m][i] = q_row[elem * VEC2_PER_LANE + i]; + } + + float acc[NUM_Q_TOKENS][VEC2_PER_LANE * 2]; + float run_max[NUM_Q_TOKENS], run_sum[NUM_Q_TOKENS]; +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) { + run_max[m] = -INFINITY; + run_sum[m] = 0.f; +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE * 2; ++i) acc[m][i] = 0.f; + } + + vec2 k_burst[UNROLL][VEC2_PER_LANE], v_burst[UNROLL][VEC2_PER_LANE]; + + // Lambdas capturing k_burst/v_burst by reference, not __device__ functions: + // only capture keeps the burst in registers. + // + // The load only issues; every s_waitcnt lives in the compute lambda, so the + // whole UNROLL-deep burst is in flight before any of it is touched. + // CheckTag and SameBlockTag are compile-time tags because a runtime test + // inside the unrolled body costs a per-step exec-mask block. + auto load_burst = [&](auto CheckTag, auto SameBlockTag, int burst_start) { + constexpr bool CHECK_BOUNDS = decltype(CheckTag)::value; + constexpr bool SAME_BLOCK = decltype(SameBlockTag)::value; + if constexpr (SAME_BLOCK) { + // Every position in the burst lives in one KV block, so the block-table + // lookup and the 64-bit base are computed once and the per-step delta + // is the loop-invariant Hkv*HEAD_DIM. block_size is a runtime + // argument, so pos/block_size and pos%block_size would otherwise + // compile to a magic-number division per step. + const int block_id = + block_table[seq * max_blocks_per_seq + burst_start / block_size]; + const size_t base = + (size_t)block_id * block_stride + + (size_t)(burst_start % block_size) * Hkv * HEAD_DIM + + (size_t)kv_head * HEAD_DIM; + const size_t row_step = (size_t)Hkv * HEAD_DIM; +#pragma unroll + for (int u = 0; u < UNROLL; ++u) { + if constexpr (CHECK_BOUNDS) { + // Zero rather than break: the two-pass form below walks every slot + // and multiplies by exp2(-inf - new_max) = 0, and 0 * uninitialised + // is NaN the moment a stale register holds one. + if (burst_start + u >= kv_end) { +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) { + k_burst[u][i] = Traits::zero2(); + v_burst[u][i] = Traits::zero2(); + } + continue; + } + } + const vec2* k_row = (const vec2*)__builtin_assume_aligned( + k_cache + base + u * row_step, 16); + const vec2* v_row = (const vec2*)__builtin_assume_aligned( + v_cache + base + u * row_step, 16); +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + k_burst[u][i] = k_row[elem * VEC2_PER_LANE + i]; +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + v_burst[u][i] = v_row[elem * VEC2_PER_LANE + i]; + } + } else { +#pragma unroll + for (int u = 0; u < UNROLL; ++u) { + const int pos = burst_start + u; + if constexpr (CHECK_BOUNDS) { + if (pos >= kv_end) { // see the SAME_BLOCK arm above +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) { + k_burst[u][i] = Traits::zero2(); + v_burst[u][i] = Traits::zero2(); + } + continue; + } + } + const int block_id = + block_table[seq * max_blocks_per_seq + pos / block_size]; + const size_t base = (size_t)block_id * block_stride + + (size_t)(pos % block_size) * Hkv * HEAD_DIM + + (size_t)kv_head * HEAD_DIM; + const vec2* k_row = + (const vec2*)__builtin_assume_aligned(k_cache + base, 16); + const vec2* v_row = + (const vec2*)__builtin_assume_aligned(v_cache + base, 16); +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + k_burst[u][i] = k_row[elem * VEC2_PER_LANE + i]; +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + v_burst[u][i] = v_row[elem * VEC2_PER_LANE + i]; + } + } + }; + + // Score the whole burst, then rescale once. Past the first few tokens the + // running max rarely moves, so the per-token `acc *= rescale` is almost + // always a multiply by 1.0 across NUM_Q_TOKENS*VEC2_PER_LANE*2 + // accumulators. Costs NUM_Q_TOKENS*UNROLL scores held live across the two + // passes, which is affordable while NUM_Q_TOKENS*VEC2_PER_LANE is small. + auto compute_burst = [&](auto CheckTag, int burst_start) { + constexpr bool CHECK_BOUNDS = decltype(CheckTag)::value; + float scores[NUM_Q_TOKENS][UNROLL]; + float burst_max[NUM_Q_TOKENS]; +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) burst_max[m] = -INFINITY; +#pragma unroll + for (int u = 0; u < UNROLL; ++u) { + const int pos = burst_start + u; + const bool in_range = !CHECK_BOUNDS || pos < kv_end; +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) { + float score = 0.f; +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + score = Traits::dot2(q_regs[m][i], k_burst[u][i], score); + score = lane_group_sum(score) * scale; + if constexpr (USE_SOFTCAP) + score = softcap_score(score, softcap, inv_softcap); + if (pos > ctx_len + m) score = -INFINITY; + if constexpr (CHECK_BOUNDS) { + if (!in_range) score = -INFINITY; + } + scores[m][u] = score; + burst_max[m] = fmaxf(burst_max[m], score); + } + } + +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) { + const float new_max = fmaxf(run_max[m], burst_max[m]); + // A non-empty segment can still be entirely masked, leaving both at + // -inf; -inf - -inf is NaN, so skip rather than rescale by it. + if (new_max == -INFINITY) continue; + const float rescale = exp2_fast(run_max[m] - new_max); + run_max[m] = new_max; + float sum = run_sum[m] * rescale; +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE * 2; ++i) acc[m][i] *= rescale; +#pragma unroll + for (int u = 0; u < UNROLL; ++u) { + const float weight = exp2_fast(scores[m][u] - new_max); + sum += weight; + // P.V through dot2 against a one-hot weight (w, 0): dot2 selects + // v's low half, and the swapped (0, w) selects the high one. The + // natural `acc += weight * float(v)` form lowers to v_fma_mix_f32, + // which is VOP3P packed math and so cannot be VOPD-paired, while + // v_dot2acc_f32_f16 pairs freely. + const vec2 w2 = Traits::pack_weight(weight); +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) + Traits::accum_pv(acc[m][2 * i], acc[m][2 * i + 1], w2, + v_burst[u][i]); + } + run_sum[m] = sum; + } + }; + + const int full_burst_end = + seg_empty ? kv_beg : kv_beg + ((kv_end - kv_beg) / UNROLL) * UNROLL; + for (int pos = kv_beg; pos < full_burst_end; pos += UNROLL) { + // A burst that lies inside one page hoists the block-table lookup and the + // 64-bit base out of the burst, leaving a loop-invariant row stride. + if (pos % block_size + UNROLL <= block_size) + load_burst(std::false_type{}, std::true_type{}, pos); + else + load_burst(std::false_type{}, std::false_type{}, pos); + compute_burst(std::false_type{}, pos); + } + if (full_burst_end < kv_end) { + if (full_burst_end % block_size + UNROLL <= block_size) + load_burst(std::true_type{}, std::true_type{}, full_burst_end); + else + load_burst(std::true_type{}, std::false_type{}, full_burst_end); + compute_burst(std::true_type{}, full_burst_end); + } + + // The wave owns this (head, segment) partial outright, so it writes + // straight out: no LDS staging, no barrier. + const size_t out_base = + (((size_t)(seq * Hq + q_head) * num_kv_segments + seg) * NUM_Q_TOKENS) * + (size_t)HEAD_DIM; +#pragma unroll + for (int m = 0; m < NUM_Q_TOKENS; ++m) { +#pragma unroll + for (int i = 0; i < VEC2_PER_LANE; ++i) { + partial_out[out_base + (size_t)m * HEAD_DIM + + (elem * VEC2_PER_LANE + i) * 2] = acc[m][2 * i]; + partial_out[out_base + (size_t)m * HEAD_DIM + + (elem * VEC2_PER_LANE + i) * 2 + 1] = acc[m][2 * i + 1]; + } + // elem == 0, NOT lane == 0: each head group owns its own stats, and with + // HEADS_PER_WAVE > 1 a single lane would write only head_slot 0's, + // leaving the rest of the group's partials uninitialised for the reduce. + if (elem == 0) { + const int stat_index = + ((seq * Hq + q_head) * num_kv_segments + seg) * NUM_Q_TOKENS + m; + partial_max[stat_index] = run_max[m]; + partial_sum[stat_index] = run_sum[m]; + } + } + } +} + +// The reduce pass that merges the per-segment partials lives in +// attn_wide_common.hip: all three split kernels write the same +// [num_seqs, Hq, num_kv_segments, NUM_Q_TOKENS, HEAD_DIM] layout, so they +// share one reduce kernel rather than three identical copies. + +} // namespace rdna35_mha + +namespace { + +// The shipped config. The op takes num_kv_segments == 0 to mean "use this", +// so callers never choose one themselves. +// +// A rule rather than a lookup table: over the tuned grid it lands within a few +// percent of the best config per cell, where the spread between the best and +// worst config of a cell reaches 37x, and a table keyed on context length would +// buy about 1% for an entry per (head count, head_dim, M, ctx). +// +// The segment count is NOT constrained to a power of two and must not be +// rounded to one. It targets a constant task count -- about 6.4 tasks per SIMD +// across the 80 SIMDs on gfx1151 -- and rounding down costs 11-43% of the +// parallelism at the head counts where it bites (Hq=40 wants 6, Hq=48 wants 5, +// Hq=71 wants 7). Both kernels take it at runtime for that reason. +bool rdna35_mha_tuned(int num_q_tokens, int head_dim, int Hkv, + int* num_kv_segments, int* unroll, int* heads_per_wave) { + const int vec2_per_lane = head_dim / 64; + if (vec2_per_lane < 1 || vec2_per_lane * 64 != head_dim) return false; + + // Head grouping first: it changes how many tasks there are. 2 heads per wave + // below 16 heads and 4 at or above -- both give a b128 lane at head_dim 64, + // and below the crossover grouping by 4 would leave too few groups (Hq=8 + // gives two). Above head_dim 64 a single head already fills the lane. + int hpw = 1; + if (vec2_per_lane == 1) { + const int want = (Hkv <= 16) ? 2 : 4; + for (int c = want; c >= 2; c >>= 1) { + if (Hkv % c == 0) { + hpw = c; + break; + } + } + } + + // The numerator is 256 at head_dim 64 and 512 above it: a grouped wave covers + // hpw heads per step, so the same task count is reached with fewer segments. + // It is conditioned rather than lowered outright because 256 costs geomean + // 1.070 (worst 1.305) at head_dim >= 128, against 1.012 / 1.086 for 512. + // + // The count stays keyed on Hkv, not on the group count. Holding the task + // count near its ungrouped value by dividing by the groups instead measures + // worse everywhere it was tried: a grouped wave does hpw tokens' worth of + // work per step, so fewer, fatter tasks saturate much like more, thinner + // ones, and the extra segments only fragment the KV walk. + const int numerator = (vec2_per_lane == 1) ? 256 : 512; + int segments = numerator / (Hkv * vec2_per_lane); + + // The cap is 32 at head_dim 64 and 16 above it, for the same reason: at + // head_dim 64 vec2_per_lane is 1, so the quotient is 8x larger and a cap of + // 16 would bind on every shape. Conditioned for the same reason as well -- + // 32 everywhere costs a little at head_dim >= 128. + const int cap = (vec2_per_lane == 1) ? 32 : 16; + segments = segments < 1 ? 1 : (segments > cap ? cap : segments); + + const int q_times_vec2 = num_q_tokens * vec2_per_lane; + *num_kv_segments = segments; + *heads_per_wave = hpw; + *unroll = (q_times_vec2 <= 4 || (Hkv * segments <= 128 && q_times_vec2 <= 8)) + ? 4 + : 2; + return true; +} + +// THREADS_PER_BLOCK for the reduce, derived from HEAD_DIM. The kernel is one +// thread per output element, so this only changes how the same total threads +// are grouped, and the value below is the fastest at each head size. +// +// At HEAD_DIM 64 it is a HARD CAP rather than a preference: with +// THREADS_PER_BLOCK > HEAD_DIM the surplus threads compute d >= HEAD_DIM and +// write past the end of the row. The kernel's static_assert does not catch it +// (256 % 64 == 0 passes), so this function is the only guard. +constexpr int reduce_tpb(int head_dim) { + return head_dim >= 256 ? 256 : (head_dim >= 128 ? 128 : 64); +} + +} // namespace + +// Dispatch over (element type, num_q_tokens, head_dim, head group, unroll, +// softcap). Only M in {1,4} is instantiated -- M=1 is plain decode and M=4 +// speculative -- and the host pads query lengths 2 and 3 up to 4 rather than +// doubling the build for them. Padding costs a few percent and keeps every +// launch on a tuned configuration. +#define RDNA35_MHA_LAUNCH(T_, M_, D_, HPW_, UNRL_, SOFTCAP_) \ + do { \ + /* One task per (seq, head GROUP, segment): a wave owns HPW_ heads. */ \ + const int num_tasks = num_seqs * (num_kv_heads / (HPW_)) * nseg; \ + rdna35_mha::rdna35_causal_mha \ + <<>>( \ + (const T_*)query.data_ptr(), (const T_*)key_cache.data_ptr(), \ + (const T_*)value_cache.data_ptr(), \ + block_table.data_ptr(), seq_lens.data_ptr(), \ + partial_out.data_ptr(), partial_max.data_ptr(), \ + partial_sum.data_ptr(), num_heads, num_kv_heads, \ + block_size, max_blocks_per_seq, nseg, scale_log2, softcap_log2, \ + inv_softcap, block_stride, num_seqs, num_tasks); \ + rdna35_mha_common::rdna35_mha_reduce \ + <<>>( \ + partial_out.data_ptr(), partial_max.data_ptr(), \ + partial_sum.data_ptr(), (T_*)out.data_ptr(), num_heads, \ + nseg); \ + return; \ + } while (0) + +#define RDNA35_MHA_CASE(T_, M_, D_, HPW_, UNRL_) \ + do { \ + if (num_q_tokens == (M_) && head_size == (D_) && unroll == (UNRL_) && \ + hpw == (HPW_) && num_kv_heads % (HPW_) == 0) { \ + if (softcap > 0.f) \ + RDNA35_MHA_LAUNCH(T_, M_, D_, HPW_, UNRL_, true); \ + else \ + RDNA35_MHA_LAUNCH(T_, M_, D_, HPW_, UNRL_, false); \ + } \ + } while (0) + +// The rule only ever emits unroll 2 or 4; 1 is instantiated because the +// dispatch is keyed on the value the rule returns and a stale rule should fail +// loudly rather than silently pick a neighbour. +#define RDNA35_MHA_GRID(T_, M_, D_, HPW_) \ + do { \ + RDNA35_MHA_CASE(T_, M_, D_, HPW_, 1); \ + RDNA35_MHA_CASE(T_, M_, D_, HPW_, 2); \ + RDNA35_MHA_CASE(T_, M_, D_, HPW_, 4); \ + } while (0) + +#define RDNA35_MHA_ALL_SHAPES(T_) \ + do { \ + /* head_dim 64 builds the grouped arms too; above it a single head fills \ + the lane, so only the ungrouped one exists. */ \ + RDNA35_MHA_GRID(T_, 1, 64, 1); \ + RDNA35_MHA_GRID(T_, 1, 64, 2); \ + RDNA35_MHA_GRID(T_, 1, 64, 4); \ + RDNA35_MHA_GRID(T_, 4, 64, 1); \ + RDNA35_MHA_GRID(T_, 4, 64, 2); \ + RDNA35_MHA_GRID(T_, 4, 64, 4); \ + RDNA35_MHA_GRID(T_, 1, 128, 1); \ + RDNA35_MHA_GRID(T_, 4, 128, 1); \ + RDNA35_MHA_GRID(T_, 1, 256, 1); \ + RDNA35_MHA_GRID(T_, 4, 256, 1); \ + RDNA35_MHA_GRID(T_, 1, 512, 1); \ + RDNA35_MHA_GRID(T_, 4, 512, 1); \ + } while (0) + +// query/out: [num_seqs, num_q_tokens, num_heads, head_size] +// key_cache/value_cache: [num_blocks, block_size, num_kv_heads, head_size] +// views into vLLM's [num_blocks, 2, ...] allocation, so block_stride is taken +// from stride(0) rather than recomputed. +// seq_lens[i] is the TOTAL length including the num_q_tokens new tokens. +// +// Every constraint is a TORCH_CHECK rather than a status code: the Python +// predicate (rdna35_causal_mha_attn.can_run) is expected to have screened the +// call, so reaching here with a bad shape is a bug, and the failure mode this +// replaces -- returning -1 with `out` left untouched -- is silently wrong +// output. +void rdna35_causal_mha_attn(torch::Tensor& out, torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& block_table, torch::Tensor& seq_lens, + torch::Tensor& partial_out, + torch::Tensor& partial_max, + torch::Tensor& partial_sum, double scale, + double softcap) { + TORCH_CHECK(query.dim() == 4, + "query must be [num_seqs, M, num_heads, head_size]"); + TORCH_CHECK(out.sizes() == query.sizes(), "out must match query shape"); + TORCH_CHECK(query.scalar_type() == out.scalar_type(), + "out dtype must match query"); + TORCH_CHECK(key_cache.scalar_type() == query.scalar_type() && + value_cache.scalar_type() == query.scalar_type(), + "KV cache dtype must match query (no quantized KV)"); + TORCH_CHECK(query.is_contiguous() && out.is_contiguous(), + "query and out must be contiguous"); + + const int num_seqs = static_cast(query.size(0)); + const int num_q_tokens = static_cast(query.size(1)); + const int num_heads = static_cast(query.size(2)); + const int head_size = static_cast(query.size(3)); + + TORCH_CHECK( + key_cache.dim() == 4 && value_cache.dim() == 4, + "kv cache must be [num_blocks, block_size, num_kv_heads, head_size]"); + const int block_size = static_cast(key_cache.size(1)); + const int num_kv_heads = static_cast(key_cache.size(2)); + TORCH_CHECK( + key_cache.size(3) == head_size && value_cache.size(3) == head_size, + "kv cache head_size must match query"); + TORCH_CHECK( + value_cache.size(1) == block_size && value_cache.size(2) == num_kv_heads, + "key and value caches must have the same layout"); + + // MHA only: the kernel maps one query head per KV head by construction. + TORCH_CHECK(num_heads == num_kv_heads, + "rdna35_causal_mha_attn is MHA-only: num_heads (", num_heads, + ") must equal num_kv_heads (", num_kv_heads, ")"); + TORCH_CHECK( + num_q_tokens == 1 || num_q_tokens == 4, + "only M=1 and M=4 are instantiated (the host pads 2 and 3 to 4), got ", + num_q_tokens); + + TORCH_CHECK(block_table.dim() == 2 && block_table.size(0) == num_seqs, + "block_table must be [num_seqs, max_blocks_per_seq]"); + TORCH_CHECK(block_table.scalar_type() == torch::kInt32 && + seq_lens.scalar_type() == torch::kInt32, + "block_table and seq_lens must be int32"); + TORCH_CHECK(seq_lens.numel() >= num_seqs, "seq_lens shorter than num_seqs"); + const int max_blocks_per_seq = static_cast(block_table.size(1)); + + // Both caches are views into one [num_blocks, 2, block_size, H, D] tensor, so + // the element stride between blocks is 2x the per-cache block extent. Taking + // it from the tensor keeps a separate-allocation layout working too. + TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0), + "key and value caches must share a block stride"); + const int64_t block_stride_64 = key_cache.stride(0); + TORCH_CHECK(block_stride_64 <= std::numeric_limits::max(), + "block stride overflows int"); + const int block_stride = static_cast(block_stride_64); + + int nseg = 0, unroll = 0, hpw = 1; + TORCH_CHECK(rdna35_mha_tuned(num_q_tokens, head_size, num_kv_heads, &nseg, + &unroll, &hpw), + "no tuned config for head_size ", head_size); + TORCH_CHECK(hpw >= 1 && num_kv_heads % hpw == 0, "head group ", hpw, + " does not divide num_kv_heads ", num_kv_heads); + + // The partials are indexed [num_seqs, num_heads, nseg, M, head_size], so a + // segment count that disagrees with the launch is a memory fault rather than + // a compile error. Check it here instead. + TORCH_CHECK(partial_out.dim() == 5 && partial_max.dim() == 4 && + partial_sum.dim() == 4, + "partial buffers have the wrong rank"); + TORCH_CHECK( + partial_out.size(0) >= num_seqs && partial_out.size(1) == num_heads && + partial_out.size(2) >= nseg && partial_out.size(3) == num_q_tokens && + partial_out.size(4) == head_size, + "partial_out must be [>=num_seqs, num_heads, >=", nseg, ", ", + num_q_tokens, ", ", head_size, "]"); + TORCH_CHECK(partial_out.size(2) == nseg, "partial buffers were sized for ", + partial_out.size(2), " KV segments but the tuned config wants ", + nseg); + TORCH_CHECK(partial_out.scalar_type() == torch::kFloat32 && + partial_max.scalar_type() == torch::kFloat32 && + partial_sum.scalar_type() == torch::kFloat32, + "partial buffers must be float32"); + // Both kernels index the partials with raw pointer arithmetic, so a strided + // view -- e.g. slicing the M dimension off a workspace allocated at M=4 -- + // would silently read the wrong elements rather than fault. + TORCH_CHECK(partial_out.is_contiguous() && partial_max.is_contiguous() && + partial_sum.is_contiguous(), + "partial buffers must be contiguous"); + // NOT is_contiguous(): the caches are k/v views into one [num_blocks, 2, ...] + // allocation, so they are strided across blocks by construction -- that gap + // is exactly what block_stride carries. What the addressing does require is + // that each block be dense internally. + TORCH_CHECK(key_cache.stride(3) == 1 && value_cache.stride(3) == 1 && + key_cache.stride(2) == head_size && + value_cache.stride(2) == head_size && + key_cache.stride(1) == (int64_t)num_kv_heads * head_size && + value_cache.stride(1) == (int64_t)num_kv_heads * head_size, + "kv cache blocks must be dense [block_size, num_kv_heads, " + "head_size] (NHD); got key strides ", + key_cache.strides()); + TORCH_CHECK(block_table.is_contiguous() && seq_lens.is_contiguous(), + "block_table and seq_lens must be contiguous"); + + // Fold log2(e) into the scale so the softmax can use the natively base-2 + // v_exp_f32. cap*tanh(s/cap) is homogeneous of degree 1, so scaling the cap + // by the same factor keeps the softcap exact. + constexpr float LOG2E = 1.44269504088896340736f; + const float scale_log2 = static_cast(scale) * LOG2E; + const float softcap_log2 = static_cast(softcap) * LOG2E; + const float inv_softcap = softcap_log2 > 0.f ? 1.f / softcap_log2 : 0.f; + + // The hipified stream accessor, as moe_gemm_w4a16_wmma.cu uses: the + // c10/cuda headers pull in a cuda_cmake_macros.h that is generated at torch + // build time and absent from the wheel. + const hipStream_t stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(); + + if (query.scalar_type() == at::ScalarType::Half) { + RDNA35_MHA_ALL_SHAPES(half); + } else if (query.scalar_type() == at::ScalarType::BFloat16) { + RDNA35_MHA_ALL_SHAPES(__hip_bfloat16); + } else { + TORCH_CHECK(false, + "rdna35_causal_mha_attn supports float16 and bfloat16, got ", + query.scalar_type()); + } + + TORCH_CHECK(false, + "rdna35_causal_mha_attn: no instantiation for M=", num_q_tokens, + " head_size=", head_size, " unroll=", unroll); +} + +#undef RDNA35_MHA_ALL_SHAPES +#undef RDNA35_MHA_GRID +#undef RDNA35_MHA_CASE +#undef RDNA35_MHA_LAUNCH diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 78709ec145b0..f51a0aac0ea2 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -184,6 +184,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "int n_valid_tokens, int top_k, int block_m, int num_blocks) -> ()"); rocm_ops.impl("moe_gemm_w4a16", torch::kCUDA, &moe_gemm_w4a16); + // Wide decode attention for MHA (num_heads == num_kv_heads) on AMD RDNA3 + // (gfx11). Always registered; the kernel body is gfx11-only (stub elsewhere) + // and Python gates calls on on_gfx1151(). Mutates out and the three fp32 + // partial buffers in place; an unsupported shape raises via TORCH_CHECK + // (callers gate via the Python rdna35_causal_mha_attn.can_run predicate). + rocm_ops.def( + "rdna35_causal_mha_attn(Tensor! out, Tensor query, Tensor key_cache, " + "Tensor value_cache, Tensor block_table, Tensor seq_lens, " + "Tensor! partial_out, Tensor! partial_max, Tensor! partial_sum, " + "float scale, float softcap) -> ()"); + rocm_ops.impl("rdna35_causal_mha_attn", torch::kCUDA, + &rdna35_causal_mha_attn); + // Custom attention op // Compute the attention between an input query and the cached // keys/values using PagedAttention. diff --git a/tests/kernels/attention/test_rdna35_causal_mha_attn.py b/tests/kernels/attention/test_rdna35_causal_mha_attn.py new file mode 100644 index 000000000000..d2ee00df40b5 --- /dev/null +++ b/tests/kernels/attention/test_rdna35_causal_mha_attn.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the gfx1151 wide MHA decode-attention kernel. + +Validates ``torch.ops._rocm_C.rdna35_causal_mha_attn`` against an explicit fp32 +attention over the paged KV cache, plus the host-side predicate and the +gather/scatter that packs a varlen decode batch into the kernel's fixed-M block. + +The reference is deliberately not fused and not fast: it exists to catch a +kernel that is quick and wrong, which is the live risk here. Two cases get +specific attention because neither is reachable from the shape grid alone: + +* **num_seqs > 1**, with the sequences given DIFFERENT lengths -- equal ones + would still pass if a batch collapsed onto one sequence's Q. +* **head counts whose tuned KV segment count is not a power of two** (40, 48, + 56 give 6, 5, 4). Hq=40 at head_size 128 is Llama-2-13B, Qwen-14B, OPT-13B + and Baichuan-13B. + +Run ``pytest tests/kernels/attention/test_rdna35_causal_mha_attn.py``. +""" + +import pytest +import torch + +# The op is compiled into _rocm_C on all ROCm builds but has a real body only on +# gfx11 (stub elsewhere). Probe op presence before importing rocm.py so non-ROCm +# platforms skip without touching ROCm-only platform code. +try: + import vllm._rocm_C # noqa: F401 + + _have_op = hasattr(torch.ops._rocm_C, "rdna35_causal_mha_attn") +except Exception: + _have_op = False +if not _have_op: + pytest.skip("_rocm_C.rdna35_causal_mha_attn not available", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx1151 + +if not on_gfx1151(): + pytest.skip("requires gfx1151 (RDNA3.5)", allow_module_level=True) + +from vllm.v1.attention.ops import rdna35_causal_mha_attn # noqa: E402 + +BLOCK_SIZE = 16 + +# fp16 keeps 11 mantissa bits, bf16 only 8, so the same algorithm lands about an +# order of magnitude further out in bf16. Both accumulate in fp32; the gap is +# input quantization, not a different computation. +TOLERANCE = {torch.float16: 5e-3, torch.bfloat16: 4e-2} + + +def _tuned_nseg(head_size: int, num_kv_heads: int) -> int: + return rdna35_causal_mha_attn.tuned_num_kv_segments(1, head_size, num_kv_heads) + + +def _reference(q, k_cache, v_cache, block_table, seq_lens, scale, softcap=0.0): + """Explicit fp32 attention over the paged KV, one sequence at a time. + + ``seq_lens[i]`` is the TOTAL length including the M new query tokens, so row + ``m`` of a block attends to ``[0, seq_len - M + m]``. + + A row can be entirely masked when ``seq_len < M`` -- with seq_len=1 and M=4 + the first three rows see nothing. ``torch.softmax`` over an all -inf row is + NaN, while the kernel defines that case as zero (an all -inf running max is + clamped before the rescale, so no NaN is produced). Zero is the useful + answer: those rows are query-padding and get discarded. So the reference + zeroes them rather than propagating NaN, otherwise this asserts against the + reference's own degenerate case instead of against the kernel. + """ + num_seqs, m, num_heads, head_size = q.shape + out = torch.empty_like(q, dtype=torch.float32) + for s in range(num_seqs): + ctx = int(seq_lens[s]) + nb = (ctx + BLOCK_SIZE - 1) // BLOCK_SIZE + blocks = block_table[s, :nb].tolist() + kf = torch.cat([k_cache[b] for b in blocks]).reshape(-1, num_heads, head_size) + vf = torch.cat([v_cache[b] for b in blocks]).reshape(-1, num_heads, head_size) + kf, vf = kf[:ctx].float(), vf[:ctx].float() + pos = torch.arange(ctx, device=q.device) + qpos = (ctx - m) + torch.arange(m, device=q.device) + for h in range(num_heads): + sc = (q[s, :, h, :].float() @ kf[:, h, :].T) * scale + if softcap > 0: + sc = softcap * torch.tanh(sc / softcap) + sc = sc.masked_fill(pos[None, :] > qpos[:, None], float("-inf")) + probs = torch.softmax(sc, dim=-1) + # Fully masked rows: match the kernel's zero instead of NaN. + probs = torch.where(qpos[:, None] >= 0, probs, probs.new_zeros(())) + out[s, :, h, :] = probs @ vf[:, h, :] + return out + + +def _run(num_heads, head_size, m, ctxs, dtype, softcap=0.0, seed=0): + """One launch over len(ctxs) sequences; returns (kernel_out, reference).""" + torch.manual_seed(seed) + dev = "cuda" + num_seqs = len(ctxs) + scale = head_size**-0.5 + nb_per = [(c + BLOCK_SIZE - 1) // BLOCK_SIZE for c in ctxs] + max_nb = max(nb_per) + + # vLLM's layout: one [num_blocks, 2, block_size, num_kv_heads, head_size] + # allocation, split into k/v views, so block_stride comes from stride(0). + kv = torch.randn( + sum(nb_per), 2, BLOCK_SIZE, num_heads, head_size, device=dev, dtype=dtype + ) + k_cache, v_cache = kv.unbind(1) + + block_table = torch.zeros(num_seqs, max_nb, device=dev, dtype=torch.int32) + off = 0 + for s, n in enumerate(nb_per): + block_table[s, :n] = torch.arange(off, off + n, device=dev, dtype=torch.int32) + off += n + seq_lens = torch.tensor(ctxs, device=dev, dtype=torch.int32) + + q = torch.randn(num_seqs, m, num_heads, head_size, device=dev, dtype=dtype) + out = torch.zeros_like(q) + + nseg = _tuned_nseg(head_size, num_heads) + po = torch.empty( + num_seqs, num_heads, nseg, m, head_size, device=dev, dtype=torch.float32 + ) + pm = torch.empty(num_seqs, num_heads, nseg, m, device=dev, dtype=torch.float32) + ps = torch.empty(num_seqs, num_heads, nseg, m, device=dev, dtype=torch.float32) + + torch.ops._rocm_C.rdna35_causal_mha_attn( + out, + q, + k_cache, + v_cache, + block_table, + seq_lens, + po, + pm, + ps, + scale, + softcap, + ) + torch.accelerator.synchronize() + return out, _reference(q, k_cache, v_cache, block_table, seq_lens, scale, softcap) + + +@pytest.mark.parametrize("num_heads", [8, 16, 32]) +@pytest.mark.parametrize("head_size", [64, 128, 256, 512]) +@pytest.mark.parametrize("m", [1, 4]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_shapes(num_heads, head_size, m, dtype): + """The tuned grid: head counts whose segment count is already a power of 2.""" + out, ref = _run(num_heads, head_size, m, [1024], dtype) + assert torch.isfinite(out).all() + assert (out.float() - ref).abs().max().item() < TOLERANCE[dtype] + + +@pytest.mark.parametrize("num_heads", [40, 48, 56]) +@pytest.mark.parametrize("head_size", [128, 256]) +@pytest.mark.parametrize("m", [1, 4]) +def test_non_power_of_two_segments(num_heads, head_size, m): + """Head counts whose tuned segment count is not a power of two. + + The reduce takes the count at runtime precisely so these run; a dispatch + that switched over instantiated counts would have to reject them. + """ + nseg = _tuned_nseg(head_size, num_heads) + out, ref = _run(num_heads, head_size, m, [2048], torch.float16) + assert torch.isfinite(out).all() + assert (out.float() - ref).abs().max().item() < TOLERANCE[torch.float16], ( + f"num_heads={num_heads} head_size={head_size} nseg={nseg}" + ) + + +@pytest.mark.parametrize("m", [1, 4]) +@pytest.mark.parametrize("num_seqs", [2, 5]) +def test_batched_unequal_lengths(m, num_seqs): + """Several sequences of DIFFERENT lengths in one launch. + + Equal lengths would still pass if the kernel collapsed the batch onto one + sequence's Q, so the lengths differ. + """ + ctxs = [512 + 137 * i for i in range(num_seqs)] + out, ref = _run(32, 128, m, ctxs, torch.float16) + assert torch.isfinite(out).all() + assert (out.float() - ref).abs().max().item() < TOLERANCE[torch.float16] + # Distinct sequences must produce distinct output; a collapsed batch would + # make these equal. + if num_seqs > 1: + assert not torch.allclose(out[0], out[1]) + + +@pytest.mark.parametrize("ctx", [1, 4, 15, 16, 17, 127, 4097]) +def test_context_edges(ctx): + """Contexts that are not multiples of the block or the burst, and ctx < M. + + At ctx < M the leading rows are entirely masked; the reduce must turn that + into zeros rather than NaN. + """ + out, ref = _run(16, 128, 4, [ctx], torch.float16) + assert torch.isfinite(out).all(), f"non-finite output at ctx={ctx}" + assert (out.float() - ref).abs().max().item() < TOLERANCE[torch.float16] + + +def test_softcap(): + out, ref = _run(16, 128, 1, [1024], torch.float16, softcap=30.0) + assert torch.isfinite(out).all() + assert (out.float() - ref).abs().max().item() < TOLERANCE[torch.float16] + + +def test_rejects_gqa(): + """MHA only: the kernel maps one query head per KV head by construction.""" + dev = "cuda" + q = torch.randn(1, 1, 16, 128, device=dev, dtype=torch.float16) + kv = torch.randn(4, 2, BLOCK_SIZE, 4, 128, device=dev, dtype=torch.float16) + k, v = kv.unbind(1) + bt = torch.zeros(1, 4, device=dev, dtype=torch.int32) + sl = torch.tensor([64], device=dev, dtype=torch.int32) + po = torch.empty(1, 16, 8, 1, 128, device=dev, dtype=torch.float32) + pm = torch.empty(1, 16, 8, 1, device=dev, dtype=torch.float32) + with pytest.raises(RuntimeError, match="MHA-only"): + torch.ops._rocm_C.rdna35_causal_mha_attn( + q.clone(), q, k, v, bt, sl, po, pm, pm.clone(), 0.088, 0.0 + ) + + +def test_rejects_mismatched_segment_count(): + """Partial buffers sized for a different segment count must raise. + + This is the failure the C++ check exists for: the partials are indexed + [num_seqs, num_heads, nseg, M, head_size], so a silent disagreement would be + a memory fault rather than an exception. + """ + dev = "cuda" + num_heads, head_size = 32, 128 + wrong = _tuned_nseg(head_size, num_heads) + 1 + q = torch.randn(1, 1, num_heads, head_size, device=dev, dtype=torch.float16) + kv = torch.randn( + 4, 2, BLOCK_SIZE, num_heads, head_size, device=dev, dtype=torch.float16 + ) + k, v = kv.unbind(1) + bt = torch.zeros(1, 4, device=dev, dtype=torch.int32) + sl = torch.tensor([64], device=dev, dtype=torch.int32) + po = torch.empty(1, num_heads, wrong, 1, head_size, device=dev, dtype=torch.float32) + pm = torch.empty(1, num_heads, wrong, 1, device=dev, dtype=torch.float32) + with pytest.raises(RuntimeError, match="KV segments"): + torch.ops._rocm_C.rdna35_causal_mha_attn( + q.clone(), q, k, v, bt, sl, po, pm, pm.clone(), 0.088, 0.0 + ) + + +def test_rejects_strided_partials(): + """A strided partial buffer must raise, not read the wrong elements. + + The workspace is allocated once at the worst case (M=4, max_num_seqs) and + reused for smaller launches. Slicing the M dimension off it -- + ``buf[:n, :, :, :1]`` -- keeps the parent's strides, and since the kernels + index the partials with raw pointer arithmetic that would read the wrong + elements silently. This pins the guard. + """ + dev = "cuda" + num_heads, head_size = 32, 128 + nseg = _tuned_nseg(head_size, num_heads) + q = torch.randn(1, 1, num_heads, head_size, device=dev, dtype=torch.float16) + kv = torch.randn( + 4, 2, BLOCK_SIZE, num_heads, head_size, device=dev, dtype=torch.float16 + ) + k, v = kv.unbind(1) + bt = torch.zeros(1, 4, device=dev, dtype=torch.int32) + sl = torch.tensor([64], device=dev, dtype=torch.int32) + + # Allocated at M=4, then sliced down to M=1: the shape is right and the + # strides are not. + po_full = torch.empty( + 2, num_heads, nseg, 4, head_size, device=dev, dtype=torch.float32 + ) + pm_full = torch.empty(2, num_heads, nseg, 4, device=dev, dtype=torch.float32) + po, pm, ps = po_full[:1, :, :, :1], pm_full[:1, :, :, :1], pm_full[:1, :, :, :1] + assert not po.is_contiguous() + with pytest.raises(RuntimeError, match="contiguous"): + torch.ops._rocm_C.rdna35_causal_mha_attn( + q.clone(), q, k, v, bt, sl, po, pm, ps, 0.088, 0.0 + ) + + +def test_kv_cache_from_unbind_is_accepted(): + """k/v views from a [num_blocks, 2, ...] allocation are strided by design. + + They must NOT be rejected as non-contiguous: the gap between blocks is what + block_stride carries. Only density *within* a block is required. + """ + out, ref = _run(16, 128, 1, [512], torch.float16) + assert (out.float() - ref).abs().max().item() < TOLERANCE[torch.float16] + + +@pytest.mark.parametrize("lens", [[1, 3, 2, 4], [1, 1, 1], [4], [2, 2], [1, 4, 1]]) +def test_gather_scatter_roundtrip(lens): + """The varlen -> fixed-M packing must be a bijection on the real tokens.""" + m = rdna35_causal_mha_attn.KERNEL_M + num_seqs, total = len(lens), sum(lens) + qsl = torch.tensor([0] + torch.tensor(lens).cumsum(0).tolist()) + rows, scatter_rows = rdna35_causal_mha_attn.build_gather_index( + qsl, num_seqs, scratch_row=total + ) + + # Fixed shape, independent of the query lengths: a data-dependent shape + # would sync and would not be cudagraph-capturable. + assert rows.shape == (num_seqs * m,) + assert scatter_rows.shape == (num_seqs * m,) + + q = torch.arange(total * 6, dtype=torch.float32).reshape(total, 2, 3) + padded = q.index_select(0, rows).view(num_seqs, m, 2, 3) + + # Real tokens are the LAST len_i rows of each block, in order; padding rows + # are aimed at the scratch slot. + dest_rows = scatter_rows.view(num_seqs, m) + for i, length in enumerate(lens): + assert (dest_rows[i, : m - length] == total).all() + assert (dest_rows[i, m - length :] != total).all() + + # Identity "kernel": scattering back must reproduce the input exactly, and + # every real row must be written (no NaN survives). + out = torch.full((total + 1, 2, 3), float("nan")) + out.index_copy_(0, scatter_rows, padded.view(num_seqs * m, 2, 3)) + assert torch.equal(out[:total], q) + + +def test_predicate_rejects_unsupported(): + """The host predicate must screen everything the kernel cannot do.""" + base = dict( + num_heads=32, + num_kv_heads=32, + head_size=128, + max_query_len=1, + dtype=torch.float16, + kv_quant_mode=rdna35_causal_mha_attn.KVQuantMode.NONE, + alibi_slopes=None, + sinks=None, + sliding_window=None, + output_scale=None, + kv_cache_layout="NHD", + ) + assert rdna35_causal_mha_attn.can_run(**base) + + slopes = torch.zeros(32) + for override in ( + {"num_kv_heads": 8}, # GQA + {"head_size": 80}, # not instantiated + {"head_size": 96}, # multiple of 32 but not of 64 + {"max_query_len": 5}, # beyond the padded M + {"max_query_len": 0}, + {"dtype": torch.float32}, + {"alibi_slopes": slopes}, + {"sinks": slopes}, + {"sliding_window": 1024}, + {"output_scale": torch.ones(1)}, + {"kv_cache_layout": "HND"}, + {"causal": False}, + ): + assert not rdna35_causal_mha_attn.can_run(**{**base, **override}), override + + # Query lengths 2 and 3 are padded up to 4 rather than rejected. + for qlen in (2, 3, 4): + assert rdna35_causal_mha_attn.can_run(**{**base, "max_query_len": qlen}) + + # Every real window is rejected, including 1. The Impl stores a window W as + # (W-1, 0), so W=1 becomes (0, 0) -- a `> 0` test on the stored extent would + # let it through and the kernel would attend to the whole context. + for window in (1, 2, 1024): + assert not rdna35_causal_mha_attn.can_run( + **{**base, "sliding_window": window} + ), window + + # head_size 64 with 32 heads is the one shape the kernel does not beat + # Triton's 2D path at M=1, but it is routed here anyway rather than carved + # out -- so it must be accepted, at every query length. + for qlen in (1, 2, 3, 4): + assert rdna35_causal_mha_attn.can_run( + **{ + **base, + "num_heads": 32, + "num_kv_heads": 32, + "head_size": 64, + "max_query_len": qlen, + } + ), qlen diff --git a/vllm/envs.py b/vllm/envs.py index 00105424a0fd..7e2f8603c438 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -121,6 +121,7 @@ VLLM_MOE_HYBRID_W4A16: bool = False VLLM_MOE_HIP: str | None = None VLLM_GDN_HIP: bool = True + VLLM_ROCM_RDNA35_CAUSAL_MHA: str | None = None VLLM_ROCM_USE_MOE_WNA16_CUDA_KERNEL: bool = False VLLM_ROCM_USE_AITER: bool = False VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False @@ -1158,6 +1159,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_GDN_HIP": lambda: ( os.getenv("VLLM_GDN_HIP", "True").lower() in ("true", "1") ), + # Wide decode attention for MHA on gfx1151. Tri-state: unset = default-on + # wherever the kernel is built and the shape qualifies, "1" forces on, "0" + # forces the Triton path. Read via rdna35_causal_mha_attn.is_enabled(). + "VLLM_ROCM_RDNA35_CAUSAL_MHA": lambda: os.environ.get( + "VLLM_ROCM_RDNA35_CAUSAL_MHA", None + ), # Optional: enable external Oink custom ops (e.g., Blackwell RMSNorm). # Disabled by default. "VLLM_USE_OINK_OPS": lambda: ( diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index e3e24a91b5e7..8024d51af783 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -34,6 +34,7 @@ compute_mm_prefix_range_tensor, get_kv_cache_layout, ) +from vllm.v1.attention.ops import rdna35_causal_mha_attn from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash, @@ -93,6 +94,10 @@ class TritonAttentionMetadata: mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None mm_prefix_range_tensor: torch.Tensor | None = None + # fp32 partials for the gfx1151 wide MHA decode kernel, or None when this + # layer's shape cannot use it. See rdna35_causal_mha_attn. + rdna35_mha_partials: tuple[torch.Tensor, ...] | None = None + class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS @@ -181,6 +186,48 @@ def __init__( device=device, ) + # Workspace for the gfx1151 wide MHA decode kernel, when this layer's + # shape can use it. Allocated here rather than in forward() so the + # buffers are stable across a cudagraph capture, and sized from the same + # helper the kernel's tuned rule uses -- the partials are indexed + # [num_seqs, num_heads, nseg, M, head_size], so a segment count that + # disagreed with the launch would be a memory fault. The op re-derives + # the count and checks it against partial_out.size(2). + # can_run() only bounds max_query_len against MAX_QUERY_LEN, so its + # verdict is the same for every query length the fast path can take; + # probing at 1 answers for all of them. + self.rdna35_mha_partials: tuple[torch.Tensor, ...] | None = None + if rdna35_causal_mha_attn.can_run( + num_heads=self.num_heads_q, + num_kv_heads=self.num_heads_kv, + head_size=self.headdim, + max_query_len=1, + dtype=model_config.dtype, + kv_quant_mode=kv_cache_spec.kv_quant_mode, + alibi_slopes=None, + sinks=None, + sliding_window=None, + output_scale=None, + kv_cache_layout=get_kv_cache_layout(), + ): + # NOT max_num_seqs alone: under cudagraph capture the request count + # is padded up to a capture size, and max_cudagraph_capture_size + # defaults to min(max_num_seqs*2, 512), so a captured batch can + # carry more rows than the scheduler's own limit. Size for the + # larger of the two or the kernel writes past the workspace. + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + all_capture_sizes = vllm_config.compilation_config.cudagraph_capture_sizes + if all_capture_sizes: + max_num_seqs = max(max_num_seqs, max(all_capture_sizes)) + out_shape, stat_shape = rdna35_causal_mha_attn.workspace_shapes( + max_num_seqs, self.num_heads_q, self.headdim, self.num_heads_kv + ) + self.rdna35_mha_partials = ( + torch.empty(out_shape, dtype=torch.float32, device=device), + torch.empty(stat_shape, dtype=torch.float32, device=device), + torch.empty(stat_shape, dtype=torch.float32, device=device), + ) + def build_for_cudagraph_capture( self, common_attn_metadata: CommonAttentionMetadata ) -> TritonAttentionMetadata: @@ -243,6 +290,7 @@ def build( softmax_segm_output=self.softmax_segm_output, softmax_segm_max=self.softmax_segm_max, softmax_segm_expsum=self.softmax_segm_expsum, + rdna35_mha_partials=self.rdna35_mha_partials, ) mm_ranges = common_attn_metadata.mm_req_doc_ranges @@ -255,6 +303,123 @@ def build( return attn_metadata +def _rdna35_causal_mha_attn_forward( + *, + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + partials: tuple[torch.Tensor, ...], + num_reqs: int, + num_actual_tokens: int, + max_query_len: int, + scale: float, + softcap: float, +) -> None: + """Run the gfx1151 wide MHA decode kernel over a decode-only batch. + + The kernel wants a dense ``[num_seqs, M, num_heads, head_size]`` block with + a compile-time M, while vLLM's query is varlen and flat. At M=1 every + request contributes exactly one token, so the reshape is a free view. Above + that the rows are gathered into an M=4 block with the padding at the front + and scattered back afterwards -- see rdna35_causal_mha_attn.build_gather_index + for why the padding goes where it does. + + ``num_reqs`` drives everything, NOT ``seq_lens.shape[0]`` or + ``num_actual_tokens``. Under cudagraph capture those three are padded + independently: gpu_model_runner builds seq_lens as ``[num_reqs_padded]`` and + num_actual_tokens as ``num_tokens_padded``. num_reqs is the one that matches + both the block_table rows and the query_start_loc entries, and at M=1 it is + the row count of the query block -- taking the token count instead would + give the view a wrong shape whenever the two paddings differ. + """ + num_seqs = num_reqs + num_heads, head_size = query.shape[1], query.shape[2] + partial_out, partial_max, partial_sum = partials + m = rdna35_causal_mha_attn.KERNEL_M + nseg = partial_out.shape[2] + + def carve(buf: torch.Tensor, *shape: int) -> torch.Tensor: + """A dense view of ``shape`` over the front of ``buf``'s storage. + + The workspace is allocated at the worst case ([max_num_seqs, ..., M=4, + ...]) and a given launch needs a smaller dense block. Slicing would not + do: ``buf[:n, :, :, :1]`` keeps the parent's strides and the kernel + indexes its partials with raw pointer arithmetic, so it would read the + wrong elements. Flatten and re-view instead, which is dense by + construction and allocates nothing. + """ + n = 1 + for s in shape: + n *= s + return buf.view(-1)[:n].view(*shape) + + if max_query_len == 1: + # One token per request, so the first num_seqs rows of the flat query + # are already the dense [num_seqs, 1, num_heads, head_size] block the + # kernel wants -- a view, no copy. Slice to num_seqs rather than to + # num_actual_tokens: with cudagraph token padding the latter can be + # larger, and .view() would then raise on the size mismatch. + q_in = query[:num_seqs].view(num_seqs, 1, num_heads, head_size) + out_view = output[:num_seqs].view(num_seqs, 1, num_heads, head_size) + torch.ops.vllm.rdna35_causal_mha_attn( + out_view, + q_in, + key_cache, + value_cache, + block_table[:num_seqs], + seq_lens[:num_seqs], + carve(partial_out, num_seqs, num_heads, nseg, 1, head_size), + carve(partial_max, num_seqs, num_heads, nseg, 1), + carve(partial_sum, num_seqs, num_heads, nseg, 1), + scale, + softcap, + ) + return + + # The scatter sends padding rows to a scratch slot one past the real output + # rather than masking them out, so both index tensors keep a fixed + # num_seqs*M shape -- a boolean mask would make the shape data-dependent, + # which syncs and cannot be captured in a cudagraph. + rows, scatter_rows = rdna35_causal_mha_attn.build_gather_index( + query_start_loc, num_seqs, scratch_row=num_actual_tokens + ) + q_pad = ( + query[:num_actual_tokens] + .index_select(0, rows) + .view(num_seqs, m, num_heads, head_size) + ) + out_pad = torch.empty_like(q_pad) + torch.ops.vllm.rdna35_causal_mha_attn( + out_pad, + q_pad, + key_cache, + value_cache, + block_table[:num_seqs], + seq_lens[:num_seqs], + carve(partial_out, num_seqs, num_heads, nseg, m, head_size), + carve(partial_max, num_seqs, num_heads, nseg, m), + carve(partial_sum, num_seqs, num_heads, nseg, m), + scale, + softcap, + ) + # Scatter every padded row back, the padding ones onto the scratch row. + # + # index_scatter over a flat [num_actual_tokens + 1] destination: the extra + # row absorbs the padding writes and is dropped. Writing into a slice of + # `output` would be neater, but only if a spare row happened to exist past + # num_actual_tokens, and branching on that would reintroduce exactly the + # data-dependent control flow this formulation exists to avoid. A one-row + # scratch tensor is unconditional and costs a single row. + flat = out_pad.view(num_seqs * m, num_heads, head_size) + dest = output.new_empty(num_actual_tokens + 1, num_heads, head_size) + dest.index_copy_(0, scatter_rows, flat) + output[:num_actual_tokens].copy_(dest[:num_actual_tokens]) + + class TritonAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [ torch.float16, @@ -622,6 +787,72 @@ def forward( mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor + # gfx1151 wide MHA decode kernel, when the whole batch is decode and the + # shape qualifies. The builder only allocates the partials for a layer + # that can use it, so a None here means "not this layer" and costs one + # attribute load. Everything else falls through to unified_attention + # below -- the fallback is literally the next statement. + # + # max_query_len <= 4 IS the decode-only test: it bounds every request in + # the batch, so no prefill can hide in it. That is the same predicate + # split_decodes_and_prefills uses (`max_query_len <= decode_threshold` + # short-circuits to "all decodes"). Do not add a check comparing against + # query_start_loc[-1] -- that is a device tensor, so the comparison + # would sync every layer of every forward and would not be capturable. + # + # A mixed batch is excluded rather than split because splitting needs + # the decodes contiguous, which needs reordering, and the reorder + # threshold is set on the builder and taken as a min across every + # attention group -- it would change batch ordering for every + # TRITON_ATTN user on every platform. Query lengths 2 and 3 are padded + # up to the M=4 kernel rather than instantiated, so no uniformity is + # required either. + if ( + attn_metadata.rdna35_mha_partials is not None + and not attn_metadata.use_cascade + and max_seqlen_q <= rdna35_causal_mha_attn.MAX_QUERY_LEN + # DECODER only. ENCODER/ENCODER_ONLY already returned above, but + # ENCODER_DECODER reaches here and its cross-attention is not + # causal, while the kernel's causal mask is unconditional. + and self.attn_type == AttentionType.DECODER + and rdna35_causal_mha_attn.can_run( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + max_query_len=max_seqlen_q, + dtype=query.dtype, + kv_quant_mode=self._kv_quant_mode, + alibi_slopes=self.alibi_slopes, + sinks=self.sinks, + # (-1, -1) is how this Impl spells "no window"; anything else is + # a real window, including the degenerate (0, 0) that a + # sliding_window of 1 produces. Pass the stored left extent + # rather than testing it against 0, or W=1 would slip through + # and the kernel would silently attend to the whole context. + sliding_window=None + if self.sliding_window == (-1, -1) + else self.sliding_window[0] + 1, + output_scale=output_scale, + kv_cache_layout=get_kv_cache_layout(), + ) + ): + _rdna35_causal_mha_attn_forward( + output=output, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=block_table, + seq_lens=seqused_k, + query_start_loc=cu_seqlens_q, + partials=attn_metadata.rdna35_mha_partials, + num_reqs=seqused_k.shape[0], + num_actual_tokens=num_actual_tokens, + max_query_len=max_seqlen_q, + scale=self.scale, + softcap=self.logits_soft_cap or 0.0, + ) + return output + with create_attention_profiler_scope( backend_name="TRITON_ATTN", batch_size=seqused_k.shape[0], diff --git a/vllm/v1/attention/ops/rdna35_causal_mha_attn.py b/vllm/v1/attention/ops/rdna35_causal_mha_attn.py new file mode 100644 index 000000000000..82e13d9bab71 --- /dev/null +++ b/vllm/v1/attention/ops/rdna35_causal_mha_attn.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dispatch for the wide MHA decode-attention kernel (gfx1151 / RDNA3.5). + +Enabled by default on gfx1151, where the kernel is built and tuned; set +``VLLM_ROCM_RDNA35_CAUSAL_MHA=0`` to force the Triton path or ``=1`` to +force-enable. When the shape qualifies, this replaces Triton's unified +attention for the decode part of the batch; anything else falls through +unchanged, so the path is A/B-able. + +The kernel is compiled into ``_rocm_C`` (``csrc/rocm/rdna35_causal_mha_attn.cu``; the +body is gfx11-only, a stub elsewhere) and exposed as +``torch.ops._rocm_C.rdna35_causal_mha_attn``. This module makes the dispatch decision +host-side and wraps the call as a registered vLLM custom op with a no-op fake, +so the path is graph-safe under torch.compile. + +Why the constraints below are what they are -- each is a property of the kernel, +not a policy choice, and routing a case that violates one produces wrong output +rather than an error: + +* **MHA only.** The kernel assigns one query head per KV head by construction + (``q_head = kv_head``), so GQA and MQA shapes would read the wrong head. +* **head_size in {64,128,256,512}.** The head dimension is a template + parameter and splits across a 32-lane wave as ``head_size/64`` vec2 per lane. +* **M in {1,2,3,4}.** Only M=1 and M=4 are instantiated; 2 and 3 are padded up + to 4 here, which costs a few percent and keeps every launch on a tuned + configuration. +* **fp16 / bf16, unquantized KV.** There is no scale parameter in the ABI, so + an fp8 cache would be read as raw bits. +* **Causal, no sliding window, no ALiBi, no sinks.** The causal mask is + unconditional and the other three do not exist in the kernel. +* **NHD cache layout.** The kernel indexes + ``block_id*block_stride + (pos%block_size)*num_kv_heads*head_size + + kv_head*head_size``; the HND stride order permutes block_size and + num_kv_heads, which silently changes what it reads. +""" + +from __future__ import annotations + +import torch + +import vllm.envs as envs +from vllm.platforms.rocm import on_gfx1151 +from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.kv_cache_interface import KVQuantMode + +# Instantiated head dimensions. head_size % 64 == 0 is a static_assert in the +# kernel; these four are what the dispatch actually builds. +SUPPORTED_HEAD_SIZES = (64, 128, 256, 512) + +# NOTE on head_size 64 with 32 heads at M=1: this is the one shape the kernel +# does not win, running 0.90-0.96x of Triton's 2D path (though 1.05-1.07x of its +# 3D one). A lane holds the least work there, and Triton feeds a 16x16x16 WMMA +# tile that reuses each K across 16 rows where this kernel does one dot2 per +# token -- at M=1 there are no query rows to amortise. At M=4 there are, and +# this kernel wins by 1.19-1.42x. It is routed here anyway rather than split: +# one predicate for every shape is easier to reason about than a table of +# exceptions, and the loss is bounded and small. + +# Only M=1 (plain decode) and M=4 (speculative) are instantiated. Query lengths +# 2 and 3 run on the M=4 kernel with the leading rows padded. +KERNEL_M = 4 +MAX_QUERY_LEN = 4 + + +def is_enabled() -> bool: + """Tri-state: unset = default-on on gfx1151, "1" forces on, "0" forces off.""" + val = envs.VLLM_ROCM_RDNA35_CAUSAL_MHA + if val is not None: + return val == "1" + return on_gfx1151() + + +def tuned_num_kv_segments(num_q_tokens: int, head_size: int, num_kv_heads: int) -> int: + """The KV segment count the kernel will pick, mirroring rdna35_mha_tuned(). + + Must agree with the C++ rule exactly: the partial buffers are indexed + ``[num_seqs, num_heads, nseg, M, head_size]``, so a disagreement would be a + memory fault rather than a compile error. The op re-derives the count and + checks it against ``partial_out.size(2)`` with TORCH_CHECK, which turns a + drift between the two copies into an exception instead of corruption. + + The result is NOT always a power of two, and must not be rounded to one: it + targets a constant task count, and rounding down costs 11-43% of the + parallelism. + + Both constants are conditioned on head_size 64, where a wave groups several + KV heads (see :func:`tuned_heads_per_wave`) and each lane therefore holds a + quarter of what it holds at 128. Applying either everywhere regresses + head_size >= 128. + """ + vec2_per_lane = head_size // 64 + numerator = 256 if vec2_per_lane == 1 else 512 + cap = 32 if vec2_per_lane == 1 else 16 + segments = numerator // (num_kv_heads * vec2_per_lane) + return max(1, min(cap, segments)) + + +def tuned_heads_per_wave(head_size: int, num_kv_heads: int) -> int: + """Adjacent KV heads one wave loads together, mirroring rdna35_mha_tuned(). + + Only at head_size 64. There a single head leaves each lane with 4 bytes and + the compiler emits ``global_load_b32``, while the whole 32-lane wave reduces + one dot2 per token -- the DPP chain becomes most of the inner loop. Adjacent + KV heads are contiguous in the paged layout, so grouping them widens the + lane to 16 bytes (``b128``) and shortens the reduction from 5 steps to 3. + + Above head_size 64 a single head already fills the lane, so grouping buys no + load width and this returns 1. + """ + if head_size // 64 != 1: + return 1 + want = 2 if num_kv_heads <= 16 else 4 + for c in (want, 2): + if num_kv_heads % c == 0: + return c + return 1 + + +def can_run( + *, + num_heads: int, + num_kv_heads: int, + head_size: int, + max_query_len: int, + dtype: torch.dtype, + kv_quant_mode: KVQuantMode, + alibi_slopes: torch.Tensor | None, + sinks: torch.Tensor | None, + sliding_window: int | None, + output_scale: torch.Tensor | None, + kv_cache_layout: str, + causal: bool = True, +) -> bool: + """True iff this attention layer's decode batch can run on the kernel. + + Takes scalars and dtypes rather than tensors so the result is cacheable by + the caller and cheap to evaluate per forward. + """ + if not is_enabled() or not on_gfx1151(): + return False + # MHA only: one query head per KV head is baked into the wave mapping. + if num_heads != num_kv_heads: + return False + if head_size not in SUPPORTED_HEAD_SIZES: + return False + if not 1 <= max_query_len <= MAX_QUERY_LEN: + return False + if dtype not in (torch.float16, torch.bfloat16): + return False + # No dequant in the kernel: there is no k_scale/v_scale in the ABI at all, + # so a quantized cache would be read as raw bits. + if kv_quant_mode != KVQuantMode.NONE: + return False + if alibi_slopes is not None or sinks is not None: + return False + # The MHA kernel has no window parameter (its GQA sibling does). + if sliding_window is not None and sliding_window > 0: + return False + # No fused output quantization. + if output_scale is not None: + return False + # HND permutes block_size and num_kv_heads inside a block. + if kv_cache_layout != "NHD": + return False + # The causal mask is unconditional. + return causal + + +def workspace_shapes( + max_num_seqs: int, num_heads: int, head_size: int, num_kv_heads: int +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Shapes for the fp32 partial buffers, sized for the worst case. + + Returns ``(partial_out_shape, partial_stat_shape)``. Sized at M=KERNEL_M + because query lengths 2 and 3 are padded up to it, and at the segment count + for that M -- the count does not depend on M, but deriving it from the same + helper keeps the two in step. + """ + nseg = tuned_num_kv_segments(KERNEL_M, head_size, num_kv_heads) + return ( + (max_num_seqs, num_heads, nseg, KERNEL_M, head_size), + (max_num_seqs, num_heads, nseg, KERNEL_M), + ) + + +def build_gather_index( + query_start_loc: torch.Tensor, num_decodes: int, scratch_row: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Row indices mapping the padded [num_decodes, KERNEL_M] block to/from Q. + + vLLM's query is varlen and flat -- ``[num_tokens, num_heads, head_size]`` + with ``query_start_loc`` marking each request -- while the kernel wants a + dense ``[num_seqs, M, num_heads, head_size]`` block with a fixed M. This + builds the indices for that copy and for the copy back. + + Padding goes at the FRONT, which is forced by the kernel's mask: it applies + ``pos > ctx_len + m`` with ``ctx_len = seq_len - M``, so row ``M-1`` is the + one that sees the real last token. Real query token ``j`` of a request with + ``len_i`` tokens therefore lands at row ``M - len_i + j``, where the mask + lets it attend to exactly ``[0, num_computed + j]`` -- the ``-M`` and the + ``+m`` cancel, which is why ``seq_lens`` is passed through unmodified and + why this reproduces Triton's ``context_len = seq_len - cur_batch_query_len`` + exactly. + + Padding rows duplicate the request's first real query rather than being + zeroed: a zero row would be a softmax over a null vector, and the values are + discarded on the way back anyway. If ``seq_len < M`` the leading rows come + out fully masked, which the reduce turns into zeros rather than NaN. + + Both returned tensors have a FIXED shape of ``num_decodes * KERNEL_M``, + independent of the query lengths. That matters: the obvious formulation + selects the real rows with a boolean mask, but a mask produces a + data-dependent output shape, which forces a device-to-host sync to + materialize and cannot be captured in a cudagraph -- and TRITON_ATTN + declares AttentionCGSupport.ALWAYS. Instead every padded row gets a + destination, and the padding ones are aimed at ``scratch_row``: a slot the + caller appends past the real output, whose contents are then discarded. + + Returns ``(gather_rows, scatter_rows)``: ``gather_rows[r]`` is the flat + query row to copy into padded row ``r``, and ``scatter_rows[r]`` is where + padded row ``r`` goes on the way back (``scratch_row`` when it is padding). + """ + starts = query_start_loc[:num_decodes] + ends = query_start_loc[1 : num_decodes + 1] + lens = ends - starts + + m_idx = torch.arange(KERNEL_M, device=query_start_loc.device) + # offset of row m within the request: m - (M - len) == m - M + len + offset = m_idx.unsqueeze(0) - KERNEL_M + lens.unsqueeze(1) + real = offset >= 0 + # Padding rows clamp to offset 0, i.e. the request's first real query. + gather_rows = starts.unsqueeze(1) + offset.clamp(min=0) + scatter_rows = torch.where( + real, gather_rows, torch.full_like(gather_rows, scratch_row) + ) + return gather_rows.reshape(-1), scatter_rows.reshape(-1) + + +def _rdna35_causal_mha_attn_impl( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + partial_out: torch.Tensor, + partial_max: torch.Tensor, + partial_sum: torch.Tensor, + scale: float, + softcap: float, +) -> None: + """Run the kernel, writing into ``out`` in place. + + Callers MUST gate on :func:`can_run` first. The kernel validates everything + it can with TORCH_CHECK, so a mismatch raises rather than silently leaving + ``out`` at its previous contents. + """ + torch.ops._rocm_C.rdna35_causal_mha_attn( + out, + query, + key_cache, + value_cache, + block_table, + seq_lens, + partial_out, + partial_max, + partial_sum, + float(scale), + float(softcap), + ) + + +def _rdna35_causal_mha_attn_fake( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + partial_out: torch.Tensor, + partial_max: torch.Tensor, + partial_sum: torch.Tensor, + scale: float, + softcap: float, +) -> None: + # Everything is pre-allocated and mutated in place; nothing to allocate. + return None + + +direct_register_custom_op( + op_name="rdna35_causal_mha_attn", + op_func=_rdna35_causal_mha_attn_impl, + mutates_args=["out", "partial_out", "partial_max", "partial_sum"], + fake_impl=_rdna35_causal_mha_attn_fake, +)