@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 223 lines 7.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 from typing import Any
5
6 import torch
7
8 from vllm.model_executor.warmup.jit_warmup import kernel_launcher
9 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10 LaunchSpec,
11 TritonWarmupTensor,
12 VllmTritonJitKernel,
13 )
14 from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16 from vllm.utils.math_utils import next_power_of_2
17
18
19 class FusedQKVRMSNormKernel(VllmTritonJitKernel["FusedQKVRMSNormKernel.CompileKey"]):
20 @dataclass(frozen=True)
21 class CompileKey:
22 dtype: torch.dtype
23 q_size: int
24 kv_size: int
25 block_size: int
26 q_in_stride: int
27 q_out_stride: int
28 kv_in_stride: int
29 kv_out_stride: int
30 eps: float
31 launch_pdl: bool
32
33 @staticmethod
34 @triton.jit
35 def kernel(
36 q_ptr,
37 q_out_ptr,
38 q_weight_ptr,
39 q_in_stride,
40 q_out_stride,
41 kv_ptr,
42 kv_out_ptr,
43 kv_weight_ptr,
44 kv_in_stride,
45 kv_out_stride,
46 eps,
47 Q_SIZE: tl.constexpr,
48 KV_SIZE: tl.constexpr,
49 BLOCK_SIZE: tl.constexpr,
50 launch_pdl: tl.constexpr,
51 ):
52 # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
53 # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
54 # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
55 # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
56 # past num_tokens ~87K under large chunked prefill.
57 token_idx = tl.program_id(0).to(tl.int64)
58 pid_task = tl.program_id(1)
59
60 if pid_task == 0:
61 SIZE = Q_SIZE
62 row_in = q_ptr + token_idx * q_in_stride
63 weight_ptr = q_weight_ptr
64 row_out = q_out_ptr + token_idx * q_out_stride
65 else:
66 SIZE = KV_SIZE
67 row_in = kv_ptr + token_idx * kv_in_stride
68 weight_ptr = kv_weight_ptr
69 row_out = kv_out_ptr + token_idx * kv_out_stride
70
71 # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
72 # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
73 # keep x, rrms, and w all in fp32 and perform a single cast at store.
74 block = tl.arange(0, BLOCK_SIZE)
75 mask = block < SIZE
76 # The weight load does not depend on the producer's output, so issue it
77 # before the PDL wait: gamma streams in while the producer finishes,
78 # leaving a single dependent global round trip (x) after the wait.
79 w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
80
81 if launch_pdl:
82 tl.extra.cuda.gdc_wait()
83 tl.extra.cuda.gdc_launch_dependents()
84
85 x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
86 variance = tl.sum(x * x, axis=0) / SIZE
87 rrms = tl.rsqrt(variance + eps)
88 y = x * rrms * w
89 tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
90
91 def dispatch( # type: ignore[override]
92 self,
93 *,
94 dtype: torch.dtype,
95 q_size: int,
96 kv_size: int,
97 eps: float,
98 launch_pdl: bool,
99 **compile_key_fields: int,
100 ) -> CompileKey:
101 max_size = q_size if q_size >= kv_size else kv_size
102 return self.CompileKey(
103 **compile_key_fields,
104 dtype=dtype,
105 q_size=q_size,
106 kv_size=kv_size,
107 block_size=next_power_of_2(max_size),
108 eps=eps,
109 launch_pdl=launch_pdl,
110 )
111
112 def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
113 hf_config = vllm_config.model_config.hf_config
114 q_size = int(getattr(hf_config, "q_lora_rank", 0) or 0)
115 kv_size = int(getattr(hf_config, "head_dim", 0) or 0)
116 if q_size <= 0 or kv_size <= 0:
117 return []
118
119 input_stride = q_size + kv_size
120 return self._trace_dispatch(self.dispatch)(
121 dtype=vllm_config.model_config.dtype,
122 q_size=q_size,
123 kv_size=kv_size,
124 q_in_stride=input_stride,
125 q_out_stride=(input_stride, q_size),
126 kv_in_stride=input_stride,
127 kv_out_stride=(input_stride, kv_size),
128 eps=float(hf_config.rms_norm_eps),
129 launch_pdl=current_platform.is_arch_support_pdl(),
130 )
131
132 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
133 return dict(
134 q=TritonWarmupTensor(
135 compile_key.dtype,
136 shape=(1, compile_key.q_size),
137 strides=(compile_key.q_in_stride, 1),
138 ),
139 kv=TritonWarmupTensor(
140 compile_key.dtype,
141 shape=(1, compile_key.kv_size),
142 strides=(compile_key.kv_in_stride, 1),
143 ),
144 q_out=TritonWarmupTensor(
145 compile_key.dtype,
146 shape=(1, compile_key.q_size),
147 strides=(compile_key.q_out_stride, 1),
148 ),
149 kv_out=TritonWarmupTensor(
150 compile_key.dtype,
151 shape=(1, compile_key.kv_size),
152 strides=(compile_key.kv_out_stride, 1),
153 ),
154 q_weight=TritonWarmupTensor(
155 compile_key.dtype,
156 shape=(compile_key.q_size,),
157 ),
158 kv_weight=TritonWarmupTensor(
159 compile_key.dtype,
160 shape=(compile_key.kv_size,),
161 ),
162 eps=compile_key.eps,
163 )
164
165 @kernel_launcher
166 def __call__(
167 self,
168 q: torch.Tensor,
169 kv: torch.Tensor,
170 q_out: torch.Tensor,
171 kv_out: torch.Tensor,
172 q_weight: torch.Tensor,
173 kv_weight: torch.Tensor,
174 eps: float,
175 ) -> LaunchSpec:
176 q_size = q.shape[1]
177 kv_size = kv.shape[1]
178 block_size = next_power_of_2(max(q_size, kv_size))
179 return (q.shape[0], 2), dict(
180 q_in_stride=q.stride(0),
181 q_out_stride=q_out.stride(0),
182 kv_in_stride=kv.stride(0),
183 kv_out_stride=kv_out.stride(0),
184 Q_SIZE=q_size,
185 KV_SIZE=kv_size,
186 BLOCK_SIZE=block_size,
187 launch_pdl=current_platform.is_arch_support_pdl(),
188 # One vectorized 128-bit access per thread for K3/DSv4-sized rows,
189 # instead of a serial per-thread chain at the default 4 warps.
190 num_warps=8 if block_size >= 2048 else 4,
191 )
192
193
194 def fused_q_kv_rmsnorm(
195 qr: torch.Tensor,
196 kv: torch.Tensor,
197 q_weight: torch.Tensor,
198 kv_weight: torch.Tensor,
199 eps: float,
200 ) -> tuple[torch.Tensor, torch.Tensor]:
201 assert qr.ndim == 2 and kv.ndim == 2
202 assert qr.shape[0] == kv.shape[0], (
203 f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
204 )
205 assert qr.stride(-1) == 1 and kv.stride(-1) == 1
206 assert q_weight.is_contiguous() and kv_weight.is_contiguous()
207
208 qr_out = torch.empty(qr.shape, dtype=qr.dtype, device=qr.device)
209 kv_out = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
210 if qr.shape[0] > 0:
211 _FUSED_Q_KV_RMSNORM_KERNEL(
212 qr,
213 kv,
214 qr_out,
215 kv_out,
216 q_weight,
217 kv_weight,
218 eps,
219 )
220 return qr_out, kv_out
221
222
223 _FUSED_Q_KV_RMSNORM_KERNEL = FusedQKVRMSNormKernel()