| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused CUDA kernels for the Inkling vision/audio towers. |
| 4 | |
| 5 | Both kernels keep the reference paths' fp32 accumulation and per-op bf16 |
| 6 | rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently |
| 7 | bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order |
| 8 | (real-checkpoint-weight cosine vs reference > 0.9999998). |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.triton_utils import tl, tldevice, triton |
| 16 | |
| 17 | from .norm import _get_num_warps_from_block_size |
| 18 | |
| 19 | |
| 20 | @triton.jit |
| 21 | def _dmel_embed_sum_norm_kernel( |
| 22 | idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) |
| 23 | w_ptr, # [NB * VOCAB, D] bf16 embedding table |
| 24 | norm_w_ptr, # [D] (unused if HAS_NORM=False) |
| 25 | out_ptr, # [T, D] bf16 |
| 26 | eps, |
| 27 | D, |
| 28 | stride_idx_t, |
| 29 | NB: tl.constexpr, |
| 30 | VOCAB: tl.constexpr, |
| 31 | D_P2: tl.constexpr, |
| 32 | HAS_NORM: tl.constexpr, |
| 33 | ): |
| 34 | t = tl.program_id(0).to(tl.int64) |
| 35 | offs = tl.arange(0, D_P2) |
| 36 | mask = offs < D |
| 37 | # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), |
| 38 | # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). |
| 39 | acc = tl.zeros([D_P2], dtype=tl.float32) |
| 40 | for b in tl.static_range(NB): |
| 41 | v = tl.load(idx_ptr + t * stride_idx_t + b) |
| 42 | row = (b * VOCAB + v).to(tl.int64) |
| 43 | acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) |
| 44 | h = acc.to(tl.bfloat16) |
| 45 | if HAS_NORM: |
| 46 | # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded |
| 47 | # bf16 multiply with the bf16 weight. |
| 48 | x32 = h.to(tl.float32) |
| 49 | var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D |
| 50 | xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) |
| 51 | w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) |
| 52 | h = xn * w |
| 53 | tl.store(out_ptr + t * D + offs, h, mask=mask) |
| 54 | |
| 55 | |
| 56 | def dmel_embed_sum_norm( |
| 57 | dmel_idx: torch.Tensor, # [T, NB] int32 |
| 58 | weight: torch.Tensor, # [NB * VOCAB, D] bf16 |
| 59 | norm_weight: torch.Tensor | None, |
| 60 | eps: float, |
| 61 | ) -> torch.Tensor: |
| 62 | """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no |
| 63 | [T, NB, D] intermediate).""" |
| 64 | T, nb = dmel_idx.shape |
| 65 | D = weight.shape[1] |
| 66 | assert weight.shape[0] % nb == 0 |
| 67 | vocab = weight.shape[0] // nb |
| 68 | out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) |
| 69 | if T == 0: |
| 70 | return out |
| 71 | d_p2 = triton.next_power_of_2(D) |
| 72 | _dmel_embed_sum_norm_kernel[(T,)]( |
| 73 | dmel_idx, |
| 74 | weight, |
| 75 | norm_weight if norm_weight is not None else weight, |
| 76 | out, |
| 77 | eps, |
| 78 | D, |
| 79 | dmel_idx.stride(0), |
| 80 | NB=nb, |
| 81 | VOCAB=vocab, |
| 82 | D_P2=d_p2, |
| 83 | HAS_NORM=norm_weight is not None, |
| 84 | # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% |
| 85 | # at large T (the 80-row gather chain is latency- not lane-bound). |
| 86 | num_warps=8, |
| 87 | ) |
| 88 | return out |
| 89 | |
| 90 | |
| 91 | @triton.jit |
| 92 | def _rmsnorm_gelu_kernel( |
| 93 | x_ptr, # [R, D] bf16 |
| 94 | w_ptr, # [D] |
| 95 | out_ptr, # [R, D] bf16 (or the folded layout when FOLD) |
| 96 | eps, |
| 97 | R, |
| 98 | D, |
| 99 | D_P2: tl.constexpr, |
| 100 | BLOCK_M: tl.constexpr, # rows per block (>1 for small D) |
| 101 | HAS_GELU: tl.constexpr, |
| 102 | FOLD: tl.constexpr, |
| 103 | # fold geometry: input rows index [N, T, H, W]; the store scatters each |
| 104 | # row to (out_row, slot) of fold_timespace_to_depth's output layout. |
| 105 | FT: tl.constexpr, |
| 106 | FH: tl.constexpr, |
| 107 | FW: tl.constexpr, |
| 108 | TF: tl.constexpr, |
| 109 | HF: tl.constexpr, |
| 110 | ): |
| 111 | pid = tl.program_id(0).to(tl.int64) |
| 112 | rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) |
| 113 | rmask = rows < R |
| 114 | offs = tl.arange(0, D_P2) |
| 115 | mask = rmask[:, None] & (offs < D)[None, :] |
| 116 | x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( |
| 117 | tl.float32 |
| 118 | ) |
| 119 | var = tl.sum(x32 * x32, axis=1) / D |
| 120 | xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) |
| 121 | w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) |
| 122 | h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm |
| 123 | if HAS_GELU: |
| 124 | # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, |
| 125 | # matching F.gelu's opmath on a bf16 tensor. |
| 126 | g32 = h.to(tl.float32) |
| 127 | h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) |
| 128 | if FOLD: |
| 129 | # Store directly into the next layer's folded layout (a pure |
| 130 | # permutation — replaces the separate fold copy pass). |
| 131 | t = (rows // (FH * FW)) % FT |
| 132 | hh = (rows // FW) % FH |
| 133 | ww = rows % FW |
| 134 | n = rows // (FT * FH * FW) |
| 135 | slot = ((t % TF) * HF + hh % HF) * HF + ww % HF |
| 136 | out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( |
| 137 | FW // HF |
| 138 | ) + ww // HF |
| 139 | base = (out_row * (TF * HF * HF) + slot) * D |
| 140 | tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) |
| 141 | else: |
| 142 | tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) |
| 143 | |
| 144 | |
| 145 | def rmsnorm_gelu( |
| 146 | x: torch.Tensor, # [..., D] bf16 contiguous |
| 147 | weight: torch.Tensor, |
| 148 | eps: float, |
| 149 | gelu: bool = True, |
| 150 | fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold |
| 151 | ) -> torch.Tensor: |
| 152 | """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block |
| 153 | when D is small. With ``fold``, x must be [N, T, H, W, D] and the output |
| 154 | comes back as ``fold_timespace_to_depth(result, *fold)``.""" |
| 155 | D = x.shape[-1] |
| 156 | flat = x.reshape(-1, D) |
| 157 | assert flat.stride(1) == 1 and flat.stride(0) == D |
| 158 | R = flat.shape[0] |
| 159 | if fold is None: |
| 160 | out = torch.empty_like(flat) |
| 161 | ft = fh = fw = tf = hf = 1 |
| 162 | out_shape = x.shape |
| 163 | else: |
| 164 | tf, hf = fold |
| 165 | N, ft, fh, fw, _ = x.shape |
| 166 | out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) |
| 167 | out = torch.empty(out_shape, dtype=x.dtype, device=x.device) |
| 168 | if R == 0: |
| 169 | return out.reshape(out_shape) |
| 170 | d_p2 = triton.next_power_of_2(D) |
| 171 | block_m = max(1, 4096 // d_p2) |
| 172 | _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( |
| 173 | flat, |
| 174 | weight, |
| 175 | out, |
| 176 | eps, |
| 177 | R, |
| 178 | D, |
| 179 | D_P2=d_p2, |
| 180 | BLOCK_M=block_m, |
| 181 | HAS_GELU=gelu, |
| 182 | FOLD=fold is not None, |
| 183 | FT=ft, |
| 184 | FH=fh, |
| 185 | FW=fw, |
| 186 | TF=tf, |
| 187 | HF=hf, |
| 188 | num_warps=_get_num_warps_from_block_size(d_p2 * block_m), |
| 189 | ) |
| 190 | return out.reshape(out_shape) |