| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch.""" |
| 4 | |
| 5 | import os |
| 6 | import subprocess |
| 7 | from functools import partial |
| 8 | from importlib.metadata import version |
| 9 | from statistics import median |
| 10 | |
| 11 | import torch |
| 12 | import torch.nn.functional as F |
| 13 | |
| 14 | import vllm |
| 15 | from vllm.models.hy_v4.nvidia.triton_ihc import ( |
| 16 | triton_ihc_post, |
| 17 | triton_ihc_pre, |
| 18 | ) |
| 19 | from vllm.triton_utils import triton |
| 20 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 21 | |
| 22 | |
| 23 | def eager_pre( |
| 24 | x: torch.Tensor, |
| 25 | weight: torch.Tensor, |
| 26 | scale: torch.Tensor, |
| 27 | base: torch.Tensor, |
| 28 | magnitude: float, |
| 29 | hc_eps: float, |
| 30 | norm_eps: float, |
| 31 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 32 | num_tokens, hc_mult, hidden_size = x.shape |
| 33 | x_flat = x.flatten(1).float() |
| 34 | reciprocal_rms = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + norm_eps) |
| 35 | mixes = F.linear(x_flat, weight) * reciprocal_rms |
| 36 | pre = torch.sigmoid(mixes[:, :hc_mult] * scale[0] + base[:hc_mult]) |
| 37 | post = torch.sigmoid(mixes[:, hc_mult:] * scale[1] + base[hc_mult:]) |
| 38 | pre = pre + hc_eps |
| 39 | post = magnitude * post + hc_eps |
| 40 | output = torch.sum(pre.unsqueeze(-1) * x.float(), dim=1) |
| 41 | return output.to(x.dtype).reshape(num_tokens, hidden_size), post |
| 42 | |
| 43 | |
| 44 | def eager_post( |
| 45 | x: torch.Tensor, |
| 46 | residual: torch.Tensor, |
| 47 | post: torch.Tensor, |
| 48 | ) -> torch.Tensor: |
| 49 | return (post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()).to( |
| 50 | x.dtype |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | def _timer(method: str): |
| 55 | if method == "cupti": |
| 56 | from flashinfer.testing import bench_gpu_time_with_cupti |
| 57 | |
| 58 | return partial( |
| 59 | bench_gpu_time_with_cupti, |
| 60 | use_cuda_graph=True, |
| 61 | cold_l2_cache=True, |
| 62 | ) |
| 63 | if method == "cudagraph": |
| 64 | from flashinfer.testing import bench_gpu_time_with_cudagraph |
| 65 | |
| 66 | return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True) |
| 67 | raise ValueError(f"unknown timing method: {method}") |
| 68 | |
| 69 | |
| 70 | @torch.inference_mode() |
| 71 | def run_benchmark( |
| 72 | token_counts: list[int], |
| 73 | hidden_size: int, |
| 74 | dtype: torch.dtype, |
| 75 | method: str, |
| 76 | ) -> None: |
| 77 | hc_mult = 4 |
| 78 | device = torch.device("cuda") |
| 79 | torch.manual_seed(0) |
| 80 | weight = torch.randn( |
| 81 | 2 * hc_mult, |
| 82 | hc_mult * hidden_size, |
| 83 | device=device, |
| 84 | dtype=torch.float32, |
| 85 | ) |
| 86 | scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01 |
| 87 | base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32) |
| 88 | timer = _timer(method) |
| 89 | |
| 90 | properties = torch.cuda.get_device_properties(device) |
| 91 | git_branch = subprocess.run( |
| 92 | ["git", "branch", "--show-current"], |
| 93 | check=False, |
| 94 | capture_output=True, |
| 95 | text=True, |
| 96 | ).stdout.strip() |
| 97 | git_commit = subprocess.run( |
| 98 | ["git", "rev-parse", "--short", "HEAD"], |
| 99 | check=False, |
| 100 | capture_output=True, |
| 101 | text=True, |
| 102 | ).stdout.strip() |
| 103 | print(f"device: {properties.name}") |
| 104 | print(f"branch: {git_branch}; commit: {git_commit}") |
| 105 | print( |
| 106 | f"vllm: {vllm.__version__}; torch: {torch.__version__}; " |
| 107 | f"CUDA: {torch.version.cuda}" |
| 108 | ) |
| 109 | print( |
| 110 | f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; " |
| 111 | f"dtype: {dtype}; method: {method}; cache: cold L2" |
| 112 | ) |
| 113 | print( |
| 114 | "env: VLLM_ENABLE_HPC_OPS=" |
| 115 | f"{os.getenv('VLLM_ENABLE_HPC_OPS', '<unset>')}; " |
| 116 | "VLLM_BATCH_INVARIANT=" |
| 117 | f"{os.getenv('VLLM_BATCH_INVARIANT', '<unset>')}" |
| 118 | ) |
| 119 | print(f"hidden_size: {hidden_size}; hc_mult: {hc_mult}") |
| 120 | print( |
| 121 | f"{'tokens':>8} {'op':>6} {'eager (us)':>12} " |
| 122 | f"{'triton (us)':>12} {'speedup':>9} {'GiB':>8} " |
| 123 | f"{'eager GB/s':>12} {'triton GB/s':>13}" |
| 124 | ) |
| 125 | |
| 126 | for num_tokens in token_counts: |
| 127 | x = torch.randn( |
| 128 | num_tokens, |
| 129 | hc_mult, |
| 130 | hidden_size, |
| 131 | device=device, |
| 132 | dtype=dtype, |
| 133 | ) |
| 134 | block_output = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype) |
| 135 | residual = torch.randn_like(x) |
| 136 | eager_output, post = eager_pre(x, weight, scale, base, 2.0, 1e-6, 1e-5) |
| 137 | triton_output, triton_post = triton_ihc_pre( |
| 138 | x, weight, scale, base, 2.0, 1e-6, 1e-5 |
| 139 | ) |
| 140 | torch.testing.assert_close(triton_output, eager_output, atol=2e-2, rtol=1e-2) |
| 141 | torch.testing.assert_close(triton_post, post, atol=2e-5, rtol=1e-5) |
| 142 | torch.testing.assert_close( |
| 143 | triton_ihc_post(block_output, residual, post), |
| 144 | eager_post(block_output, residual, post), |
| 145 | atol=0, |
| 146 | rtol=0, |
| 147 | ) |
| 148 | |
| 149 | benchmarks = ( |
| 150 | ( |
| 151 | "pre", |
| 152 | partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5), |
| 153 | partial(triton_ihc_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5), |
| 154 | x.nbytes |
| 155 | + weight.nbytes |
| 156 | + scale.nbytes |
| 157 | + base.nbytes |
| 158 | + eager_output.nbytes |
| 159 | + post.nbytes, |
| 160 | ), |
| 161 | ( |
| 162 | "post", |
| 163 | partial(eager_post, block_output, residual, post), |
| 164 | partial(triton_ihc_post, block_output, residual, post), |
| 165 | block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes, |
| 166 | ), |
| 167 | ) |
| 168 | for op_name, eager_fn, triton_fn, logical_bytes in benchmarks: |
| 169 | eager_us = median(timer(eager_fn)) * 1e3 |
| 170 | triton_us = median(timer(triton_fn)) * 1e3 |
| 171 | speedup = eager_us / triton_us |
| 172 | logical_gib = logical_bytes / 2**30 |
| 173 | eager_gbps = logical_bytes / eager_us / 1e3 |
| 174 | triton_gbps = logical_bytes / triton_us / 1e3 |
| 175 | print( |
| 176 | f"{num_tokens:>8} {op_name:>6} {eager_us:>12.1f} " |
| 177 | f"{triton_us:>12.1f} {speedup:>8.2f}x {logical_gib:>8.3f} " |
| 178 | f"{eager_gbps:>12.1f} {triton_gbps:>13.1f}" |
| 179 | ) |
| 180 | |
| 181 | |
| 182 | if __name__ == "__main__": |
| 183 | parser = FlexibleArgumentParser(description=__doc__) |
| 184 | parser.add_argument( |
| 185 | "--token-counts", |
| 186 | type=int, |
| 187 | nargs="+", |
| 188 | default=[1, 2, 4, 8, 16, 64, 256, 1024, 4096, 8192], |
| 189 | ) |
| 190 | parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144) |
| 191 | parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16") |
| 192 | parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti") |
| 193 | args = parser.parse_args() |
| 194 | run_benchmark( |
| 195 | args.token_counts, |
| 196 | args.hidden_size, |
| 197 | getattr(torch, args.dtype), |
| 198 | args.method, |
| 199 | ) |