| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compatibility wrapper for HPC API changes. |
| 4 | |
| 5 | Users of vLLM should always import **only** these wrappers. |
| 6 | """ |
| 7 | |
| 8 | import functools |
| 9 | import importlib |
| 10 | import importlib.util |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.logger import init_logger |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | @functools.cache |
| 20 | def has_hpc() -> bool: |
| 21 | """Return `True` if hpc package is available.""" |
| 22 | # Use find_spec to check if the module exists without importing it |
| 23 | # This avoids potential CUDA initialization side effects |
| 24 | if importlib.util.find_spec("hpc") is None: |
| 25 | logger.warning_once( |
| 26 | "HPC attention requires the hpc module to be installed. " |
| 27 | "Please install it from https://github.com/Tencent/hpc-ops" |
| 28 | ) |
| 29 | return False |
| 30 | return True |
| 31 | |
| 32 | |
| 33 | # Remove 'torch._library.custom_ops': |
| 34 | # The output of this custom operator (1) must not also be an input to |
| 35 | # this custom operator and (2) may not alias any inputs to this custom |
| 36 | # operator or other returns. The most common way to trigger this error |
| 37 | # is if we have y = custom_op(x) and y and x are the same Tensor. |
| 38 | # Please instead return a clone of the offending output tensor(s) (e.g. |
| 39 | # return x.clone()) or refactor the custom operator to not return y. |
| 40 | # @torch.library.custom_op( |
| 41 | # "vllm::fuse_moe_impl", |
| 42 | # mutates_args=[], |
| 43 | # device_types="cuda", |
| 44 | # ) |
| 45 | def fuse_moe_impl( |
| 46 | x: torch.Tensor, |
| 47 | gate_up_weight: torch.Tensor, |
| 48 | down_weight: torch.Tensor, |
| 49 | gate_up_scale: torch.Tensor, |
| 50 | down_scale: torch.Tensor, |
| 51 | act_and_mul_scale: torch.Tensor, |
| 52 | topk_ids: torch.Tensor, |
| 53 | topk_scale: torch.Tensor, |
| 54 | rank_ep: int, |
| 55 | num_expert_total: int, |
| 56 | use_bf16_mul: bool = True, |
| 57 | shared_output: torch.Tensor = None, |
| 58 | output: torch.Tensor = None, |
| 59 | ) -> torch.Tensor: |
| 60 | from hpc import fuse_moe as fuse_moe_ |
| 61 | |
| 62 | return fuse_moe_( |
| 63 | x, |
| 64 | gate_up_weight, |
| 65 | down_weight, |
| 66 | gate_up_scale, |
| 67 | down_scale, |
| 68 | act_and_mul_scale, |
| 69 | topk_ids, |
| 70 | topk_scale, |
| 71 | rank_ep, |
| 72 | num_expert_total, |
| 73 | use_bf16_mul, |
| 74 | shared_output, |
| 75 | output=output, |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def hpc_fuse_moe( |
| 80 | x: torch.Tensor, |
| 81 | gate_up_weight: torch.Tensor, |
| 82 | down_weight: torch.Tensor, |
| 83 | gate_up_scale: torch.Tensor, |
| 84 | down_scale: torch.Tensor, |
| 85 | act_and_mul_scale: torch.Tensor, |
| 86 | topk_ids: torch.Tensor, |
| 87 | topk_scale: torch.Tensor, |
| 88 | rank_ep: int, |
| 89 | num_expert_total: int, |
| 90 | use_bf16_mul: bool = True, |
| 91 | shared_output: torch.Tensor = None, |
| 92 | output: torch.Tensor = None, |
| 93 | ) -> torch.Tensor: |
| 94 | return fuse_moe_impl( |
| 95 | x, |
| 96 | gate_up_weight, |
| 97 | down_weight, |
| 98 | gate_up_scale, |
| 99 | down_scale, |
| 100 | act_and_mul_scale, |
| 101 | topk_ids, |
| 102 | topk_scale, |
| 103 | rank_ep, |
| 104 | num_expert_total, |
| 105 | use_bf16_mul, |
| 106 | shared_output, |
| 107 | output=output, |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | # @torch.library.custom_op( |
| 112 | # "vllm::fuse_moe_blockwise_impl", |
| 113 | # mutates_args=[], |
| 114 | # device_types="cuda", |
| 115 | # ) |
| 116 | def fuse_moe_blockwise_impl( |
| 117 | x: torch.Tensor, |
| 118 | x_scale: torch.Tensor, |
| 119 | gate_up_weight: torch.Tensor, |
| 120 | gate_up_weight_scale: torch.Tensor, |
| 121 | down_weight: torch.Tensor, |
| 122 | down_weight_scale: torch.Tensor, |
| 123 | topk_ids: torch.Tensor, |
| 124 | topk_scale: torch.Tensor, |
| 125 | rank_ep: int, |
| 126 | num_expert_total: int, |
| 127 | shared_output: torch.Tensor = None, |
| 128 | output: torch.Tensor = None, |
| 129 | ) -> torch.Tensor: |
| 130 | from hpc import fuse_moe_blockwise as fuse_moe_blockwise_ |
| 131 | |
| 132 | return fuse_moe_blockwise_( |
| 133 | x, |
| 134 | x_scale, |
| 135 | gate_up_weight, |
| 136 | gate_up_weight_scale, |
| 137 | down_weight, |
| 138 | down_weight_scale, |
| 139 | topk_ids, |
| 140 | topk_scale, |
| 141 | rank_ep, |
| 142 | num_expert_total, |
| 143 | shared_output, |
| 144 | output=output, |
| 145 | ) |
| 146 | |
| 147 | |
| 148 | def hpc_fuse_moe_blockwise( |
| 149 | x: torch.Tensor, |
| 150 | x_scale: torch.Tensor, |
| 151 | gate_up_weight: torch.Tensor, |
| 152 | gate_up_weight_scale: torch.Tensor, |
| 153 | down_weight: torch.Tensor, |
| 154 | down_weight_scale: torch.Tensor, |
| 155 | topk_ids: torch.Tensor, |
| 156 | topk_scale: torch.Tensor, |
| 157 | rank_ep: int, |
| 158 | num_expert_total: int, |
| 159 | shared_output: torch.Tensor = None, |
| 160 | output: torch.Tensor = None, |
| 161 | ) -> torch.Tensor: |
| 162 | return fuse_moe_blockwise_impl( |
| 163 | x, |
| 164 | x_scale, |
| 165 | gate_up_weight, |
| 166 | gate_up_weight_scale, |
| 167 | down_weight, |
| 168 | down_weight_scale, |
| 169 | topk_ids, |
| 170 | topk_scale, |
| 171 | rank_ep, |
| 172 | num_expert_total, |
| 173 | shared_output, |
| 174 | output=output, |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | __all__ = [ |
| 179 | "has_hpc", |
| 180 | "hpc_fuse_moe", |
| 181 | "hpc_fuse_moe_blockwise", |
| 182 | ] |