| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from functools import lru_cache |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.triton_utils import tl, triton |
| 10 | |
| 11 | _MAX_FUSED_SIZE = 65536 |
| 12 | |
| 13 | |
| 14 | def _get_num_warps_from_block_size(block_size: int) -> int: |
| 15 | if block_size >= 32768: |
| 16 | return 32 |
| 17 | if block_size >= 8192: |
| 18 | return 16 |
| 19 | if block_size >= 2048: |
| 20 | return 8 |
| 21 | return 4 |
| 22 | |
| 23 | |
| 24 | def _largest_power_of_2(n: int) -> int: |
| 25 | assert n > 0, f"{n=}" |
| 26 | return 1 << (n.bit_length() - 1) |
| 27 | |
| 28 | |
| 29 | @lru_cache(maxsize=128) |
| 30 | def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int: |
| 31 | num_sms = torch.cuda.get_device_properties(device).multi_processor_count |
| 32 | return _largest_power_of_2(num_sms) * factor |
| 33 | |
| 34 | |
| 35 | @triton.jit |
| 36 | def _rmsnorm_fwd_kernel( |
| 37 | x_ptr, |
| 38 | weight_ptr, |
| 39 | y_ptr, |
| 40 | rstd_ptr, |
| 41 | eps, |
| 42 | x_stride_0, |
| 43 | y_stride_0, |
| 44 | n_cols, |
| 45 | block_size_n: tl.constexpr, |
| 46 | ): |
| 47 | pid_m = tl.program_id(0).to(tl.int64) |
| 48 | offs_n = tl.arange(0, block_size_n) |
| 49 | mask_n = offs_n < n_cols |
| 50 | |
| 51 | weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 52 | x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to( |
| 53 | tl.float32 |
| 54 | ) |
| 55 | row_var = tl.sum(x * x, axis=0) / n_cols |
| 56 | rstd = tl.math.rsqrt(row_var + eps) |
| 57 | tl.store(rstd_ptr + pid_m, rstd) |
| 58 | |
| 59 | y = x * rstd * weight |
| 60 | tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) |
| 61 | |
| 62 | |
| 63 | @triton.jit(do_not_specialize=["n_rows"]) |
| 64 | def _rmsnorm_fwd_kernel_block_m( |
| 65 | x_ptr, |
| 66 | weight_ptr, |
| 67 | y_ptr, |
| 68 | rstd_ptr, |
| 69 | eps, |
| 70 | x_stride_0, |
| 71 | y_stride_0, |
| 72 | n_rows, |
| 73 | n_cols, |
| 74 | block_size_m: tl.constexpr, |
| 75 | block_size_n: tl.constexpr, |
| 76 | ): |
| 77 | pid_m = tl.program_id(0).to(tl.int64) |
| 78 | offs_n = tl.arange(0, block_size_n) |
| 79 | mask_n = offs_n < n_cols |
| 80 | |
| 81 | num_blocks_m = tl.cdiv(n_rows, block_size_m) |
| 82 | blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0)) |
| 83 | block_id_start = pid_m * blocks_per_pid |
| 84 | block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m) |
| 85 | |
| 86 | weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 87 | |
| 88 | for block_id in range(block_id_start, block_id_end): |
| 89 | offs_m = block_id * block_size_m + tl.arange(0, block_size_m) |
| 90 | mask_m = offs_m < n_rows |
| 91 | mask_mn = mask_m[:, None] & mask_n[None, :] |
| 92 | x = tl.load( |
| 93 | x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :], |
| 94 | mask=mask_mn, |
| 95 | other=0.0, |
| 96 | ).to(tl.float32) |
| 97 | row_var = tl.sum(x * x, axis=1) / n_cols |
| 98 | rstd = tl.math.rsqrt(row_var + eps) |
| 99 | tl.store(rstd_ptr + offs_m, rstd, mask=mask_m) |
| 100 | |
| 101 | y = x * rstd[:, None] * weight |
| 102 | tl.store( |
| 103 | y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :], |
| 104 | y, |
| 105 | mask=mask_mn, |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | @triton.jit |
| 110 | def _add_rmsnorm_fwd_kernel( |
| 111 | res_ptr, # [T, N] residual (read) |
| 112 | delta_ptr, # [T, N] delta to add (read) |
| 113 | weight_ptr, |
| 114 | y_ptr, # [T, N] normed output |
| 115 | res_out_ptr, # [T, N] updated residual output |
| 116 | eps, |
| 117 | res_stride_0, |
| 118 | delta_stride_0, |
| 119 | y_stride_0, |
| 120 | res_out_stride_0, |
| 121 | n_cols, |
| 122 | block_size_n: tl.constexpr, |
| 123 | ): |
| 124 | pid_m = tl.program_id(0).to(tl.int64) |
| 125 | offs_n = tl.arange(0, block_size_n) |
| 126 | mask_n = offs_n < n_cols |
| 127 | |
| 128 | weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 129 | r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to( |
| 130 | tl.float32 |
| 131 | ) |
| 132 | d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to( |
| 133 | tl.float32 |
| 134 | ) |
| 135 | # Round the sum to the residual dtype first (matches the eager |
| 136 | # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit). |
| 137 | s = (r + d).to(res_out_ptr.dtype.element_ty) |
| 138 | tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n) |
| 139 | x = s.to(tl.float32) |
| 140 | row_var = tl.sum(x * x, axis=0) / n_cols |
| 141 | rstd = tl.math.rsqrt(row_var + eps) |
| 142 | y = x * rstd * weight |
| 143 | tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) |
| 144 | |
| 145 | |
| 146 | def add_rmsnorm( |
| 147 | residual: torch.Tensor, |
| 148 | delta: torch.Tensor, |
| 149 | weight: torch.Tensor, |
| 150 | eps: float, |
| 151 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 152 | """Fused ``res = residual + delta; y = rmsnorm(res)``. |
| 153 | |
| 154 | Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no |
| 155 | in-place update of the inputs). |
| 156 | """ |
| 157 | assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape) |
| 158 | n_rows, n_cols = residual.shape |
| 159 | assert weight.shape[0] == n_cols |
| 160 | y = torch.empty_like(residual) |
| 161 | res_out = torch.empty_like(residual) |
| 162 | if n_rows == 0: |
| 163 | return y, res_out |
| 164 | |
| 165 | block_size_n = triton.next_power_of_2(n_cols) |
| 166 | max_block_size_n = _MAX_FUSED_SIZE // residual.element_size() |
| 167 | if max_block_size_n < block_size_n: |
| 168 | raise RuntimeError(f"Large {n_cols=} is not supported") |
| 169 | num_warps = _get_num_warps_from_block_size(block_size_n) |
| 170 | _add_rmsnorm_fwd_kernel[(n_rows,)]( |
| 171 | residual, |
| 172 | delta, |
| 173 | weight, |
| 174 | y, |
| 175 | res_out, |
| 176 | eps, |
| 177 | residual.stride(0), |
| 178 | delta.stride(0), |
| 179 | y.stride(0), |
| 180 | res_out.stride(0), |
| 181 | n_cols, |
| 182 | block_size_n, |
| 183 | num_warps=num_warps, |
| 184 | ) |
| 185 | return y, res_out |
| 186 | |
| 187 | |
| 188 | @triton.jit |
| 189 | def _embed_rmsnorm_kernel( |
| 190 | ids_ptr, # [T] token ids |
| 191 | table_ptr, # [V, N] embedding table |
| 192 | weight_ptr, # [N] (HAS_NORM only) |
| 193 | chain_weight_ptr, # [N] (HAS_CHAIN only) |
| 194 | out_ptr, # [T, N] rmsnorm(table[ids], weight) |
| 195 | chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only) |
| 196 | eps, |
| 197 | table_stride_0, |
| 198 | n_cols, |
| 199 | block_size_n: tl.constexpr, |
| 200 | HAS_NORM: tl.constexpr, |
| 201 | HAS_CHAIN: tl.constexpr, |
| 202 | ): |
| 203 | pid_m = tl.program_id(0).to(tl.int64) |
| 204 | offs_n = tl.arange(0, block_size_n) |
| 205 | mask_n = offs_n < n_cols |
| 206 | row = tl.load(ids_ptr + pid_m).to(tl.int64) |
| 207 | x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0) |
| 208 | if HAS_NORM: |
| 209 | xf = x.to(tl.float32) |
| 210 | rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) |
| 211 | w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 212 | # Round to the output dtype so the chained norm is bit-exact vs the |
| 213 | # unfused pair (which stores bf16 in between). |
| 214 | x = (xf * rstd * w).to(out_ptr.dtype.element_ty) |
| 215 | tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n) |
| 216 | if HAS_CHAIN: |
| 217 | xf = x.to(tl.float32) |
| 218 | rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) |
| 219 | w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 220 | tl.store( |
| 221 | chain_out_ptr + pid_m * n_cols + offs_n, |
| 222 | (xf * rstd * w).to(chain_out_ptr.dtype.element_ty), |
| 223 | mask=mask_n, |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | def embed_rmsnorm( |
| 228 | input_ids: torch.Tensor, |
| 229 | embed_table: torch.Tensor, |
| 230 | weight: torch.Tensor | None, |
| 231 | eps: float, |
| 232 | chain_weight: torch.Tensor | None = None, |
| 233 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 234 | """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm. |
| 235 | |
| 236 | Requires the full vocab on-rank (replicated or tp_size == 1). |
| 237 | ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure |
| 238 | embedding-table gather. |
| 239 | ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the |
| 240 | first decoder layer's pre-attention norm) as a second output, still one |
| 241 | launch. Bit-exact vs the unfused module sequence.""" |
| 242 | ids = input_ids.view(-1) |
| 243 | (T,) = ids.shape |
| 244 | n = embed_table.shape[1] |
| 245 | out = torch.empty( |
| 246 | (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device |
| 247 | ) |
| 248 | chain_out = torch.empty_like(out) if chain_weight is not None else None |
| 249 | if T > 0: |
| 250 | block_size_n = triton.next_power_of_2(n) |
| 251 | _embed_rmsnorm_kernel[(T,)]( |
| 252 | ids, |
| 253 | embed_table, |
| 254 | weight if weight is not None else embed_table, |
| 255 | chain_weight if chain_weight is not None else embed_table, |
| 256 | out, |
| 257 | chain_out if chain_out is not None else out, |
| 258 | eps, |
| 259 | embed_table.stride(0), |
| 260 | n, |
| 261 | block_size_n, |
| 262 | HAS_NORM=weight is not None, |
| 263 | HAS_CHAIN=chain_weight is not None, |
| 264 | num_warps=_get_num_warps_from_block_size(block_size_n), |
| 265 | ) |
| 266 | if chain_out is not None: |
| 267 | return out, chain_out |
| 268 | return out |
| 269 | |
| 270 | |
| 271 | @triton.jit |
| 272 | def _embed_dual_rmsnorm_cat_kernel( |
| 273 | hidden_ptr, # [T, N] |
| 274 | emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER |
| 275 | ids_ptr, # [T] token ids (GATHER only) |
| 276 | w_hidden_ptr, # [N] |
| 277 | w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only) |
| 278 | w_embed_ptr, # [N] |
| 279 | out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))] |
| 280 | eps, |
| 281 | hidden_stride_0, |
| 282 | emb_stride_0, |
| 283 | n_cols, |
| 284 | block_size_n: tl.constexpr, |
| 285 | GATHER: tl.constexpr, |
| 286 | HAS_PRE_NORM: tl.constexpr, |
| 287 | ): |
| 288 | pid_m = tl.program_id(0).to(tl.int64) |
| 289 | which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N) |
| 290 | offs_n = tl.arange(0, block_size_n) |
| 291 | mask_n = offs_n < n_cols |
| 292 | if which == 0: |
| 293 | x = tl.load( |
| 294 | hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0 |
| 295 | ).to(tl.float32) |
| 296 | w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 297 | else: |
| 298 | row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m |
| 299 | x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to( |
| 300 | tl.float32 |
| 301 | ) |
| 302 | if HAS_PRE_NORM: |
| 303 | w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 304 | rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) |
| 305 | # Round-trip through the output dtype so the chained norm is |
| 306 | # bit-exact vs the unfused pair (which stores bf16 in between). |
| 307 | x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32) |
| 308 | w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) |
| 309 | rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) |
| 310 | tl.store( |
| 311 | out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n, |
| 312 | (x * rstd * w).to(out_ptr.dtype.element_ty), |
| 313 | mask=mask_n, |
| 314 | ) |
| 315 | |
| 316 | |
| 317 | def embed_dual_rmsnorm_cat( |
| 318 | hidden: torch.Tensor, |
| 319 | hidden_weight: torch.Tensor, |
| 320 | embed_weight: torch.Tensor, |
| 321 | eps: float, |
| 322 | *, |
| 323 | embeds: torch.Tensor | None = None, |
| 324 | input_ids: torch.Tensor | None = None, |
| 325 | embed_table: torch.Tensor | None = None, |
| 326 | pre_norm_weight: torch.Tensor | None = None, |
| 327 | ) -> torch.Tensor: |
| 328 | """The MTP depth-layer input in one launch: |
| 329 | ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``. |
| 330 | |
| 331 | The embed side is either a fused row gather ``embed_table[input_ids]`` |
| 332 | (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged |
| 333 | multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the |
| 334 | backbone embed_norm in front of the depth embed_norm (bit-exact vs the |
| 335 | unfused sequence). The concat copies collapse into direct writes.""" |
| 336 | T, n = hidden.shape |
| 337 | if embeds is not None: |
| 338 | assert embeds.shape == hidden.shape |
| 339 | src, ids, src_stride = embeds, embeds, embeds.stride(0) |
| 340 | gather = False |
| 341 | else: |
| 342 | assert input_ids is not None and embed_table is not None |
| 343 | assert input_ids.shape == (T,) and embed_table.shape[1] == n |
| 344 | src, ids, src_stride = embed_table, input_ids, embed_table.stride(0) |
| 345 | gather = True |
| 346 | out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device) |
| 347 | if T == 0: |
| 348 | return out |
| 349 | block_size_n = triton.next_power_of_2(n) |
| 350 | _embed_dual_rmsnorm_cat_kernel[(T, 2)]( |
| 351 | hidden, |
| 352 | src, |
| 353 | ids, |
| 354 | hidden_weight, |
| 355 | pre_norm_weight if pre_norm_weight is not None else embed_weight, |
| 356 | embed_weight, |
| 357 | out, |
| 358 | eps, |
| 359 | hidden.stride(0), |
| 360 | src_stride, |
| 361 | n, |
| 362 | block_size_n, |
| 363 | GATHER=gather, |
| 364 | HAS_PRE_NORM=pre_norm_weight is not None, |
| 365 | num_warps=_get_num_warps_from_block_size(block_size_n), |
| 366 | ) |
| 367 | return out |
| 368 | |
| 369 | |
| 370 | def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: |
| 371 | assert x.ndim == 2, f"{x.shape=}" |
| 372 | assert weight.ndim == 1, f"{weight.shape=}" |
| 373 | n_rows, n_cols = x.shape |
| 374 | assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}" |
| 375 | y = torch.empty_like(x) |
| 376 | rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device) |
| 377 | |
| 378 | block_size_n = triton.next_power_of_2(n_cols) |
| 379 | max_block_size_n = _MAX_FUSED_SIZE // x.element_size() |
| 380 | if max_block_size_n < block_size_n: |
| 381 | raise RuntimeError(f"Large {n_cols=} is not supported") |
| 382 | block_size_m = max(1, 4096 // block_size_n) |
| 383 | num_warps = _get_num_warps_from_block_size(block_size_n) |
| 384 | |
| 385 | if block_size_m == 1: |
| 386 | _rmsnorm_fwd_kernel[(n_rows,)]( |
| 387 | x, |
| 388 | weight, |
| 389 | y, |
| 390 | rstd, |
| 391 | eps, |
| 392 | x.stride(0), |
| 393 | y.stride(0), |
| 394 | n_cols, |
| 395 | block_size_n, |
| 396 | num_warps=num_warps, |
| 397 | ) |
| 398 | else: |
| 399 | grid_size = _get_grid_size_for_mem_bw_kernel(x.device) |
| 400 | _rmsnorm_fwd_kernel_block_m[(grid_size,)]( |
| 401 | x, |
| 402 | weight, |
| 403 | y, |
| 404 | rstd, |
| 405 | eps, |
| 406 | x.stride(0), |
| 407 | y.stride(0), |
| 408 | n_rows, |
| 409 | n_cols, |
| 410 | block_size_m, |
| 411 | block_size_n, |
| 412 | num_warps=num_warps, |
| 413 | ) |
| 414 | return y |