@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 155 lines 4.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Fused Gemma-style RMSNorm for AMD ROCm via Triton.
4
5 Gemma RMSNorm = normalize(x) * (1 + weight), computed in fp32. FlashInfer's
6 ``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` CUDA kernels are unavailable on
7 ROCm, so the AMD path previously used a ~8-op PyTorch sequence (float cast, add,
8 pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing
9 fp32 intermediates. These kernels collapse that into a single pass per row.
10
11 Two entry points:
12 * ``gemma_rmsnorm(x, w, eps)`` -> normalized tensor
13 * ``gemma_fused_add_rmsnorm(x, res, w, eps)`` -> (normalized, x + res)
14
15 Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
16 so they serve both the full-hidden norms (input/post-attn/final) and the
17 per-head q_norm/k_norm (N == head_dim). Inputs may be non-contiguous views
18 (e.g. ``qkv.split`` slices); strides are passed through and outputs are written
19 contiguous.
20 """
21
22 import torch
23
24 from vllm.triton_utils import tl, triton
25
26
27 @triton.jit
28 def _gemma_rmsnorm_kernel(
29 x_ptr,
30 w_ptr,
31 out_ptr,
32 n_cols,
33 stride_row,
34 stride_col,
35 eps,
36 BLOCK_N: tl.constexpr,
37 ):
38 row = tl.program_id(0)
39 cols = tl.arange(0, BLOCK_N)
40 mask = cols < n_cols
41 x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
42 tl.float32
43 )
44 var = tl.sum(x * x, axis=0) / n_cols
45 rstd = 1.0 / tl.sqrt(var + eps)
46 w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
47 out = x * rstd * (1.0 + w)
48 tl.store(
49 out_ptr + row * n_cols + cols,
50 out.to(out_ptr.dtype.element_ty),
51 mask=mask,
52 )
53
54
55 @triton.jit
56 def _gemma_fused_add_rmsnorm_kernel(
57 x_ptr,
58 res_ptr,
59 w_ptr,
60 out_ptr,
61 res_out_ptr,
62 n_cols,
63 stride_xrow,
64 stride_xcol,
65 stride_rrow,
66 stride_rcol,
67 eps,
68 BLOCK_N: tl.constexpr,
69 ):
70 row = tl.program_id(0)
71 cols = tl.arange(0, BLOCK_N)
72 mask = cols < n_cols
73 x = tl.load(
74 x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
75 ).to(tl.float32)
76 r = tl.load(
77 res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
78 ).to(tl.float32)
79 s = x + r
80 # residual_out is the pre-norm sum (consumed by the next layer's add).
81 tl.store(
82 res_out_ptr + row * n_cols + cols,
83 s.to(res_out_ptr.dtype.element_ty),
84 mask=mask,
85 )
86 var = tl.sum(s * s, axis=0) / n_cols
87 rstd = 1.0 / tl.sqrt(var + eps)
88 w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
89 out = s * rstd * (1.0 + w)
90 tl.store(
91 out_ptr + row * n_cols + cols,
92 out.to(out_ptr.dtype.element_ty),
93 mask=mask,
94 )
95
96
97 def _num_warps(block_n: int) -> int:
98 if block_n >= 4096:
99 return 16
100 if block_n >= 1024:
101 return 8
102 return 4
103
104
105 def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
106 orig_shape = x.shape
107 n = orig_shape[-1]
108 x2 = x.reshape(-1, n)
109 m = x2.shape[0]
110 out = torch.empty((m, n), dtype=x.dtype, device=x.device)
111 block_n = triton.next_power_of_2(n)
112 _gemma_rmsnorm_kernel[(m,)](
113 x2,
114 weight,
115 out,
116 n,
117 x2.stride(0),
118 x2.stride(1),
119 eps,
120 BLOCK_N=block_n,
121 num_warps=_num_warps(block_n),
122 )
123 return out.reshape(orig_shape)
124
125
126 def gemma_fused_add_rmsnorm(
127 x: torch.Tensor,
128 residual: torch.Tensor,
129 weight: torch.Tensor,
130 eps: float,
131 ) -> tuple[torch.Tensor, torch.Tensor]:
132 orig_shape = x.shape
133 n = orig_shape[-1]
134 x2 = x.reshape(-1, n)
135 r2 = residual.reshape(-1, n)
136 m = x2.shape[0]
137 out = torch.empty((m, n), dtype=x.dtype, device=x.device)
138 res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
139 block_n = triton.next_power_of_2(n)
140 _gemma_fused_add_rmsnorm_kernel[(m,)](
141 x2,
142 r2,
143 weight,
144 out,
145 res_out,
146 n,
147 x2.stride(0),
148 x2.stride(1),
149 r2.stride(0),
150 r2.stride(1),
151 eps,
152 BLOCK_N=block_n,
153 num_warps=_num_warps(block_n),
154 )
155 return out.reshape(orig_shape), res_out.reshape(orig_shape)