| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import random |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( |
| 11 | persistent_masked_m_silu_mul_quant, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 14 | get_fp8_min_max, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.utils.deep_gemm import ( |
| 18 | DeepGemmQuantScaleFMT, |
| 19 | has_deep_gemm, |
| 20 | transform_sf_into_required_layout, |
| 21 | ) |
| 22 | from vllm.utils.math_utils import cdiv, round_up |
| 23 | from vllm.utils.torch_utils import set_random_seed |
| 24 | |
| 25 | fp8_dtype = current_platform.fp8_dtype() |
| 26 | DEVICE = current_platform.device_type |
| 27 | |
| 28 | pytestmark = pytest.mark.skipif( |
| 29 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 30 | reason="persistent_masked_m_silu_mul_quant requires a CUDA-alike or XPU device.", |
| 31 | ) |
| 32 | |
| 33 | CASES = [ |
| 34 | (1, 1, 128, fp8_dtype), |
| 35 | (1, 4, 128 * 1, fp8_dtype), |
| 36 | (2, 4, 128 * 2, fp8_dtype), |
| 37 | (1, 4, 128 * 3, fp8_dtype), |
| 38 | (8, 16, 128 * 4, fp8_dtype), |
| 39 | (8, 16, 128 * 5, fp8_dtype), |
| 40 | (8, 16, 128 * 6, fp8_dtype), |
| 41 | (8, 16, 128 * 7, fp8_dtype), |
| 42 | (8, 16, 128 * 8, fp8_dtype), |
| 43 | (8, 16, 128 * 9, fp8_dtype), |
| 44 | (8, 64, 7168, fp8_dtype), |
| 45 | (8, 128, 128 * 33, fp8_dtype), |
| 46 | (1, 4, 128 * 10, fp8_dtype), |
| 47 | (8, 128, 7168, fp8_dtype), |
| 48 | (8, 512, 7168, fp8_dtype), |
| 49 | (8, 1024, 7168, fp8_dtype), |
| 50 | (17, 31, 768, fp8_dtype), |
| 51 | (32, 64, 256, fp8_dtype), |
| 52 | (256, 8, 7168, fp8_dtype), |
| 53 | (256, 32, 7168, fp8_dtype), |
| 54 | (256, 64, 7168, fp8_dtype), |
| 55 | # Only add a few fnuz tests to help with long CI times. |
| 56 | (8, 512, 7168, torch.float8_e4m3fnuz), |
| 57 | (8, 1024, 7168, torch.float8_e4m3fnuz), |
| 58 | ] |
| 59 | |
| 60 | |
| 61 | def as_uint8(x) -> torch.Tensor: |
| 62 | return ( |
| 63 | torch.empty(x.shape, dtype=x.dtype, device=x.device).copy_(x).view(torch.uint8) |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def silu(x: torch.Tensor) -> torch.Tensor: |
| 68 | x_f32 = x.to(torch.float32) |
| 69 | act_f32 = x_f32 / (1.0 + torch.exp(-x_f32)) |
| 70 | if current_platform.is_cuda(): |
| 71 | # C++ kernel returns bf16 |
| 72 | return act_f32.to(torch.bfloat16) |
| 73 | # Triton fallback stays in f32 |
| 74 | return act_f32 |
| 75 | |
| 76 | |
| 77 | def do_quant(x: torch.Tensor, group_size: int, ceil_ue8m0: bool): |
| 78 | fp8_min_val, fp8_max_val = get_fp8_min_max() |
| 79 | eps_bf16 = torch.tensor([1e-10], device=x.device, dtype=torch.bfloat16) |
| 80 | one_bf16 = torch.tensor([1.0], device=x.device, dtype=torch.bfloat16) |
| 81 | fp8_max_bf16 = torch.tensor([fp8_max_val], device=x.device, dtype=torch.bfloat16) |
| 82 | fp8_min_bf16 = torch.tensor([fp8_min_val], device=x.device, dtype=torch.bfloat16) |
| 83 | fp8_max_inv = one_bf16 / fp8_max_bf16 |
| 84 | assert fp8_max_inv.dtype == torch.bfloat16 |
| 85 | |
| 86 | assert x.size(-1) % group_size == 0 |
| 87 | num_groups = x.numel() // group_size |
| 88 | x_og_shape = x.shape |
| 89 | |
| 90 | if current_platform.is_cuda(): |
| 91 | # C++ kernel computes entirely in bf16 |
| 92 | x = x.to(torch.bfloat16) |
| 93 | x = x.view((-1, group_size)) |
| 94 | amax = x.abs().amax(dim=1).clamp(min=eps_bf16) |
| 95 | assert amax.dtype == torch.bfloat16 |
| 96 | s = amax * fp8_max_inv |
| 97 | |
| 98 | if ceil_ue8m0: |
| 99 | s = torch.exp2( |
| 100 | torch.ceil(torch.log2(s).to(torch.bfloat16)).to(torch.bfloat16) |
| 101 | ).to(torch.bfloat16) |
| 102 | |
| 103 | inv_s = one_bf16 / s |
| 104 | inv_s = inv_s.view((num_groups, 1)) |
| 105 | xq = torch.clamp( |
| 106 | x * inv_s, min=fp8_min_bf16.item(), max=fp8_max_bf16.item() |
| 107 | ).to(fp8_dtype) |
| 108 | else: |
| 109 | # Triton fallback computes in f32. Use multiply-by-reciprocal |
| 110 | # to match Triton's constexpr evaluation of 1.0/fp8_max. |
| 111 | fp8_min_f, fp8_max_f = get_fp8_min_max() |
| 112 | |
| 113 | x = x.to(torch.float32).view((-1, group_size)) |
| 114 | amax = x.abs().amax(dim=1).clamp(min=1e-10) |
| 115 | s = amax * (1.0 / fp8_max_f) |
| 116 | if ceil_ue8m0: |
| 117 | s = torch.exp2(torch.ceil(torch.log2(s))) |
| 118 | inv_s = (1.0 / s).view((num_groups, 1)) |
| 119 | xq = torch.clamp(x * inv_s, min=fp8_min_f, max=fp8_max_f).to(fp8_dtype) |
| 120 | |
| 121 | xq = xq.view(x_og_shape) |
| 122 | xs = s.view((-1, xq.size(-1) // group_size)) |
| 123 | return xq, xs |
| 124 | |
| 125 | |
| 126 | def silu_mul_quant( |
| 127 | gate: torch.Tensor, up: torch.Tensor, group_size: int, ceil_ue8m0: bool |
| 128 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 129 | assert gate.size(-1) % group_size == 0 |
| 130 | assert up.size(-1) % group_size == 0 |
| 131 | |
| 132 | assert gate.dtype == torch.bfloat16 |
| 133 | assert up.dtype == torch.bfloat16 |
| 134 | |
| 135 | act = silu(gate) |
| 136 | |
| 137 | # act & mul |
| 138 | a_m = act * up |
| 139 | |
| 140 | q, s = do_quant(a_m, group_size, ceil_ue8m0) |
| 141 | return q, s |
| 142 | |
| 143 | |
| 144 | def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: |
| 145 | """Pack float32 scales into a int32 tensor.""" |
| 146 | assert x.dtype == torch.float32 |
| 147 | E, T, G = x.size() |
| 148 | |
| 149 | # Add i32_padding here so we can view it as a i32 tensor later on. |
| 150 | i32_padding = round_up(G, 4) - G |
| 151 | ref_s_i8 = torch.empty((E, T, G + i32_padding), dtype=torch.uint8, device=DEVICE) |
| 152 | for e in range(E): |
| 153 | nt = tokens_per_expert[e].item() |
| 154 | ref_s_i8[e, :nt, :G] = x[e, :nt].view(torch.int32) >> 23 |
| 155 | |
| 156 | ref_s_i32 = ref_s_i8.view(torch.int32) |
| 157 | |
| 158 | return ref_s_i32 |
| 159 | |
| 160 | |
| 161 | def ref_with_scale_fmt( |
| 162 | E: int, |
| 163 | T: int, |
| 164 | H: int, |
| 165 | group_size: int, |
| 166 | tokens_per_expert: torch.Tensor, |
| 167 | gate: torch.Tensor, |
| 168 | up: torch.Tensor, |
| 169 | scale_fmt: DeepGemmQuantScaleFMT, |
| 170 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 171 | """The precision types of the operations triggered by this function |
| 172 | match closely with the kernel implementation so we compare more |
| 173 | accurately. |
| 174 | """ |
| 175 | scale_dtype = ( |
| 176 | torch.int32 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0 else torch.float32 |
| 177 | ) |
| 178 | ceil_ue8m0 = scale_fmt in [ |
| 179 | DeepGemmQuantScaleFMT.UE8M0, |
| 180 | DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, |
| 181 | ] |
| 182 | |
| 183 | ref_q = torch.empty((E, T, H), dtype=fp8_dtype, device=DEVICE) |
| 184 | ref_s_f32 = torch.empty( |
| 185 | (E, T, cdiv(H, group_size)), dtype=torch.float32, device=DEVICE |
| 186 | ) |
| 187 | |
| 188 | for e in range(E): |
| 189 | nt = tokens_per_expert[e].item() |
| 190 | if nt == 0: |
| 191 | continue |
| 192 | ref_q[e, :nt], ref_s_f32[e, :nt] = silu_mul_quant( |
| 193 | gate[e, :nt], up[e, :nt], group_size, ceil_ue8m0=ceil_ue8m0 |
| 194 | ) |
| 195 | |
| 196 | if scale_dtype == torch.float32: |
| 197 | return ref_q, ref_s_f32 |
| 198 | |
| 199 | assert scale_dtype == torch.int32 |
| 200 | return ref_q, pack_scales(ref_s_f32, tokens_per_expert) |
| 201 | |
| 202 | |
| 203 | def token_random(E, T, H2, tokens_per_expert): |
| 204 | """Initialize each token in a random range so we test a range of |
| 205 | scale values. |
| 206 | """ |
| 207 | y = torch.empty((E, T, H2), dtype=torch.bfloat16, device=DEVICE) |
| 208 | for e in range(E): |
| 209 | for t in range(tokens_per_expert[e].item()): |
| 210 | exp = random.choice(range(1, 20)) |
| 211 | y[e, t].uniform_(-(2**exp), 2**exp) |
| 212 | return y |
| 213 | |
| 214 | |
| 215 | @pytest.mark.parametrize("E,T,H,fp8_type", CASES) |
| 216 | @torch.inference_mode() |
| 217 | def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dtype): |
| 218 | group_size = 128 |
| 219 | set_random_seed(42) |
| 220 | |
| 221 | tokens_per_expert = torch.randint( |
| 222 | low=0, |
| 223 | high=T, |
| 224 | size=(E,), |
| 225 | dtype=torch.int32, |
| 226 | device=DEVICE, |
| 227 | ) |
| 228 | |
| 229 | # Input tensor of shape (E, T, 2*H) |
| 230 | y = token_random(E, T, 2 * H, tokens_per_expert) |
| 231 | |
| 232 | gate = y[..., :H].to(torch.bfloat16) |
| 233 | up = y[..., H:].to(torch.bfloat16) |
| 234 | |
| 235 | scale_fmts = [ |
| 236 | DeepGemmQuantScaleFMT.FLOAT32, |
| 237 | DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, |
| 238 | DeepGemmQuantScaleFMT.UE8M0, |
| 239 | ] |
| 240 | |
| 241 | # Run the SiLU V2 kernel |
| 242 | for scale_fmt in scale_fmts: |
| 243 | y_q, y_s = persistent_masked_m_silu_mul_quant( |
| 244 | y, |
| 245 | tokens_per_expert, |
| 246 | group_size=group_size, |
| 247 | quant_scale_fmt=scale_fmt, |
| 248 | ) |
| 249 | |
| 250 | ref_y_q, ref_y_s = ref_with_scale_fmt( |
| 251 | E, T, H, group_size, tokens_per_expert, gate, up, scale_fmt=scale_fmt |
| 252 | ) |
| 253 | |
| 254 | # deepgemm scales transform |
| 255 | dg_scales = None |
| 256 | if ( |
| 257 | has_deep_gemm() |
| 258 | and current_platform.has_device_capability(100) |
| 259 | and scale_fmt == DeepGemmQuantScaleFMT.UE8M0 |
| 260 | ): |
| 261 | _q, _s = ref_with_scale_fmt( |
| 262 | E, |
| 263 | T, |
| 264 | H, |
| 265 | group_size, |
| 266 | tokens_per_expert, |
| 267 | gate, |
| 268 | up, |
| 269 | scale_fmt=DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, |
| 270 | ) |
| 271 | dg_scales = transform_sf_into_required_layout( |
| 272 | sf=_s, |
| 273 | mn=_q.size(1), |
| 274 | k=_q.size(2), |
| 275 | recipe=(1, 128, 128), |
| 276 | num_groups=_q.size(0), |
| 277 | is_sfa=True, |
| 278 | ) |
| 279 | |
| 280 | expected_scale_dtype = ( |
| 281 | torch.int32 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0 else torch.float32 |
| 282 | ) |
| 283 | assert y_s.dtype == expected_scale_dtype |
| 284 | assert ref_y_s.dtype == expected_scale_dtype |
| 285 | |
| 286 | for e in range(E): |
| 287 | nt = tokens_per_expert[e].item() |
| 288 | |
| 289 | if current_platform.is_rocm() or current_platform.is_xpu(): |
| 290 | # ROCm and XPU both run the Triton fallback kernel, whose f32 |
| 291 | # math intrinsics (tl.exp) may differ from PyTorch's torch.exp |
| 292 | # by 1 ULP. At FP8 quantization boundaries this can flip one |
| 293 | # representable value. Allow 1 FP8 quantum of tolerance. |
| 294 | torch.testing.assert_close( |
| 295 | y_q[e, :nt].to(torch.float32), |
| 296 | ref_y_q[e, :nt].to(torch.float32), |
| 297 | atol=32.0, |
| 298 | rtol=0.2, |
| 299 | ) |
| 300 | else: |
| 301 | torch.testing.assert_close( |
| 302 | y_q[e, :nt].to(torch.float32), |
| 303 | ref_y_q[e, :nt].to(torch.float32), |
| 304 | ) |
| 305 | |
| 306 | if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: |
| 307 | G = H // group_size |
| 308 | y_s_sliced = as_uint8(y_s[e]) |
| 309 | ref_s_sliced = as_uint8(ref_y_s[e]) |
| 310 | torch.testing.assert_close(y_s_sliced[:nt, :G], ref_s_sliced[:nt, :G]) |
| 311 | if dg_scales is not None: |
| 312 | dg_sliced = as_uint8(dg_scales[e]) |
| 313 | torch.testing.assert_close(y_s_sliced[:nt, :G], dg_sliced[:nt, :G]) |
| 314 | else: |
| 315 | torch.testing.assert_close( |
| 316 | y_s[e, :nt], |
| 317 | ref_y_s[e, :nt], |
| 318 | ) |