@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 193 lines 6.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 from typing import Any
6
7 import torch
8
9 from vllm.model_executor.warmup.jit_warmup import kernel_launcher
10 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
11 LaunchSpec,
12 TritonWarmupTensor,
13 VllmTritonJitKernel,
14 )
15 from vllm.triton_utils import tl, triton
16 from vllm.utils.math_utils import next_power_of_2
17
18
19 class SavePartialStatesKernel(
20 VllmTritonJitKernel["SavePartialStatesKernel.CompileKey"]
21 ):
22 @dataclass(frozen=True)
23 class CompileKey:
24 head_size: int
25 triton_block_size: int
26 state_width: int
27 compress_ratio: int
28 kv_stride: int
29 score_stride: int
30 ape_stride: int
31 state_cache_stride0: int
32 state_cache_stride1: int
33 block_size: int
34
35 @staticmethod
36 @triton.jit
37 def kernel(
38 kv_ptr,
39 kv_stride,
40 score_ptr,
41 score_stride,
42 ape_ptr,
43 ape_stride,
44 positions_ptr,
45 state_cache_ptr,
46 state_cache_stride0,
47 state_cache_stride1,
48 slot_mapping_ptr,
49 block_size,
50 HEAD_SIZE: tl.constexpr,
51 TRITON_BLOCK_SIZE: tl.constexpr,
52 # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
53 STATE_WIDTH: tl.constexpr,
54 COMPRESS_RATIO: tl.constexpr,
55 ):
56 token_idx = tl.program_id(0)
57 slot_id = tl.load(slot_mapping_ptr + token_idx)
58
59 # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
60 # by vLLM). During CUDA graph replay the batch may contain padding
61 # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
62 # illegal memory access.
63 if slot_id < 0:
64 return
65
66 block_idx = slot_id // block_size
67 pos_in_block = slot_id % block_size
68 base_ptr = (
69 state_cache_ptr
70 + block_idx * state_cache_stride0
71 + pos_in_block * state_cache_stride1
72 )
73
74 block = tl.arange(0, TRITON_BLOCK_SIZE)
75 mask = block < HEAD_SIZE
76
77 kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
78 tl.store(base_ptr + block, kv, mask=mask)
79
80 # Fused: score += ape[position % compress_ratio]
81 position = tl.load(positions_ptr + token_idx)
82 ape_row = position % COMPRESS_RATIO
83 ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
84 score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
85 tl.store(
86 base_ptr + STATE_WIDTH + block,
87 score + ape,
88 mask=mask,
89 )
90
91 def dispatch( # type: ignore[override]
92 self,
93 *,
94 head_size: int,
95 **compile_key_fields: int,
96 ) -> CompileKey:
97 return self.CompileKey(
98 **compile_key_fields,
99 head_size=head_size,
100 triton_block_size=next_power_of_2(head_size),
101 )
102
103 def get_warmup_keys(
104 self,
105 *,
106 head_dim: int,
107 compress_ratio: int,
108 ) -> list[CompileKey]:
109 if head_dim <= 0 or compress_ratio not in (4, 128):
110 return []
111
112 coefficient = 2 if compress_ratio == 4 else 1
113 return self._trace_dispatch(self.dispatch)(
114 head_size=coefficient * head_dim,
115 state_width=coefficient * head_dim,
116 compress_ratio=compress_ratio,
117 kv_stride=2 * coefficient * head_dim,
118 score_stride=2 * coefficient * head_dim,
119 ape_stride=coefficient * head_dim,
120 state_cache_stride0=8 * coefficient * head_dim,
121 state_cache_stride1=2 * coefficient * head_dim,
122 block_size=4 if compress_ratio == 4 else 8,
123 )
124
125 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
126 return dict(
127 kv=TritonWarmupTensor(
128 torch.float32,
129 shape=(1, compile_key.head_size),
130 strides=(compile_key.kv_stride, 1),
131 ),
132 score=TritonWarmupTensor(
133 torch.float32,
134 shape=(1, 1),
135 strides=(compile_key.score_stride, 1),
136 ),
137 ape=TritonWarmupTensor(
138 torch.float32,
139 shape=(1, 1),
140 strides=(compile_key.ape_stride, 1),
141 ),
142 positions=TritonWarmupTensor(torch.int64),
143 state_cache=TritonWarmupTensor(
144 torch.float32,
145 shape=(1, 1, compile_key.state_width),
146 strides=(
147 compile_key.state_cache_stride0,
148 compile_key.state_cache_stride1,
149 1,
150 ),
151 ),
152 slot_mapping=TritonWarmupTensor(torch.int64),
153 block_size=compile_key.block_size,
154 state_width=compile_key.state_width,
155 compress_ratio=compile_key.compress_ratio,
156 )
157
158 @kernel_launcher
159 def __call__(
160 self,
161 kv: torch.Tensor,
162 score: torch.Tensor,
163 ape: torch.Tensor,
164 positions: torch.Tensor,
165 state_cache: torch.Tensor,
166 slot_mapping: torch.Tensor,
167 block_size: int,
168 state_width: int,
169 compress_ratio: int,
170 pdl_kwargs: dict | None = None,
171 ) -> LaunchSpec:
172 """Write packed [kv, score+ape] partial states into the compressor cache.
173
174 One program per token; pads (slot_id == -1) are skipped.
175 """
176 num_actual = slot_mapping.shape[0]
177 head_size = kv.shape[-1]
178 return (num_actual,), dict(
179 kv_stride=kv.stride(0),
180 score_stride=score.stride(0),
181 ape_stride=ape.stride(0),
182 state_cache_stride0=state_cache.stride(0),
183 state_cache_stride1=state_cache.stride(1),
184 block_size=block_size,
185 HEAD_SIZE=head_size,
186 TRITON_BLOCK_SIZE=next_power_of_2(head_size),
187 STATE_WIDTH=state_width,
188 COMPRESS_RATIO=compress_ratio,
189 **(pdl_kwargs or {}),
190 )
191
192
193 _SAVE_PARTIAL_STATES_KERNEL = SavePartialStatesKernel()