A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Kernel] support merge_attn_states CUDA kernel, 3x speedup (#16173)
Signed-off-by: DefTruth <qiustudent_r@163.com>
DefTruth committed
Apr 11, 2025 at 20:50 UTC
e9528f6dc614879952aa871fae8df296adcfc559
10 files changed
+519
-4
CMakeLists.txt
+1
@@ -230,6 +230,7 @@ set(VLLM_EXT_SRC
230
"csrc/cache_kernels.cu"
231
"csrc/attention/paged_attention_v1.cu"
232
"csrc/attention/paged_attention_v2.cu"
233
+ "csrc/attention/merge_attn_states.cu"
234
"csrc/pos_encoding_kernels.cu"
235
"csrc/activation_kernels.cu"
236
"csrc/layernorm_kernels.cu"
csrc/attention/merge_attn_states.cu
new
+173
@@ -0,0 +1,173 @@
1
+#include <optional>
2
+#include <torch/all.h>
3
+#include <ATen/cuda/CUDAContext.h>
4
+#include <c10/cuda/CUDAGuard.h>
5
+#include <algorithm>
6
+
7
+#include "attention_dtypes.h"
8
+#include "attention_utils.cuh"
9
+
10
+namespace vllm {
11
+
12
+// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
13
+// can be used to combine partial attention results (in the split-KV case)
14
+template <typename scalar_t, const uint NUM_THREADS>
15
+__global__ void merge_attn_states_kernel(
16
+ scalar_t* output, float* output_lse, const scalar_t* prefix_output,
17
+ const float* prefix_lse, const scalar_t* suffix_output,
18
+ const float* suffix_lse, const uint num_tokens, const uint num_heads,
19
+ const uint head_size) {
20
+ using pack_128b_t = uint4;
21
+ const uint pack_size = 16 / sizeof(scalar_t);
22
+ const uint threads_per_head = head_size / pack_size;
23
+
24
+ const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x;
25
+ const uint token_head_threads = num_tokens * num_heads * threads_per_head;
26
+
27
+ if (global_idx >= token_head_threads) return;
28
+
29
+ // global_idx -> token_idx + head_idx + pack_idx
30
+ const uint token_head_idx = global_idx / threads_per_head;
31
+ const uint pack_idx = global_idx % threads_per_head;
32
+
33
+ const uint token_idx = token_head_idx / num_heads;
34
+ const uint head_idx = token_head_idx % num_heads;
35
+
36
+ const uint pack_offset = pack_idx * pack_size; // (0~15)*8, etc.
37
+ const uint head_offset =
38
+ token_idx * num_heads * head_size + head_idx * head_size;
39
+ const scalar_t* prefix_head_ptr = prefix_output + head_offset;
40
+ const scalar_t* suffix_head_ptr = suffix_output + head_offset;
41
+ scalar_t* output_head_ptr = output + head_offset;
42
+
43
+ float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
44
+ float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
45
+ p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
46
+ s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
47
+
48
+ const float max_lse = fmaxf(p_lse, s_lse);
49
+ p_lse = p_lse - max_lse;
50
+ s_lse = s_lse - max_lse;
51
+ const float p_se = expf(p_lse);
52
+ const float s_se = expf(s_lse);
53
+ const float out_se = p_se + s_se;
54
+ const float p_scale = p_se / out_se;
55
+ const float s_scale = s_se / out_se;
56
+
57
+ if (pack_offset < head_size) {
58
+ // Pack 128b load
59
+ pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
60
+ prefix_head_ptr)[pack_offset / pack_size];
61
+ pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
62
+ suffix_head_ptr)[pack_offset / pack_size];
63
+ pack_128b_t o_out_pack;
64
+
65
+#pragma unroll
66
+ for (uint i = 0; i < pack_size; ++i) {
67
+ // Always use float for FMA to keep high precision.
68
+ // half(uint16_t), bfloat16, float -> float.
69
+ const float p_out_f =
70
+ vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
71
+ const float s_out_f =
72
+ vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
73
+ // fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale)
74
+ const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale);
75
+ // float -> half(uint16_t), bfloat16, float.
76
+ vllm::from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i], o_out_f);
77
+ }
78
+
79
+ // Pack 128b storage
80
+ reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
81
+ o_out_pack;
82
+ }
83
+ // We only need to write to output_lse once per head.
84
+ if (output_lse != nullptr && pack_idx == 0) {
85
+ float out_lse = logf(out_se) + max_lse;
86
+ output_lse[head_idx * num_tokens + token_idx] = out_lse;
87
+ }
88
+}
89
+
90
+} // namespace vllm
91
+
92
+// The following macro is used to dispatch the conversion function based on
93
+// the output data type. The FN is a macro that calls a function with
94
+// template<typename scalar_t>.
95
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
96
+ { \
97
+ if (scalar_dtype == at::ScalarType::Float) { \
98
+ fn(float); \
99
+ } else if (scalar_dtype == at::ScalarType::Half) { \
100
+ fn(uint16_t); \
101
+ } else if (scalar_dtype == at::ScalarType::BFloat16) { \
102
+ fn(__nv_bfloat16); \
103
+ } else { \
104
+ TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
105
+ } \
106
+ }
107
+
108
+#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS) \
109
+ { \
110
+ vllm::merge_attn_states_kernel<scalar_t, NUM_THREADS><<<grid, block>>>( \
111
+ reinterpret_cast<scalar_t*>(output.data_ptr()), output_lse_ptr, \
112
+ reinterpret_cast<scalar_t*>(prefix_output.data_ptr()), \
113
+ reinterpret_cast<float*>(prefix_lse.data_ptr()), \
114
+ reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \
115
+ reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens, \
116
+ num_heads, head_size); \
117
+ }
118
+
119
+/*@brief Merges the attention states from prefix and suffix
120
+ * into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d
121
+ *
122
+ * @param output [n,h,d] The output tensor to store the merged attention states.
123
+ * @param output_lse [h,d] Optional tensor to store the log-sum-exp values.
124
+ * @param prefix_output [n,h,d] The prefix attention states.
125
+ * @param prefix_lse [h,d] The log-sum-exp values for the prefix attention
126
+ * states.
127
+ * @param suffix_output [n,h,d] The suffix attention states.
128
+ * @param suffix_lse [h,d] The log-sum-exp values for the suffix attention
129
+ * states.
130
+ */
131
+template <typename scalar_t>
132
+void merge_attn_states_launcher(torch::Tensor& output,
133
+ std::optional<torch::Tensor> output_lse,
134
+ const torch::Tensor& prefix_output,
135
+ const torch::Tensor& prefix_lse,
136
+ const torch::Tensor& suffix_output,
137
+ const torch::Tensor& suffix_lse) {
138
+ constexpr uint NUM_THREADS = 128;
139
+ const uint num_tokens = output.size(0);
140
+ const uint num_heads = output.size(1);
141
+ const uint head_size = output.size(2);
142
+ const uint pack_size = 16 / sizeof(scalar_t);
143
+ TORCH_CHECK(head_size % pack_size == 0,
144
+ "headsize must be multiple of pack_size:", pack_size);
145
+ float* output_lse_ptr = nullptr;
146
+ if (output_lse.has_value()) {
147
+ output_lse_ptr = output_lse.value().data_ptr<float>();
148
+ }
149
+ // process one pack elements per thread. float -> 4, half/bf16 -> 8
150
+ const uint threads_per_head = head_size / pack_size;
151
+ const uint total_threads = num_tokens * num_heads * threads_per_head;
152
+
153
+ dim3 block(NUM_THREADS);
154
+ dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
155
+
156
+ LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);
157
+}
158
+
159
+#define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \
160
+ { \
161
+ merge_attn_states_launcher<scalar_t>(output, output_lse, prefix_output, \
162
+ prefix_lse, suffix_output, \
163
+ suffix_lse); \
164
+ }
165
+
166
+void merge_attn_states(torch::Tensor& output,
167
+ std::optional<torch::Tensor> output_lse,
168
+ const torch::Tensor& prefix_output,
169
+ const torch::Tensor& prefix_lse,
170
+ const torch::Tensor& suffix_output,
171
+ const torch::Tensor& suffix_lse) {
172
+ DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);
173
+}
csrc/ops.h
+9
@@ -52,6 +52,15 @@ void paged_attention_v2(
52
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
53
const int64_t blocksparse_head_sliding_step);
54
55
+#ifndef USE_ROCM
56
+void merge_attn_states(torch::Tensor& output,
57
+ std::optional<torch::Tensor> output_lse,
58
+ const torch::Tensor& prefix_output,
59
+ const torch::Tensor& prefix_lse,
60
+ const torch::Tensor& suffix_output,
61
+ const torch::Tensor& suffix_lse);
62
+#endif
63
+
64
void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
65
double epsilon);
66
csrc/torch_bindings.cpp
+15
@@ -64,6 +64,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
64
" int blocksparse_head_sliding_step) -> ()");
65
ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
66
67
+#ifndef USE_ROCM
68
+ // Merge attn states
69
+ // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
70
+ // can be used to combine partial attention results (in the split-KV case)
71
+ ops.def(
72
+ "merge_attn_states("
73
+ " Tensor! output,"
74
+ " Tensor!? output_lse,"
75
+ " Tensor prefix_output,"
76
+ " Tensor prefix_lse,"
77
+ " Tensor suffix_output,"
78
+ " Tensor suffix_lse) -> ()");
79
+ ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
80
+#endif
81
+
82
// Activation ops
83
// Activation function used in SwiGLU.
84
ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
tests/kernels/test_merge_attn_states.py
new
+265
@@ -0,0 +1,265 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
8
+from vllm.attention.ops.triton_merge_attn_states import (
9
+ merge_attn_states as merge_attn_states_triton)
10
+from vllm.platforms import current_platform
11
+
12
+
13
+# Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
14
+# can be used to combine partial attention results (in the split-KV case)
15
+def merge_attn_states_torch(
16
+ output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
17
+ prefix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
18
+ prefix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]
19
+ suffix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
20
+ suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]
21
+ output_lse: Optional[torch.Tensor] = None, # [NUM_HEADS, NUM_TOKENS]
22
+):
23
+ p_lse = prefix_lse
24
+ s_lse = suffix_lse
25
+ # inf -> -inf
26
+ p_lse[p_lse == torch.inf] = -torch.inf
27
+ s_lse[s_lse == torch.inf] = -torch.inf
28
+ # max_lse [NUM_HEADS, NUM_TOKENS]
29
+ max_lse = torch.maximum(p_lse, s_lse)
30
+ p_lse = p_lse - max_lse
31
+ s_lse = s_lse - max_lse
32
+ p_lse_exp = torch.exp(p_lse)
33
+ s_lse_exp = torch.exp(s_lse)
34
+ out_se = (p_lse_exp + s_lse_exp)
35
+ if output_lse is not None:
36
+ output_lse = torch.log(out_se) + max_lse
37
+ p_scale = p_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]
38
+ s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]
39
+ p_scale = torch.transpose(p_scale, 0,
40
+ 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
41
+ s_scale = torch.transpose(s_scale, 0,
42
+ 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
43
+ output = prefix_output * p_scale + suffix_output * s_scale
44
+ return output, output_lse
45
+
46
+
47
+NUM_BATCH_TOKENS = [256, 512, 613, 1024, 1536, 4096]
48
+NUM_QUERY_HEADS = [4, 8, 16, 32, 48, 64]
49
+HEAD_SIZES = [32, 48, 64, 96, 128, 256]
50
+DTYPES = [torch.float32, torch.half, torch.bfloat16]
51
+
52
+all_case_info: list[tuple] = []
53
+
54
+
55
+def generate_markdown_table():
56
+ global all_case_info
57
+ table_header = ("| tokens | heads | headsize | dtype "
58
+ "| device | torch | triton | cuda | speedup |")
59
+ table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"
60
+
61
+ def shortly_dtype(dtype: torch.dtype) -> str:
62
+ return str(dtype).removeprefix("torch.")
63
+
64
+ def shortly_device(device: str) -> str:
65
+ return device.removeprefix("NVIDIA").strip()
66
+
67
+ print(table_header)
68
+ print(table_separator)
69
+ for info in all_case_info:
70
+ (num_tokens, num_heads, head_size, dtype, device,
71
+ avg_time_torch_kernel, avg_time_triton_kernel, avg_time_cuda_kernel,
72
+ performance_improved) = info
73
+ dtype = shortly_dtype(dtype)
74
+ device = shortly_device(device)
75
+ print(f"| {num_tokens} | {num_heads} | {head_size} "
76
+ f"| {dtype} | {device} | {avg_time_torch_kernel:.5f}ms "
77
+ f"| {avg_time_triton_kernel:.5f}ms "
78
+ f"| {avg_time_cuda_kernel:.5f}ms "
79
+ f"| {performance_improved:.4f}x |")
80
+
81
+
82
+@pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)
83
+@pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)
84
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
85
+@pytest.mark.parametrize("output_dtype", DTYPES)
86
+@torch.inference_mode()
87
+def test_merge_attn_states(num_tokens: int, num_query_heads: int,
88
+ head_size: int, output_dtype: torch.dtype):
89
+ if not current_platform.is_cuda():
90
+ pytest.skip('Currently only support compare triton merge_attn_states '
91
+ 'with custom cuda merge_attn_states kernel')
92
+
93
+ NUM_TOKENS = num_tokens
94
+ NUM_HEADS = num_query_heads
95
+ HEAD_SIZE = head_size
96
+
97
+ print(f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "
98
+ f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, "
99
+ f"Device: {current_platform.get_device_name()}")
100
+
101
+ # prefix_lse and suffix_lse contain inf and normal values
102
+ prefix_lse = torch.randn(NUM_HEADS,
103
+ NUM_TOKENS,
104
+ dtype=torch.float32,
105
+ device="cuda")
106
+ suffix_lse = torch.randn(NUM_HEADS,
107
+ NUM_TOKENS,
108
+ dtype=torch.float32,
109
+ device="cuda")
110
+
111
+ # Generate boolean masks
112
+ mask_prefix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1
113
+ mask_suffix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1
114
+ # Ensure that the same position is not True at the same time
115
+ combined_mask = torch.logical_and(mask_prefix, mask_suffix)
116
+ mask_prefix = torch.logical_and(mask_prefix, ~combined_mask)
117
+ mask_suffix = torch.logical_and(mask_suffix, ~combined_mask)
118
+
119
+ prefix_lse[mask_prefix] = float('inf')
120
+ suffix_lse[mask_suffix] = float('inf')
121
+
122
+ # Other input tensors (need to be initialized but
123
+ # no actual calculation needed)
124
+ output = torch.zeros((NUM_TOKENS, NUM_HEADS, HEAD_SIZE),
125
+ dtype=output_dtype,
126
+ device="cuda")
127
+ output_lse = torch.zeros((NUM_HEADS, NUM_TOKENS),
128
+ dtype=torch.float32,
129
+ device="cuda")
130
+ prefix_output = torch.randn((NUM_TOKENS, NUM_HEADS, HEAD_SIZE),
131
+ dtype=output_dtype,
132
+ device="cuda")
133
+ suffix_output = torch.randn((NUM_TOKENS, NUM_HEADS, HEAD_SIZE),
134
+ dtype=output_dtype,
135
+ device="cuda")
136
+
137
+ warmup_times = 2
138
+ repeat_times = 20
139
+
140
+ output_torch = output.clone()
141
+ output_lse_torch = output_lse.clone()
142
+ total_time_torch_kernel = 0
143
+ start = torch.cuda.Event(enable_timing=True)
144
+ end = torch.cuda.Event(enable_timing=True)
145
+
146
+ # 0. Run the Torch kernel
147
+ prefix_lse_torch = prefix_lse.clone()
148
+ suffix_lse_torch = suffix_lse.clone()
149
+ for _ in range(warmup_times):
150
+ output_torch, output_lse_torch = merge_attn_states_torch(
151
+ output_torch, prefix_output, prefix_lse_torch, suffix_output,
152
+ suffix_lse_torch, output_lse_torch)
153
+ torch.cuda.synchronize()
154
+
155
+ for _ in range(repeat_times):
156
+ start.record()
157
+ output_torch, output_lse_torch = merge_attn_states_torch(
158
+ output_torch, prefix_output, prefix_lse_torch, suffix_output,
159
+ suffix_lse_torch, output_lse_torch)
160
+ end.record()
161
+ torch.cuda.synchronize()
162
+ total_time_torch_kernel += start.elapsed_time(end)
163
+
164
+ avg_time_torch_kernel = total_time_torch_kernel / repeat_times
165
+
166
+ # 1. Run the Triton kernel
167
+ output_ref_triton = output.clone()
168
+ output_lse_ref_triton = output_lse.clone()
169
+
170
+ total_time_triton_kernel = 0
171
+ start = torch.cuda.Event(enable_timing=True)
172
+ end = torch.cuda.Event(enable_timing=True)
173
+
174
+ for _ in range(warmup_times):
175
+ merge_attn_states_triton(output_ref_triton, prefix_output, prefix_lse,
176
+ suffix_output, suffix_lse,
177
+ output_lse_ref_triton)
178
+ torch.cuda.synchronize()
179
+
180
+ for _ in range(repeat_times):
181
+ start.record()
182
+ merge_attn_states_triton(output_ref_triton, prefix_output, prefix_lse,
183
+ suffix_output, suffix_lse,
184
+ output_lse_ref_triton)
185
+ end.record()
186
+ torch.cuda.synchronize()
187
+ total_time_triton_kernel += start.elapsed_time(end)
188
+
189
+ avg_time_triton_kernel = total_time_triton_kernel / repeat_times
190
+
191
+ # 2. Run the CUDA kernel
192
+ total_time_cuda_kernel = 0
193
+ output_cuda = output.clone()
194
+ output_lse_cuda = output_lse.clone()
195
+
196
+ for _ in range(warmup_times):
197
+ merge_attn_states_cuda(output_cuda, prefix_output, prefix_lse,
198
+ suffix_output, suffix_lse, output_lse_cuda)
199
+ torch.cuda.synchronize()
200
+
201
+ for _ in range(repeat_times):
202
+ start.record()
203
+ merge_attn_states_cuda(output_cuda, prefix_output, prefix_lse,
204
+ suffix_output, suffix_lse, output_lse_cuda)
205
+ end.record()
206
+ torch.cuda.synchronize()
207
+ total_time_cuda_kernel += start.elapsed_time(end)
208
+
209
+ avg_time_cuda_kernel = total_time_cuda_kernel / repeat_times
210
+
211
+ # 3. Performance compare
212
+ performance_improved = avg_time_triton_kernel / avg_time_cuda_kernel
213
+ print(f" Torch time: {avg_time_torch_kernel:.6f}ms")
214
+ print(f"Triton time: {avg_time_triton_kernel:.6f}ms")
215
+ print(f" CUDA time: {avg_time_cuda_kernel:.6f}ms, "
216
+ f"Performance: {performance_improved:.5f}x")
217
+ print("-" * 100)
218
+
219
+ # 4. Correctness compare
220
+ # Liger Kernel: Efficient Triton Kernels for LLM Training
221
+ # https://arxiv.org/pdf/2410.10989, 3.3 Correctness
222
+ # use rtol = 1e-2 for bfloat16.
223
+ rtol = 1e-2 if output_dtype == torch.bfloat16 else 1e-3
224
+
225
+ def diff(a: torch.Tensor, b: torch.Tensor):
226
+ max_diff = torch.max(torch.abs(a.float() - b.float()))
227
+ return max_diff
228
+
229
+ # Use Triton output as reference because we want to replace
230
+ # the Triton kernel with custom CUDA kernel for merge attn
231
+ # states operation.
232
+ output_ref = output_ref_triton
233
+ output_lse_ref = output_lse_ref_triton
234
+ torch.testing.assert_close(output_cuda.float(),
235
+ output_ref.float(),
236
+ atol=1e-3,
237
+ rtol=rtol)
238
+ print("Output all match, max abs diff:")
239
+ print(f"(Triton vs Torch) : {diff(output_torch, output_ref)}")
240
+ print(f" (CUDA vs Torch) : {diff(output_torch, output_cuda)}")
241
+ print(f" (CUDA vs Triton): {diff(output_ref, output_cuda)}")
242
+ print("-" * 100)
243
+
244
+ torch.testing.assert_close(output_lse_cuda.float(),
245
+ output_lse_ref.float(),
246
+ atol=1e-3,
247
+ rtol=rtol)
248
+ print("Output LSE all match, max abs diff:")
249
+ print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}")
250
+ print(f" (CUDA vs Torch) : {diff(output_lse_torch, output_lse_cuda)}")
251
+ print(f" (CUDA vs Triton): {diff(output_lse_ref, output_lse_cuda)}")
252
+ print("-" * 100)
253
+
254
+ print("All output values test passed! All inf values "
255
+ "are correctly replaced with -inf.")
256
+ print("-" * 100)
257
+
258
+ device = current_platform.get_device_name()
259
+ all_case_info.append(
260
+ (NUM_TOKENS, NUM_HEADS, HEAD_SIZE, output_dtype, device,
261
+ avg_time_torch_kernel, avg_time_triton_kernel, avg_time_cuda_kernel,
262
+ performance_improved))
263
+ if len(all_case_info) == (len(NUM_BATCH_TOKENS) * len(HEAD_SIZES) *
264
+ len(NUM_QUERY_HEADS) * len(DTYPES)):
265
+ generate_markdown_table()
vllm/_custom_ops.py
+11
@@ -138,6 +138,17 @@ def mla_decode_kvcache_cpu(
138
block_tables, seq_lens)
139
140
141
+# merge attn states ops
142
+def merge_attn_states(output: torch.Tensor,
143
+ prefix_output: torch.Tensor,
144
+ prefix_lse: torch.Tensor,
145
+ suffix_output: torch.Tensor,
146
+ suffix_lse: torch.Tensor,
147
+ output_lse: Optional[torch.Tensor] = None) -> None:
148
+ torch.ops._C.merge_attn_states(output, output_lse, prefix_output,
149
+ prefix_lse, suffix_output, suffix_lse)
150
+
151
+
152
# pos encoding ops
153
def rotary_embedding(
154
positions: torch.Tensor,
vllm/attention/backends/mla/common.py
+1
-2
@@ -204,6 +204,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
204
from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
205
compute_slot_mapping_start_idx,
206
is_block_tables_empty)
207
+from vllm.attention.ops.merge_attn_states import merge_attn_states
208
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
209
LinearBase, RowParallelLinear,
210
UnquantizedLinearMethod)
@@ -217,9 +218,7 @@ from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version
218
219
if HAS_TRITON:
220
from vllm.attention.ops.triton_flash_attention import triton_attention
220
- from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
221
else:
222
- merge_attn_states = None
222
triton_attention = None
223
224
try:
vllm/attention/ops/merge_attn_states.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import torch
5
+
6
+from vllm.platforms import current_platform
7
+
8
+
9
+def merge_attn_states(
10
+ output: torch.Tensor,
11
+ prefix_output: torch.Tensor,
12
+ prefix_lse: torch.Tensor,
13
+ suffix_output: torch.Tensor,
14
+ suffix_lse: torch.Tensor,
15
+ output_lse: Optional[torch.Tensor] = None,
16
+) -> None:
17
+
18
+ # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel
19
+ # is not support for FP8 dtype, fallback to use Triton kernel.
20
+ def supported_dtypes(o: torch.Tensor) -> bool:
21
+ return o.dtype in [torch.float32, torch.half, torch.bfloat16]
22
+
23
+ # NOTE(DefTruth): Currently, custom merge_attn_states CUDA
24
+ # kernel load/store 128b(16 bytes) per memory issue within
25
+ # thread. Namely, the headsize(headdim) must be multiple of
26
+ # pack_size (float32 -> 4, half/bfloat16 -> 8).
27
+ def supported_headdim(o: torch.Tensor) -> bool:
28
+ headdim = o.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
29
+ if o.dtype == torch.float32:
30
+ return headdim % 4 == 0
31
+ return headdim % 8 == 0
32
+
33
+ if (current_platform.is_cuda() and supported_dtypes(output)
34
+ and supported_headdim(output)):
35
+ from vllm._custom_ops import merge_attn_states
36
+ return merge_attn_states(output, prefix_output, prefix_lse,
37
+ suffix_output, suffix_lse, output_lse)
38
+ else:
39
+ from vllm.attention.ops.triton_merge_attn_states import (
40
+ merge_attn_states)
41
+ return merge_attn_states(output, prefix_output, prefix_lse,
42
+ suffix_output, suffix_lse, output_lse)
vllm/v1/attention/backends/flash_attn.py
+1
-1
@@ -10,7 +10,7 @@ from vllm import _custom_ops as ops
10
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionMetadata, AttentionType,
12
is_quantized_kv_cache)
13
-from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
13
+from vllm.attention.ops.merge_attn_states import merge_attn_states
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
from vllm.utils import cdiv
vllm/v1/attention/backends/mla/common.py
+1
-1
@@ -195,7 +195,7 @@ from vllm import _custom_ops as ops
195
from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
196
AttentionMetadata,
197
MLAAttentionImpl)
198
-from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
198
+from vllm.attention.ops.merge_attn_states import merge_attn_states
199
from vllm.logger import init_logger
200
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
201
LinearBase, RowParallelLinear,