@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 193 lines 6.39 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4 #
5 # This file contains code adapted from the flash-linear-attention project.
6 # The original source code was licensed under the MIT license and included
7 # the following copyright notice:
8 # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
9
10
11 import torch
12
13 from vllm.triton_utils import tl, triton
14
15
16 @triton.jit
17 def _attn_res_kernel(
18 prefix_ptr,
19 delta_ptr,
20 blocks_ptr,
21 norm_weight_ptr,
22 qk_weight_ptr,
23 output_norm_weight_ptr,
24 output_ptr,
25 stride_prefix_m: tl.constexpr,
26 stride_delta_m: tl.constexpr,
27 stride_block_m: tl.constexpr,
28 stride_block_r: tl.constexpr,
29 stride_output_m: tl.constexpr,
30 num_blocks: tl.constexpr,
31 hidden_size: tl.constexpr,
32 block_write_idx: tl.constexpr,
33 eps: tl.constexpr,
34 output_norm_eps: tl.constexpr,
35 HAS_DELTA: tl.constexpr,
36 WRITE_BLOCK: tl.constexpr,
37 APPLY_OUTPUT_NORM: tl.constexpr,
38 BLOCK_L: tl.constexpr,
39 BLOCK_D: tl.constexpr,
40 ):
41 row_idx = tl.program_id(0).to(tl.int64)
42 d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D)
43 d_mask = d_offsets < hidden_size
44
45 updated_prefix = tl.load(
46 prefix_ptr + row_idx * stride_prefix_m + d_offsets,
47 mask=d_mask,
48 other=0.0,
49 ).to(tl.float32)
50 if HAS_DELTA:
51 delta = tl.load(
52 delta_ptr + row_idx * stride_delta_m + d_offsets,
53 mask=d_mask,
54 other=0.0,
55 ).to(tl.float32)
56 updated_prefix += delta
57 # Match the BF16 prefix-add result before using it as a residual source.
58 updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32)
59 tl.store(
60 prefix_ptr + row_idx * stride_prefix_m + d_offsets,
61 updated_prefix,
62 mask=d_mask,
63 )
64 if WRITE_BLOCK:
65 tl.store(
66 blocks_ptr
67 + row_idx * stride_block_m
68 + block_write_idx * stride_block_r
69 + d_offsets,
70 updated_prefix,
71 mask=d_mask,
72 )
73 # With only the prefix source, the AttnRes softmax is exactly one.
74 if num_blocks == 0:
75 mixed = updated_prefix
76 else:
77 # Reloading avoids keeping the full prefix vector live across the loop.
78 if HAS_DELTA:
79 tl.debug_barrier()
80 input_qk_weight = tl.load(
81 norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
82 ).to(tl.float32) * tl.load(
83 qk_weight_ptr + d_offsets, mask=d_mask, other=0.0
84 ).to(tl.float32)
85 max_logit = tl.full((), -float("inf"), tl.float32)
86 denominator = tl.zeros((), tl.float32)
87 mixed = tl.zeros((BLOCK_D,), tl.float32)
88
89 num_sources = num_blocks + 1
90 for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
91 source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
92 source_mask = source_offsets < num_sources
93 is_prefix = source_offsets == num_blocks
94 block_ptrs = (
95 blocks_ptr
96 + row_idx * stride_block_m
97 + source_offsets[:, None] * stride_block_r
98 + d_offsets[None, :]
99 )
100 block_values = tl.load(
101 block_ptrs,
102 mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]),
103 other=0.0,
104 eviction_policy="evict_first",
105 ).to(tl.float32)
106 values = tl.where(is_prefix[:, None], updated_prefix[None, :], block_values)
107 reciprocal_std = tl.rsqrt(
108 tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
109 )
110 logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
111 scores = tl.where(source_mask, logits, -float("inf"))
112
113 new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
114 old_scale = tl.exp(max_logit - new_max_logit)
115 block_scales = tl.exp(scores - new_max_logit)
116 denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
117 mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
118 max_logit = new_max_logit
119
120 mixed /= denominator
121 output = mixed
122
123 if APPLY_OUTPUT_NORM:
124 output_reciprocal_std = tl.rsqrt(
125 tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size)
126 + output_norm_eps
127 )
128 output_norm_weight = tl.load(
129 output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
130 ).to(tl.float32)
131 output = mixed * output_reciprocal_std * output_norm_weight
132 tl.store(
133 output_ptr + row_idx * stride_output_m + d_offsets,
134 output,
135 mask=d_mask,
136 )
137
138
139 def attn_res(
140 prefix: torch.Tensor,
141 delta: torch.Tensor | None,
142 blocks: torch.Tensor,
143 norm_weight: torch.Tensor,
144 qk_weight: torch.Tensor,
145 output_norm_weight: torch.Tensor | None,
146 num_blocks: int,
147 block_write_idx: int,
148 eps: float,
149 output_norm_eps: float,
150 ) -> torch.Tensor:
151 num_tokens, hidden_size = prefix.shape
152 assert prefix.stride(-1) == 1
153 assert delta is None or delta.stride(-1) == 1
154 assert blocks.stride(-1) == 1
155 assert norm_weight.stride(-1) == 1
156 assert qk_weight.stride(-1) == 1
157 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
158 output = prefix.new_empty(prefix.shape)
159 if num_tokens == 0:
160 return output
161
162 # Source tiling helps decode, while one-source tiles scale better for prefill.
163 if num_tokens >= 256 or num_blocks <= 1:
164 block_l, num_warps = 1, 4
165 else:
166 block_l, num_warps = 4, 8
167 _attn_res_kernel[(num_tokens,)](
168 prefix,
169 delta,
170 blocks,
171 norm_weight,
172 qk_weight,
173 output_norm_weight,
174 output,
175 prefix.stride(0),
176 0 if delta is None else delta.stride(0),
177 blocks.stride(0),
178 blocks.stride(1),
179 output.stride(0),
180 num_blocks,
181 hidden_size,
182 block_write_idx,
183 eps,
184 output_norm_eps,
185 HAS_DELTA=delta is not None,
186 WRITE_BLOCK=block_write_idx >= 0,
187 APPLY_OUTPUT_NORM=output_norm_weight is not None,
188 BLOCK_L=block_l,
189 BLOCK_D=triton.next_power_of_2(hidden_size),
190 num_warps=num_warps,
191 num_stages=2,
192 )
193 return output