| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | import torch.nn.functional as F |
| 5 | from torch import Tensor |
| 6 | |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | from vllm.utils.torch_utils import direct_register_custom_op |
| 9 | |
| 10 | |
| 11 | @triton.jit |
| 12 | def _rmsnorm_nw_kernel( |
| 13 | x_ptr, |
| 14 | out_ptr, |
| 15 | stride_row, |
| 16 | D, |
| 17 | eps, |
| 18 | RBLOCK: tl.constexpr, |
| 19 | ): |
| 20 | """Weight-free RMSNorm Triton kernel: out = x * rsqrt(mean(x², -1) + eps).""" |
| 21 | row = tl.program_id(0) |
| 22 | cols = tl.arange(0, RBLOCK) |
| 23 | mask = cols < D |
| 24 | |
| 25 | x = tl.load( |
| 26 | x_ptr + row * stride_row + cols, |
| 27 | mask=mask, |
| 28 | other=0.0, |
| 29 | eviction_policy="evict_first", |
| 30 | ).to(tl.float32) |
| 31 | |
| 32 | var = tl.sum(x * x, 0) / D |
| 33 | rstd = tl.rsqrt(var + eps) |
| 34 | |
| 35 | out = (x * rstd).to(out_ptr.dtype.element_ty) |
| 36 | tl.store(out_ptr + row * D + cols, out, mask=mask, eviction_policy="evict_first") |
| 37 | |
| 38 | |
| 39 | def rmsnorm_nw(x: Tensor, eps: float) -> Tensor: |
| 40 | """Weight-free RMSNorm over the last dimension. |
| 41 | |
| 42 | Treats *x* as ``[num_rows, D]`` where ``num_rows = product(shape[:-1])``. |
| 43 | Returns a contiguous tensor with the same shape and dtype as *x*. |
| 44 | """ |
| 45 | orig_shape = x.shape |
| 46 | D = orig_shape[-1] |
| 47 | x_2d = x.reshape(-1, D) |
| 48 | num_rows = x_2d.shape[0] |
| 49 | |
| 50 | out = torch.empty_like(x_2d) |
| 51 | RBLOCK = triton.next_power_of_2(D) |
| 52 | |
| 53 | _rmsnorm_nw_kernel[(num_rows,)]( |
| 54 | x_2d, |
| 55 | out, |
| 56 | x_2d.stride(0), |
| 57 | D, |
| 58 | eps, |
| 59 | RBLOCK=RBLOCK, |
| 60 | num_warps=1 if RBLOCK <= 512 else (4 if RBLOCK <= 4096 else 8), |
| 61 | ) |
| 62 | return out.view(orig_shape) |
| 63 | |
| 64 | |
| 65 | @triton.jit |
| 66 | def _hc_head_reduce_store_kernel( |
| 67 | pre_ptr, |
| 68 | x_ptr, |
| 69 | out_ptr, |
| 70 | hidden_size: tl.constexpr, |
| 71 | hc_mult: tl.constexpr, |
| 72 | pre_stride_t: tl.constexpr, |
| 73 | pre_stride_m: tl.constexpr, |
| 74 | x_stride_t: tl.constexpr, |
| 75 | x_stride_m: tl.constexpr, |
| 76 | x_stride_h: tl.constexpr, |
| 77 | out_stride_t: tl.constexpr, |
| 78 | out_stride_h: tl.constexpr, |
| 79 | BLOCK_H: tl.constexpr, |
| 80 | ): |
| 81 | token_idx = tl.program_id(0).to(tl.int64) |
| 82 | block_idx = tl.program_id(1) |
| 83 | offsets = block_idx * BLOCK_H + tl.arange(0, BLOCK_H) |
| 84 | mask = offsets < hidden_size |
| 85 | |
| 86 | acc = tl.zeros((BLOCK_H,), dtype=tl.float32) |
| 87 | for mix_idx in tl.static_range(0, hc_mult): |
| 88 | pre = tl.load(pre_ptr + token_idx * pre_stride_t + mix_idx * pre_stride_m).to( |
| 89 | tl.float32 |
| 90 | ) |
| 91 | x = tl.load( |
| 92 | x_ptr |
| 93 | + token_idx * x_stride_t |
| 94 | + mix_idx * x_stride_m |
| 95 | + offsets * x_stride_h, |
| 96 | mask=mask, |
| 97 | other=0.0, |
| 98 | ).to(tl.float32) |
| 99 | acc += pre * x |
| 100 | |
| 101 | tl.store( |
| 102 | out_ptr + token_idx * out_stride_t + offsets * out_stride_h, |
| 103 | acc, |
| 104 | mask=mask, |
| 105 | ) |
| 106 | |
| 107 | |
| 108 | def hc_collapse_triton(x: Tensor, pre_mix: Tensor) -> Tensor: |
| 109 | """Collapse BF16 residual streams with FP32 pre-mix coefficients.""" |
| 110 | assert x.ndim == 3 and x.dtype == torch.bfloat16 |
| 111 | num_tokens, hc_mult, hidden_size = x.shape |
| 112 | assert pre_mix.shape == (num_tokens, hc_mult) |
| 113 | assert pre_mix.dtype == torch.float32 |
| 114 | out = torch.empty(num_tokens, hidden_size, dtype=x.dtype, device=x.device) |
| 115 | if num_tokens == 0: |
| 116 | return out |
| 117 | |
| 118 | block_h = 1024 |
| 119 | _hc_head_reduce_store_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( |
| 120 | pre_mix, |
| 121 | x, |
| 122 | out, |
| 123 | hidden_size, |
| 124 | hc_mult, |
| 125 | pre_mix.stride(0), |
| 126 | pre_mix.stride(1), |
| 127 | x.stride(0), |
| 128 | x.stride(1), |
| 129 | x.stride(2), |
| 130 | out.stride(0), |
| 131 | out.stride(1), |
| 132 | BLOCK_H=block_h, |
| 133 | num_warps=4, |
| 134 | # Preserve the separate FP32 multiply and sum in the Torch reference. |
| 135 | enable_fp_fusion=False, |
| 136 | ) |
| 137 | return out |
| 138 | |
| 139 | |
| 140 | def _hc_collapse_triton_fake(x: Tensor, pre_mix: Tensor) -> Tensor: |
| 141 | return torch.empty(x.shape[0], x.shape[2], dtype=x.dtype, device=x.device) |
| 142 | |
| 143 | |
| 144 | @triton.jit |
| 145 | def _mhc_pre_mix_kernel( |
| 146 | gemm_ptr, |
| 147 | sqrsum_ptr, |
| 148 | hc_scale_ptr, |
| 149 | hc_base_ptr, |
| 150 | out_ptr, |
| 151 | splitk, |
| 152 | hc_mult: tl.constexpr, |
| 153 | gemm_stride_k, |
| 154 | gemm_stride_t, |
| 155 | gemm_stride_j, |
| 156 | sqrsum_stride_k, |
| 157 | sqrsum_stride_t, |
| 158 | out_stride_t, |
| 159 | out_stride_j, |
| 160 | inv_hc_hidden, |
| 161 | rms_eps, |
| 162 | hc_pre_eps, |
| 163 | SPLITK_BLOCK: tl.constexpr, |
| 164 | HC_BLOCK: tl.constexpr, |
| 165 | ): |
| 166 | """Recover the pre-mix gate from a split-k mHC pre GEMM output.""" |
| 167 | token_idx = tl.program_id(0).to(tl.int64) |
| 168 | ks = tl.arange(0, SPLITK_BLOCK) |
| 169 | js = tl.arange(0, HC_BLOCK) |
| 170 | kmask = ks < splitk |
| 171 | jmask = js < hc_mult |
| 172 | |
| 173 | # Only the first hc_mult GEMM columns feed the pre gate. |
| 174 | gemm = tl.load( |
| 175 | gemm_ptr |
| 176 | + ks[:, None] * gemm_stride_k |
| 177 | + token_idx * gemm_stride_t |
| 178 | + js[None, :] * gemm_stride_j, |
| 179 | mask=kmask[:, None] & jmask[None, :], |
| 180 | other=0.0, |
| 181 | ).to(tl.float32) |
| 182 | mixes = tl.sum(gemm, 0) |
| 183 | |
| 184 | sqrsum = tl.load( |
| 185 | sqrsum_ptr + ks * sqrsum_stride_k + token_idx * sqrsum_stride_t, |
| 186 | mask=kmask, |
| 187 | other=0.0, |
| 188 | ).to(tl.float32) |
| 189 | rstd = tl.rsqrt(tl.sum(sqrsum, 0) * inv_hc_hidden + rms_eps) |
| 190 | |
| 191 | scale = tl.load(hc_scale_ptr).to(tl.float32) |
| 192 | base = tl.load(hc_base_ptr + js, mask=jmask, other=0.0).to(tl.float32) |
| 193 | |
| 194 | pre = tl.sigmoid(mixes * rstd * scale + base) + hc_pre_eps |
| 195 | tl.store(out_ptr + token_idx * out_stride_t + js * out_stride_j, pre, mask=jmask) |
| 196 | |
| 197 | |
| 198 | def mhc_pre_mix_triton( |
| 199 | gemm_out: Tensor, |
| 200 | sqrsum: Tensor, |
| 201 | hc_scale: Tensor, |
| 202 | hc_base: Tensor, |
| 203 | hc_mult: int, |
| 204 | hc_hidden_size: int, |
| 205 | rms_eps: float, |
| 206 | hc_pre_eps: float, |
| 207 | ) -> Tensor: |
| 208 | """Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output. |
| 209 | |
| 210 | AITER's ``mhc_pre_big_fuse`` consumes the unreduced ``[splitk, tokens, |
| 211 | hc_mult3]`` GEMM output and the matching row square-sums, but only returns |
| 212 | the post and comb gates. The delayed formulation also needs the pre gate, |
| 213 | to carry into the next sublayer seam. It is the same slice of the same |
| 214 | numbers, so recover it here rather than repeating the projection. |
| 215 | """ |
| 216 | assert gemm_out.ndim == 3 and gemm_out.dtype == torch.float32 |
| 217 | assert sqrsum.ndim == 2 and sqrsum.dtype == torch.float32 |
| 218 | splitk, num_tokens = gemm_out.shape[0], gemm_out.shape[1] |
| 219 | assert sqrsum.shape == (splitk, num_tokens) |
| 220 | |
| 221 | out = torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=gemm_out.device) |
| 222 | if num_tokens == 0: |
| 223 | return out |
| 224 | |
| 225 | _mhc_pre_mix_kernel[(num_tokens,)]( |
| 226 | gemm_out, |
| 227 | sqrsum, |
| 228 | hc_scale, |
| 229 | hc_base, |
| 230 | out, |
| 231 | splitk, |
| 232 | hc_mult, |
| 233 | gemm_out.stride(0), |
| 234 | gemm_out.stride(1), |
| 235 | gemm_out.stride(2), |
| 236 | sqrsum.stride(0), |
| 237 | sqrsum.stride(1), |
| 238 | out.stride(0), |
| 239 | out.stride(1), |
| 240 | 1.0 / hc_hidden_size, |
| 241 | rms_eps, |
| 242 | hc_pre_eps, |
| 243 | SPLITK_BLOCK=triton.next_power_of_2(splitk), |
| 244 | HC_BLOCK=triton.next_power_of_2(hc_mult), |
| 245 | num_warps=1, |
| 246 | # Match the separate FP32 multiply and add in the Torch reference. |
| 247 | enable_fp_fusion=False, |
| 248 | ) |
| 249 | return out |
| 250 | |
| 251 | |
| 252 | def _mhc_pre_mix_triton_fake( |
| 253 | gemm_out: Tensor, |
| 254 | sqrsum: Tensor, |
| 255 | hc_scale: Tensor, |
| 256 | hc_base: Tensor, |
| 257 | hc_mult: int, |
| 258 | hc_hidden_size: int, |
| 259 | rms_eps: float, |
| 260 | hc_pre_eps: float, |
| 261 | ) -> Tensor: |
| 262 | return torch.empty( |
| 263 | gemm_out.shape[1], hc_mult, dtype=torch.float32, device=gemm_out.device |
| 264 | ) |
| 265 | |
| 266 | |
| 267 | def hc_head_reduce_triton_kernel( |
| 268 | x: torch.Tensor, |
| 269 | hc_fn: torch.Tensor, |
| 270 | hc_scale: torch.Tensor, |
| 271 | hc_base: torch.Tensor, |
| 272 | out: torch.Tensor, |
| 273 | norm_eps: float, |
| 274 | hc_eps: float, |
| 275 | ) -> None: |
| 276 | x_flat = x.flatten(-2) |
| 277 | x_normed = rmsnorm_nw(x_flat, norm_eps) |
| 278 | mixes = F.linear(x_normed.float(), hc_fn) |
| 279 | pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps |
| 280 | |
| 281 | hidden_size = x.shape[-1] |
| 282 | hc_mult = x.shape[-2] |
| 283 | block_h = 1024 |
| 284 | _hc_head_reduce_store_kernel[(x.shape[0], (hidden_size + block_h - 1) // block_h)]( |
| 285 | pre, |
| 286 | x, |
| 287 | out, |
| 288 | hidden_size, |
| 289 | hc_mult, |
| 290 | pre.stride(0), |
| 291 | pre.stride(1), |
| 292 | x.stride(0), |
| 293 | x.stride(1), |
| 294 | x.stride(2), |
| 295 | out.stride(0), |
| 296 | out.stride(1), |
| 297 | BLOCK_H=block_h, |
| 298 | num_warps=4, |
| 299 | ) |
| 300 | |
| 301 | |
| 302 | def _hc_head_triton( |
| 303 | hs_flat: torch.Tensor, |
| 304 | fn: torch.Tensor, |
| 305 | hc_scale: torch.Tensor, |
| 306 | hc_base: torch.Tensor, |
| 307 | out: torch.Tensor, |
| 308 | hidden_size: int, |
| 309 | rms_eps: float, |
| 310 | hc_eps: float, |
| 311 | hc_mult: int, |
| 312 | ) -> None: |
| 313 | """Fill pre-allocated `out` (T, H) in-place with the hc_head result.""" |
| 314 | if hs_flat.shape[0] == 0: |
| 315 | return |
| 316 | |
| 317 | hc_head_reduce_triton_kernel( |
| 318 | hs_flat, |
| 319 | fn, |
| 320 | hc_scale, |
| 321 | hc_base, |
| 322 | out, |
| 323 | rms_eps, |
| 324 | hc_eps, |
| 325 | ) |
| 326 | return |
| 327 | |
| 328 | |
| 329 | direct_register_custom_op( |
| 330 | op_name="hc_head_triton", |
| 331 | op_func=_hc_head_triton, |
| 332 | mutates_args=["out"], |
| 333 | ) |
| 334 | |
| 335 | |
| 336 | direct_register_custom_op( |
| 337 | op_name="hc_collapse_triton", |
| 338 | op_func=hc_collapse_triton, |
| 339 | mutates_args=[], |
| 340 | fake_impl=_hc_collapse_triton_fake, |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | direct_register_custom_op( |
| 345 | op_name="mhc_pre_mix_triton", |
| 346 | op_func=mhc_pre_mix_triton, |
| 347 | mutates_args=[], |
| 348 | fake_impl=_mhc_pre_mix_triton_fake, |
| 349 | ) |