diff --git a/benchmarks/kernels/benchmark_hy_v4_ihc.py b/benchmarks/kernels/benchmark_hy_v4_ihc.py new file mode 100644 index 0000000000..96220a66cb --- /dev/null +++ b/benchmarks/kernels/benchmark_hy_v4_ihc.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch.""" + +import os +import subprocess +from functools import partial +from importlib.metadata import version +from statistics import median + +import torch +import torch.nn.functional as F + +import vllm +from vllm.models.hy_v4.nvidia.triton_ihc import ( + triton_ihc_post, + triton_ihc_pre, +) +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def eager_pre( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + magnitude: float, + hc_eps: float, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens, hc_mult, hidden_size = x.shape + x_flat = x.flatten(1).float() + reciprocal_rms = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + norm_eps) + mixes = F.linear(x_flat, weight) * reciprocal_rms + pre = torch.sigmoid(mixes[:, :hc_mult] * scale[0] + base[:hc_mult]) + post = torch.sigmoid(mixes[:, hc_mult:] * scale[1] + base[hc_mult:]) + pre = pre + hc_eps + post = magnitude * post + hc_eps + output = torch.sum(pre.unsqueeze(-1) * x.float(), dim=1) + return output.to(x.dtype).reshape(num_tokens, hidden_size), post + + +def eager_post( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, +) -> torch.Tensor: + return (post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()).to( + x.dtype + ) + + +def _timer(method: str): + if method == "cupti": + from flashinfer.testing import bench_gpu_time_with_cupti + + return partial( + bench_gpu_time_with_cupti, + use_cuda_graph=True, + cold_l2_cache=True, + ) + if method == "cudagraph": + from flashinfer.testing import bench_gpu_time_with_cudagraph + + return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True) + raise ValueError(f"unknown timing method: {method}") + + +@torch.inference_mode() +def run_benchmark( + token_counts: list[int], + hidden_size: int, + dtype: torch.dtype, + method: str, +) -> None: + hc_mult = 4 + device = torch.device("cuda") + torch.manual_seed(0) + weight = torch.randn( + 2 * hc_mult, + hc_mult * hidden_size, + device=device, + dtype=torch.float32, + ) + scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01 + base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32) + timer = _timer(method) + + properties = torch.cuda.get_device_properties(device) + git_branch = subprocess.run( + ["git", "branch", "--show-current"], + check=False, + capture_output=True, + text=True, + ).stdout.strip() + git_commit = subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + check=False, + capture_output=True, + text=True, + ).stdout.strip() + print(f"device: {properties.name}") + print(f"branch: {git_branch}; commit: {git_commit}") + print( + f"vllm: {vllm.__version__}; torch: {torch.__version__}; " + f"CUDA: {torch.version.cuda}" + ) + print( + f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; " + f"dtype: {dtype}; method: {method}; cache: cold L2" + ) + print( + "env: VLLM_ENABLE_HPC_OPS=" + f"{os.getenv('VLLM_ENABLE_HPC_OPS', '')}; " + "VLLM_BATCH_INVARIANT=" + f"{os.getenv('VLLM_BATCH_INVARIANT', '')}" + ) + print(f"hidden_size: {hidden_size}; hc_mult: {hc_mult}") + print( + f"{'tokens':>8} {'op':>6} {'eager (us)':>12} " + f"{'triton (us)':>12} {'speedup':>9} {'GiB':>8} " + f"{'eager GB/s':>12} {'triton GB/s':>13}" + ) + + for num_tokens in token_counts: + x = torch.randn( + num_tokens, + hc_mult, + hidden_size, + device=device, + dtype=dtype, + ) + block_output = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype) + residual = torch.randn_like(x) + eager_output, post = eager_pre(x, weight, scale, base, 2.0, 1e-6, 1e-5) + triton_output, triton_post = triton_ihc_pre( + x, weight, scale, base, 2.0, 1e-6, 1e-5 + ) + torch.testing.assert_close(triton_output, eager_output, atol=2e-2, rtol=1e-2) + torch.testing.assert_close(triton_post, post, atol=2e-5, rtol=1e-5) + torch.testing.assert_close( + triton_ihc_post(block_output, residual, post), + eager_post(block_output, residual, post), + atol=0, + rtol=0, + ) + + benchmarks = ( + ( + "pre", + partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5), + partial(triton_ihc_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5), + x.nbytes + + weight.nbytes + + scale.nbytes + + base.nbytes + + eager_output.nbytes + + post.nbytes, + ), + ( + "post", + partial(eager_post, block_output, residual, post), + partial(triton_ihc_post, block_output, residual, post), + block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes, + ), + ) + for op_name, eager_fn, triton_fn, logical_bytes in benchmarks: + eager_us = median(timer(eager_fn)) * 1e3 + triton_us = median(timer(triton_fn)) * 1e3 + speedup = eager_us / triton_us + logical_gib = logical_bytes / 2**30 + eager_gbps = logical_bytes / eager_us / 1e3 + triton_gbps = logical_bytes / triton_us / 1e3 + print( + f"{num_tokens:>8} {op_name:>6} {eager_us:>12.1f} " + f"{triton_us:>12.1f} {speedup:>8.2f}x {logical_gib:>8.3f} " + f"{eager_gbps:>12.1f} {triton_gbps:>13.1f}" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser(description=__doc__) + parser.add_argument( + "--token-counts", + type=int, + nargs="+", + default=[1, 2, 4, 8, 16, 64, 256, 1024, 4096, 8192], + ) + parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144) + parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16") + parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti") + args = parser.parse_args() + run_benchmark( + args.token_counts, + args.hidden_size, + getattr(torch, args.dtype), + args.method, + ) diff --git a/vllm/models/hy_v4/nvidia/hc.py b/vllm/models/hy_v4/nvidia/hc.py index a89deed0f1..130d3842ec 100644 --- a/vllm/models/hy_v4/nvidia/hc.py +++ b/vllm/models/hy_v4/nvidia/hc.py @@ -9,9 +9,8 @@ scatters the result back over the channels (``HYV4HCPostLayer``). The final ``HYV4HCHeadLayer`` merges the channels before the model's output norm. NOTE: Each of the three steps has an optional single-kernel HPC replacement -(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). They are only constructed -when the hpc package is installed, ``VLLM_ENABLE_HPC_OPS=1`` and the shape / -device constraints hold; otherwise the eager path below runs unchanged. +(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). Pre and post fall back to +in-tree Triton kernels on CUDA when HPC is unavailable, then to the eager path. TODO: port the cross-layer post+pre fusion (``HpcIHCPostPre``) as well; it requires restructuring the decoder-layer forward scheduling. """ @@ -22,6 +21,11 @@ from transformers import PretrainedConfig from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.models.hy_v4.nvidia.triton_ihc import ( + triton_ihc_post, + triton_ihc_pre, + triton_ihc_supported, +) class HYV4HCPreLayer(nn.Module): @@ -106,6 +110,16 @@ class HYV4HCPreLayer(nn.Module): """ if self.hpc_op is not None: return self.hpc_op(x) + if triton_ihc_supported(x): + return triton_ihc_pre( + x, + self.hc_fn.weight, + self.hc_scale, + self.hc_base, + self.magnitude, + self.hc_eps, + self.layernorm_epsilon, + ) shape = x.size() # [num_tokens, hc, d] hc = self.hc_mult @@ -175,6 +189,8 @@ class HYV4HCPostLayer(nn.Module): """ if self.hpc_op is not None: return self.hpc_op(x, residual, post) + if triton_ihc_supported(x): + return triton_ihc_post(x, residual, post) dtype = x.dtype x = x.float() diff --git a/vllm/models/hy_v4/nvidia/triton_ihc.py b/vllm/models/hy_v4/nvidia/triton_ihc.py new file mode 100644 index 0000000000..54d4fd7dca --- /dev/null +++ b/vllm/models/hy_v4/nvidia/triton_ihc.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton iHC pre/post kernels for HY V4. + +Adapted from the SGLang HY V4 implementation: +https://github.com/sgl-project/sglang/pull/36805 +""" + +import torch + +from vllm import envs +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_BLOCK_K = 1024 +_BLOCK_D = 1024 + + +def triton_ihc_supported(x: torch.Tensor) -> bool: + """Return whether the in-tree Triton path can run for this input.""" + return ( + HAS_TRITON + and current_platform.is_cuda() + and x.is_cuda + and x.dtype in (torch.float16, torch.bfloat16) + and not envs.VLLM_BATCH_INVARIANT + ) + + +@triton.jit +def _ihc_pre_stage1( + x_ptr, + weight_ptr, + partial_ptr, + K_TOTAL: tl.constexpr, + HC_MULT: tl.constexpr, + HC_POW2: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_K: tl.constexpr, + PARTIAL_STRIDE: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + split_idx = tl.program_id(1) + hc_idx = tl.arange(0, HC_POW2) + hc_mask = hc_idx < HC_MULT + + k_offsets = split_idx * BLOCK_K + tl.arange(0, BLOCK_K) + k_mask = k_offsets < K_TOTAL + x = tl.load( + x_ptr + token_idx * K_TOTAL + k_offsets, + mask=k_mask, + other=0.0, + ).to(tl.float32) + sum_squares = tl.sum(x * x, axis=0) + + weight_offsets = hc_idx[:, None] * K_TOTAL + k_offsets[None, :] + weight_mask = hc_mask[:, None] & k_mask[None, :] + pre_mix = tl.sum( + tl.load(weight_ptr + weight_offsets, mask=weight_mask, other=0.0) * x[None, :], + axis=1, + ) + post_mix = tl.sum( + tl.load( + weight_ptr + HC_MULT * K_TOTAL + weight_offsets, + mask=weight_mask, + other=0.0, + ) + * x[None, :], + axis=1, + ) + + partial = partial_ptr + (token_idx * NUM_SPLITS + split_idx) * PARTIAL_STRIDE + tl.store(partial, sum_squares) + tl.store(partial + 1 + hc_idx, pre_mix, mask=hc_mask) + tl.store(partial + 1 + HC_POW2 + hc_idx, post_mix, mask=hc_mask) + + +@triton.jit +def _ihc_pre_stage2( + x_ptr, + partial_ptr, + scale_ptr, + base_ptr, + output_ptr, + post_ptr, + HIDDEN_SIZE: tl.constexpr, + K_TOTAL: tl.constexpr, + HC_MULT: tl.constexpr, + HC_POW2: tl.constexpr, + NUM_SPLITS: tl.constexpr, + PARTIAL_STRIDE: tl.constexpr, + BLOCK_D: tl.constexpr, + MAGNITUDE: tl.constexpr, + NORM_EPS: tl.constexpr, + HC_EPS: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + hidden_block_idx = tl.program_id(1) + hc_idx = tl.arange(0, HC_POW2) + hc_mask = hc_idx < HC_MULT + + partial_row = partial_ptr + token_idx * NUM_SPLITS * PARTIAL_STRIDE + sum_squares = tl.zeros((), dtype=tl.float32) + pre_mix = tl.zeros((HC_POW2,), dtype=tl.float32) + post_mix = tl.zeros((HC_POW2,), dtype=tl.float32) + for split_idx in tl.static_range(NUM_SPLITS): + partial = partial_row + split_idx * PARTIAL_STRIDE + sum_squares += tl.load(partial) + pre_mix += tl.load(partial + 1 + hc_idx, mask=hc_mask, other=0.0) + post_mix += tl.load( + partial + 1 + HC_POW2 + hc_idx, + mask=hc_mask, + other=0.0, + ) + + reciprocal_rms = tl.rsqrt(sum_squares / K_TOTAL + NORM_EPS) + pre_scale = tl.load(scale_ptr) + post_scale = tl.load(scale_ptr + 1) + pre_base = tl.load(base_ptr + hc_idx, mask=hc_mask, other=0.0) + post_base = tl.load(base_ptr + HC_MULT + hc_idx, mask=hc_mask, other=0.0) + pre = tl.sigmoid(pre_mix * reciprocal_rms * pre_scale + pre_base) + HC_EPS + + if hidden_block_idx == 0: + post = ( + MAGNITUDE * tl.sigmoid(post_mix * reciprocal_rms * post_scale + post_base) + + HC_EPS + ) + tl.store( + post_ptr + token_idx * HC_MULT + hc_idx, + post, + mask=hc_mask, + ) + + hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D) + hidden_mask = hidden_offsets < HIDDEN_SIZE + x_row = x_ptr + token_idx * K_TOTAL + output = tl.zeros((BLOCK_D,), dtype=tl.float32) + for channel_idx in tl.static_range(HC_MULT): + channel = tl.load( + x_row + channel_idx * HIDDEN_SIZE + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + channel_gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0) + output += channel_gate * channel + + tl.store( + output_ptr + token_idx * HIDDEN_SIZE + hidden_offsets, + output.to(output_ptr.dtype.element_ty), + mask=hidden_mask, + ) + + +@triton.jit +def _ihc_post_kernel( + x_ptr, + residual_ptr, + post_ptr, + output_ptr, + HIDDEN_SIZE: tl.constexpr, + HC_MULT: tl.constexpr, + HC_POW2: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + hidden_block_idx = tl.program_id(1) + hc_idx = tl.arange(0, HC_POW2) + hc_mask = hc_idx < HC_MULT + post = tl.load( + post_ptr + token_idx * HC_MULT + hc_idx, + mask=hc_mask, + other=0.0, + ) + + hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D) + hidden_mask = hidden_offsets < HIDDEN_SIZE + x = tl.load( + x_ptr + token_idx * HIDDEN_SIZE + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_row = residual_ptr + token_idx * HC_MULT * HIDDEN_SIZE + output_row = output_ptr + token_idx * HC_MULT * HIDDEN_SIZE + for channel_idx in tl.static_range(HC_MULT): + residual = tl.load( + residual_row + channel_idx * HIDDEN_SIZE + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + channel_gate = tl.sum(tl.where(hc_idx == channel_idx, post, 0.0), axis=0) + output = channel_gate * x + residual + tl.store( + output_row + channel_idx * HIDDEN_SIZE + hidden_offsets, + output.to(output_ptr.dtype.element_ty), + mask=hidden_mask, + ) + + +def triton_ihc_pre( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + magnitude: float, + hc_eps: float, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reduce iHC channels and produce the post gates.""" + assert x.dim() == 3 + assert x.is_cuda and weight.is_cuda and scale.is_cuda and base.is_cuda + assert weight.dtype == torch.float32 + assert scale.dtype == torch.float32 and base.dtype == torch.float32 + + x = x.contiguous() + assert weight.is_contiguous() + scale = scale.contiguous() + base = base.contiguous() + num_tokens, hc_mult, hidden_size = x.shape + k_total = hc_mult * hidden_size + assert weight.shape == (2 * hc_mult, k_total) + assert scale.shape == (2,) + assert base.shape == (2 * hc_mult,) + + if num_tokens == 0: + return ( + torch.empty((0, hidden_size), dtype=x.dtype, device=x.device), + torch.empty((0, hc_mult), dtype=torch.float32, device=x.device), + ) + + output = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device) + post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device) + hc_pow2 = triton.next_power_of_2(hc_mult) + num_splits = triton.cdiv(k_total, _BLOCK_K) + partial_stride = 1 + 2 * hc_pow2 + partial = torch.empty( + (num_tokens, num_splits, partial_stride), + dtype=torch.float32, + device=x.device, + ) + + _ihc_pre_stage1[(num_tokens, num_splits)]( + x, + weight, + partial, + K_TOTAL=k_total, + HC_MULT=hc_mult, + HC_POW2=hc_pow2, + NUM_SPLITS=num_splits, + BLOCK_K=_BLOCK_K, + PARTIAL_STRIDE=partial_stride, + num_warps=8, + enable_fp_fusion=False, + ) + _ihc_pre_stage2[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))]( + x, + partial, + scale, + base, + output, + post, + HIDDEN_SIZE=hidden_size, + K_TOTAL=k_total, + HC_MULT=hc_mult, + HC_POW2=hc_pow2, + NUM_SPLITS=num_splits, + PARTIAL_STRIDE=partial_stride, + BLOCK_D=_BLOCK_D, + MAGNITUDE=magnitude, + NORM_EPS=norm_eps, + HC_EPS=hc_eps, + num_warps=4, + enable_fp_fusion=False, + ) + return output, post + + +def triton_ihc_post( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, +) -> torch.Tensor: + """Scatter a sub-block output back over the iHC residual channels.""" + assert x.dim() == 2 + assert x.is_cuda and residual.is_cuda and post.is_cuda + assert post.dtype == torch.float32 + + x = x.contiguous() + residual = residual.contiguous() + post = post.contiguous() + num_tokens, hidden_size = x.shape + hc_mult = post.shape[-1] + assert residual.shape == (num_tokens, hc_mult, hidden_size) + assert post.shape == (num_tokens, hc_mult) + + if num_tokens == 0: + return torch.empty((0, hc_mult, hidden_size), dtype=x.dtype, device=x.device) + + output = torch.empty( + (num_tokens, hc_mult, hidden_size), dtype=x.dtype, device=x.device + ) + _ihc_post_kernel[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))]( + x, + residual, + post, + output, + HIDDEN_SIZE=hidden_size, + HC_MULT=hc_mult, + HC_POW2=triton.next_power_of_2(hc_mult), + BLOCK_D=_BLOCK_D, + num_warps=4, + enable_fp_fusion=False, + ) + return output