@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 223 lines 7.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from vllm.platforms import current_platform
7 from vllm.triton_utils import tl, triton
8
9 float8_info = torch.finfo(current_platform.fp8_dtype())
10
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 def merge_attn_states(
15 output: torch.Tensor,
16 prefix_output: torch.Tensor,
17 prefix_lse: torch.Tensor,
18 suffix_output: torch.Tensor,
19 suffix_lse: torch.Tensor,
20 output_lse: torch.Tensor | None = None,
21 prefill_tokens_with_context: int | None = None,
22 output_scale: torch.Tensor | None = None,
23 ) -> None:
24 num_tokens = output.shape[0]
25 num_query_heads = output.shape[1]
26 head_size = output.shape[2]
27 padded_head_size = triton.next_power_of_2(head_size)
28 # We assume the output stride on num_head is not always as same as the
29 # `suffix_output` and `prefix_output`, as them might be padded by the
30 # attention backend.
31 prefix_head_stride = prefix_output.stride(1)
32 output_head_stride = output.stride(1)
33 # lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views
34 # (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index
35 # them by their actual strides rather than assuming a contiguous layout.
36 prefix_lse_head_stride = prefix_lse.stride(0)
37 prefix_lse_token_stride = prefix_lse.stride(1)
38 suffix_lse_head_stride = suffix_lse.stride(0)
39 suffix_lse_token_stride = suffix_lse.stride(1)
40 output_lse_head_stride = output_lse.stride(0) if output_lse is not None else 0
41 output_lse_token_stride = output_lse.stride(1) if output_lse is not None else 0
42
43 # If prefill_tokens_with_context is None, all tokens should use prefix context
44 if prefill_tokens_with_context is None:
45 prefill_tokens_with_context = num_tokens
46
47 # TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead.
48 merge_attn_states_kernel[(num_tokens, num_query_heads)](
49 output,
50 output_lse,
51 prefix_output,
52 prefix_lse,
53 suffix_output,
54 suffix_lse,
55 prefix_head_stride,
56 output_head_stride,
57 prefix_lse_head_stride,
58 prefix_lse_token_stride,
59 suffix_lse_head_stride,
60 suffix_lse_token_stride,
61 output_lse_head_stride,
62 output_lse_token_stride,
63 output_scale,
64 prefill_tokens_with_context,
65 head_size,
66 padded_head_size,
67 output_lse is not None,
68 output_scale is not None,
69 )
70
71
72 @triton.jit
73 def merge_attn_states_kernel(
74 output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
75 output_lse, # [NUM_HEADS, NUM_TOKENS]
76 prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
77 prefix_lse, # [NUM_HEADS, NUM_TOKENS]
78 suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
79 suffix_lse, # [NUM_HEADS, NUM_TOKENS]
80 prefix_head_stride,
81 output_head_stride,
82 prefix_lse_head_stride,
83 prefix_lse_token_stride,
84 suffix_lse_head_stride,
85 suffix_lse_token_stride,
86 output_lse_head_stride,
87 output_lse_token_stride,
88 output_scale, # scale tensor or None
89 prefill_tokens_with_context,
90 HEAD_SIZE: tl.constexpr,
91 PADDED_HEAD_SIZE: tl.constexpr,
92 OUTPUT_LSE: tl.constexpr,
93 USE_FP8: tl.constexpr,
94 FP8_MIN: tl.constexpr = float8_info.min,
95 FP8_MAX: tl.constexpr = float8_info.max,
96 ):
97 token_idx = tl.program_id(0)
98 head_idx = tl.program_id(1)
99 num_heads = tl.num_programs(1)
100
101 prefix_mask = token_idx < prefill_tokens_with_context
102
103 head_arange = tl.arange(0, PADDED_HEAD_SIZE)
104 head_mask = head_arange < HEAD_SIZE
105
106 # For tokens without context (token_idx >= prefill_tokens_with_context),
107 # directly copy from suffix_output
108 if not prefix_mask:
109 s_lse = tl.load(
110 suffix_lse
111 + head_idx * suffix_lse_head_stride
112 + token_idx * suffix_lse_token_stride
113 )
114 if OUTPUT_LSE:
115 tl.store(
116 output_lse
117 + head_idx * output_lse_head_stride
118 + token_idx * output_lse_token_stride,
119 s_lse,
120 )
121
122 s_out = tl.load(
123 suffix_output
124 + token_idx * num_heads * prefix_head_stride
125 + head_idx * prefix_head_stride
126 + head_arange,
127 mask=head_mask,
128 )
129
130 if USE_FP8:
131 s_out = s_out * (1.0 / tl.load(output_scale))
132 s_out = tl.clamp(s_out, FP8_MIN, FP8_MAX)
133 s_out = s_out.to(output.dtype.element_ty)
134
135 tl.store(
136 output
137 + token_idx * num_heads * output_head_stride
138 + head_idx * output_head_stride
139 + head_arange,
140 s_out,
141 mask=head_mask,
142 )
143 return
144
145 # For tokens with context (token_idx < prefill_tokens_with_context),
146 # perform normal merge operation
147 p_lse = tl.load(
148 prefix_lse
149 + head_idx * prefix_lse_head_stride
150 + token_idx * prefix_lse_token_stride
151 )
152 s_lse = tl.load(
153 suffix_lse
154 + head_idx * suffix_lse_head_stride
155 + token_idx * suffix_lse_token_stride
156 )
157
158 # FA2 and FA3 have different behavior for when the sum-exp is 0, this namely
159 # arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf.
160 # If we see an inf assume FA2 and convert inf to -inf for consistency
161 # and correctness. Inf generally doesn't make sense in this context outside
162 # of undefined-behavior/FA2-case, so I think this a safe assumption.
163 p_lse = float("-inf") if p_lse == float("inf") else p_lse
164 s_lse = float("-inf") if s_lse == float("inf") else s_lse
165
166 max_lse = tl.maximum(p_lse, s_lse)
167 p_lse = p_lse - max_lse
168 s_lse = s_lse - max_lse
169 # Will reuse precomputed Exp values for scale factor computation.
170 p_se = tl.exp(p_lse)
171 s_se = tl.exp(s_lse)
172 out_se = p_se + s_se
173
174 if OUTPUT_LSE:
175 out_lse = tl.log(out_se) + max_lse
176 # Both sides empty (max_lse == -inf) => undefined merge; keep -inf so
177 # downstream merges continue to treat the token as empty.
178 out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse)
179 tl.store(
180 output_lse
181 + head_idx * output_lse_head_stride
182 + token_idx * output_lse_token_stride,
183 out_lse,
184 )
185
186 p_out = tl.load(
187 prefix_output
188 + token_idx * num_heads * prefix_head_stride
189 + head_idx * prefix_head_stride
190 + head_arange,
191 mask=head_mask,
192 )
193 s_out = tl.load(
194 suffix_output
195 + token_idx * num_heads * prefix_head_stride
196 + head_idx * prefix_head_stride
197 + head_arange,
198 mask=head_mask,
199 )
200
201 # NOTE(woosuk): Be careful with the numerical stability.
202 # We should compute the scale first, and then multiply it with the output.
203 # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
204 p_scale = p_se / out_se
205 s_scale = s_se / out_se
206 out = p_out * p_scale + s_out * s_scale
207 # If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit
208 # zeros rather than NaN.
209 out = tl.where(max_lse == float("-inf"), 0.0, out)
210
211 if USE_FP8:
212 out = out * (1.0 / tl.load(output_scale))
213 out = tl.clamp(out, FP8_MIN, FP8_MAX)
214 out = out.to(output.dtype.element_ty)
215
216 tl.store(
217 output
218 + token_idx * num_heads * output_head_stride
219 + head_idx * output_head_stride
220 + head_arange,
221 out,
222 mask=head_mask,
223 )