| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CUTLASS based Fused MoE kernels.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 11 | ApplyMoEActivationConfig, |
| 12 | MoEActivation, |
| 13 | apply_moe_activation, |
| 14 | apply_moe_activation_supported, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe.config import ( |
| 17 | FusedMoEConfig, |
| 18 | FusedMoEParallelConfig, |
| 19 | FusedMoEQuantConfig, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( |
| 22 | MoEPermuteScratch, |
| 23 | moe_permute, |
| 24 | moe_permute_unpermute_supported, |
| 25 | moe_unpermute, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 28 | TopKWeightAndReduceDelegate, |
| 29 | TopKWeightAndReduceNoOP, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 32 | _resize_cache, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 35 | QuantKey, |
| 36 | kFp8DynamicTensorSym, |
| 37 | kFp8DynamicTokenSym, |
| 38 | kFp8StaticChannelSym, |
| 39 | kFp8StaticTensorSym, |
| 40 | kInt4Static, |
| 41 | kMxfp4Dynamic, |
| 42 | kMxfp4Static, |
| 43 | kNvfp4Dynamic, |
| 44 | kNvfp4Static, |
| 45 | ) |
| 46 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 47 | cutlass_group_gemm_supported, |
| 48 | ) |
| 49 | from vllm.platforms import current_platform |
| 50 | from vllm.scalar_type import scalar_types |
| 51 | |
| 52 | logger = init_logger(__name__) |
| 53 | |
| 54 | |
| 55 | def run_cutlass_moe_fp8( |
| 56 | output: torch.Tensor, |
| 57 | hidden_states: torch.Tensor, |
| 58 | w1: torch.Tensor, |
| 59 | w2: torch.Tensor, |
| 60 | topk_ids: torch.Tensor, |
| 61 | activation: MoEActivation, |
| 62 | global_num_experts: int, |
| 63 | expert_map: torch.Tensor | None, |
| 64 | w1_scale: torch.Tensor | None, |
| 65 | w2_scale: torch.Tensor | None, |
| 66 | a1q_scale: torch.Tensor | None, |
| 67 | a2_scale: torch.Tensor | None, |
| 68 | ab_strides1: torch.Tensor, |
| 69 | ab_strides2: torch.Tensor, |
| 70 | c_strides1: torch.Tensor, |
| 71 | c_strides2: torch.Tensor, |
| 72 | workspace13: torch.Tensor, |
| 73 | workspace2: torch.Tensor, |
| 74 | expert_num_tokens: torch.Tensor | None, |
| 75 | out_dtype: torch.dtype, |
| 76 | per_act_token: bool, |
| 77 | per_out_ch: bool, |
| 78 | use_batched_format: bool, |
| 79 | topk_weights: torch.Tensor | None, |
| 80 | permute_scratch: MoEPermuteScratch | None, |
| 81 | *, |
| 82 | activation_config: ApplyMoEActivationConfig | None = None, |
| 83 | ): |
| 84 | a1q = hidden_states |
| 85 | |
| 86 | assert activation.is_gated, "Only gated activation is supported" |
| 87 | assert w1_scale is not None |
| 88 | assert w2_scale is not None |
| 89 | assert w1.dtype == torch.float8_e4m3fn |
| 90 | assert w2.dtype == torch.float8_e4m3fn |
| 91 | assert a1q.size(-1) == w1.size(2), "Hidden size mismatch w1" |
| 92 | assert w1.size(1) == w2.size(2) * 2, "Hidden size mismatch w2" |
| 93 | assert ( |
| 94 | w1_scale.dim() == 1 or w1_scale.size(1) == 1 or w1_scale.shape[1] == w1.size(1) |
| 95 | ), "W1 scale shape mismatch" |
| 96 | assert ( |
| 97 | w2_scale.dim() == 1 or w2_scale.size(1) == 1 or w2_scale.shape[1] == w2.size(1) |
| 98 | ), "W2 scale shape mismatch" |
| 99 | assert w1.size(0) == w2.size(0), "Expert number mismatch" |
| 100 | assert ( |
| 101 | a1q_scale is None |
| 102 | or a1q_scale.dim() == 0 |
| 103 | or a1q_scale.size(0) == 1 |
| 104 | or a1q_scale.size(0) == a1q.shape[0] |
| 105 | ), "Input scale shape mismatch" |
| 106 | assert w1.size(0) == w2.size(0), "Weights expert number mismatch" |
| 107 | assert w1.size(0) == w1_scale.size(0), "w1 scales expert number mismatch" |
| 108 | assert w1.size(0) == w2_scale.size(0), "w2 scales expert number mismatch" |
| 109 | assert ( |
| 110 | a2_scale is None |
| 111 | or a2_scale.dim() == 0 |
| 112 | or a2_scale.size(0) == 1 |
| 113 | or a2_scale.size(0) == a1q.shape[0] |
| 114 | ), "Intermediate scale shape mismatch" |
| 115 | assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype" |
| 116 | |
| 117 | # NOTE(rob): the expert_map is used for the STANDARD case and |
| 118 | # the batched format is used by the BATCHED case. |
| 119 | # TODO(rob): update the MK interface to only pass the expert_map |
| 120 | # during the STANDARD case to make this clearer across all kernels. |
| 121 | if use_batched_format: |
| 122 | assert expert_num_tokens is not None |
| 123 | else: |
| 124 | assert expert_num_tokens is None |
| 125 | |
| 126 | # We have two modes: batched experts and non-batched experts. |
| 127 | # In the non-batched mode, the input tokens are not padded: thus, the shape |
| 128 | # of the input is [total_num_tokens, hidden_size]. The input and output |
| 129 | # require shuffling by a_map and c_map such that the tokens assigned to |
| 130 | # each expert are contiguous. |
| 131 | # In the batched mode, the input tokens are padded per expert to ensure that |
| 132 | # the batched dispatch and combine functions work correctly: thus, the shape |
| 133 | # of the input is [num_experts, max_num_tokens_per_expert, hidden_size]. |
| 134 | # The batched input and output require no shuffling by a_map and c_map since |
| 135 | # their tokens are already contiguous for each expert as a result of |
| 136 | # the dispatch function. |
| 137 | |
| 138 | M = a1q.size(0) # non batched expert M |
| 139 | padded_M = a1q.size(1) # batched expert M |
| 140 | _, K, N = w2.shape |
| 141 | device = a1q.device |
| 142 | |
| 143 | assert w1.size(2) == K |
| 144 | assert global_num_experts != -1 |
| 145 | assert a1q_scale is not None |
| 146 | |
| 147 | topk = topk_ids.size(1) |
| 148 | local_E = w1.size(0) |
| 149 | |
| 150 | if use_batched_format: |
| 151 | mm1_out = _resize_cache(workspace13, (local_E * padded_M, N * 2)) |
| 152 | act_out = _resize_cache(workspace2, (local_E * padded_M, N)) |
| 153 | quant_out = _resize_cache( |
| 154 | workspace13.view(dtype=torch.float8_e4m3fn), (local_E * padded_M, N) |
| 155 | ) |
| 156 | mm2_out = _resize_cache(workspace2, (local_E * padded_M, K)) |
| 157 | else: |
| 158 | a1q_perm = _resize_cache( |
| 159 | workspace2.view(dtype=torch.float8_e4m3fn), (M * topk, K) |
| 160 | ) |
| 161 | mm1_out = _resize_cache(workspace13, (M * topk, N * 2)) |
| 162 | act_out = _resize_cache(workspace2, (M * topk, N)) |
| 163 | # original workspace are based on input hidden_states dtype (bf16) |
| 164 | quant_out = _resize_cache( |
| 165 | workspace13.view(dtype=torch.float8_e4m3fn), (M * topk, N) |
| 166 | ) |
| 167 | mm2_out = _resize_cache(workspace2, (M * topk, K)) |
| 168 | |
| 169 | if use_batched_format: |
| 170 | assert expert_num_tokens is not None |
| 171 | |
| 172 | expert_offsets = torch.empty((local_E), dtype=torch.int32, device=device) |
| 173 | problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 174 | problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 175 | |
| 176 | ops.get_cutlass_batched_moe_mm_data( |
| 177 | expert_offsets, |
| 178 | problem_sizes1, |
| 179 | problem_sizes2, |
| 180 | expert_num_tokens, |
| 181 | local_E, |
| 182 | padded_M, |
| 183 | N, |
| 184 | K, |
| 185 | ) |
| 186 | |
| 187 | w1_scale = w1_scale.reshape(w1_scale.size(0), -1) |
| 188 | w2_scale = w2_scale.reshape(w2_scale.size(0), -1) |
| 189 | a1q = a1q.reshape(-1, a1q.size(2)) |
| 190 | a1q_scale = a1q_scale.reshape(-1, a1q_scale.size(2)).contiguous() |
| 191 | # c3x get_group_gemm_starts expects int64 to avoid overflow |
| 192 | # during offset calculations |
| 193 | expert_offsets = expert_offsets.to(torch.int64) |
| 194 | else: |
| 195 | problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 196 | problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 197 | |
| 198 | num_expert = global_num_experts if expert_map is None else expert_map.size(0) |
| 199 | # permuted a1q reuses workspace2 |
| 200 | a1q, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute( |
| 201 | a1q, |
| 202 | a1q_scale, |
| 203 | topk_ids, |
| 204 | num_expert, |
| 205 | local_E, |
| 206 | expert_map, |
| 207 | permuted_hidden_states=a1q_perm, |
| 208 | scratch=permute_scratch, |
| 209 | ) |
| 210 | # swap_ab is a CUTLASS grouped-GEMM optimization (M <= 64 reduces padding). |
| 211 | swap_ab = a1q.size(0) <= 64 |
| 212 | ops.get_cutlass_moe_mm_problem_sizes_from_expert_offsets( |
| 213 | expert_first_token_offset, problem_sizes1, problem_sizes2, N, K, swap_ab |
| 214 | ) |
| 215 | expert_offsets = expert_first_token_offset[:-1] |
| 216 | |
| 217 | if not per_act_token and (expert_map is not None or use_batched_format): |
| 218 | # this is necessary to avoid imprecise scale calculation caused by |
| 219 | # random data in the unused workspace. The workspace is unused when |
| 220 | # this rank handles only partial tokens, or when it is batched . |
| 221 | mm1_out.fill_(0) |
| 222 | |
| 223 | ops.cutlass_moe_mm( |
| 224 | mm1_out, |
| 225 | a1q, |
| 226 | w1, |
| 227 | a1q_scale, |
| 228 | w1_scale, |
| 229 | expert_offsets, |
| 230 | problem_sizes1, |
| 231 | ab_strides1, |
| 232 | ab_strides1, |
| 233 | c_strides1, |
| 234 | per_act_token, |
| 235 | per_out_ch, |
| 236 | ) |
| 237 | |
| 238 | apply_moe_activation( |
| 239 | activation, |
| 240 | act_out, |
| 241 | mm1_out, |
| 242 | activation_config=activation_config, |
| 243 | ) |
| 244 | |
| 245 | a2q, a2q_scale = ops.scaled_fp8_quant( |
| 246 | act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out |
| 247 | ) |
| 248 | |
| 249 | ops.cutlass_moe_mm( |
| 250 | mm2_out, |
| 251 | a2q, |
| 252 | w2, |
| 253 | a2q_scale, |
| 254 | w2_scale, |
| 255 | expert_offsets, |
| 256 | problem_sizes2, |
| 257 | ab_strides2, |
| 258 | ab_strides2, |
| 259 | c_strides2, |
| 260 | per_act_token, |
| 261 | per_out_ch, |
| 262 | ) |
| 263 | |
| 264 | if use_batched_format: |
| 265 | output.copy_(mm2_out.reshape(local_E, padded_M, K), non_blocking=True) |
| 266 | else: |
| 267 | # for non-chunking mode the output is resized from workspace13 |
| 268 | # so we need to make sure mm2_out uses workspace2. |
| 269 | moe_unpermute( |
| 270 | out=output, |
| 271 | permuted_hidden_states=mm2_out, |
| 272 | topk_weights=topk_weights, |
| 273 | inv_permuted_idx=inv_perm, |
| 274 | expert_first_token_offset=expert_first_token_offset, |
| 275 | ) |
| 276 | |
| 277 | |
| 278 | class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular): |
| 279 | def __init__( |
| 280 | self, |
| 281 | moe_config: FusedMoEConfig, |
| 282 | quant_config: FusedMoEQuantConfig, |
| 283 | max_num_tokens: int | None = None, |
| 284 | num_dispatchers: int | None = None, |
| 285 | ): |
| 286 | super().__init__( |
| 287 | moe_config=moe_config, |
| 288 | quant_config=quant_config, |
| 289 | max_num_tokens=max_num_tokens, |
| 290 | num_dispatchers=num_dispatchers, |
| 291 | ) |
| 292 | assert quant_config.use_fp8_w8a8 |
| 293 | |
| 294 | e = moe_config.num_local_experts |
| 295 | n = moe_config.intermediate_size_per_partition |
| 296 | k = moe_config.hidden_dim |
| 297 | device = moe_config.device |
| 298 | ab_strides1_c_strides2 = torch.full((e,), k, device=device, dtype=torch.int64) |
| 299 | ab_strides2 = torch.full((e,), n, device=device, dtype=torch.int64) |
| 300 | c_strides1 = torch.full((e,), 2 * n, device=device, dtype=torch.int64) |
| 301 | |
| 302 | self.out_dtype = moe_config.in_dtype |
| 303 | self.ab_strides1 = ab_strides1_c_strides2 |
| 304 | self.ab_strides2 = ab_strides2 |
| 305 | self.c_strides1 = c_strides1 |
| 306 | self.c_strides2 = ab_strides1_c_strides2 |
| 307 | self._permute_scratch: MoEPermuteScratch | None = None |
| 308 | |
| 309 | @staticmethod |
| 310 | def _supports_current_device() -> bool: |
| 311 | return cutlass_group_gemm_supported() |
| 312 | |
| 313 | @staticmethod |
| 314 | def _supports_no_act_and_mul() -> bool: |
| 315 | return False |
| 316 | |
| 317 | @staticmethod |
| 318 | def _supports_quant_scheme( |
| 319 | weight_key: QuantKey | None, |
| 320 | activation_key: QuantKey | None, |
| 321 | ) -> bool: |
| 322 | SUPPORTED_W_A = [ |
| 323 | (kFp8StaticChannelSym, kFp8DynamicTokenSym), |
| 324 | (kFp8StaticTensorSym, kFp8DynamicTensorSym), |
| 325 | (kFp8StaticTensorSym, kFp8StaticTensorSym), |
| 326 | ] |
| 327 | return (weight_key, activation_key) in SUPPORTED_W_A |
| 328 | |
| 329 | @staticmethod |
| 330 | def _supports_activation(activation: MoEActivation) -> bool: |
| 331 | return activation.is_gated and apply_moe_activation_supported(activation) |
| 332 | |
| 333 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 334 | # Let PrepareAndFinalize::finalize() decide the impl. |
| 335 | return TopKWeightAndReduceDelegate() |
| 336 | |
| 337 | def _get_permute_scratch(self) -> MoEPermuteScratch | None: |
| 338 | if self._permute_scratch is None and moe_permute_unpermute_supported(): |
| 339 | self._permute_scratch = MoEPermuteScratch( |
| 340 | max_num_tokens=self.moe_config.max_num_tokens, |
| 341 | topk=self.moe_config.experts_per_token, |
| 342 | num_experts=self.moe_config.num_experts, |
| 343 | num_local_experts=self.moe_config.num_local_experts, |
| 344 | device=torch.device(self.moe_config.device), |
| 345 | ) |
| 346 | return self._permute_scratch |
| 347 | |
| 348 | def apply( |
| 349 | self, |
| 350 | output: torch.Tensor, |
| 351 | hidden_states: torch.Tensor, |
| 352 | w1: torch.Tensor, |
| 353 | w2: torch.Tensor, |
| 354 | topk_weights: torch.Tensor, |
| 355 | topk_ids: torch.Tensor, |
| 356 | activation: MoEActivation, |
| 357 | global_num_experts: int, |
| 358 | expert_map: torch.Tensor | None, |
| 359 | a1q_scale: torch.Tensor | None, |
| 360 | a2_scale: torch.Tensor | None, |
| 361 | workspace13: torch.Tensor, |
| 362 | workspace2: torch.Tensor, |
| 363 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 364 | apply_router_weight_on_input: bool, |
| 365 | ): |
| 366 | assert self.w1_zp is None, "w1_zp is not supported in CUTLASS MoE" |
| 367 | assert self.w2_zp is None, "w2_zp is not supported in CUTLASS MoE" |
| 368 | |
| 369 | expert_num_tokens = None |
| 370 | if expert_tokens_meta is not None: |
| 371 | expert_num_tokens = expert_tokens_meta.expert_num_tokens |
| 372 | |
| 373 | use_batched_format = ( |
| 374 | self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts |
| 375 | ) |
| 376 | |
| 377 | in_dtype = hidden_states.dtype |
| 378 | run_cutlass_moe_fp8( |
| 379 | output, |
| 380 | hidden_states, |
| 381 | w1, |
| 382 | w2, |
| 383 | topk_ids, |
| 384 | activation, |
| 385 | global_num_experts, |
| 386 | expert_map, |
| 387 | self.w1_scale, |
| 388 | self.w2_scale, |
| 389 | a1q_scale, |
| 390 | a2_scale, |
| 391 | self.ab_strides1, |
| 392 | self.ab_strides2, |
| 393 | self.c_strides1, |
| 394 | self.c_strides2, |
| 395 | workspace13, |
| 396 | workspace2, |
| 397 | expert_num_tokens, |
| 398 | self.out_dtype if self.out_dtype is not None else in_dtype, |
| 399 | self.per_act_token_quant, |
| 400 | self.per_out_ch_quant, |
| 401 | use_batched_format, |
| 402 | topk_weights, |
| 403 | self._get_permute_scratch(), |
| 404 | activation_config=self.activation_config, |
| 405 | ) |
| 406 | |
| 407 | |
| 408 | class CutlassExpertsFp8(CutlassExpertsFp8Base): |
| 409 | """CUTLASS FP8 fused MoE expert implementation.""" |
| 410 | |
| 411 | @staticmethod |
| 412 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 413 | return mk.FusedMoEActivationFormat.Standard |
| 414 | |
| 415 | @staticmethod |
| 416 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 417 | # CutlassExpertsFp8 does not support expert map, which is |
| 418 | # needed for STANDARD activation format kernels in DP/EP mode. |
| 419 | # Note that the BATCHED activation format does not use |
| 420 | # the expert map for identifying experts. |
| 421 | return not ( |
| 422 | moe_parallel_config.use_fi_nvl_two_sided_kernels |
| 423 | or moe_parallel_config.use_deepep_ht_kernels |
| 424 | or moe_parallel_config.use_fi_nvl_one_sided_kernels |
| 425 | ) |
| 426 | |
| 427 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 428 | # topk weights and reduction are fused in moe_unpermute cuda kernel |
| 429 | return TopKWeightAndReduceNoOP() |
| 430 | |
| 431 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 432 | return self.out_dtype if self.out_dtype is not None else act_dtype |
| 433 | |
| 434 | def workspace_shapes( |
| 435 | self, |
| 436 | M: int, |
| 437 | N: int, |
| 438 | K: int, |
| 439 | topk: int, |
| 440 | global_num_experts: int, |
| 441 | local_num_experts: int, |
| 442 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 443 | activation: MoEActivation, |
| 444 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 445 | activation_out_dim = self.adjust_N_for_activation(N, activation) |
| 446 | workspace1 = (M * topk, max(N, K)) |
| 447 | workspace2 = (M * topk, max(activation_out_dim, K)) |
| 448 | output = (M, K) |
| 449 | return (workspace1, workspace2, output) |
| 450 | |
| 451 | |
| 452 | class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base): |
| 453 | """Batched CUTLASS FP8 fused MoE expert implementation.""" |
| 454 | |
| 455 | @staticmethod |
| 456 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 457 | # BATCHED activation format works with EP because |
| 458 | # expert_map is not used to identify experts (the |
| 459 | # info is encoded/managed by the P/F logic). |
| 460 | return True |
| 461 | |
| 462 | @staticmethod |
| 463 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 464 | return mk.FusedMoEActivationFormat.BatchedExperts |
| 465 | |
| 466 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 467 | return self.out_dtype if self.out_dtype is not None else act_dtype |
| 468 | |
| 469 | def workspace_shapes( |
| 470 | self, |
| 471 | M: int, |
| 472 | N: int, |
| 473 | K: int, |
| 474 | topk: int, |
| 475 | global_num_experts: int, |
| 476 | local_num_experts: int, |
| 477 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 478 | activation: MoEActivation, |
| 479 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 480 | experts_per_worker = self.moe_config.num_local_experts |
| 481 | activation_out_dim = self.adjust_N_for_activation(N, activation) |
| 482 | workspace1 = (experts_per_worker, M, max(N, K)) |
| 483 | workspace2 = ( |
| 484 | experts_per_worker, |
| 485 | M, |
| 486 | max(activation_out_dim, K), |
| 487 | ) |
| 488 | output = (experts_per_worker, M, K) |
| 489 | return (workspace1, workspace2, output) |
| 490 | |
| 491 | |
| 492 | FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() |
| 493 | FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 494 | |
| 495 | |
| 496 | def run_cutlass_moe_fp4( |
| 497 | output: torch.Tensor, |
| 498 | a: torch.Tensor, |
| 499 | a1_gscale: torch.Tensor, |
| 500 | w1_fp4: torch.Tensor, |
| 501 | w1_blockscale: torch.Tensor, |
| 502 | w1_alphas: torch.Tensor, |
| 503 | a2_gscale: torch.Tensor, |
| 504 | w2_fp4: torch.Tensor, |
| 505 | w2_blockscale: torch.Tensor, |
| 506 | w2_alphas: torch.Tensor, |
| 507 | topk_weights: torch.Tensor, |
| 508 | topk_ids: torch.Tensor, |
| 509 | activation: MoEActivation, |
| 510 | workspace13: torch.Tensor, |
| 511 | workspace2: torch.Tensor, |
| 512 | m: int, |
| 513 | n: int, |
| 514 | k: int, |
| 515 | e: int, |
| 516 | device: torch.device, |
| 517 | apply_router_weight_on_input: bool = False, |
| 518 | *, |
| 519 | activation_config: ApplyMoEActivationConfig | None = None, |
| 520 | ) -> None: |
| 521 | """MoE implementation for FP4 Inputs. |
| 522 | |
| 523 | # Gemm 1 |
| 524 | a: Input tensor: [m, k] (half/bfloat16) |
| 525 | a1_gscale: Activation scale per expert: [e] (float32) |
| 526 | w1 (not an argument to cutlass_moe_fp4): [e, w1_n, k] |
| 527 | w1_fp4: [e, w1_n, k // 2], dtype: torch.uint8 (stacked fp4: E2M1) |
| 528 | where w1_n = 2*n for gated activations (gate+up), n for non-gated (up only). |
| 529 | (Note: `n` is the up projection output dim, `k` is the input dim in |
| 530 | full precision) |
| 531 | w1_blockscale: [e, w1_n, k // block_size] (float8_e4m3) |
| 532 | (Block size = 16 for NVFP4) |
| 533 | |
| 534 | # Gemm 2 |
| 535 | a2_gscale: Activation scale per expert: [e] |
| 536 | w2(down projection) (not an argument to cutlass_moe_fp4): [e, k, n] |
| 537 | w2_fp4: [e, k, n // 2], dtype: torch.uint8 (stacked E2M1) |
| 538 | w2_blockscale: [e, k, n // block_size], dtype: float8_e4m3 |
| 539 | |
| 540 | topk_weights: [m, topk] dtype: float8 |
| 541 | topk_ids: [m, topk] dtype: float8 |
| 542 | |
| 543 | m, n, k: Unquantized weight shapes, dtype: int |
| 544 | e: number of experts, dtype: int |
| 545 | |
| 546 | assumes that topk < k < n to satisfy - up/down projection expectations. |
| 547 | """ |
| 548 | is_gated = activation.is_gated |
| 549 | # For gated activations (e.g. SiLU), w1 output is 2*n (gate + up). |
| 550 | # For non-gated activations (e.g. SiLU_NO_MUL), w1 output is n (up only). |
| 551 | w1_n = n * 2 if is_gated else n |
| 552 | |
| 553 | assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" |
| 554 | assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8" |
| 555 | assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8" |
| 556 | assert ( |
| 557 | w1_fp4.ndim == 3 |
| 558 | and w2_fp4.ndim == 3 |
| 559 | and w1_blockscale.ndim == 3 |
| 560 | and w2_blockscale.ndim == 3 |
| 561 | ), "All Weights must be of rank 3 for cutlass_moe_fp4" |
| 562 | m_a, k_a = a.shape |
| 563 | e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape |
| 564 | e_w2, k_w2, half_n_w2 = w2_fp4.shape |
| 565 | |
| 566 | assert e_w1 == e_w2 and e_w1 == e, ( |
| 567 | "Number of experts must match", |
| 568 | f" between weights. {e_w1}, {e_w2}, {e}", |
| 569 | ) |
| 570 | assert k_a == half_k_w1 * 2 and k == k_w2, ( |
| 571 | "Hidden size mismatch between a, w1 and w2" |
| 572 | ) |
| 573 | assert w1_n_actual == w1_n and half_n_w2 * 2 == n, "mismatch in expected `n`" |
| 574 | assert m == m_a, "input shape mismatch" |
| 575 | assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1" |
| 576 | assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype" |
| 577 | assert topk_weights.size(0) == m and topk_ids.size(0) == m, ( |
| 578 | "topk must be provided for each row of a" |
| 579 | ) |
| 580 | topk = topk_ids.size(1) |
| 581 | out_dtype = a.dtype |
| 582 | num_topk = topk_ids.size(1) |
| 583 | |
| 584 | expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) |
| 585 | blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) |
| 586 | # Problem size: (num_experts, (m,2n,k)) |
| 587 | problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device) |
| 588 | # Problem size: (num_experts, (m,n,k)) |
| 589 | problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device) |
| 590 | |
| 591 | a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) |
| 592 | c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) |
| 593 | |
| 594 | if apply_router_weight_on_input: |
| 595 | # TODO: this only works for topK=1, will need to update for topK>1 |
| 596 | assert num_topk == 1, ( |
| 597 | "apply_router_weight_on_input is only implemented for topk=1" |
| 598 | ) |
| 599 | a.mul_(topk_weights.to(out_dtype)) |
| 600 | |
| 601 | # problem shapes should have [m, n, k] |
| 602 | # Note that problem sizes are based on logical number of elements. |
| 603 | ops.get_cutlass_moe_mm_data( |
| 604 | topk_ids, |
| 605 | expert_offsets, |
| 606 | problem_sizes1, |
| 607 | problem_sizes2, |
| 608 | a_map, |
| 609 | c_map, |
| 610 | e, |
| 611 | n, |
| 612 | k, |
| 613 | blockscale_offsets, |
| 614 | is_gated=is_gated, |
| 615 | ) |
| 616 | |
| 617 | a = ops.shuffle_rows(a, a_map) |
| 618 | rep_a_fp4, rep_a_blockscale = ops.scaled_fp4_experts_quant( |
| 619 | a, |
| 620 | a1_gscale, |
| 621 | expert_offsets, |
| 622 | blockscale_offsets, |
| 623 | num_topk, |
| 624 | ) |
| 625 | c1 = _resize_cache(workspace13, (m * topk, w1_n)) |
| 626 | c2 = _resize_cache(workspace2, (m * topk, n)) |
| 627 | c3 = _resize_cache(workspace13, (m * topk, k)) |
| 628 | ops.cutlass_fp4_moe_mm( |
| 629 | c1, |
| 630 | rep_a_fp4, |
| 631 | w1_fp4, |
| 632 | rep_a_blockscale, |
| 633 | w1_blockscale, |
| 634 | w1_alphas, |
| 635 | problem_sizes1, |
| 636 | expert_offsets[:-1], |
| 637 | blockscale_offsets[:-1], |
| 638 | ) |
| 639 | del rep_a_fp4, rep_a_blockscale |
| 640 | if activation == MoEActivation.SILU and ( |
| 641 | activation_config is None or activation_config.clamp_limit is None |
| 642 | ): |
| 643 | # Fused SiLU+Mul+NVFP4 quantization |
| 644 | # Note: c2 workspace is no longer needed since SiLU is fused with quantization. |
| 645 | # c3 reuses workspace13 after c1 is consumed. |
| 646 | int_fp4, int_blockscale = ops.silu_and_mul_scaled_fp4_experts_quant( |
| 647 | c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk |
| 648 | ) |
| 649 | else: |
| 650 | apply_moe_activation( |
| 651 | activation, |
| 652 | c2, |
| 653 | c1, |
| 654 | activation_config=activation_config, |
| 655 | ) |
| 656 | int_fp4, int_blockscale = ops.scaled_fp4_experts_quant( |
| 657 | c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk |
| 658 | ) |
| 659 | |
| 660 | ops.cutlass_fp4_moe_mm( |
| 661 | c3, |
| 662 | int_fp4, |
| 663 | w2_fp4, |
| 664 | int_blockscale, |
| 665 | w2_blockscale, |
| 666 | w2_alphas, |
| 667 | problem_sizes2, |
| 668 | expert_offsets[:-1], |
| 669 | blockscale_offsets[:-1], |
| 670 | ) |
| 671 | del int_fp4, int_blockscale |
| 672 | |
| 673 | c3 = ops.shuffle_rows(c3, c_map) |
| 674 | |
| 675 | assert output.dtype == out_dtype |
| 676 | if not apply_router_weight_on_input: |
| 677 | output.copy_( |
| 678 | ( |
| 679 | c3.view(m, num_topk, k) |
| 680 | * topk_weights.view(m, num_topk, 1).to(out_dtype) |
| 681 | ).sum(dim=1), |
| 682 | non_blocking=True, |
| 683 | ) |
| 684 | else: |
| 685 | output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True) |
| 686 | return |
| 687 | |
| 688 | |
| 689 | class CutlassExpertsFp4(mk.FusedMoEExpertsModular): |
| 690 | """CUTLASS FP4 fused MoE expert implementation.""" |
| 691 | |
| 692 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 693 | # Fuse activation scales into w_scale_2 in-place so that |
| 694 | # g1/g2_alphas (which reference the same tensor) stay in sync |
| 695 | # when EPLB rearranges the parameter. |
| 696 | layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) |
| 697 | layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) |
| 698 | |
| 699 | @property |
| 700 | def expects_unquantized_inputs(self) -> bool: |
| 701 | return True |
| 702 | |
| 703 | @staticmethod |
| 704 | def _supports_current_device() -> bool: |
| 705 | p = current_platform |
| 706 | return p.is_cuda() and ( |
| 707 | p.is_device_capability_family(100) |
| 708 | or p.is_device_capability_family(110) |
| 709 | or p.is_device_capability_family(120) |
| 710 | ) |
| 711 | |
| 712 | @staticmethod |
| 713 | def _supports_no_act_and_mul() -> bool: |
| 714 | return True |
| 715 | |
| 716 | @staticmethod |
| 717 | def _supports_quant_scheme( |
| 718 | weight_key: QuantKey | None, |
| 719 | activation_key: QuantKey | None, |
| 720 | ) -> bool: |
| 721 | return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) |
| 722 | |
| 723 | @staticmethod |
| 724 | def _supports_activation(activation: MoEActivation) -> bool: |
| 725 | # SILU uses a fused silu+mul+fp4_quant kernel path. |
| 726 | # Other gated activations use the generic apply_moe_activation() |
| 727 | # fallback + separate fp4 quantization in run_cutlass_moe_fp4(). |
| 728 | # Non-gated activations (_NO_MUL) are also supported for models |
| 729 | # like Nemotron-Nano that don't use gated MLP. |
| 730 | return apply_moe_activation_supported(activation) |
| 731 | |
| 732 | @staticmethod |
| 733 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 734 | # CutlassExpertsFp4 does not support expert map, which is |
| 735 | # needed for STANDARD activation format kernels in EP mode. |
| 736 | return moe_parallel_config.ep_size == 1 |
| 737 | |
| 738 | @staticmethod |
| 739 | def _supports_batch_invariance() -> bool: |
| 740 | return True |
| 741 | |
| 742 | @staticmethod |
| 743 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 744 | return mk.FusedMoEActivationFormat.Standard |
| 745 | |
| 746 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 747 | return TopKWeightAndReduceNoOP() |
| 748 | |
| 749 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 750 | return act_dtype |
| 751 | |
| 752 | def workspace_shapes( |
| 753 | self, |
| 754 | M: int, |
| 755 | N: int, |
| 756 | K: int, |
| 757 | topk: int, |
| 758 | global_num_experts: int, |
| 759 | local_num_experts: int, |
| 760 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 761 | activation: MoEActivation, |
| 762 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 763 | workspace1 = (M * topk, max(2 * N, K)) |
| 764 | workspace2 = (M * topk, N) |
| 765 | output = (M, K) |
| 766 | return (workspace1, workspace2, output) |
| 767 | |
| 768 | def apply( |
| 769 | self, |
| 770 | output: torch.Tensor, |
| 771 | hidden_states: torch.Tensor, |
| 772 | w1: torch.Tensor, |
| 773 | w2: torch.Tensor, |
| 774 | topk_weights: torch.Tensor, |
| 775 | topk_ids: torch.Tensor, |
| 776 | activation: MoEActivation, |
| 777 | global_num_experts: int, |
| 778 | expert_map: torch.Tensor | None, |
| 779 | a1q_scale: torch.Tensor | None, # unused |
| 780 | a2_scale: torch.Tensor | None, # unused |
| 781 | workspace13: torch.Tensor | None, |
| 782 | workspace2: torch.Tensor | None, |
| 783 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 784 | apply_router_weight_on_input: bool, |
| 785 | ): |
| 786 | e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids) |
| 787 | n = w2.shape[2] * 2 |
| 788 | |
| 789 | run_cutlass_moe_fp4( |
| 790 | output=output, |
| 791 | a=hidden_states, |
| 792 | a1_gscale=self.a1_gscale, |
| 793 | w1_fp4=w1, |
| 794 | w1_blockscale=self.w1_scale, |
| 795 | w1_alphas=self.g1_alphas, |
| 796 | a2_gscale=self.a2_gscale, |
| 797 | w2_fp4=w2, |
| 798 | w2_blockscale=self.w2_scale, |
| 799 | w2_alphas=self.g2_alphas, |
| 800 | topk_weights=topk_weights, |
| 801 | topk_ids=topk_ids, |
| 802 | activation=activation, |
| 803 | workspace13=workspace13, |
| 804 | workspace2=workspace2, |
| 805 | m=m, |
| 806 | n=n, |
| 807 | k=k, |
| 808 | e=e, |
| 809 | device=hidden_states.device, |
| 810 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 811 | activation_config=self.activation_config, |
| 812 | ) |
| 813 | |
| 814 | |
| 815 | def run_cutlass_moe_mxfp4( |
| 816 | output: torch.Tensor, |
| 817 | a: torch.Tensor, |
| 818 | w1_fp4: torch.Tensor, |
| 819 | w1_blockscale: torch.Tensor, |
| 820 | w2_fp4: torch.Tensor, |
| 821 | w2_blockscale: torch.Tensor, |
| 822 | topk_weights: torch.Tensor, |
| 823 | topk_ids: torch.Tensor, |
| 824 | activation: MoEActivation, |
| 825 | workspace13: torch.Tensor, |
| 826 | workspace2: torch.Tensor, |
| 827 | m: int, |
| 828 | n: int, |
| 829 | k: int, |
| 830 | e: int, |
| 831 | device: torch.device, |
| 832 | apply_router_weight_on_input: bool = False, |
| 833 | *, |
| 834 | activation_config: ApplyMoEActivationConfig | None = None, |
| 835 | ) -> None: |
| 836 | """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM.""" |
| 837 | is_gated = activation.is_gated |
| 838 | w1_n = n * 2 if is_gated else n |
| 839 | |
| 840 | assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" |
| 841 | assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8" |
| 842 | assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8" |
| 843 | assert ( |
| 844 | w1_fp4.ndim == 3 |
| 845 | and w2_fp4.ndim == 3 |
| 846 | and w1_blockscale.ndim == 3 |
| 847 | and w2_blockscale.ndim == 3 |
| 848 | ), "All Weights must be of rank 3 for cutlass_moe_mxfp4" |
| 849 | m_a, k_a = a.shape |
| 850 | e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape |
| 851 | e_w2, k_w2, half_n_w2 = w2_fp4.shape |
| 852 | |
| 853 | assert e_w1 == e_w2 and e_w1 == e |
| 854 | assert k_a == half_k_w1 * 2 and k == k_w2 |
| 855 | assert w1_n_actual == w1_n and half_n_w2 * 2 == n |
| 856 | assert m == m_a |
| 857 | assert 2 * half_k_w1 == k_w2 |
| 858 | assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype" |
| 859 | assert topk_weights.size(0) == m and topk_ids.size(0) == m |
| 860 | |
| 861 | topk = topk_ids.size(1) |
| 862 | out_dtype = a.dtype |
| 863 | num_topk = topk_ids.size(1) |
| 864 | |
| 865 | expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) |
| 866 | blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) |
| 867 | problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device) |
| 868 | problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device) |
| 869 | |
| 870 | a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) |
| 871 | c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) |
| 872 | |
| 873 | if apply_router_weight_on_input: |
| 874 | assert num_topk == 1, ( |
| 875 | "apply_router_weight_on_input is only implemented for topk=1" |
| 876 | ) |
| 877 | a.mul_(topk_weights.to(out_dtype)) |
| 878 | |
| 879 | ops.get_cutlass_moe_mm_data( |
| 880 | topk_ids, |
| 881 | expert_offsets, |
| 882 | problem_sizes1, |
| 883 | problem_sizes2, |
| 884 | a_map, |
| 885 | c_map, |
| 886 | e, |
| 887 | n, |
| 888 | k, |
| 889 | blockscale_offsets, |
| 890 | is_gated=is_gated, |
| 891 | ) |
| 892 | |
| 893 | a = ops.shuffle_rows(a, a_map) |
| 894 | rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant( |
| 895 | a, |
| 896 | expert_offsets, |
| 897 | blockscale_offsets, |
| 898 | e, |
| 899 | num_topk, |
| 900 | ) |
| 901 | c1 = _resize_cache(workspace13, (m * topk, w1_n)) |
| 902 | c2 = _resize_cache(workspace2, (m * topk, n)) |
| 903 | c3 = _resize_cache(workspace13, (m * topk, k)) |
| 904 | |
| 905 | ops.cutlass_mxfp4_moe_mm( |
| 906 | c1, |
| 907 | rep_a_fp4, |
| 908 | w1_fp4, |
| 909 | rep_a_blockscale, |
| 910 | w1_blockscale, |
| 911 | problem_sizes1, |
| 912 | expert_offsets[:-1], |
| 913 | blockscale_offsets[:-1], |
| 914 | ) |
| 915 | del rep_a_fp4, rep_a_blockscale |
| 916 | if activation == MoEActivation.SILU and ( |
| 917 | activation_config is None or activation_config.clamp_limit is None |
| 918 | ): |
| 919 | int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant( |
| 920 | c1, expert_offsets, blockscale_offsets, e, num_topk |
| 921 | ) |
| 922 | else: |
| 923 | apply_moe_activation( |
| 924 | activation, |
| 925 | c2, |
| 926 | c1, |
| 927 | activation_config=activation_config, |
| 928 | ) |
| 929 | int_fp4, int_blockscale = ops.mxfp4_experts_quant( |
| 930 | c2, expert_offsets, blockscale_offsets, e, num_topk |
| 931 | ) |
| 932 | |
| 933 | ops.cutlass_mxfp4_moe_mm( |
| 934 | c3, |
| 935 | int_fp4, |
| 936 | w2_fp4, |
| 937 | int_blockscale, |
| 938 | w2_blockscale, |
| 939 | problem_sizes2, |
| 940 | expert_offsets[:-1], |
| 941 | blockscale_offsets[:-1], |
| 942 | ) |
| 943 | del int_fp4, int_blockscale |
| 944 | |
| 945 | c3 = ops.shuffle_rows(c3, c_map) |
| 946 | |
| 947 | assert output.dtype == out_dtype |
| 948 | if not apply_router_weight_on_input: |
| 949 | output.copy_( |
| 950 | ( |
| 951 | c3.view(m, num_topk, k) |
| 952 | * topk_weights.view(m, num_topk, 1).to(out_dtype) |
| 953 | ).sum(dim=1), |
| 954 | non_blocking=True, |
| 955 | ) |
| 956 | else: |
| 957 | output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True) |
| 958 | return |
| 959 | |
| 960 | |
| 961 | def swizzle_mxfp4_scales( |
| 962 | scales: torch.Tensor, |
| 963 | N: int, |
| 964 | K: int, |
| 965 | ) -> torch.Tensor: |
| 966 | """Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout. |
| 967 | |
| 968 | CUTLASS expects MX scale factors in a tiled layout: |
| 969 | [numMTiles, numKTiles, 32, 4, 4] |
| 970 | where numMTiles = ceil(N/128), numKTiles = ceil(K/128), |
| 971 | and the inner dimensions correspond to the swizzle pattern: |
| 972 | mTileIdx = mIdx / 128 |
| 973 | outerMIdx = mIdx % 32 |
| 974 | innerMIdx = (mIdx / 32) % 4 |
| 975 | kTileIdx = kIdx / 4 |
| 976 | innerKIdx = kIdx % 4 |
| 977 | with kIdx = col_in_scale_space (i.e., index into K//32). |
| 978 | """ |
| 979 | assert scales.dtype == torch.uint8 |
| 980 | num_scale_cols = K // 32 # number of E8M0 scale values per row |
| 981 | |
| 982 | num_m_tiles = (N + 127) // 128 |
| 983 | num_k_tiles = (num_scale_cols + 3) // 4 |
| 984 | |
| 985 | # Pad N to multiple of 128 and scale_cols to multiple of 4 |
| 986 | padded_N = num_m_tiles * 128 |
| 987 | padded_scale_cols = num_k_tiles * 4 |
| 988 | |
| 989 | # Start with flat scales, pad if needed |
| 990 | padded = torch.zeros( |
| 991 | padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device |
| 992 | ) |
| 993 | padded[:N, :num_scale_cols] = scales |
| 994 | |
| 995 | # Reshape to tile structure: |
| 996 | # [numMTiles, 4, 32, numKTiles, 4] |
| 997 | # mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx |
| 998 | tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4) |
| 999 | # Permute to [numMTiles, numKTiles, 32, 4, 4] |
| 1000 | # (outerMIdx, innerMIdx, innerKIdx) |
| 1001 | tiled = tiled.permute(0, 3, 2, 1, 4).contiguous() |
| 1002 | return tiled.reshape(-1) |
| 1003 | |
| 1004 | |
| 1005 | class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular): |
| 1006 | """CUTLASS MXFP4 x MXFP4 fused MoE expert implementation.""" |
| 1007 | |
| 1008 | @property |
| 1009 | def expects_unquantized_inputs(self) -> bool: |
| 1010 | return True |
| 1011 | |
| 1012 | @staticmethod |
| 1013 | def _supports_current_device() -> bool: |
| 1014 | p = current_platform |
| 1015 | capability = p.get_device_capability() |
| 1016 | return ( |
| 1017 | p.is_cuda() |
| 1018 | and capability is not None |
| 1019 | and ops.mxfp4_experts_quant_supported(capability.to_int()) |
| 1020 | ) |
| 1021 | |
| 1022 | @staticmethod |
| 1023 | def _supports_no_act_and_mul() -> bool: |
| 1024 | return True |
| 1025 | |
| 1026 | @staticmethod |
| 1027 | def _supports_quant_scheme( |
| 1028 | weight_key: QuantKey | None, |
| 1029 | activation_key: QuantKey | None, |
| 1030 | ) -> bool: |
| 1031 | return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic) |
| 1032 | |
| 1033 | @staticmethod |
| 1034 | def _supports_activation(activation: MoEActivation) -> bool: |
| 1035 | return apply_moe_activation_supported(activation) |
| 1036 | |
| 1037 | @staticmethod |
| 1038 | def _supports_parallel_config( |
| 1039 | moe_parallel_config: FusedMoEParallelConfig, |
| 1040 | ) -> bool: |
| 1041 | return moe_parallel_config.ep_size == 1 |
| 1042 | |
| 1043 | @staticmethod |
| 1044 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 1045 | return mk.FusedMoEActivationFormat.Standard |
| 1046 | |
| 1047 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 1048 | return TopKWeightAndReduceNoOP() |
| 1049 | |
| 1050 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 1051 | return act_dtype |
| 1052 | |
| 1053 | def workspace_shapes( |
| 1054 | self, |
| 1055 | M: int, |
| 1056 | N: int, |
| 1057 | K: int, |
| 1058 | topk: int, |
| 1059 | global_num_experts: int, |
| 1060 | local_num_experts: int, |
| 1061 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 1062 | activation: MoEActivation, |
| 1063 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 1064 | workspace1 = (M * topk, max(2 * N, K)) |
| 1065 | workspace2 = (M * topk, N) |
| 1066 | output = (M, K) |
| 1067 | return (workspace1, workspace2, output) |
| 1068 | |
| 1069 | def apply( |
| 1070 | self, |
| 1071 | output: torch.Tensor, |
| 1072 | hidden_states: torch.Tensor, |
| 1073 | w1: torch.Tensor, |
| 1074 | w2: torch.Tensor, |
| 1075 | topk_weights: torch.Tensor, |
| 1076 | topk_ids: torch.Tensor, |
| 1077 | activation: MoEActivation, |
| 1078 | global_num_experts: int, |
| 1079 | expert_map: torch.Tensor | None, |
| 1080 | a1q_scale: torch.Tensor | None, |
| 1081 | a2_scale: torch.Tensor | None, |
| 1082 | workspace13: torch.Tensor | None, |
| 1083 | workspace2: torch.Tensor | None, |
| 1084 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 1085 | apply_router_weight_on_input: bool, |
| 1086 | ): |
| 1087 | e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids) |
| 1088 | n = w2.shape[2] * 2 |
| 1089 | |
| 1090 | run_cutlass_moe_mxfp4( |
| 1091 | output=output, |
| 1092 | a=hidden_states, |
| 1093 | w1_fp4=w1, |
| 1094 | w1_blockscale=self.w1_scale, |
| 1095 | w2_fp4=w2, |
| 1096 | w2_blockscale=self.w2_scale, |
| 1097 | topk_weights=topk_weights, |
| 1098 | topk_ids=topk_ids, |
| 1099 | activation=activation, |
| 1100 | workspace13=workspace13, |
| 1101 | workspace2=workspace2, |
| 1102 | m=m, |
| 1103 | n=n, |
| 1104 | k=k, |
| 1105 | e=e, |
| 1106 | device=hidden_states.device, |
| 1107 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1108 | activation_config=self.activation_config, |
| 1109 | ) |
| 1110 | |
| 1111 | |
| 1112 | # W4A8 |
| 1113 | def run_cutlass_moe_w4a8_fp8( |
| 1114 | output: torch.Tensor, |
| 1115 | hidden_states: torch.Tensor, |
| 1116 | w1: torch.Tensor, |
| 1117 | w2: torch.Tensor, |
| 1118 | topk_ids: torch.Tensor, |
| 1119 | activation: MoEActivation, |
| 1120 | global_num_experts: int, |
| 1121 | expert_map: torch.Tensor | None, |
| 1122 | w1_scale: torch.Tensor | None, |
| 1123 | w2_scale: torch.Tensor | None, |
| 1124 | a1q_scale: torch.Tensor | None, |
| 1125 | a2_scale: torch.Tensor | None, |
| 1126 | w1_chan_scale: torch.Tensor, |
| 1127 | w2_chan_scale: torch.Tensor, |
| 1128 | a_strides1: torch.Tensor, |
| 1129 | a_strides2: torch.Tensor, |
| 1130 | b_strides1: torch.Tensor, |
| 1131 | b_strides2: torch.Tensor, |
| 1132 | c_strides1: torch.Tensor, |
| 1133 | c_strides2: torch.Tensor, |
| 1134 | s_strides1: torch.Tensor, |
| 1135 | s_strides2: torch.Tensor, |
| 1136 | workspace13: torch.Tensor, |
| 1137 | workspace2: torch.Tensor, |
| 1138 | expert_num_tokens: torch.Tensor | None, |
| 1139 | out_dtype: torch.dtype, |
| 1140 | per_act_token: bool, |
| 1141 | per_out_ch: bool, |
| 1142 | use_batched_format: bool, |
| 1143 | topk_weights: torch.Tensor | None, |
| 1144 | group_size: int, |
| 1145 | permute_scratch: MoEPermuteScratch | None, |
| 1146 | *, |
| 1147 | activation_config: ApplyMoEActivationConfig | None = None, |
| 1148 | ): |
| 1149 | a1q = hidden_states |
| 1150 | M = a1q.size(0) |
| 1151 | local_E = w1.size(0) |
| 1152 | device = a1q.device |
| 1153 | _, K, N_packed = w2.shape |
| 1154 | N = N_packed * 8 # logical N, pack 8 int4 into 1 int32 |
| 1155 | |
| 1156 | assert per_act_token, "W4A8 must use per-token scales" |
| 1157 | assert per_out_ch, "W4A8 must use per-channel scales" |
| 1158 | assert w1_scale is not None |
| 1159 | assert w2_scale is not None |
| 1160 | assert w1_scale.dtype == torch.float8_e4m3fn |
| 1161 | assert w2_scale.dtype == torch.float8_e4m3fn |
| 1162 | assert w1.dtype == torch.int32 |
| 1163 | assert w2.dtype == torch.int32 |
| 1164 | assert w1_chan_scale.dtype == torch.float32 |
| 1165 | assert w2_chan_scale.dtype == torch.float32 |
| 1166 | assert w1.size(0) == w2.size(0), "Weights expert number mismatch" |
| 1167 | assert a1q_scale is not None |
| 1168 | assert a2_scale is None |
| 1169 | assert out_dtype in [torch.bfloat16], f"Invalid output dtype: {out_dtype}" |
| 1170 | if expert_map is not None: |
| 1171 | assert expert_num_tokens is None |
| 1172 | assert not use_batched_format, "batched format not supported yet" |
| 1173 | assert group_size == 128, f"Only group size 128 supported but got {group_size=}" |
| 1174 | |
| 1175 | assert global_num_experts != -1 |
| 1176 | assert w1.size(2) * 8 == K, ( |
| 1177 | f"w1 hidden size mismatch: got {w1.size(2) * 8}, expected {K=}" |
| 1178 | ) |
| 1179 | |
| 1180 | topk = topk_ids.size(1) |
| 1181 | a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn), (M * topk, K)) |
| 1182 | mm1_out = _resize_cache(workspace13, (M * topk, N * 2)) |
| 1183 | act_out = _resize_cache(workspace2, (M * topk, N)) |
| 1184 | # original workspace are based on input hidden_states dtype (bf16) |
| 1185 | quant_out = _resize_cache( |
| 1186 | workspace13.view(dtype=torch.float8_e4m3fn), (M * topk, N) |
| 1187 | ) |
| 1188 | mm2_out = _resize_cache(workspace2, (M * topk, K)) |
| 1189 | |
| 1190 | problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 1191 | problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device) |
| 1192 | |
| 1193 | num_expert = global_num_experts if expert_map is None else expert_map.size(0) |
| 1194 | # permuted a1q reuses workspace2 |
| 1195 | a1q, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute( |
| 1196 | a1q, |
| 1197 | a1q_scale, |
| 1198 | topk_ids, |
| 1199 | num_expert, |
| 1200 | local_E, |
| 1201 | expert_map, |
| 1202 | permuted_hidden_states=a1q_perm, |
| 1203 | scratch=permute_scratch, |
| 1204 | ) |
| 1205 | # for RS gemm SwapAB is always enabled (swap logical M, N in the problem shape). |
| 1206 | ops.get_cutlass_moe_mm_problem_sizes_from_expert_offsets( |
| 1207 | expert_first_token_offset, problem_sizes1, problem_sizes2, N, K, True |
| 1208 | ) |
| 1209 | expert_offsets = expert_first_token_offset[:-1] |
| 1210 | |
| 1211 | ops.cutlass_w4a8_moe_mm( |
| 1212 | mm1_out, |
| 1213 | a1q, |
| 1214 | w1, |
| 1215 | a1q_scale, |
| 1216 | w1_chan_scale, |
| 1217 | w1_scale, |
| 1218 | group_size, |
| 1219 | expert_offsets, |
| 1220 | problem_sizes1, |
| 1221 | a_strides1, |
| 1222 | b_strides1, |
| 1223 | c_strides1, |
| 1224 | s_strides1, |
| 1225 | ) |
| 1226 | |
| 1227 | apply_moe_activation( |
| 1228 | activation, |
| 1229 | act_out, |
| 1230 | mm1_out, |
| 1231 | activation_config=activation_config, |
| 1232 | ) |
| 1233 | |
| 1234 | a2q, a2q_scale = ops.scaled_fp8_quant( |
| 1235 | act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out |
| 1236 | ) |
| 1237 | |
| 1238 | ops.cutlass_w4a8_moe_mm( |
| 1239 | mm2_out, |
| 1240 | a2q, |
| 1241 | w2, |
| 1242 | a2q_scale, |
| 1243 | w2_chan_scale, |
| 1244 | w2_scale, |
| 1245 | group_size, |
| 1246 | expert_offsets, |
| 1247 | problem_sizes2, |
| 1248 | a_strides2, |
| 1249 | b_strides2, |
| 1250 | c_strides2, |
| 1251 | s_strides2, |
| 1252 | ) |
| 1253 | |
| 1254 | # for non-chunking mode the output is resized from workspace13 |
| 1255 | # so we need to make sure mm2_out uses workspace2. |
| 1256 | moe_unpermute( |
| 1257 | out=output, |
| 1258 | permuted_hidden_states=mm2_out, |
| 1259 | topk_weights=topk_weights, |
| 1260 | inv_permuted_idx=inv_perm, |
| 1261 | expert_first_token_offset=expert_first_token_offset, |
| 1262 | ) |
| 1263 | |
| 1264 | |
| 1265 | class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular): |
| 1266 | def __init__( |
| 1267 | self, |
| 1268 | moe_config: FusedMoEConfig, |
| 1269 | quant_config: FusedMoEQuantConfig, |
| 1270 | b_strides1: torch.Tensor, |
| 1271 | b_strides2: torch.Tensor, |
| 1272 | group_size: int, |
| 1273 | ): |
| 1274 | super().__init__(moe_config=moe_config, quant_config=quant_config) |
| 1275 | |
| 1276 | e = moe_config.num_local_experts |
| 1277 | n = moe_config.intermediate_size_per_partition |
| 1278 | k = moe_config.hidden_dim |
| 1279 | device = moe_config.device |
| 1280 | |
| 1281 | self.out_dtype = moe_config.in_dtype |
| 1282 | |
| 1283 | a_strides1_c_strides2 = torch.full((e,), k, device=device, dtype=torch.int64) |
| 1284 | self.a_strides1 = a_strides1_c_strides2 |
| 1285 | self.a_strides2 = torch.full((e,), n, device=device, dtype=torch.int64) |
| 1286 | self.c_strides1 = torch.full((e,), 2 * n, device=device, dtype=torch.int64) |
| 1287 | self.c_strides2 = a_strides1_c_strides2 |
| 1288 | |
| 1289 | self.b_strides1 = b_strides1 |
| 1290 | self.b_strides2 = b_strides2 |
| 1291 | |
| 1292 | # sizeof(StrideS) = 16 bytes, encoded as 2xint64. |
| 1293 | self.s_strides1 = torch.zeros((e, 2), device=device, dtype=torch.int64) |
| 1294 | self.s_strides1[:, 0] = 2 * n |
| 1295 | self.s_strides2 = torch.zeros((e, 2), device=device, dtype=torch.int64) |
| 1296 | self.s_strides2[:, 0] = k |
| 1297 | |
| 1298 | self.group_size = group_size |
| 1299 | self._permute_scratch: MoEPermuteScratch | None = None |
| 1300 | |
| 1301 | @staticmethod |
| 1302 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 1303 | return mk.FusedMoEActivationFormat.Standard |
| 1304 | |
| 1305 | @staticmethod |
| 1306 | def is_supported_config( |
| 1307 | cls: type[mk.FusedMoEExperts], |
| 1308 | moe_config: FusedMoEConfig, |
| 1309 | weight_key: QuantKey | None, |
| 1310 | activation_key: QuantKey | None, |
| 1311 | activation_format: mk.FusedMoEActivationFormat, |
| 1312 | ) -> tuple[bool, str | None]: |
| 1313 | if moe_config.in_dtype != torch.bfloat16: |
| 1314 | return ( |
| 1315 | False, |
| 1316 | f"kernel does not support {moe_config.in_dtype} input/output dtype", |
| 1317 | ) |
| 1318 | |
| 1319 | return mk.FusedMoEExperts.is_supported_config( |
| 1320 | cls, |
| 1321 | moe_config, |
| 1322 | weight_key, |
| 1323 | activation_key, |
| 1324 | activation_format, |
| 1325 | ) |
| 1326 | |
| 1327 | @staticmethod |
| 1328 | def _supports_current_device() -> bool: |
| 1329 | return ( |
| 1330 | current_platform.is_device_capability(90) and cutlass_group_gemm_supported() |
| 1331 | ) |
| 1332 | |
| 1333 | @staticmethod |
| 1334 | def _supports_no_act_and_mul() -> bool: |
| 1335 | return False |
| 1336 | |
| 1337 | @staticmethod |
| 1338 | def _supports_quant_scheme( |
| 1339 | weight_key: QuantKey | None, |
| 1340 | activation_key: QuantKey | None, |
| 1341 | ) -> bool: |
| 1342 | return (weight_key, activation_key) == (kInt4Static, kFp8DynamicTokenSym) |
| 1343 | |
| 1344 | @staticmethod |
| 1345 | def _supports_activation(activation: MoEActivation) -> bool: |
| 1346 | return activation.is_gated and apply_moe_activation_supported(activation) |
| 1347 | |
| 1348 | @staticmethod |
| 1349 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 1350 | return True |
| 1351 | |
| 1352 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 1353 | # topk weights and reduction are fused in moe_unpermute cuda kernel |
| 1354 | return TopKWeightAndReduceNoOP() |
| 1355 | |
| 1356 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 1357 | return self.out_dtype if self.out_dtype is not None else act_dtype |
| 1358 | |
| 1359 | def _get_permute_scratch(self) -> MoEPermuteScratch | None: |
| 1360 | if self._permute_scratch is None and moe_permute_unpermute_supported(): |
| 1361 | self._permute_scratch = MoEPermuteScratch( |
| 1362 | max_num_tokens=self.moe_config.max_num_tokens, |
| 1363 | topk=self.moe_config.experts_per_token, |
| 1364 | num_experts=self.moe_config.num_experts, |
| 1365 | num_local_experts=self.moe_config.num_local_experts, |
| 1366 | device=torch.device(self.moe_config.device), |
| 1367 | ) |
| 1368 | return self._permute_scratch |
| 1369 | |
| 1370 | def workspace_shapes( |
| 1371 | self, |
| 1372 | M: int, |
| 1373 | N: int, |
| 1374 | K: int, |
| 1375 | topk: int, |
| 1376 | global_num_experts: int, |
| 1377 | local_num_experts: int, |
| 1378 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 1379 | activation: MoEActivation, |
| 1380 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 1381 | activation_out_dim = self.adjust_N_for_activation(N, activation) |
| 1382 | workspace1 = (M * topk, max(N, K)) |
| 1383 | workspace2 = (M * topk, max(activation_out_dim, K)) |
| 1384 | output = (M, K) |
| 1385 | return (workspace1, workspace2, output) |
| 1386 | |
| 1387 | def apply( |
| 1388 | self, |
| 1389 | output: torch.Tensor, |
| 1390 | hidden_states: torch.Tensor, |
| 1391 | w1: torch.Tensor, |
| 1392 | w2: torch.Tensor, |
| 1393 | topk_weights: torch.Tensor, |
| 1394 | topk_ids: torch.Tensor, |
| 1395 | activation: MoEActivation, |
| 1396 | global_num_experts: int, |
| 1397 | expert_map: torch.Tensor | None, |
| 1398 | a1q_scale: torch.Tensor | None, |
| 1399 | a2_scale: torch.Tensor | None, |
| 1400 | workspace13: torch.Tensor | None, |
| 1401 | workspace2: torch.Tensor | None, |
| 1402 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 1403 | apply_router_weight_on_input: bool, |
| 1404 | ): |
| 1405 | assert self.w1_zp is None, "w1_zp is not supported in CUTLASS MoE" |
| 1406 | assert self.w2_zp is None, "w2_zp is not supported in CUTLASS MoE" |
| 1407 | |
| 1408 | expert_num_tokens = None |
| 1409 | |
| 1410 | use_batched_format = ( |
| 1411 | self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts |
| 1412 | ) |
| 1413 | assert not use_batched_format, "batched format not supported" |
| 1414 | |
| 1415 | in_dtype = hidden_states.dtype |
| 1416 | |
| 1417 | run_cutlass_moe_w4a8_fp8( |
| 1418 | output, |
| 1419 | hidden_states, |
| 1420 | w1, |
| 1421 | w2, |
| 1422 | topk_ids, |
| 1423 | activation, |
| 1424 | global_num_experts, |
| 1425 | expert_map, |
| 1426 | self.w1_scale, |
| 1427 | self.w2_scale, |
| 1428 | a1q_scale, |
| 1429 | a2_scale, |
| 1430 | self.g1_alphas, # per-channel scales |
| 1431 | self.g2_alphas, # per-channel scales |
| 1432 | self.a_strides1, |
| 1433 | self.a_strides2, |
| 1434 | self.b_strides1, |
| 1435 | self.b_strides2, |
| 1436 | self.c_strides1, |
| 1437 | self.c_strides2, |
| 1438 | self.s_strides1, |
| 1439 | self.s_strides2, |
| 1440 | workspace13, |
| 1441 | workspace2, |
| 1442 | expert_num_tokens, |
| 1443 | self.out_dtype if self.out_dtype is not None else in_dtype, |
| 1444 | self.per_act_token_quant, |
| 1445 | self.per_out_ch_quant, |
| 1446 | use_batched_format, |
| 1447 | topk_weights, |
| 1448 | self.group_size, |
| 1449 | self._get_permute_scratch(), |
| 1450 | activation_config=self.activation_config, |
| 1451 | ) |