iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model][Hardware][AMD][Kernel] Part 2/2 -> Enable e2e QK Norm + RoPE + KV Cache runtime fusion for Qwen3-30B-A3B on ROCM_ATTN (#42755)

Browse Source 
Signed-off-by: Jack Hu <Jack.Hu@amd.com>
Signed-off-by: Jack Hu <jingxujaster@hotmail.com>
Signed-off-by: Jack Hu <jack.hu@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
This commit is contained in:
Jack Hu
2026-10-15 17:10:14 -07:00
committed by GitHub
co-authored by Douglas Lehr
parent af28c171f5
commit ef38e8d006
23 changed files with 701 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cache_kernels.cu
+21 -9
View File
@@ -273,7 +273,7 @@ __global__ void reshape_and_cache_kernel(
const int64_t* __restrict__ slot_mapping, // [num_tokens]
const int key_stride, const int value_stride, const int num_heads,
const int head_size, const int block_size, const int x,
const float* k_scale, const float* v_scale) {
const float* k_scale, const float* v_scale, const bool interleaved_v) {
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
if (slot_idx < 0) {
@@ -301,10 +301,20 @@ __global__ void reshape_and_cache_kernel(
key_cache + block_idx * num_heads * h_block_count * block_size * x +
head_idx * h_block_count * block_size * x + h_block * block_size * x +
block_offset * x;
const int64_t tgt_value_start =
block_idx * num_heads * h_block_count * x * block_size +
head_idx * h_block_count * x * block_size + h_block * x * block_size +
block_offset;
// V cache target: standard [.., head_size, block_size], or the interleaved
// [.., block_size/x, head_size, x] the fused AITER writer/reader use.
const int64_t v_head_base = block_idx * num_heads * head_size * block_size +
head_idx * head_size * block_size;
int64_t tgt_value_start;
int v_dim_stride;
if (interleaved_v) {
tgt_value_start = v_head_base + (block_offset / x) * (head_size * x) +
(h_block * x) * x + (block_offset % x);
v_dim_stride = x;
} else {
tgt_value_start = v_head_base + h_block * x * block_size + block_offset;
v_dim_stride = block_size;
}
constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
@@ -318,7 +328,7 @@ __global__ void reshape_and_cache_kernel(
cache_t* __restrict__ value_dst = value_cache + tgt_value_start;
#pragma unroll
for (int i = 0; i < x; i++) {
v_op(value_dst[i * block_size], value_src[i]);
v_op(value_dst[i * v_dim_stride], value_src[i]);
}
}
@@ -775,7 +785,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
slot_mapping.const_data_ptr<int64_t>(), key_stride, value_stride, \
num_heads, head_size, block_size, x, \
reinterpret_cast<const float*>(k_scale.data_ptr()), \
reinterpret_cast<const float*>(v_scale.data_ptr()));
reinterpret_cast<const float*>(v_scale.data_ptr()), interleaved_v);
void reshape_and_cache(
torch::stable::Tensor& key, // [num_tokens, num_heads, head_size]
@@ -783,10 +793,12 @@ void reshape_and_cache(
torch::stable::Tensor&
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
torch::stable::Tensor&
value_cache, // [num_blocks, num_heads, head_size, block_size]
value_cache, // [num_blocks, num_heads, head_size, block_size], or the
// interleaved [.., block_size/x, head_size, x] when
// interleaved_v is set
torch::stable::Tensor& slot_mapping, // [num_tokens]
const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale) {
torch::stable::Tensor& v_scale, bool interleaved_v) {
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
csrc/ops.h
+1 -1
View File
@@ -650,7 +650,7 @@ void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
torch::stable::Tensor& slot_mapping,
const std::string& kv_cache_dtype,
torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale);
torch::stable::Tensor& v_scale, bool interleaved_v);
void reshape_and_cache_flash(
torch::stable::Tensor& key, torch::stable::Tensor& value,
csrc/rocm/attention.cu
+154 -67
View File
@@ -355,7 +355,8 @@ __device__ float warpReduceMax(float val) {
// clang-format off
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE>
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -593,26 +594,59 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
_B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too
const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
// KX_V = number of contiguous token elements per interleaved chunk
// (e.g. 8 for bf16, 16 for fp8). This is also the 'x' dimension size.
constexpr int KX_V = 16 / sizeof(cache_t);
// v fetches are 16head elems across lanes x 16 tokens per lane
for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
if constexpr (USE_INTERLEAVED_V_CACHE) {
// Interleaved V-cache layout: [num_blocks, num_heads, block_size/x,
// head_dim, x] Address for v[block][head][token t][dim d]:
// chunk = t / x, sub = t % x
// addr = base + chunk * head_dim * x + d * x + sub
// Each 16B fetch reads KX_V contiguous elements within a single chunk, so
// a lane's token offset must land on a chunk boundary. This holds only
// when BLOCK_SIZE is a multiple of KX_V.
static_assert(BLOCK_SIZE % KX_V == 0,
"Interleaved V-cache read requires BLOCK_SIZE % KX_V == 0");
const cache_t* v_ptr_head = v_cache + wg_start_kv_head_idx * kv_head_stride;
for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
const int vblock_depth = 0;
const int64_t vblock_number = static_cast<int64_t>(
vphysical_block_number[vtoken_depth][vblock_depth]);
const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
const cache_t* v_fetch_ptr =
v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
const _B16x8* v_fetch_ptr_16B =
reinterpret_cast<const _B16x8*>(v_fetch_ptr);
Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
const cache_t* v_block_ptr =
v_ptr_head + vblock_number * kv_block_stride;
for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
const int bo = ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE) +
vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
const int bo_chunk = bo / KX_V; // which chunk of x tokens
const cache_t* v_fetch_ptr =
v_block_ptr + (bo_chunk * HEAD_SIZE + vhead_elem) * KX_V +
(bo % KX_V);
Vlocal[vtoken_depth][vhe_depth][vfetch_depth] =
*reinterpret_cast<const _B16x8*>(v_fetch_ptr);
}
}
}
} else {
const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
const int vblock_depth = 0;
const int64_t vblock_number = static_cast<int64_t>(
vphysical_block_number[vtoken_depth][vblock_depth]);
const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
const cache_t* v_fetch_ptr =
v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
const _B16x8* v_fetch_ptr_16B =
reinterpret_cast<const _B16x8*>(v_fetch_ptr);
Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
}
}
}
}
@@ -996,7 +1030,8 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
int GQA_RATIO>
int GQA_RATIO,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -1177,25 +1212,41 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
}
}
// KX_V = number of contiguous token elements per interleaved chunk
constexpr int KX_V = 16 / sizeof(cache_t);
const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride;
// fetch vcache in kv cache auto case
// fetch vcache in kv cache auto case (bf16/fp16)
if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
const _B16x8* v_ptrh8 = reinterpret_cast<const _B16x8*>(v_ptr);
// iterate over each v block
for (int b = 0; b < VBLOCKS; b++) {
// int32 physical_block_number leads to overflow when multiplied with
// kv_block_stride
const int64_t vphysical_block_number =
static_cast<int64_t>(vphysical_blocks[b]);
const _B16x8* v_ptrh8b =
v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
// iterate over each head elem (within head_size)
for (int h = 0; h < VHELOOP; h++) {
const int head_size_elem = h * WARP_SIZE + laneid;
const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
// iterate over all velems within block
for (int d = 0; d < BLOCK_SIZE / 8; d++) {
Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
if constexpr (USE_INTERLEAVED_V_CACHE) {
// Interleaved layout: [num_blocks, num_heads, block_size/x,
// head_dim, x] Iterate over chunks of x tokens; within each chunk,
// x values for the same head_size_elem are contiguous (16-byte
// aligned).
for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) {
const _B16x8* v_chunk =
v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8;
for (int xd = 0; xd < KX_V / 8; xd++) {
Vlocal[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] =
v_chunk[xd];
}
}
} else {
// Standard PagedAttention layout: [num_blocks, num_heads, head_dim,
// block_size] All block_size tokens for one head_size_elem are
// contiguous.
const _B16x8* v_ptrh8be =
v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
for (int d = 0; d < BLOCK_SIZE / 8; d++) {
Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
}
}
}
}
@@ -1203,21 +1254,29 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
// fetch vcache in fp8 case
else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto)
const _B8x8* v_ptrh8 = reinterpret_cast<const _B8x8*>(v_ptr);
// iterate over each v block
for (int b = 0; b < VBLOCKS; b++) {
// int32 physical_block_number leads to overflow when multiplied with
// kv_block_stride
const int64_t vphysical_block_number =
static_cast<int64_t>(vphysical_blocks[b]);
const _B8x8* v_ptrh8b =
v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
// iterate over each head elem (within head_size)
for (int h = 0; h < VHELOOP; h++) {
const int head_size_elem = h * WARP_SIZE + laneid;
const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
// iterate over all velems within block
for (int d = 0; d < BLOCK_SIZE / 8; d++) {
Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
if constexpr (USE_INTERLEAVED_V_CACHE) {
// Same interleaved addressing as bf16 above, but with fp8
// elements (KX_V=16, using _B8x8 loads instead of _B16x8).
for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) {
const _B8x8* v_chunk =
v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8;
for (int xd = 0; xd < KX_V / 8; xd++) {
Vlocalb8[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] =
v_chunk[xd];
}
}
} else {
const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
for (int d = 0; d < BLOCK_SIZE / 8; d++) {
Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
}
}
}
}
@@ -1831,7 +1890,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
MFMAType MFMA_TYPE>
MFMAType MFMA_TYPE,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -2285,8 +2345,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
int GQA_RATIO>
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -2611,7 +2671,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
MFMAType MFMA_TYPE>
MFMAType MFMA_TYPE,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -3030,8 +3091,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
int GQA_RATIO>
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -3261,7 +3322,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
int GQA_RATIO, MFMAType MFMA_TYPE>
int GQA_RATIO, MFMAType MFMA_TYPE,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -3288,7 +3350,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
template <typename scalar_t, typename cache_t,
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
int GQA_RATIO>
int GQA_RATIO,
bool USE_INTERLEAVED_V_CACHE>
__global__
__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -3330,10 +3393,11 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
#endif
// clang-format off
#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
HEAD_SIZE, NTHR, ALIBI_ENABLED, \
GQA_RATIO, MFMA_TYPE> \
paged_attention_ll4mi_QKV_mfma16_kernel< \
T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED, \
GQA_RATIO, MFMA_TYPE, USE_INTERLEAVED_V_CACHE> \
<<<grid, block, 0, stream>>>( \
query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \
@@ -3342,15 +3406,16 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
max_ctx_blocks, k_scale_ptr, v_scale_ptr);
#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \
paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
HEAD_SIZE, NTHR, ALIBI_ENABLED, \
GQA_RATIO> \
paged_attention_ll4mi_QKV_mfma4_kernel< \
T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED, \
GQA_RATIO, USE_INTERLEAVED_V_CACHE> \
<<<grid, block, 0, stream>>>( \
query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \
max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
max_ctx_blocks, k_scale_ptr, v_scale_ptr);
// clang-format on
#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
@@ -3361,7 +3426,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
void paged_attention_custom_launcher(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3517,7 +3582,7 @@ void paged_attention_custom_launcher(
template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
void paged_attention_custom_launcher_navi(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3688,21 +3753,42 @@ void paged_attention_custom_launcher_navi(
}
}
#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
PSIZE, ALIBI_ENABLED, MFMA_TYPE) \
if (!is_navi) { \
paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE>( \
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
} else { \
paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE, \
HEAD_SIZE, OUTT, PSIZE, \
ALIBI_ENABLED, MFMA_TYPE>( \
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
max_seq_len, alibi_slopes, k_scale, v_scale); \
// USE_INTERLEAVED_V_CACHE gates the V-cache read path at compile time.
// Currently only the GFX9 kernels (MI300x/MI325x/MI350x/MI355x) have the
// actual interleaved V-fetch logic implemented. The GFX11 (RDNA 3) and
// GFX12 (RDNA 4) mfma16 kernels do NOT support interleaved V-cache addressing.
// The template parameter is carried on all variants for compilation
// compatibility; a TORCH_CHECK prevents dispatch to NAVI with interleaved V.
// TODO: implement interleaved V-cache support in the NAVI (GFX11/GFX12)
// kernels if shuffle KV cache layout is needed on RDNA GPUs.
#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
PSIZE, ALIBI_ENABLED, MFMA_TYPE) \
if (use_interleaved_v_cache) { \
TORCH_CHECK(!is_navi, \
"Reading interleaved V-cache layout from paged decode " \
"attention not supported on NAVI GPUs"); \
paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE, \
true>( \
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
} else { \
if (!is_navi) { \
paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE, \
false>( \
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
} else { \
paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE, \
HEAD_SIZE, OUTT, PSIZE, \
ALIBI_ENABLED, MFMA_TYPE, false>( \
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
max_seq_len, alibi_slopes, k_scale, v_scale); \
} \
}
#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
@@ -3799,7 +3885,8 @@ void paged_attention(
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale,
const std::optional<torch::Tensor>& fp8_out_scale,
const std::string& mfma_type) {
const std::string& mfma_type,
bool use_interleaved_v_cache) {
// clang-format on
bool is_navi = is_navi_gpu();
const int head_size = query.size(2);
csrc/rocm/ops.h
+1 -1
View File
@@ -50,4 +50,4 @@ void paged_attention(
int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale,
const std::string& mfma_type);
const std::string& mfma_type, bool use_interleaved_v_cache);
csrc/rocm/torch_bindings.cpp
+2 -1
View File
@@ -92,7 +92,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale,"
" Tensor? fp8_out_scale,"
" str mfma_type) -> ()");
" str mfma_type,"
" bool use_interleaved_v_cache=False) -> ()");
rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention);
}
csrc/torch_bindings.cpp
+2 -1
View File
@@ -955,7 +955,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
" Tensor! key_cache, Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
" Tensor k_scale, Tensor v_scale,"
" bool interleaved_v=False) -> ()");
// Reshape the key and value tensors and cache them.
ops.def(
tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
+67 -25
View File
@@ -157,10 +157,17 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
)
if self.kv_cache_dtype != self.dtype:
self.attn._k_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
self.attn._v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
self.attn._k_scale_float = 1.0
self.attn._v_scale_float = 1.0
# Non-unit scales so the fused and unfused writers must agree on them.
self.attn._k_scale = torch.tensor(0.5, dtype=torch.float32, device=device)
self.attn._v_scale = torch.tensor(0.25, dtype=torch.float32, device=device)
self.attn._k_scale_cpu = torch.tensor(
0.5, dtype=torch.float32, device="cpu"
)
self.attn._v_scale_cpu = torch.tensor(
0.25, dtype=torch.float32, device="cpu"
)
self.attn._k_scale_float = 0.5
self.attn._v_scale_float = 0.25
else:
self.attn._k_scale = self.attn._k_scale.to(device)
self.attn._v_scale = self.attn._v_scale.to(device)
@@ -527,16 +534,6 @@ def _run_qk_norm_rope_kvcache_fusion_test(
else:
torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
if not expect_fusion or (not is_fp8_cache and mrope_section is None):
# The AITER PTS kernel populates k_out only for non-FP8 caches.
# With FP8, the kernel writes quantized K directly to the cache
# and may leave k_out uninitialised. In production this is fine
# because downstream attention reads K from the cache.
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
# Should be bit exact since no processing had been done on v for both paths
torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
# Fused and unfused arithmetic can straddle an FP8 rounding boundary.
# Allow one E4M3 quantization step while still comparing every block.
if is_fp8_cache:
@@ -544,6 +541,18 @@ def _run_qk_norm_rope_kvcache_fusion_test(
else:
cache_atol, cache_rtol = ATOL, RTOL
# The fused k_out is the dequantized cache write, so it carries the same
# rounding; FA and ROCM_ATTN prefill consume it directly. Unified reads K
# from the cache and asks the kernel not to write k_out at all.
if attn_backend != AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
torch.testing.assert_close(
k_unfused, k_fused, atol=cache_atol, rtol=cache_rtol
)
# Should be bit exact since no processing had been done on v for both paths
torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
# Whole cache, so the interleaved V write (ROCM_ATTN kv_cache[1]) is checked.
torch.testing.assert_close(
kv_cache_unfused.float(),
kv_cache_fused.float(),
@@ -570,20 +579,52 @@ _FUSION_CONFIGS = [
_FUSION_CONFIGS,
)
@pytest.mark.parametrize(
"attn_backend",
"attn_backend, use_shuffle_kv_layout, kv_layout",
[
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
# ROCM_ATTN always writes the shuffled, interleaved-V layout and needs each
# K/V side contiguous per block, so head-major only.
pytest.param(
AttentionBackendEnum.ROCM_ATTN,
"0",
KVCacheLayout.LBHNC,
id="rocm_attn",
),
# Unified never reads the shuffle env; its strided views take any layout.
pytest.param(
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
"0",
KVCacheLayout.LBHNC,
id="unified-head_major",
),
pytest.param(
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
"0",
KVCacheLayout.LBNHC,
id="unified-token_major",
),
# FA takes any layout unshuffled; its shuffle writer needs each K/V side
# contiguous per block, so head-major only.
pytest.param(
AttentionBackendEnum.ROCM_AITER_FA,
"0",
KVCacheLayout.LBHNC,
id="fa-head_major",
),
pytest.param(
AttentionBackendEnum.ROCM_AITER_FA,
"0",
KVCacheLayout.LBNHC,
id="fa-token_major",
),
pytest.param(
AttentionBackendEnum.ROCM_AITER_FA,
"1",
KVCacheLayout.LBHNC,
id="fa-shuffle",
),
],
)
@pytest.mark.parametrize("num_tokens", [5, 2048])
@pytest.mark.parametrize(
"kv_layout",
[
pytest.param(KVCacheLayout.LBHNC, id="head_major"),
pytest.param(KVCacheLayout.LBNHC, id="token_major"),
],
)
@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
@pytest.mark.parametrize("block_size", [16, 32, 64])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@@ -615,6 +656,7 @@ def test_qk_norm_rope_kvcache_fusion(
kv_cache_dtype: CacheDType,
rms_norm_eps: float,
custom_op: str,
use_shuffle_kv_layout: str,
monkeypatch: pytest.MonkeyPatch,
):
_run_qk_norm_rope_kvcache_fusion_test(
@@ -627,7 +669,7 @@ def test_qk_norm_rope_kvcache_fusion(
rotary_dim=rotary_dim,
block_size=block_size,
is_neox=is_neox,
use_shuffle_kv_layout="0",
use_shuffle_kv_layout=use_shuffle_kv_layout,
kv_layout=kv_layout,
dtype=dtype,
kv_cache_dtype=kv_cache_dtype,
tests/kernels/attention/test_attention.py
+25 -1
View File
@@ -122,6 +122,7 @@ def ref_single_query_cached_kv_attention(
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("use_interleaved_v_cache", [False, True])
def test_paged_attention(
kv_cache_factory,
num_seqs: int,
@@ -133,9 +134,14 @@ def test_paged_attention(
kv_cache_dtype: str,
seed: int,
device: str,
use_interleaved_v_cache: bool,
) -> None:
if current_platform.is_navi() and (
kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
kv_cache_dtype == "fp8"
or head_size != 128
or block_size != 16
or use_alibi
or use_interleaved_v_cache
):
pytest.skip()
@@ -201,6 +207,20 @@ def test_paged_attention(
last_block_idx = block_tables_lst[seq_idx][seq_len // block_size]
value_cache[last_block_idx, :, :, padding_start:] = padding_nan
standard_value_cache = value_cache
if use_interleaved_v_cache:
# Hand the kernel V in the interleaved layout (same 4D shape); the
# reference below gets the standard layout back.
x_v = 16 // value_cache.element_size()
value_cache = (
value_cache.view(
NUM_BLOCKS, num_kv_heads, head_size, block_size // x_v, x_v
)
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(NUM_BLOCKS, num_kv_heads, head_size, block_size)
)
# Call the paged attention kernel.
output = torch.empty_like(query)
num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
@@ -234,6 +254,7 @@ def test_paged_attention(
kv_cache_dtype,
k_scale,
v_scale,
use_interleaved_v_cache=use_interleaved_v_cache,
)
opcheck(
@@ -259,10 +280,13 @@ def test_paged_attention(
v_scale,
None,
"f16",
use_interleaved_v_cache,
),
cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
)
value_cache = standard_value_cache
# Run the reference implementation.
if kv_cache_dtype == "fp8":
# Convert cache data back to dtype.
tests/kernels/attention/test_prefix_prefill.py
+90 -2
View File
@@ -109,6 +109,7 @@ def create_alibi_causal_mask(
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
@pytest.mark.parametrize("op", OPS)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention(
num_heads: int,
@@ -119,6 +120,7 @@ def test_contexted_kv_attention(
kv_cache_dtype: str,
device: str,
op: Callable,
interleave_v: bool,
block_size: int = 32,
) -> None:
if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
@@ -225,6 +227,23 @@ def test_contexted_kv_attention(
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
extra_op_kwargs = {}
if interleave_v:
# Rewrite the V cache into the interleaved layout (same 4D shape). Both
# context_attention_fwd (prefill) and kernel_paged_attention_2d (decode)
# read it back interleaved.
x_v = 16 // v_cache.element_size()
v_cache = (
v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(cache_size, num_kv_heads, head_size, block_size)
)
if op is context_attention_fwd:
extra_op_kwargs["interleaved_v_pack_factor"] = x_v
else:
extra_op_kwargs["use_interleaved_v_cache"] = True
# Warm up the Triton kernel by calling it once before actually measuring
# generation time
op(
@@ -243,6 +262,7 @@ def test_contexted_kv_attention(
k_scale,
v_scale,
sliding_window=sliding_window,
**extra_op_kwargs,
)
torch.accelerator.synchronize()
start_time = time.time()
@@ -262,6 +282,7 @@ def test_contexted_kv_attention(
k_scale,
v_scale,
sliding_window=sliding_window,
**extra_op_kwargs,
)
torch.accelerator.synchronize()
end_time = time.time()
@@ -386,6 +407,7 @@ def test_rocm_backend_sliding_window_includes_boundary_token(
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention_cached_kv(
num_heads: int,
@@ -394,6 +416,7 @@ def test_contexted_kv_attention_cached_kv(
dtype: torch.dtype,
kv_cache_dtype: str,
device: str,
interleave_v: bool,
block_size: int = 32,
) -> None:
# Exercises the KV_FROM_CACHE path of context_attention_fwd: the current
@@ -487,6 +510,18 @@ def test_contexted_kv_attention_cached_kv(
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
interleaved_v_pack_factor = 0
if interleave_v:
# Rewrite the V cache into the interleaved layout (same 4D shape).
interleaved_v_pack_factor = 16 // v_cache.element_size()
x_v = interleaved_v_pack_factor
v_cache = (
v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(cache_size, num_kv_heads, head_size, block_size)
)
# Cached-K/V path: current-chunk k and v are None.
context_attention_fwd(
query,
@@ -504,6 +539,7 @@ def test_contexted_kv_attention_cached_kv(
k_scale,
v_scale,
sliding_window=0,
interleaved_v_pack_factor=interleaved_v_pack_factor,
)
torch.accelerator.synchronize()
@@ -544,8 +580,11 @@ def test_contexted_kv_attention_cached_kv(
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None:
def test_contexted_kv_attention_cached_kv_block_table_boundary(
device: str, interleave_v: bool
) -> None:
# Boundary guard for the KV_FROM_CACHE block-table load. With an
# exact-sized block table (row length == number of blocks for the
# sequence) and a query that ends the sequence, the last K/V tile has
@@ -611,6 +650,18 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
interleaved_v_pack_factor = 0
if interleave_v:
# Rewrite the V cache into the interleaved layout (same 4D shape).
interleaved_v_pack_factor = 16 // v_cache.element_size()
x_v = interleaved_v_pack_factor
v_cache = (
v_cache.view(num_blocks, num_kv_heads, head_size, block_size // x_v, x_v)
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(num_blocks, num_kv_heads, head_size, block_size)
)
# Cached-K/V path: current-chunk k and v are None.
context_attention_fwd(
query,
@@ -628,6 +679,7 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N
k_scale,
v_scale,
sliding_window=0,
interleaved_v_pack_factor=interleaved_v_pack_factor,
)
torch.accelerator.synchronize()
@@ -723,6 +775,7 @@ def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("op", OPS)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention_alibi(
num_heads: int,
@@ -732,6 +785,7 @@ def test_contexted_kv_attention_alibi(
kv_cache_dtype: str,
device: str,
op: Callable,
interleave_v: bool,
block_size: int = 32,
) -> None:
if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
@@ -860,6 +914,23 @@ def test_contexted_kv_attention_alibi(
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
extra_op_kwargs = {}
if interleave_v:
# Rewrite the V cache into the interleaved layout (same 4D shape). Both
# context_attention_fwd (prefill) and kernel_paged_attention_2d (decode)
# read it back interleaved.
x_v = 16 // v_cache.element_size()
v_cache = (
v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(cache_size, num_kv_heads, head_size, block_size)
)
if op is context_attention_fwd:
extra_op_kwargs["interleaved_v_pack_factor"] = x_v
else:
extra_op_kwargs["use_interleaved_v_cache"] = True
# Warm up the Triton kernel by calling it once before actually measuring
# generation time
op(
@@ -878,6 +949,7 @@ def test_contexted_kv_attention_alibi(
k_scale,
v_scale,
alibi_slopes=alibi_slopes,
**extra_op_kwargs,
)
torch.accelerator.synchronize()
start_time = time.time()
@@ -897,6 +969,7 @@ def test_contexted_kv_attention_alibi(
k_scale,
v_scale,
alibi_slopes=alibi_slopes,
**extra_op_kwargs,
)
torch.accelerator.synchronize()
end_time = time.time()
@@ -992,6 +1065,7 @@ def test_contexted_kv_attention_alibi(
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
@pytest.mark.parametrize("op", OPS)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention_f32(
num_heads: int,
@@ -1002,6 +1076,7 @@ def test_contexted_kv_attention_f32(
kv_cache_dtype: str,
device: str,
op: Callable,
interleave_v: bool,
) -> None:
test_contexted_kv_attention(
num_heads,
@@ -1012,6 +1087,7 @@ def test_contexted_kv_attention_f32(
kv_cache_dtype,
device,
op,
interleave_v=interleave_v,
)
@@ -1023,6 +1099,7 @@ def test_contexted_kv_attention_f32(
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("op", OPS)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_contexted_kv_attention_alibi_f32(
num_heads: int,
@@ -1032,9 +1109,17 @@ def test_contexted_kv_attention_alibi_f32(
kv_cache_dtype: str,
device: str,
op: Callable,
interleave_v: bool,
) -> None:
test_contexted_kv_attention_alibi(
num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op
num_heads,
num_queries_per_kv,
head_size,
dtype,
kv_cache_dtype,
device,
op,
interleave_v=interleave_v,
)
@@ -1052,6 +1137,7 @@ NONSTANDARD_BLOCK_SIZE_SHAPES = [
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("op", OPS)
@pytest.mark.parametrize("interleave_v", [False, True])
@torch.inference_mode()
def test_qwen3_nonstandard_block_size(
num_heads: int,
@@ -1061,6 +1147,7 @@ def test_qwen3_nonstandard_block_size(
dtype: torch.dtype,
device: str,
op: Callable,
interleave_v: bool,
) -> None:
"""Non-power-of-2 pages must match, even when a tile straddles a page."""
if not current_platform.is_rocm():
@@ -1076,4 +1163,5 @@ def test_qwen3_nonstandard_block_size(
kv_cache_dtype="auto",
device=device,
op=op,
interleave_v=interleave_v,
)
tests/kernels/attention/test_rocm_aiter_fa.py
+22
View File
@@ -1061,3 +1061,25 @@ def test_aiter_mha_varlen_fp8_kv(dtype):
rtol=rtol,
)
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
def test_asm_kv_scale_tables_shared_across_builders():
"""Builders share one AsmKvScaleTables; equal (shape, scale) requests must
return the same dense table, a different scale a different one."""
from vllm.v1.attention.backends.rocm_aiter_fa import (
AiterFlashAttentionMetadataBuilder,
AsmKvScaleTables,
)
tables = AsmKvScaleTables()
device = torch.device("cuda")
table = tables.get(device, 8, 2, 16, 0.5)
assert tables.get(device, 8, 2, 16, 0.5) is table
assert tables.get(device, 8, 2, 16, 1.0) is not table
assert table.shape == (8, 2, 16)
assert table.dtype == torch.float32
assert torch.all(table == 0.5)
# The builder class holds the single process-wide instance.
assert isinstance(
AiterFlashAttentionMetadataBuilder.asm_kv_scale_tables, AsmKvScaleTables
)
tests/v1/distributed/test_async_llm_dp.py
+3 -1
View File
@@ -368,9 +368,11 @@ async def test_dp_pause_late_request_does_not_block_drain():
# down first. The front-end sets its own copy optimistically when it
# forwards the wake, so sample only after several coordinator
# publishes (every 100ms while stats change) have overwritten it.
# Rank 0 reports the wave only after its first step, which can outlast
# the 2s on a cold device, so wait for the report rather than sample it.
await asyncio.sleep(2)
assert not long_request.done(), "the warmup request was too short to sample"
assert engine.engine_core.dp_engines_running(), (
assert await _poll_flag(engine, True, timeout=30), (
"the coordinator does not report the engines as running while they are"
)
vllm/_aiter_ops.py
+24 -9
View File
@@ -11,6 +11,7 @@ import torch.distributed as dist
from torch._ops import OpOverload
import vllm.envs as envs
from vllm import _custom_ops as ops
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.torch_utils import direct_register_custom_op
@@ -3887,17 +3888,29 @@ class rocm_aiter_ops:
v_scale: torch.Tensor,
kv_cache_dtype: str,
use_shuffle_layout: bool,
return_kv: bool,
) -> None:
"""Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches.
Shared by the AITER FA and unified-attention impls. The caller splits
Shared by the AITER FA, ROCM ATTN, and unified-attention impls. The caller splits
kv_cache, since the unbind dim depends on the layout (e.g. the unified
encoder-decoder path is K/V-first), and passes use_shuffle_layout
(unified reads NHD and must pass False).
(unified reads NHD and must pass False). k_scale/v_scale are the
per-tensor CPU scalars. return_kv says whether the impl consumes
k_out (FA and ROCM_ATTN prefill do; unified reads K from the cache and
passes False, so the kernel skips that store). For fp8 caches the
kernel emits k_out in the cache dtype, so a returned k_out is
converted to the activation dtype.
"""
if kv_cache_dtype.startswith("fp8"):
key_cache = key_cache.view(current_platform.fp8_dtype())
value_cache = value_cache.view(current_platform.fp8_dtype())
is_fp8_cache = kv_cache_dtype.startswith("fp8")
dequant_k_out = return_kv and is_fp8_cache
kernel_k_out = k_out
if is_fp8_cache:
fp8_dtype = current_platform.fp8_dtype()
key_cache = key_cache.view(fp8_dtype)
value_cache = value_cache.view(fp8_dtype)
if dequant_k_out:
kernel_k_out = torch.empty_like(k_out, dtype=fp8_dtype)
# Partial-rotary support (e.g. GLM-4.7 applies rotary to only a prefix
# of each head's channel dim).
rotary_dim = cos_sin_cache.shape[-1]
@@ -3920,14 +3933,16 @@ class rocm_aiter_ops:
slot_mapping=slot_mapping,
k_scale=k_scale,
v_scale=v_scale,
k_out=k_out,
k_out=kernel_k_out,
v_out=None,
return_kv=True,
return_kv=return_kv,
use_shuffle_layout=use_shuffle_layout,
block_size=key_cache.shape[1],
x=16 // key_cache.element_size(),
rotary_dim=kernel_rotary_dim,
)
if dequant_k_out:
ops.convert_fp8(k_out, kernel_k_out, float(k_scale), kv_cache_dtype)
@staticmethod
def fused_qk_norm_mrope_and_cache(
@@ -4506,8 +4521,8 @@ class rocm_aiter_ops:
scale: float,
K_QScale_hip: torch.Tensor,
V_QScale_hip: torch.Tensor,
K_QScale_asm: torch.Tensor,
V_QScale_asm: torch.Tensor,
K_QScale_asm: torch.Tensor | None,
V_QScale_asm: torch.Tensor | None,
out_: torch.Tensor,
kv_cache_dtype: str,
):
vllm/_custom_ops.py
+4
View File
@@ -122,6 +122,7 @@ def paged_attention_rocm(
v_scale: torch.Tensor,
fp8_out_scale: torch.Tensor | None = None,
mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16",
use_interleaved_v_cache: bool = False,
) -> None:
torch.ops._rocm_C.paged_attention(
out,
@@ -144,6 +145,7 @@ def paged_attention_rocm(
v_scale,
fp8_out_scale,
mfma_type,
use_interleaved_v_cache,
)
@@ -2554,6 +2556,7 @@ def reshape_and_cache(
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
interleaved_v: bool = False,
) -> None:
torch.ops._C_cache_ops.reshape_and_cache(
key,
@@ -2564,6 +2567,7 @@ def reshape_and_cache(
kv_cache_dtype,
k_scale,
v_scale,
interleaved_v,
)
vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py
+1 -1
View File
@@ -893,6 +893,7 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
layer.head_size,
)
continue
layer.impl.set_interleaved_v_cache()
if supports_rope:
for epsilon in [1e-5, 1e-6]:
for neox in [True, False]:
@@ -903,7 +904,6 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
is_neox=neox,
quant_query=quant_q,
).register(self.patterns)
if supports_mrope:
for section, is_interleaved in self.mrope_configs:
rotary_dim = 2 * sum(section)
vllm/model_executor/layers/fused_moe/modular_kernel.py
-5
View File
@@ -531,11 +531,6 @@ class FusedMoEExperts(ABC):
def is_monolithic() -> bool:
raise NotImplementedError("Implemented by subclasses.")
@staticmethod
def _supports_hash_routing() -> bool:
"""Whether a monolithic kernel implements hash routing."""
return False
@property
def expects_unquantized_inputs(self) -> bool:
"""Whether or not the PrepareFinalize should defer input quantization
vllm/v1/attention/backend.py
+10
View File
@@ -950,6 +950,16 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
"""Whether this implementation supports fused QKNorm+MRoPE+KVCache."""
return False
def set_interleaved_v_cache(self):
"""Opt this implementation into the interleaved V-cache layout for all of
its cache writers and readers. QkNormRopeKvCacheFusionPass calls this
for every eligible layer when it is constructed, i.e. before any KV
cache is written, so fused and unfused paths stay consistent whether or
not the pattern matches. Default is a no-op; backends whose kernels
read an interleaved V-cache under this fusion (e.g. ROCM_ATTN) override.
"""
return False
def fused_rope_kvcache_supported(self):
"""Does this attention implementation support RoPE+KVCache fusion.
This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
vllm/v1/attention/backends/rocm_aiter_fa.py
+95 -54
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Attention layer with AiterFlashAttention."""
from dataclasses import dataclass, replace
from dataclasses import dataclass, field, replace
from typing import ClassVar
import torch
@@ -152,8 +152,8 @@ if current_platform.is_rocm():
k_reg = tl.load(key_cache_ptr_offset + k_reg_offset)
v_reg = tl.load(value_cache_ptr_offset + v_reg_offset)
if DEQUANT:
k_scale = 1.0
v_scale = 1.0
k_scale = tl.load(k_scale_ptr)
v_scale = tl.load(v_scale_ptr)
k_reg = k_reg.to(tl.float32) * k_scale
v_reg = v_reg.to(tl.float32) * v_scale
tl.store(key_ptr_offset + col_offsets, k_reg)
@@ -270,8 +270,8 @@ if current_platform.is_rocm():
k_val = tl.load(key_ptr + src_offset_k + offset)
v_val = tl.load(value_ptr + src_offset_v + offset)
if QUANT:
k_scale = 1.0
v_scale = 1.0
k_scale = tl.load(k_scale_ptr)
v_scale = tl.load(v_scale_ptr)
k_dtype = key_cache_ptr.type.element_ty
v_dtype = value_cache_ptr.type.element_ty
k_val = (k_val.to(tl.float32) / k_scale).to(k_dtype)
@@ -410,11 +410,48 @@ class AiterFlashAttentionMetadata:
kv_sharing_metadata: AiterKVSharingMetadata | None = None
@dataclass
class AsmKvScaleTables:
"""Dense per-slot scale tables the asm decode kernel reads for an fp8 cache.
One table per distinct (device, shape, scale), shared by every layer and
every metadata builder in the process, so the target and draft models with
equal scales reuse a single allocation. Tables are constants: never write
to one in place; a per-token or dynamic scale needs its own buffer.
"""
tables: dict[tuple[str, int, int, int, float], torch.Tensor] = field(
default_factory=dict
)
def get(
self,
device: torch.device,
num_blocks: int,
num_kv_heads: int,
block_size: int,
scale: float,
) -> torch.Tensor:
key = (str(device), num_blocks, num_kv_heads, block_size, scale)
table = self.tables.get(key)
if table is None:
table = torch.full(
(num_blocks, num_kv_heads, block_size),
scale,
dtype=torch.float32,
device=device,
)
self.tables[key] = table
return table
class AiterFlashAttentionMetadataBuilder(
AttentionMetadataBuilder[AiterFlashAttentionMetadata]
):
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
asm_kv_scale_tables: ClassVar[AsmKvScaleTables] = AsmKvScaleTables()
def __init__(
self,
kv_cache_spec: AttentionSpec,
@@ -466,7 +503,8 @@ class AiterFlashAttentionMetadataBuilder(
dtype=self.model_config.dtype,
device=device,
)
self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device)
self.k_scale: dict[str, torch.Tensor] | None = None
self.v_scale: dict[str, torch.Tensor] | None = None
self.kv_sharing_workspace = (
torch.empty(
(
@@ -488,6 +526,38 @@ class AiterFlashAttentionMetadataBuilder(
common_prefix_len=0, common_attn_metadata=common_attn_metadata
)
def _build_asm_kv_scale_tables(self) -> None:
# Only for fp8 shuffle kv cache: the asm decode kernel reads a dense
# [num_blocks, num_kv_heads, block_size] scale table per layer. Map each
# layer to the shared table for its scale once the KV cache is bound.
if (
self.k_scale is not None
or not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
or not is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)
):
return
attn_layers = self.vllm_config.compilation_config.static_forward_context
k_scale_tables: dict[str, torch.Tensor] = {}
v_scale_tables: dict[str, torch.Tensor] = {}
for layer_name in self.layer_names:
layer = attn_layers[layer_name]
num_blocks = layer.kv_cache.shape[0]
k_scale_tables[layer_name] = self.asm_kv_scale_tables.get(
self.device,
num_blocks,
self.num_heads_kv,
self.block_size,
layer._k_scale_float,
)
v_scale_tables[layer_name] = self.asm_kv_scale_tables.get(
self.device,
num_blocks,
self.num_heads_kv,
self.block_size,
layer._v_scale_float,
)
self.k_scale, self.v_scale = k_scale_tables, v_scale_tables
def build(
self,
common_prefix_len: int,
@@ -499,24 +569,7 @@ class AiterFlashAttentionMetadataBuilder(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
)
# Allocate scales for fp8 shuffle kv cache with shuffle_kv_cache enabled
if (
rocm_aiter_ops.is_shuffle_kv_cache_enabled()
and self.scale.numel() == 1
and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)
):
# Size the scales from a layer this builder owns. The draft model
# runs its own builder over its own KV cache, so the first layer of
# the whole config can carry an unrelated block count.
kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[
self.layer_names[0]
].kv_cache.shape
num_blocks = kv_cache_shape[0]
self.scale = torch.ones(
[num_blocks, self.num_heads_kv, self.block_size],
dtype=torch.float32,
device=self.device,
)
self._build_asm_kv_scale_tables()
(
num_decodes,
num_extends,
@@ -709,8 +762,8 @@ class AiterFlashAttentionMetadataBuilder(
decode_metadata=decode_metadata,
prefill_metadata=prefill_metadata,
extend_metadata=extend_metadata,
k_scale=self.scale,
v_scale=self.scale,
k_scale=self.k_scale,
v_scale=self.v_scale,
kv_sharing_metadata=kv_sharing_metadata,
)
return attn_metadata
@@ -726,6 +779,7 @@ class AiterFlashAttentionMetadataBuilder(
skip split_decodes_prefills_and_extends() and avoid all .cpu() /
.item() calls that would otherwise break CUDA graph capture.
"""
self._build_asm_kv_scale_tables()
num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens
@@ -749,8 +803,8 @@ class AiterFlashAttentionMetadataBuilder(
decode_metadata=decode_metadata,
prefill_metadata=None,
extend_metadata=None,
k_scale=self.scale,
v_scale=self.scale,
k_scale=self.k_scale,
v_scale=self.v_scale,
)
@@ -1249,9 +1303,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
extend_outputs = output[extend_tokens_slice]
k_scale = layer._k_scale
v_scale = layer._v_scale
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
k_scale = attn_metadata.k_scale
v_scale = attn_metadata.v_scale
self.extend_forward(
attn_metadata=attn_metadata,
query=extend_queries,
@@ -1428,16 +1479,13 @@ class AiterFlashAttentionImpl(AttentionImpl):
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
k_qscale = (
layer._k_scale
if attn_metadata.k_scale is None
else attn_metadata.k_scale
)
v_qscale = (
layer._v_scale
if attn_metadata.v_scale is None
else attn_metadata.v_scale
)
# asm indexes this layer's dense scale buffer from build();
# the HIP kernel reads the per-tensor scale directly.
k_qscale_asm = v_qscale_asm = None
if attn_metadata.k_scale is not None:
assert attn_metadata.v_scale is not None
k_qscale_asm = attn_metadata.k_scale[layer.layer_name]
v_qscale_asm = attn_metadata.v_scale[layer.layer_name]
rocm_aiter_ops.paged_attention_common(
Q=query[:num_decode_tokens],
K=new_key_cache,
@@ -1452,10 +1500,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
0
),
scale=self.scale,
K_QScale_hip=k_qscale,
V_QScale_hip=v_qscale,
K_QScale_asm=k_qscale,
V_QScale_asm=v_qscale,
K_QScale_hip=layer._k_scale,
V_QScale_hip=layer._v_scale,
K_QScale_asm=k_qscale_asm,
V_QScale_asm=v_qscale_asm,
out_=output[:num_decode_tokens],
kv_cache_dtype=self.kv_cache_dtype,
)
@@ -1542,9 +1590,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
# the reshape_and_cache_flash op uses the slot_mapping's shape
# to determine the number of actual tokens.
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
# We may calculate per token quant scale in
# reshape_and_cache_shuffle_triton which might differ from
# vllm's style when shuffle layout is used.
k_scale = layer._k_scale
v_scale = layer._v_scale
assert k_scale is not None and v_scale is not None, (
@@ -1581,12 +1626,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
)
def fused_qk_norm_rope_kvcache_supported(self):
# Only fuse when shuffle layout is off; the shuffle write path uses a
# dedicated cache update, mirroring fused_rope_kvcache_supported.
return (
rocm_aiter_ops.is_enabled()
and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
)
return rocm_aiter_ops.is_enabled()
def do_qk_norm_rope_kvcache_update(
self,
@@ -1624,6 +1664,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
v_scale=layer._v_scale_cpu,
kv_cache_dtype=self.kv_cache_dtype,
use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(),
return_kv=True,
)
def do_rope_and_kv_cache_update(
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+6
View File
@@ -413,6 +413,11 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
def fused_qk_norm_mrope_kvcache_supported(self) -> bool:
return IS_AITER_FOUND
def set_interleaved_v_cache(self):
# No-op override of RocmAttentionImpl: this backend
# decodes with the AITER triton unified-attention kernel
pass
def do_qk_norm_rope_kvcache_update(
self,
layer: AttentionLayer,
@@ -449,6 +454,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
v_scale=layer._v_scale_cpu,
kv_cache_dtype=self.kv_cache_dtype,
use_shuffle_layout=self._use_shuffle_kv,
return_kv=False,
)
def do_qk_norm_mrope_kvcache_update(
vllm/v1/attention/backends/rocm_attn.py
+63
View File
@@ -292,6 +292,8 @@ class RocmAttentionImpl(AttentionImpl):
f"num_heads: {num_heads}."
)
self._use_interleaved_v_cache = False
def _forward_encoder_attention(
self,
query: torch.Tensor,
@@ -450,6 +452,7 @@ class RocmAttentionImpl(AttentionImpl):
sm_scale=self.scale,
output_scale=output_scale,
sinks=self.sinks,
use_interleaved_v_cache=self._use_interleaved_v_cache,
causal=attn_metadata.causal,
)
@@ -475,6 +478,12 @@ class RocmAttentionImpl(AttentionImpl):
block_size = value_cache.shape[3]
has_native_layout = has_native_kv_cache_layout(key_cache, value_cache)
# Match the fused writer's interleaved V-cache when fusion is active.
interleaved_v = self._use_interleaved_v_cache
interleaved_v_pack_factor = (
(16 // value_cache.element_size()) if interleaved_v else 0
)
if block_size in (16, 32) and has_native_layout:
# Normal 16, 32 with contiguous blocks: use vLLM native HIP C++ logic.
PagedAttention.write_to_paged_cache(
@@ -486,6 +495,7 @@ class RocmAttentionImpl(AttentionImpl):
self.kv_cache_dtype,
layer._k_scale,
layer._v_scale,
interleaved_v,
)
else:
# Non-standard blocks and hybrid attention/Mamba layouts need the
@@ -501,8 +511,61 @@ class RocmAttentionImpl(AttentionImpl):
self.kv_cache_dtype,
layer._k_scale,
layer._v_scale,
interleaved_v_pack_factor,
)
def fused_qk_norm_rope_kvcache_supported(self):
return rocm_aiter_ops.is_enabled()
def set_interleaved_v_cache(self):
# Decode/prefill kernels must read the interleaved V-cache the AITER
# fused write emits when use_shuffle_layout=True.
self._use_interleaved_v_cache = True
def do_qk_norm_rope_kvcache_update(
self,
layer: AttentionLayer,
qkv: torch.Tensor,
q_out: torch.Tensor,
k_out: torch.Tensor,
positions: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
rms_norm_eps: float,
cos_sin_cache: torch.Tensor,
is_neox: bool,
kv_cache: torch.Tensor,
layer_slot_mapping: torch.Tensor,
):
# kv_cache is logical [num_blocks, 2, block_size, num_kv_heads * head_size]
# (see forward); the AITER shuffle writer takes one (B, N, H, hs) per side.
key_cache, value_cache = kv_cache.unflatten(
-1, (self.num_kv_heads, self.head_size)
).unbind(1)
use_shuffle_layout = self._use_interleaved_v_cache
rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
qkv=qkv,
q_weight=q_weight,
k_weight=k_weight,
cos_sin_cache=cos_sin_cache,
positions=positions,
num_heads_q=self.num_heads,
num_heads_k=self.num_kv_heads,
head_dim=self.head_size,
is_neox=is_neox,
rms_norm_eps=rms_norm_eps,
q_out=q_out,
k_out=k_out,
key_cache=key_cache,
value_cache=value_cache,
slot_mapping=layer_slot_mapping,
k_scale=layer._k_scale_cpu,
v_scale=layer._v_scale_cpu,
kv_cache_dtype=self.kv_cache_dtype,
use_shuffle_layout=use_shuffle_layout,
return_kv=True,
)
def fused_rope_kvcache_supported(self):
return rocm_aiter_ops.is_enabled()
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
+27 -7
View File
@@ -84,6 +84,7 @@ def kernel_paged_attention_2d(
query_start_len_ptr, # [num_seqs+1]
USE_SINKS: tl.constexpr, # bool
USE_FP8: tl.constexpr,
INTERLEAVED_V_PACK_FACTOR: tl.constexpr,
FP8_MIN: tl.constexpr = float8_info.min,
FP8_MAX: tl.constexpr = float8_info.max,
):
@@ -170,13 +171,24 @@ def kernel_paged_attention_2d(
+ (offs_d[:, None] % x) * stride_k_cache_4
)
# 4D addressing logic of V (Slot is innermost)
v_offset = (
p_block_idx[:, None] * stride_v_cache_0
+ kv_head_idx * stride_v_cache_1
+ offs_d[None, :] * stride_v_cache_2
+ internal_offsets[:, None] * stride_v_cache_3
)
# V addressing: standard 4D or interleaved 5D
if INTERLEAVED_V_PACK_FACTOR > 0:
v_offset = (
p_block_idx[:, None] * stride_v_cache_0
+ kv_head_idx * stride_v_cache_1
+ (internal_offsets[:, None] // INTERLEAVED_V_PACK_FACTOR)
* HEAD_SIZE
* INTERLEAVED_V_PACK_FACTOR
+ offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR
+ (internal_offsets[:, None] % INTERLEAVED_V_PACK_FACTOR)
)
else:
v_offset = (
p_block_idx[:, None] * stride_v_cache_0
+ kv_head_idx * stride_v_cache_1
+ offs_d[None, :] * stride_v_cache_2
+ internal_offsets[:, None] * stride_v_cache_3
)
# Only the final tile can straddle seq_len. Slots >= seq_len are
# unwritten KV cache that may hold NaN/garbage; they are score-masked
@@ -304,6 +316,7 @@ def chunked_prefill_paged_decode(
# Optional tensor for sinks
sinks=None,
is_block_table_ptr: bool = False,
use_interleaved_v_cache: bool = False,
causal: bool = True,
):
if sm_scale is None:
@@ -314,6 +327,10 @@ def chunked_prefill_paged_decode(
if sliding_window is None or sliding_window <= 0:
sliding_window = 0
interleaved_v_pack_factor = (
(16 // value_cache.element_size()) if use_interleaved_v_cache else 0
)
if max_query_len > 1:
context_attention_fwd(
q=query,
@@ -336,6 +353,7 @@ def chunked_prefill_paged_decode(
skip_decode=True,
fp8_out_scale=output_scale,
sinks=sinks,
interleaved_v_pack_factor=interleaved_v_pack_factor,
causal=causal,
)
@@ -428,6 +446,7 @@ def chunked_prefill_paged_decode(
k_scale=k_scale,
v_scale=v_scale,
fp8_out_scale=output_scale,
use_interleaved_v_cache=use_interleaved_v_cache,
)
else:
logger.warning_once(
@@ -504,4 +523,5 @@ def chunked_prefill_paged_decode(
query_start_len_ptr=query_start_loc,
USE_SINKS=sinks is not None,
USE_FP8=output_scale is not None,
INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor,
)
vllm/v1/attention/ops/paged_attn.py
+2
View File
@@ -38,6 +38,7 @@ class PagedAttention:
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
interleaved_v: bool,
) -> None:
ops.reshape_and_cache(
key,
@@ -48,4 +49,5 @@ class PagedAttention:
kv_cache_dtype,
k_scale,
v_scale,
interleaved_v,
)
vllm/v1/attention/ops/prefix_prefill.py
+51 -13
View File
@@ -56,6 +56,8 @@ def _paged_kv_cache_offsets(
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE: tl.constexpr,
BLOCK_DMODEL: tl.constexpr,
INTERLEAVED_V_PACK_FACTOR: tl.constexpr = 0,
MASK_BLOCK_TABLE: tl.constexpr = False,
):
"""Compute paged K/V cache element offsets for a tile of token positions.
@@ -65,7 +67,9 @@ def _paged_kv_cache_offsets(
tiles), then the physical element offsets are built with the same layout the
kernel uses everywhere:
K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
V cache: [num_blocks, num_kv_heads, head_size, block_size]
V cache: [num_blocks, num_kv_heads, head_size, block_size], or the
interleaved layout [num_blocks, num_kv_heads, block_size // x,
head_size, x] when `INTERLEAVED_V_PACK_FACTOR` (= x) > 0.
Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively.
When `MASK_BLOCK_TABLE` is set, the block-table load is masked with
@@ -94,12 +98,25 @@ def _paged_kv_cache_offsets(
+ internal[None, :] * stride_k_cache_bl
+ (offs_d[:, None] % x) * stride_k_cache_x
)
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ internal[:, None] * stride_v_cache_bl
)
if INTERLEAVED_V_PACK_FACTOR > 0:
# Interleaved V-cache: [num_blocks, num_kv_heads, block_size/x,
# head_size, x] flattened onto the 4D stride layout.
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ (internal[:, None] // INTERLEAVED_V_PACK_FACTOR)
* BLOCK_DMODEL
* INTERLEAVED_V_PACK_FACTOR
+ offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR
+ (internal[:, None] % INTERLEAVED_V_PACK_FACTOR)
)
else:
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ internal[:, None] * stride_v_cache_bl
)
return off_k, off_v
@@ -159,6 +176,7 @@ def _fwd_kernel(
USE_FP8: tl.constexpr,
CAUSAL: tl.constexpr = True,
MAX_Q_LEN: tl.constexpr = 0,
INTERLEAVED_V_PACK_FACTOR: tl.constexpr = 0,
MAX_CTX_LEN: tl.constexpr = 0,
FP8_MIN: tl.constexpr = float8_info.min,
FP8_MAX: tl.constexpr = float8_info.max,
@@ -263,6 +281,8 @@ def _fwd_kernel(
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE,
BLOCK_DMODEL,
INTERLEAVED_V_PACK_FACTOR,
)
if (
@@ -403,6 +423,8 @@ def _fwd_kernel(
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE,
BLOCK_DMODEL,
INTERLEAVED_V_PACK_FACTOR,
MASK_BLOCK_TABLE=True,
)
k_cur_load = tl.load(
@@ -540,6 +562,7 @@ def _fwd_kernel_alibi(
BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
BLOCK_N: tl.constexpr,
SKIP_DECODE: tl.constexpr,
INTERLEAVED_V_PACK_FACTOR: tl.constexpr,
):
# attn_bias[]
cur_batch = tl.program_id(0)
@@ -608,12 +631,24 @@ def _fwd_kernel_alibi(
+ ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl
+ (offs_d[:, None] % x) * stride_k_cache_x
)
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl
)
if INTERLEAVED_V_PACK_FACTOR > 0:
v_internal = (start_n + offs_n[:, None]) % block_size
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ (v_internal // INTERLEAVED_V_PACK_FACTOR)
* BLOCK_DMODEL
* INTERLEAVED_V_PACK_FACTOR
+ offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR
+ (v_internal % INTERLEAVED_V_PACK_FACTOR)
)
else:
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl
)
k_load = tl.load(
K_cache + off_k,
mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
@@ -796,6 +831,7 @@ def context_attention_fwd(
fp8_out_scale=None,
sinks=None,
is_block_table_ptr: bool = False,
interleaved_v_pack_factor: int = 0,
causal: bool = True,
):
q_dtype_is_f32 = q.dtype is torch.float32
@@ -937,6 +973,7 @@ def context_attention_fwd(
BLOCK_DMODEL_PADDED=Lk_padded,
BLOCK_N=BLOCK,
SKIP_DECODE=skip_decode,
INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor,
num_warps=NUM_WARPS,
num_stages=1,
)
@@ -1010,6 +1047,7 @@ def context_attention_fwd(
num_warps=4,
num_stages=1,
USE_SINKS=sinks is not None,
INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor,
CAUSAL=causal,
KV_FROM_CACHE=kv_from_cache,
**extra_kargs,
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+30 -7
View File
@@ -52,6 +52,9 @@ def reshape_and_cache_kernel_flash(
block_size: tl.constexpr,
x: tl.constexpr,
USE_HEAD_MAJOR_LAYOUT: tl.constexpr,
# >0 writes V interleaved as [Block, Head, Slot//kx, Dim, kx] (head-major
# layout only); 0 keeps the standard [Block, Head, Dim, Slot] V layout.
INTERLEAVED_V_PACK_FACTOR: tl.constexpr,
# FP8 flags
FP8_KV_CACHE: tl.constexpr,
# tune parameters
@@ -76,13 +79,25 @@ def reshape_and_cache_kernel_flash(
# Decompose the tile index back into head and dim coordinates.
cur_head = tile_pos // head_size
cur_dim = tile_pos % head_size
# Value addressing (4D): [Block, Head, Dim, Slot]
tgt_idx_v = (
block_idx * block_stride
+ cur_head * head_stride
+ cur_dim * dim_stride_v
+ block_offset * 1
)
# Value addressing: standard 4D [Block, Head, Dim, Slot], or the
# interleaved 5D [Block, Head, Slot//kx, Dim, kx] the fused AITER
# writer emits and the interleaved read path expects.
if INTERLEAVED_V_PACK_FACTOR > 0:
tgt_idx_v = (
block_idx * block_stride
+ cur_head * head_stride
+ (block_offset // INTERLEAVED_V_PACK_FACTOR)
* (head_size * INTERLEAVED_V_PACK_FACTOR)
+ cur_dim * INTERLEAVED_V_PACK_FACTOR
+ (block_offset % INTERLEAVED_V_PACK_FACTOR)
)
else:
tgt_idx_v = (
block_idx * block_stride
+ cur_head * head_stride
+ cur_dim * dim_stride_v
+ block_offset * 1
)
# Key addressing (5D): [Block, Head, Dim//8, Slot, 8]
tgt_idx_k = (
block_idx * block_stride
@@ -371,11 +386,18 @@ def triton_reshape_and_cache_flash(
kv_cache_dtype: str, # "auto", "fp8"
k_scale: torch.Tensor, # float32
v_scale: torch.Tensor, # float32
interleaved_v_pack_factor: int = 0, # >0: interleave V (head-major); 0: standard
):
num_heads = key.shape[1]
head_size = key.shape[2]
use_head_major_layout = key_cache.ndim == 5
if interleaved_v_pack_factor and not use_head_major_layout:
raise NotImplementedError(
"Interleaved V-cache write is only implemented for the head-major "
"(5D) cache layout; got a non-head-major cache with "
f"interleaved_v_pack_factor={interleaved_v_pack_factor}."
)
if use_head_major_layout:
block_size = key_cache.shape[3]
x = key_cache.shape[4]
@@ -452,6 +474,7 @@ def triton_reshape_and_cache_flash(
block_size=block_size,
x=x,
USE_HEAD_MAJOR_LAYOUT=use_head_major_layout,
INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor,
FP8_KV_CACHE=FP8_KV_CACHE,
# autotune parameters
TILE_SIZE=TILE_SIZE,