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 12, 13:30 UTC
[Performance][MLA][ROCm] Remove redundant D2D copy in deepseek (#27457)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Nov 26, 2025 at 12:45 UTC
d9d342d214b8c13f71215318a6d9252cc4a5ca47
5 files changed
+49
-41
csrc/attention/merge_attn_states.cu
+12
-15
@@ -16,7 +16,8 @@ __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) {
19
+ const uint head_size, const uint prefix_head_stride,
20
+ const uint output_head_stride) {
21
using pack_128b_t = uint4;
22
const uint pack_size = 16 / sizeof(scalar_t);
23
const uint threads_per_head = head_size / pack_size;
@@ -34,11 +35,13 @@ __global__ void merge_attn_states_kernel(
35
const uint head_idx = token_head_idx % num_heads;
36
37
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;
38
+ const uint src_head_offset = token_idx * num_heads * prefix_head_stride +
39
+ head_idx * prefix_head_stride;
40
+ const uint dst_head_offset = token_idx * num_heads * output_head_stride +
41
+ head_idx * output_head_stride;
42
+ const scalar_t* prefix_head_ptr = prefix_output + src_head_offset;
43
+ const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;
44
+ scalar_t* output_head_ptr = output + dst_head_offset;
45
46
float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
47
float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
@@ -140,7 +143,7 @@ __global__ void merge_attn_states_kernel(
143
reinterpret_cast<float*>(prefix_lse.data_ptr()), \
144
reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \
145
reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens, \
143
- num_heads, head_size); \
146
+ num_heads, head_size, prefix_head_stride, output_head_stride); \
147
}
148
149
/*@brief Merges the attention states from prefix and suffix
@@ -166,17 +169,11 @@ void merge_attn_states_launcher(torch::Tensor& output,
169
const uint num_tokens = output.size(0);
170
const uint num_heads = output.size(1);
171
const uint head_size = output.size(2);
172
+ const uint prefix_head_stride = prefix_output.stride(1);
173
+ const uint output_head_stride = output.stride(1);
174
const uint pack_size = 16 / sizeof(scalar_t);
175
TORCH_CHECK(head_size % pack_size == 0,
176
"headsize must be multiple of pack_size:", pack_size);
172
- TORCH_CHECK(output.stride(-2) == head_size && output.stride(-1) == 1,
173
- "output heads must be contiguous in memory");
174
- TORCH_CHECK(
175
- prefix_output.stride(-2) == head_size && prefix_output.stride(-1) == 1,
176
- "prefix_output heads must be contiguous in memory");
177
- TORCH_CHECK(
178
- suffix_output.stride(-2) == head_size && suffix_output.stride(-1) == 1,
179
- "suffix_output heads must be contiguous in memory");
177
float* output_lse_ptr = nullptr;
178
if (output_lse.has_value()) {
179
output_lse_ptr = output_lse.value().data_ptr<float>();
csrc/ops.h
+1
-2
@@ -52,14 +52,13 @@ 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
55
void merge_attn_states(torch::Tensor& output,
56
std::optional<torch::Tensor> output_lse,
57
const torch::Tensor& prefix_output,
58
const torch::Tensor& prefix_lse,
59
const torch::Tensor& suffix_output,
60
const torch::Tensor& suffix_lse);
62
-
61
+#ifndef USE_ROCM
62
void convert_vertical_slash_indexes(
63
torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
64
torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
csrc/torch_bindings.cpp
+1
-2
@@ -63,7 +63,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
63
" int blocksparse_head_sliding_step) -> ()");
64
ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
65
66
-#ifndef USE_ROCM
66
// Merge attn states
67
// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
68
// can be used to combine partial attention results (in the split-KV case)
@@ -76,7 +75,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
75
" Tensor suffix_output,"
76
" Tensor suffix_lse) -> ()");
77
ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
79
-
78
+#ifndef USE_ROCM
79
ops.def(
80
"convert_vertical_slash_indexes("
81
" Tensor! block_count, Tensor! block_offset, "
vllm/attention/ops/triton_merge_attn_states.py
+17
-6
@@ -20,7 +20,11 @@ def merge_attn_states(
20
num_query_heads = output.shape[1]
21
head_size = output.shape[2]
22
padded_head_size = triton.next_power_of_2(head_size)
23
-
23
+ # We assume the output stride on num_head is not always as same as the
24
+ # `suffix_output` and `prefix_output`, as them might be padded by the attention
25
+ # backend.
26
+ prefix_head_stride = prefix_output.stride(1)
27
+ output_head_stride = output.stride(1)
28
# TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead.
29
merge_attn_states_kernel[(num_tokens, num_query_heads)](
30
output,
@@ -29,6 +33,8 @@ def merge_attn_states(
33
prefix_lse,
34
suffix_output,
35
suffix_lse,
36
+ prefix_head_stride,
37
+ output_head_stride,
38
head_size,
39
padded_head_size,
40
output_lse is not None,
@@ -43,6 +49,8 @@ def merge_attn_states_kernel(
49
prefix_lse, # [NUM_HEADS, NUM_TOKENS]
50
suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
51
suffix_lse, # [NUM_HEADS, NUM_TOKENS]
52
+ prefix_head_stride,
53
+ output_head_stride,
54
HEAD_SIZE: tl.constexpr,
55
PADDED_HEAD_SIZE: tl.constexpr,
56
OUTPUT_LSE: tl.constexpr,
@@ -79,15 +87,15 @@ def merge_attn_states_kernel(
87
head_mask = head_arange < HEAD_SIZE
88
p_out = tl.load(
89
prefix_output
82
- + token_idx * num_heads * HEAD_SIZE
83
- + head_idx * HEAD_SIZE
90
+ + token_idx * num_heads * prefix_head_stride
91
+ + head_idx * prefix_head_stride
92
+ head_arange,
93
mask=head_mask,
94
)
95
s_out = tl.load(
96
suffix_output
89
- + token_idx * num_heads * HEAD_SIZE
90
- + head_idx * HEAD_SIZE
97
+ + token_idx * num_heads * prefix_head_stride
98
+ + head_idx * prefix_head_stride
99
+ head_arange,
100
mask=head_mask,
101
)
@@ -99,7 +107,10 @@ def merge_attn_states_kernel(
107
s_scale = s_se / out_se
108
out = p_out * p_scale + s_out * s_scale
109
tl.store(
102
- output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange,
110
+ output
111
+ + token_idx * num_heads * output_head_stride
112
+ + head_idx * output_head_stride
113
+ + head_arange,
114
out,
115
mask=head_mask,
116
)
vllm/v1/attention/backends/mla/common.py
+18
-16
@@ -1238,15 +1238,13 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1238
def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1239
# Convert from (B, N, L) to (N, B, L)
1240
x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1241
+
1242
if self.is_aiter_triton_fp8_bmm_enabled:
1243
+ out = out.view(-1, self.num_heads, self.v_head_dim)
1244
# Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1245
x = rocm_aiter_ops.triton_fp8_bmm(
1244
- x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1246
+ x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
1247
)
1246
- # Convert from (B, N, V) to (B, N * V)
1247
- x = x.reshape(-1, self.num_heads * self.v_head_dim)
1248
- # Copy result
1249
- out.copy_(x)
1248
else:
1249
# Convert from (B, N * V) to (N, B, V)
1250
out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)
@@ -1824,7 +1822,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1822
kv_c_and_k_pe_cache: torch.Tensor,
1823
attn_metadata: MLACommonMetadata,
1824
k_scale: torch.Tensor,
1827
- ) -> torch.Tensor:
1825
+ output: torch.Tensor,
1826
+ ) -> None:
1827
# TODO (zyongye): Prefill function here
1828
assert attn_metadata.prefill is not None
1829
assert self.dcp_world_size is not None
@@ -1837,7 +1836,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1836
1837
k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1838
1840
- output = self._run_prefill_new_tokens(
1839
+ output_prefill = self._run_prefill_new_tokens(
1840
prefill=attn_metadata.prefill,
1841
q=q,
1842
k=k,
@@ -1846,7 +1845,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1845
)
1846
1847
if has_context:
1849
- suffix_output, suffix_lse = output
1848
+ suffix_output, suffix_lse = output_prefill
1849
if self.dcp_world_size > 1:
1850
context_output, context_lse = (
1851
self._context_parallel_compute_prefill_context(
@@ -1862,7 +1861,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1861
q, kv_c_and_k_pe_cache, attn_metadata, k_scale
1862
)
1863
1865
- output = torch.empty_like(suffix_output)
1864
+ # unpad if necessary
1865
+ if self._pad_v:
1866
+ context_output = context_output[..., : v.shape[-1]]
1867
+ suffix_output = suffix_output[..., : v.shape[-1]]
1868
+
1869
+ output = output.view(-1, self.num_heads, self.v_head_dim)
1870
merge_attn_states(
1871
output=output,
1872
prefix_output=context_output,
@@ -1870,12 +1874,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1874
suffix_output=suffix_output,
1875
suffix_lse=suffix_lse,
1876
)
1873
-
1874
- # unpad if necessary
1875
- if self._pad_v:
1876
- output = output[..., : v.shape[-1]]
1877
-
1878
- return output.flatten(start_dim=-2)
1877
+ else:
1878
+ output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2)
1879
+ output.copy_(output_prefill)
1880
1881
@abstractmethod
1882
def _forward_decode(
@@ -1970,13 +1971,14 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1971
kv_cache = kv_cache.view(current_platform.fp8_dtype())
1972
1973
if has_prefill:
1973
- output[num_decode_tokens:] = self._forward_prefill(
1974
+ self._forward_prefill(
1975
prefill_q,
1976
prefill_k_c_normed,
1977
prefill_k_pe,
1978
kv_cache,
1979
attn_metadata,
1980
layer._k_scale,
1981
+ output=output[num_decode_tokens:],
1982
)
1983
1984
if has_decode: