diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 24453ff1c1..351bcf24f9 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -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(), key_stride, value_stride, \ num_heads, head_size, block_size, x, \ reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr())); + reinterpret_cast(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); diff --git a/csrc/ops.h b/csrc/ops.h index 3266d28b1e..549dc84a1c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -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, diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index caf6ee9a14..7106840964 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -355,7 +355,8 @@ __device__ float warpReduceMax(float val) { // clang-format off template + 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( 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(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(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( + 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(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 + 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(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(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(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(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 + 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 + 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 + 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 + 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 + 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 + 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 \ + 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> \ <<>>( \ 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 \ + paged_attention_ll4mi_QKV_mfma4_kernel< \ + T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED, \ + GQA_RATIO, USE_INTERLEAVED_V_CACHE> \ <<>>( \ 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 + 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 + 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( \ - 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( \ - 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( \ + 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( \ + 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( \ + 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& 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); diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index 78751841ca..de963913ea 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -50,4 +50,4 @@ void paged_attention( int64_t max_seq_len, const std::optional& alibi_slopes, const std::string& kv_cache_dtype, torch::Tensor& k_scale, torch::Tensor& v_scale, const std::optional& fp8_out_scale, - const std::string& mfma_type); + const std::string& mfma_type, bool use_interleaved_v_cache); diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 14d1822707..077c3f35ce 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -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); } diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3d4f9ab273..a92d9cc114 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -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( diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index c32469bfe2..cb421687a4 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -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, diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py index e3c3ac0b24..a81c973852 100644 --- a/tests/kernels/attention/test_attention.py +++ b/tests/kernels/attention/test_attention.py @@ -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. diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index 3f675e15e6..c93bb9259d 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -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, ) diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index c3b2c0754f..f6b3fb92cd 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -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 + ) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index f4b4ca2a79..7b350fc55e 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -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" ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 73f8f4d513..0f57c9ecb4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -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, ): diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 1af864bc74..70fed78573 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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, ) diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py index f418e8a7ef..35803e1347 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py @@ -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) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index ef3319bf12..fccbd1a9d8 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 466bba586d..d3c8d7ea24 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -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 diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 0ec0d5dd36..04c614a1d2 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -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( diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 4a9f67a637..b05482c007 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -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( diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 471bb9f058..16b4c38903 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -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() diff --git a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py b/vllm/v1/attention/ops/chunked_prefill_paged_decode.py index f48018e8fc..5f9e28d7ea 100644 --- a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py +++ b/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @@ -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, ) diff --git a/vllm/v1/attention/ops/paged_attn.py b/vllm/v1/attention/ops/paged_attn.py index 896e929b54..3215a2f04b 100644 --- a/vllm/v1/attention/ops/paged_attn.py +++ b/vllm/v1/attention/ops/paged_attn.py @@ -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, ) diff --git a/vllm/v1/attention/ops/prefix_prefill.py b/vllm/v1/attention/ops/prefix_prefill.py index 2577f52ca5..2fd9df983a 100644 --- a/vllm/v1/attention/ops/prefix_prefill.py +++ b/vllm/v1/attention/ops/prefix_prefill.py @@ -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, diff --git a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py index 882b6787ae..47261bbc73 100644 --- a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -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,