| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | |
| 9 | if not current_platform.is_device_capability_family(100): |
| 10 | pytest.skip( |
| 11 | reason="Nvfp4 Requires compute capability of 10 or above.", |
| 12 | allow_module_level=True, |
| 13 | ) |
| 14 | |
| 15 | import torch |
| 16 | from flashinfer import fp4_quantize |
| 17 | from torch.nn import functional as F |
| 18 | |
| 19 | from vllm.model_executor.layers.activation import SiluAndMul |
| 20 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501 |
| 21 | flashinfer_cutedsl_moe_masked, |
| 22 | ) |
| 23 | from vllm.utils.flashinfer import ( |
| 24 | flashinfer_cutedsl_grouped_gemm_nt_masked as cutedsl_gmm_masked, |
| 25 | ) |
| 26 | from vllm.utils.flashinfer import ( |
| 27 | scaled_fp4_grouped_quantize, |
| 28 | ) |
| 29 | |
| 30 | kE2M1ToFloat = torch.tensor( |
| 31 | [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 |
| 32 | ) |
| 33 | |
| 34 | FLOAT8_E4M3_MAX = 448.0 |
| 35 | FLOAT4_E2M1_MAX = 6.0 |
| 36 | |
| 37 | |
| 38 | def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size): |
| 39 | m_tiles = (m + 128 - 1) // 128 |
| 40 | f = block_size * 4 |
| 41 | k_tiles = (k + f - 1) // f |
| 42 | tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) |
| 43 | tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) |
| 44 | out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) |
| 45 | return out[0:m, 0:k] |
| 46 | |
| 47 | |
| 48 | def dequantize_nvfp4_to_dtype( |
| 49 | tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16 |
| 50 | ): |
| 51 | """Dequantize the fp4 tensor back to high precision.""" |
| 52 | # Two fp4 values are packed into one uint8. |
| 53 | assert tensor_fp4.dtype == torch.uint8 |
| 54 | m, packed_k = tensor_fp4.shape |
| 55 | k = packed_k * 2 |
| 56 | tensor_f32 = break_fp4_bytes(tensor_fp4, dtype) |
| 57 | tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) |
| 58 | tensor_sf = tensor_sf.view(torch.float8_e4m3fn) |
| 59 | tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) |
| 60 | tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale |
| 61 | |
| 62 | # scale the tensor |
| 63 | out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) |
| 64 | return out.to(dtype=dtype) |
| 65 | |
| 66 | |
| 67 | def break_fp4_bytes(a, dtype): |
| 68 | assert a.dtype == torch.uint8 |
| 69 | m, n = a.shape |
| 70 | |
| 71 | # Vectorized nibble processing |
| 72 | a_flat = a.flatten() |
| 73 | high = (a_flat & 0xF0) >> 4 # Upper nibbles |
| 74 | low = a_flat & 0x0F # Lower nibbles |
| 75 | |
| 76 | # Combine nibbles for batch processing |
| 77 | combined = torch.stack((low, high), dim=1).flatten() |
| 78 | |
| 79 | # Vectorized sign and magnitude extraction |
| 80 | signs = (combined & 0x08).to(torch.bool) # Sign bits |
| 81 | abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices |
| 82 | |
| 83 | # Device-aware lookup and sign application |
| 84 | kE2M1 = kE2M1ToFloat.to(device=a.device) |
| 85 | values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) |
| 86 | |
| 87 | # Reshape to final form |
| 88 | return values.reshape(m, n * 2).to(dtype=dtype) |
| 89 | |
| 90 | |
| 91 | def generate_balanced_routing( |
| 92 | hidden_states: torch.Tensor, num_experts: int, top_k: int |
| 93 | ): |
| 94 | """Generate routing weights and topk indices such that every expert is active. |
| 95 | Returns routing_weights, topk_idx |
| 96 | """ |
| 97 | num_tokens, hidden_dim = hidden_states.shape |
| 98 | # num_tokens = batch_size * seq_len |
| 99 | |
| 100 | # First, assign at least one token per expert |
| 101 | tokens_per_expert = torch.arange(num_tokens) % num_experts |
| 102 | tokens_per_expert = tokens_per_expert[torch.randperm(num_tokens)] # shuffle |
| 103 | |
| 104 | # Each token has top_k experts — start with one guaranteed expert |
| 105 | topk_idx = torch.full((num_tokens, top_k), -1, dtype=torch.long) |
| 106 | topk_idx[:, 0] = tokens_per_expert |
| 107 | |
| 108 | # For remaining top_k - 1 experts, pick randomly (allowing repeats) |
| 109 | if top_k > 1: |
| 110 | random_choices = torch.randint(0, num_experts, (num_tokens, top_k - 1)) |
| 111 | topk_idx[:, 1:] = random_choices |
| 112 | |
| 113 | # Normalize routing weights so each token's weights sum to 1 |
| 114 | routing_weights = torch.rand(num_tokens, top_k) |
| 115 | routing_weights /= routing_weights.sum(dim=-1, keepdim=True) |
| 116 | |
| 117 | # Reshape back if needed |
| 118 | routing_weights = routing_weights.view(num_tokens, top_k) |
| 119 | topk_idx = topk_idx.view(num_tokens, top_k) |
| 120 | |
| 121 | return routing_weights, topk_idx |
| 122 | |
| 123 | |
| 124 | def prepare_inputs( |
| 125 | hidden_states: torch.Tensor, |
| 126 | router_logits: torch.Tensor, |
| 127 | num_experts: int, |
| 128 | topk: int, |
| 129 | ): |
| 130 | routing_weights, topk_idx = generate_balanced_routing( |
| 131 | router_logits, num_experts, topk |
| 132 | ) |
| 133 | |
| 134 | masked_m = [] |
| 135 | for i in range(num_experts): |
| 136 | mask = topk_idx.view(-1) == i |
| 137 | masked_m.append(mask.sum()) |
| 138 | |
| 139 | masked_m = torch.tensor(masked_m, dtype=torch.int32) |
| 140 | # Initialize the hidden_states_3d with ones instead of empty to avoid nan |
| 141 | # issue. |
| 142 | hidden_states_3d = torch.ones( |
| 143 | (num_experts, max(masked_m), hidden_states.shape[1]), |
| 144 | dtype=hidden_states.dtype, |
| 145 | device=hidden_states.device, |
| 146 | ) |
| 147 | for i in range(num_experts): |
| 148 | hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i] |
| 149 | |
| 150 | return hidden_states_3d, masked_m, topk_idx, routing_weights |
| 151 | |
| 152 | |
| 153 | MNK_FACTORS = [ |
| 154 | (2, 1024, 1024), |
| 155 | (2, 1024, 1536), |
| 156 | (2, 3072, 1024), |
| 157 | (2, 3072, 1536), |
| 158 | (64, 1024, 1024), |
| 159 | (64, 1024, 1536), |
| 160 | (64, 3072, 1024), |
| 161 | (64, 2048, 1024), |
| 162 | (224, 1024, 1024), |
| 163 | (224, 1024, 1536), |
| 164 | ] |
| 165 | |
| 166 | |
| 167 | # Reference implementation of torch_moe |
| 168 | def torch_moe(a, w1, w2, score, topk, expert_map): |
| 169 | B, D = a.shape |
| 170 | a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) |
| 171 | out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) |
| 172 | score = torch.softmax(score, dim=-1, dtype=torch.float32) |
| 173 | topk_weight, topk_ids = torch.topk(score, topk) |
| 174 | topk_weight = topk_weight.view(-1) |
| 175 | topk_ids = topk_ids.view(-1) |
| 176 | if expert_map is not None: |
| 177 | topk_ids = expert_map[topk_ids] |
| 178 | for i in range(w1.shape[0]): |
| 179 | mask = topk_ids == i |
| 180 | if mask.sum(): |
| 181 | out[mask] = SiluAndMul()(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose( |
| 182 | 0, 1 |
| 183 | ) |
| 184 | return ( |
| 185 | out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype) |
| 186 | ).sum(dim=1) |
| 187 | |
| 188 | |
| 189 | def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids): |
| 190 | B, D = a.shape |
| 191 | a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) |
| 192 | out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) |
| 193 | |
| 194 | topk_weight = topk_weight.view(-1) |
| 195 | topk_ids = topk_ids.view(-1) |
| 196 | |
| 197 | for i in range(w1.shape[0]): |
| 198 | mask = topk_ids == i |
| 199 | if mask.sum(): |
| 200 | m = w1[i].shape[0] |
| 201 | assert m % 2 == 0 |
| 202 | # Note: w1 and w3 are swapped! |
| 203 | w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :] |
| 204 | inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t()) |
| 205 | inter_gs = torch.tensor(1.0).cuda() |
| 206 | inter_q, inter_blockscale = fp4_quantize(inter, inter_gs) |
| 207 | inter = dequantize_nvfp4_to_dtype( |
| 208 | inter_q, |
| 209 | inter_blockscale, |
| 210 | inter_gs, |
| 211 | dtype=inter.dtype, |
| 212 | device=inter.device, |
| 213 | block_size=16, |
| 214 | ).cuda() |
| 215 | out[mask] = inter @ w2[i].transpose(0, 1) |
| 216 | return ( |
| 217 | out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype) |
| 218 | ).sum(dim=1) |
| 219 | |
| 220 | |
| 221 | def grouped_gemm_ref( |
| 222 | hidden_states_expanded: torch.Tensor, |
| 223 | hidden_states_3d: torch.Tensor, |
| 224 | weights: torch.Tensor, |
| 225 | topk_idx: torch.Tensor, |
| 226 | masked_m: torch.Tensor, |
| 227 | B: int, |
| 228 | topk: int, |
| 229 | num_experts: int, |
| 230 | *, |
| 231 | block_size: int = 16, |
| 232 | ) -> torch.Tensor: |
| 233 | """Computes the reference grouped GEMM (fp4 quantized per-expert loop), |
| 234 | computes flashinfer grouped GEMM (for scale consistency), |
| 235 | and returns ONLY the repacked reference output: out_ref. |
| 236 | |
| 237 | Returns: |
| 238 | out_ref: Tensor [num_experts, max_m, n_out] |
| 239 | |
| 240 | """ |
| 241 | device_hs = hidden_states_expanded.device |
| 242 | device_w = weights.device |
| 243 | out_dtype = weights.dtype |
| 244 | n_out = weights.shape[1] |
| 245 | |
| 246 | # Flattened reference output (B*topk, n_out) |
| 247 | out = torch.zeros((B * topk, n_out), dtype=out_dtype, device=device_w) |
| 248 | |
| 249 | # Per-expert reference compute loop |
| 250 | for i in range(num_experts): |
| 251 | mask = topk_idx.view(-1) == i |
| 252 | if mask.any(): |
| 253 | lhs = hidden_states_expanded[mask] |
| 254 | rhs = weights[i] |
| 255 | |
| 256 | a_amax = lhs.abs().max().to(torch.float32).to(device_hs) |
| 257 | b_amax = rhs.abs().max().to(torch.float32).to(device_w) |
| 258 | |
| 259 | a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax |
| 260 | b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax |
| 261 | |
| 262 | lhsq, lhsq_sf = fp4_quantize(lhs, a_gs) |
| 263 | rhsq, rhsq_sf = fp4_quantize(rhs, b_gs) |
| 264 | |
| 265 | lhs_in_dtype = dequantize_nvfp4_to_dtype( |
| 266 | lhsq, |
| 267 | lhsq_sf, |
| 268 | a_gs, |
| 269 | dtype=lhs.dtype, |
| 270 | device=device_hs, |
| 271 | block_size=block_size, |
| 272 | ) |
| 273 | rhs_in_dtype = dequantize_nvfp4_to_dtype( |
| 274 | rhsq, |
| 275 | rhsq_sf, |
| 276 | b_gs, |
| 277 | dtype=rhs.dtype, |
| 278 | device=device_w, |
| 279 | block_size=block_size, |
| 280 | ) |
| 281 | |
| 282 | out[mask] = lhs_in_dtype @ rhs_in_dtype.t() |
| 283 | |
| 284 | # Determine per-expert max_m |
| 285 | max_m_val = int(masked_m.max().item()) |
| 286 | |
| 287 | # Repack into [num_experts, max_m, n_out] |
| 288 | out_ref = torch.zeros( |
| 289 | (num_experts, max_m_val, n_out), |
| 290 | dtype=out.dtype, |
| 291 | device=out.device, |
| 292 | ) |
| 293 | expert_slot = [0] * num_experts |
| 294 | |
| 295 | for i, expert_id in enumerate(topk_idx.view(-1).tolist()): |
| 296 | slot = expert_slot[expert_id] |
| 297 | if slot < max_m_val: |
| 298 | out_ref[expert_id, slot, :] = out[i] |
| 299 | expert_slot[expert_id] += 1 |
| 300 | else: |
| 301 | raise IndexError( |
| 302 | f"Expert {expert_id} exceeded max slots ({max_m_val}). " |
| 303 | "Increase max_m or check masked_m." |
| 304 | ) |
| 305 | |
| 306 | return out_ref |
| 307 | |
| 308 | |
| 309 | def flashinfer_cutedsl_grouped_gemm_nt_masked( |
| 310 | hidden_states: torch.Tensor, # 3d |
| 311 | input_global_scale: torch.Tensor, # (l,) |
| 312 | weights: torch.Tensor, |
| 313 | w_global_scale: torch.Tensor, # (l,) |
| 314 | masked_m: torch.Tensor, |
| 315 | ): |
| 316 | # hidden_states: [l, m, k] |
| 317 | # weights: [l, n, k] |
| 318 | aq, aq_sf = scaled_fp4_grouped_quantize( |
| 319 | hidden_states, |
| 320 | masked_m.to(hidden_states.device), |
| 321 | input_global_scale, |
| 322 | ) |
| 323 | num_experts, n, k = weights.shape |
| 324 | bq, bq_sf = scaled_fp4_grouped_quantize( |
| 325 | weights, |
| 326 | torch.full((num_experts,), n, device=weights.device, dtype=torch.int32), |
| 327 | w_global_scale, |
| 328 | ) |
| 329 | |
| 330 | out = torch.zeros( |
| 331 | (num_experts, max(masked_m), n), dtype=weights.dtype, device=aq.device |
| 332 | ) |
| 333 | out = out.permute(1, 2, 0) # requirement of kernel |
| 334 | sf_vec_size = 16 |
| 335 | ab_dtype = "float4_e2m1fn" |
| 336 | sf_dtype = "float8_e4m3fn" |
| 337 | c_dtype = "bfloat16" |
| 338 | alpha = 1.0 / (input_global_scale * w_global_scale).to(out.dtype).view( |
| 339 | 1, 1, num_experts |
| 340 | ) |
| 341 | |
| 342 | def get_cute_dtype(input: torch.Tensor) -> str: |
| 343 | if input.dtype == torch.bfloat16: |
| 344 | return "bfloat16" |
| 345 | elif input.dtype == torch.float16: |
| 346 | return "float16" |
| 347 | elif input.dtype == torch.float32: |
| 348 | return "float32" |
| 349 | else: |
| 350 | raise ValueError(f"Unsupported cute dtype {input.dtype}") |
| 351 | |
| 352 | cutedsl_gmm_masked( |
| 353 | (aq, aq_sf), |
| 354 | (bq, bq_sf), |
| 355 | out, |
| 356 | masked_m.to(aq.device), |
| 357 | ab_dtype=ab_dtype, |
| 358 | sf_dtype=sf_dtype, |
| 359 | c_dtype=c_dtype, |
| 360 | sf_vec_size=sf_vec_size, |
| 361 | alpha=alpha, |
| 362 | alpha_dtype=get_cute_dtype(alpha), |
| 363 | ) |
| 364 | |
| 365 | return out |
| 366 | |
| 367 | |
| 368 | @pytest.mark.parametrize("bs, hidden_dim, inter_dim", [(2, 128, 256), (16, 128, 512)]) |
| 369 | @pytest.mark.parametrize("topk", [1, 2, 4]) |
| 370 | @torch.inference_mode() |
| 371 | def test_flashinfer_cutedsl_moe_masked( |
| 372 | bs: int, hidden_dim: int, inter_dim: int, topk: int |
| 373 | ): |
| 374 | torch.manual_seed(42) |
| 375 | device = "cuda" |
| 376 | num_experts = 8 |
| 377 | hidden_states = ( |
| 378 | torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) / 5.0 |
| 379 | ) |
| 380 | w1 = ( |
| 381 | torch.randn( |
| 382 | num_experts, 2 * inter_dim, hidden_dim, dtype=torch.bfloat16, device=device |
| 383 | ) |
| 384 | / 10.0 |
| 385 | ) |
| 386 | w2 = ( |
| 387 | torch.randn( |
| 388 | num_experts, hidden_dim, inter_dim, dtype=torch.bfloat16, device=device |
| 389 | ) |
| 390 | / 10.0 |
| 391 | ) |
| 392 | router_logits = torch.randn(bs, num_experts, dtype=torch.float32) |
| 393 | |
| 394 | hidden_states_expanded = ( |
| 395 | hidden_states.view(bs, -1, hidden_dim) |
| 396 | .repeat(1, topk, 1) |
| 397 | .reshape(-1, hidden_dim) |
| 398 | ) |
| 399 | hidden_states_3d, masked_m, topk_idx, routing_weights = prepare_inputs( |
| 400 | hidden_states_expanded, router_logits, num_experts, topk |
| 401 | ) |
| 402 | |
| 403 | w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32).to(w1.device) |
| 404 | w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32).to(w2.device) |
| 405 | input_global_scale = torch.ones( |
| 406 | (num_experts,), dtype=torch.float32, device=hidden_states.device |
| 407 | ) |
| 408 | |
| 409 | w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax |
| 410 | w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax |
| 411 | a2_global_scale = torch.ones( |
| 412 | (num_experts,), dtype=torch.float32, device=hidden_states.device |
| 413 | ) # assume intermediate scale is 1.0 |
| 414 | |
| 415 | w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize( |
| 416 | w1, |
| 417 | torch.ones(num_experts, dtype=torch.int32, device=w1.device) * 2 * inter_dim, |
| 418 | w1_global_scale, |
| 419 | ) |
| 420 | w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize( |
| 421 | w2, |
| 422 | torch.ones(num_experts, dtype=torch.int32, device=w2.device) * hidden_dim, |
| 423 | w2_global_scale, |
| 424 | ) |
| 425 | |
| 426 | w1_alpha = 1.0 / (input_global_scale * w1_global_scale) |
| 427 | w2_alpha = 1.0 / (a2_global_scale * w2_global_scale) |
| 428 | |
| 429 | out = torch.empty_like(hidden_states_3d, device=hidden_states.device) |
| 430 | # Note: the 1st dim shouldn't be bs |
| 431 | wk = torch.empty( |
| 432 | num_experts, |
| 433 | hidden_states_3d.shape[1], |
| 434 | inter_dim * 2, |
| 435 | dtype=hidden_states_3d.dtype, |
| 436 | device=hidden_states.device, |
| 437 | ) |
| 438 | flashinfer_cutedsl_moe_masked( |
| 439 | hidden_states_3d.to(hidden_states.device), |
| 440 | input_global_scale, |
| 441 | w1_fp4.permute(2, 0, 1), |
| 442 | w1_blockscale, |
| 443 | w1_alpha, |
| 444 | w2_fp4.permute(2, 0, 1), |
| 445 | a2_global_scale, |
| 446 | w2_blockscale, |
| 447 | w2_alpha, |
| 448 | masked_m.to(hidden_states.device), |
| 449 | wk, |
| 450 | out, |
| 451 | ) |
| 452 | |
| 453 | # reference |
| 454 | # input_global_scale is per-expert ([num_experts]); fp4_quantize and |
| 455 | # dequantize_nvfp4_to_dtype are non-grouped APIs that expect [1] or |
| 456 | # [num_tokens]. Use a single element since all values are uniform here. |
| 457 | a_global = input_global_scale[:1].contiguous() |
| 458 | a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, a_global) |
| 459 | a_in_dtype = dequantize_nvfp4_to_dtype( |
| 460 | a_fp4, |
| 461 | a_scale_interleaved, |
| 462 | a_global, |
| 463 | dtype=hidden_states.dtype, |
| 464 | device=hidden_states.device, |
| 465 | block_size=16, |
| 466 | ) |
| 467 | w1_d = torch.empty( |
| 468 | (num_experts, 2 * inter_dim, hidden_dim), device=w1.device, dtype=w1.dtype |
| 469 | ) |
| 470 | w2_d = torch.empty( |
| 471 | (num_experts, hidden_dim, inter_dim), device=w2.device, dtype=w2.dtype |
| 472 | ) |
| 473 | |
| 474 | for idx in range(0, num_experts): |
| 475 | w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize( |
| 476 | w1[idx], w1_global_scale[idx] |
| 477 | ) |
| 478 | w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize( |
| 479 | w2[idx], w2_global_scale[idx] |
| 480 | ) |
| 481 | w1_d[idx] = dequantize_nvfp4_to_dtype( |
| 482 | w1_fp4_sliced, |
| 483 | w1_blockscale_sliced, |
| 484 | w1_global_scale[idx], |
| 485 | dtype=w1.dtype, |
| 486 | device=w1.device, |
| 487 | block_size=16, |
| 488 | ) |
| 489 | w2_d[idx] = dequantize_nvfp4_to_dtype( |
| 490 | w2_fp4_sliced, |
| 491 | w2_blockscale_sliced, |
| 492 | w2_global_scale[idx], |
| 493 | dtype=w2.dtype, |
| 494 | device=w2.device, |
| 495 | block_size=16, |
| 496 | ) |
| 497 | |
| 498 | ref_output = torch_moe_nvfp4( |
| 499 | a_in_dtype, |
| 500 | w1_d, |
| 501 | w2_d, |
| 502 | topk, |
| 503 | routing_weights.to(a_in_dtype.device), |
| 504 | topk_idx.to(a_in_dtype.device), |
| 505 | ) |
| 506 | out_weighted = torch.zeros_like(ref_output, device=out.device, dtype=out.dtype) |
| 507 | |
| 508 | positions = torch.nonzero(masked_m[topk_idx], as_tuple=False) |
| 509 | rows, cols = positions[:, 0], positions[:, 1] |
| 510 | experts = topk_idx[rows, cols] |
| 511 | for i in range(num_experts): |
| 512 | mask = experts == i |
| 513 | if mask.any(): |
| 514 | idx = torch.nonzero(mask, as_tuple=False).squeeze(-1) |
| 515 | r, c = rows[idx], cols[idx] |
| 516 | out_weighted[r] += out[i, : len(r), :] * routing_weights[r, c].to( |
| 517 | out.device |
| 518 | ).unsqueeze(-1) |
| 519 | torch.testing.assert_close( |
| 520 | out_weighted.cpu(), ref_output.cpu(), atol=2e-1, rtol=2e-1 |
| 521 | ) |
| 522 | |
| 523 | |
| 524 | @pytest.mark.parametrize( |
| 525 | "bs, hidden_dim, inter_dim, topk", [(2, 128, 256, 2), (16, 128, 512, 5)] |
| 526 | ) |
| 527 | @torch.inference_mode() |
| 528 | def test_grouped_gemm_nt_masked( |
| 529 | bs: int, hidden_dim: int, inter_dim: int, topk: int |
| 530 | ) -> None: |
| 531 | torch.manual_seed(42) |
| 532 | B = bs |
| 533 | D = hidden_dim |
| 534 | N = inter_dim |
| 535 | # CuteDSL group gemm has issue when not all experts are active. |
| 536 | # i.e. masked = [2, 3, 0, 0, 1] where the 2nd and 3rd experts are inactive |
| 537 | # see https://github.com/flashinfer-ai/flashinfer/issues/1856 |
| 538 | num_experts = bs |
| 539 | hidden_states = torch.randn(B, D, dtype=torch.bfloat16, device="cuda") |
| 540 | weights = torch.randn(num_experts, N, D, dtype=torch.bfloat16, device="cuda") |
| 541 | router_logits = torch.randn(B, num_experts, dtype=torch.float32) |
| 542 | |
| 543 | hidden_states_expanded = ( |
| 544 | hidden_states.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) |
| 545 | ) |
| 546 | hidden_states_3d, masked_m, topk_idx, _ = prepare_inputs( |
| 547 | hidden_states_expanded, router_logits, num_experts, topk |
| 548 | ) |
| 549 | |
| 550 | a_amax = ( |
| 551 | hidden_states_3d.abs() |
| 552 | .amax(dim=(1, 2)) |
| 553 | .to(torch.float32) |
| 554 | .to(hidden_states.device) |
| 555 | ) |
| 556 | b_amax = weights.abs().amax(dim=(1, 2)).to(torch.float32).to(weights.device) |
| 557 | a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax |
| 558 | b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax |
| 559 | out_flashinfer = flashinfer_cutedsl_grouped_gemm_nt_masked( |
| 560 | hidden_states_3d.to(hidden_states.device), a_gs, weights, b_gs, masked_m |
| 561 | ) |
| 562 | # reference |
| 563 | out_ref = grouped_gemm_ref( |
| 564 | hidden_states_expanded=hidden_states_expanded, |
| 565 | hidden_states_3d=hidden_states_3d, |
| 566 | weights=weights, |
| 567 | topk_idx=topk_idx, |
| 568 | masked_m=masked_m, |
| 569 | B=B, |
| 570 | topk=topk, |
| 571 | num_experts=num_experts, |
| 572 | ) |
| 573 | # Note: just to compare the masked position due to cutedsl may write nan |
| 574 | # into unmasked position. |
| 575 | for i in range(num_experts): |
| 576 | torch.testing.assert_close( |
| 577 | out_flashinfer.permute(2, 0, 1)[i, : masked_m[i]], |
| 578 | out_ref.to(out_flashinfer.device)[i, : masked_m[i]], |
| 579 | atol=1e-1, |
| 580 | rtol=1e-1, |
| 581 | ) |
| 582 | |
| 583 | |
| 584 | if __name__ == "__main__": |
| 585 | test_flashinfer_cutedsl_moe_masked(16, 128, 512, 4) |
| 586 | test_grouped_gemm_nt_masked(16, 128, 512, 4) |