| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.deep_gemm import ( |
| 10 | _import_deep_gemm, |
| 11 | is_deep_gemm_supported, |
| 12 | mega_mhc, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | @functools.cache |
| 17 | def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: |
| 18 | """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel.""" |
| 19 | if not ( |
| 20 | is_deep_gemm_supported() |
| 21 | and current_platform.is_device_capability_family(100) |
| 22 | and hidden_size > 0 |
| 23 | and hidden_size % 1024 == 0 |
| 24 | and hc_mult == 4 |
| 25 | ): |
| 26 | return False |
| 27 | deep_gemm = _import_deep_gemm() |
| 28 | return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) |
| 29 | |
| 30 | |
| 31 | def can_use_mega_mhc( |
| 32 | x: torch.Tensor, |
| 33 | residual: torch.Tensor, |
| 34 | pre_mix: torch.Tensor | None, |
| 35 | norm_weight: torch.Tensor | None, |
| 36 | capture_aux: bool, |
| 37 | ) -> bool: |
| 38 | return ( |
| 39 | pre_mix is not None |
| 40 | and norm_weight is not None |
| 41 | and not capture_aux |
| 42 | and x.shape[0] <= 1 << 20 |
| 43 | and is_mega_mhc_supported(x.shape[1], residual.shape[1]) |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | def mhc_shifted_post_pre_deep_gemm( |
| 48 | x: torch.Tensor, |
| 49 | residual: torch.Tensor, |
| 50 | shifted_prev_mix: torch.Tensor, |
| 51 | post_mix: torch.Tensor, |
| 52 | comb_res_mix: torch.Tensor, |
| 53 | fn: torch.Tensor, |
| 54 | mix_scales: torch.Tensor, |
| 55 | mix_bases: torch.Tensor, |
| 56 | hc_norm_eps: float, |
| 57 | hc_pre_eps: float, |
| 58 | hc_post_scale: float, |
| 59 | sinkhorn_eps: float, |
| 60 | num_sinkhorn_iters: int, |
| 61 | rmsnorm_weight: torch.Tensor, |
| 62 | rmsnorm_eps: float, |
| 63 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 64 | """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC.""" |
| 65 | num_tokens, hidden_size = x.shape |
| 66 | hc_mult = residual.shape[1] |
| 67 | new_residual = torch.empty_like(residual) |
| 68 | new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32) |
| 69 | new_post_mix = torch.empty_like(post_mix) |
| 70 | new_comb_res_mix = torch.empty_like(comb_res_mix) |
| 71 | y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16) |
| 72 | mega_mhc( |
| 73 | x=x, |
| 74 | residual=residual, |
| 75 | shifted_prev_mix=shifted_prev_mix.unsqueeze(-1), |
| 76 | post_mix=post_mix, |
| 77 | comb_res_mix=comb_res_mix, |
| 78 | fn=fn, |
| 79 | mix_scales=mix_scales, |
| 80 | mix_bases=mix_bases, |
| 81 | hc_mult=hc_mult, |
| 82 | hc_norm_eps=hc_norm_eps, |
| 83 | hc_pre_eps=hc_pre_eps, |
| 84 | hc_post_scale=hc_post_scale, |
| 85 | sinkhorn_eps=sinkhorn_eps, |
| 86 | num_sinkhorn_iters=num_sinkhorn_iters, |
| 87 | rmsnorm_weight=rmsnorm_weight, |
| 88 | rmsnorm_eps=rmsnorm_eps, |
| 89 | rmsnorm_scale=1.0, |
| 90 | new_residual=new_residual, |
| 91 | new_prev_mix=new_prev_mix, |
| 92 | new_post_mix=new_post_mix, |
| 93 | new_comb_res_mix=new_comb_res_mix, |
| 94 | y_bf16=y_bf16, |
| 95 | ) |
| 96 | return ( |
| 97 | new_residual, |
| 98 | new_post_mix, |
| 99 | new_comb_res_mix, |
| 100 | y_bf16, |
| 101 | new_prev_mix.squeeze(-1), |
| 102 | ) |