@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 110 lines 3.41 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5 from torch import Tensor
6 from torch.library import triton_op, wrap_triton
7
8 from vllm import ir
9 from vllm.platforms import current_platform
10 from vllm.triton_utils import HAS_TRITON, tl, triton
11
12 _MAX_INTERMEDIATE_SIZE = 32768
13 _SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
14
15
16 @triton.jit
17 def _gelu_and_mul_sparse_kernel(
18 x_ptr,
19 output_ptr,
20 d: tl.constexpr,
21 std_multiplier,
22 BLOCK_SIZE: tl.constexpr,
23 ):
24 row = tl.program_id(0).to(tl.int64)
25 offsets = tl.arange(0, BLOCK_SIZE)
26 mask = offsets < d
27 row_start = row * 2 * d
28
29 gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32)
30 # TODO: Investigate delaying this load to reduce spills for large reductions
31 # without regressing smaller shapes.
32 up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32)
33
34 input_dtype = x_ptr.dtype.element_ty
35 mean = tl.sum(gate, axis=0) / d
36 centered = tl.where(mask, gate - mean, 0.0)
37 variance = tl.sum(centered * centered, axis=0) / d
38 # Match the input-dtype rounding at each native PyTorch op boundary.
39 mean = mean.to(input_dtype).to(tl.float32)
40 std = tl.sqrt(variance).to(input_dtype).to(tl.float32)
41 scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32)
42 cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32)
43 sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32)
44 sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate)
45
46 inner = 0.7978845608028654 * (
47 sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate
48 )
49 tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0
50 activated = 0.5 * sparse_gate * (1.0 + tanh)
51 activated = activated.to(input_dtype).to(tl.float32)
52 tl.store(output_ptr + row * d + offsets, activated * up, mask=mask)
53
54
55 def _supports_gelu_and_mul_sparse(
56 x: Tensor, std_multiplier: float, approximate: str = "none"
57 ) -> bool:
58 del std_multiplier
59 if (
60 x.device.type != "cuda"
61 or x.dtype not in _SUPPORTED_DTYPES
62 or not x.is_contiguous()
63 or x.ndim == 0
64 or approximate != "tanh"
65 ):
66 return False
67 last_dim = x.shape[-1]
68 return (
69 last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE
70 )
71
72
73 @triton_op(
74 "vllm::gelu_and_mul_sparse_triton",
75 mutates_args=(),
76 )
77 def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor:
78 d = x.shape[-1] // 2
79 output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device)
80 if output.numel() == 0:
81 return output
82
83 n_rows = x.numel() // x.shape[-1]
84
85 def grid(meta):
86 return (n_rows,)
87
88 block_size = triton.next_power_of_2(d)
89 num_warps = min(max(block_size // 256, 1), 8)
90 wrap_triton(_gelu_and_mul_sparse_kernel)[grid](
91 x,
92 output,
93 d=d,
94 std_multiplier=std_multiplier,
95 BLOCK_SIZE=block_size,
96 num_warps=num_warps,
97 )
98 return output
99
100
101 @ir.ops.gelu_and_mul_sparse.register_impl(
102 "triton",
103 supported=HAS_TRITON and current_platform.is_cuda(),
104 supports_args=_supports_gelu_and_mul_sparse,
105 )
106 def gelu_and_mul_sparse_triton(
107 x: Tensor, std_multiplier: float, approximate: str = "none"
108 ) -> Tensor:
109 assert approximate == "tanh"
110 return _gelu_and_mul_sparse_triton_op(x, std_multiplier)