| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused MoE utilities for GPTQ.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Callable |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | import vllm._custom_ops as ops |
| 11 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 12 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 13 | ApplyMoEActivationConfig, |
| 14 | MoEActivation, |
| 15 | apply_moe_activation, |
| 16 | apply_moe_activation_masked_supported, |
| 17 | apply_moe_activation_supported, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.config import ( |
| 20 | FusedMoEConfig, |
| 21 | FusedMoEParallelConfig, |
| 22 | FusedMoEQuantConfig, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( |
| 25 | LoRAExpertsMixin, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( |
| 28 | batched_moe_align_block_size, |
| 29 | moe_align_block_size, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 32 | TopKWeightAndReduceDelegate, |
| 33 | TopKWeightAndReduceNoOP, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.fused_moe.utils import _resize_cache |
| 36 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 37 | get_marlin_input_dtype, |
| 38 | get_marlin_workspace, |
| 39 | marlin_moe_intermediate_size, |
| 40 | marlin_quant_input, |
| 41 | ) |
| 42 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 43 | QuantKey, |
| 44 | kFp8Static128BlockSym, |
| 45 | kFp8StaticChannelSym, |
| 46 | kFp8StaticTensorSym, |
| 47 | kInt4Static, |
| 48 | kInt4Static32, |
| 49 | kInt4Static32Asym, |
| 50 | kInt4StaticAsym, |
| 51 | kInt8Static, |
| 52 | kMxfp4Static, |
| 53 | kMxfp8Static, |
| 54 | kNvfp4Static, |
| 55 | ) |
| 56 | from vllm.platforms import current_platform |
| 57 | from vllm.scalar_type import ScalarType, scalar_types |
| 58 | |
| 59 | |
| 60 | def _fused_marlin_moe( |
| 61 | hidden_states: torch.Tensor, |
| 62 | w1: torch.Tensor, |
| 63 | w2: torch.Tensor, |
| 64 | bias1: torch.Tensor | None, |
| 65 | bias2: torch.Tensor | None, |
| 66 | w1_scale: torch.Tensor, |
| 67 | w2_scale: torch.Tensor, |
| 68 | topk_weights: torch.Tensor, |
| 69 | num_topk: int, |
| 70 | quant_type: ScalarType, |
| 71 | apply_router_weight_on_input: bool, |
| 72 | expert_map: torch.Tensor | None, |
| 73 | block_size_m: int, |
| 74 | sorted_token_ids: torch.Tensor, |
| 75 | expert_ids: torch.Tensor, |
| 76 | num_tokens_post_padded: torch.Tensor, |
| 77 | activation: MoEActivation = MoEActivation.SILU, |
| 78 | activation_func: Callable[..., None] | None = None, |
| 79 | topk_ids: torch.Tensor | None = None, |
| 80 | input_global_scale1: torch.Tensor | None = None, |
| 81 | input_global_scale2: torch.Tensor | None = None, |
| 82 | global_scale1: torch.Tensor | None = None, |
| 83 | global_scale2: torch.Tensor | None = None, |
| 84 | w1_zeros: torch.Tensor | None = None, |
| 85 | w2_zeros: torch.Tensor | None = None, |
| 86 | workspace: torch.Tensor | None = None, |
| 87 | intermediate_cache13: torch.Tensor | None = None, |
| 88 | intermediate_cache2: torch.Tensor | None = None, |
| 89 | output: torch.Tensor | None = None, |
| 90 | input_dtype: torch.dtype | None = None, |
| 91 | activation_config: ApplyMoEActivationConfig | None = None, |
| 92 | ) -> torch.Tensor: |
| 93 | assert hidden_states.ndim == 2 |
| 94 | M, K = hidden_states.size() |
| 95 | N = marlin_moe_intermediate_size(w1, w2) |
| 96 | w13_num_shards = 2 if activation.is_gated else 1 |
| 97 | if workspace is None: |
| 98 | workspace = get_marlin_workspace(hidden_states.device) |
| 99 | |
| 100 | if intermediate_cache13 is None: |
| 101 | intermediate_cache13 = torch.empty( |
| 102 | (M * num_topk * max(w13_num_shards * N, K),), |
| 103 | device=hidden_states.device, |
| 104 | dtype=hidden_states.dtype, |
| 105 | ) |
| 106 | |
| 107 | if intermediate_cache2 is None: |
| 108 | intermediate_cache2 = torch.empty( |
| 109 | (M * num_topk, N), |
| 110 | device=hidden_states.device, |
| 111 | dtype=hidden_states.dtype, |
| 112 | ) |
| 113 | |
| 114 | intermediate_cache1 = _resize_cache( |
| 115 | intermediate_cache13, (M * num_topk, w13_num_shards * N) |
| 116 | ) |
| 117 | |
| 118 | intermediate_cache3 = _resize_cache(intermediate_cache13, (M * num_topk, K)) |
| 119 | |
| 120 | intermediate_cache2 = _resize_cache(intermediate_cache2, (M * num_topk, N)) |
| 121 | |
| 122 | a_scales1 = None |
| 123 | gate_up_input = hidden_states |
| 124 | if input_dtype == torch.int8: |
| 125 | gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype) |
| 126 | if input_global_scale1 is not None: |
| 127 | a_scales1 = a_scales1 * input_global_scale1 |
| 128 | elif input_dtype == torch.float8_e4m3fn: |
| 129 | gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype) |
| 130 | |
| 131 | intermediate_cache1 = ops.moe_wna16_marlin_gemm( |
| 132 | gate_up_input, |
| 133 | intermediate_cache1, |
| 134 | w1, |
| 135 | bias1, |
| 136 | w1_scale, |
| 137 | a_scales1, |
| 138 | global_scale1, |
| 139 | w1_zeros, |
| 140 | workspace, |
| 141 | sorted_token_ids, |
| 142 | expert_ids, |
| 143 | num_tokens_post_padded, |
| 144 | topk_weights, |
| 145 | moe_block_size=block_size_m, |
| 146 | top_k=num_topk, |
| 147 | mul_topk_weights=apply_router_weight_on_input, |
| 148 | b_q_type=quant_type, |
| 149 | size_m=M, |
| 150 | size_n=w13_num_shards * N, |
| 151 | size_k=K, |
| 152 | use_atomic_add=False, |
| 153 | use_fp32_reduce=True, |
| 154 | is_zp_float=False, |
| 155 | ) |
| 156 | activation_input = intermediate_cache1.view(-1, w13_num_shards * N) |
| 157 | if activation_func is None: |
| 158 | config = ( |
| 159 | ApplyMoEActivationConfig() |
| 160 | if activation_config is None |
| 161 | else activation_config |
| 162 | ) |
| 163 | apply_moe_activation( |
| 164 | activation, |
| 165 | intermediate_cache2, |
| 166 | activation_input, |
| 167 | activation_config=config, |
| 168 | topk_ids=topk_ids, |
| 169 | expert_map=expert_map, |
| 170 | ) |
| 171 | else: |
| 172 | activation_func( |
| 173 | activation, |
| 174 | intermediate_cache2, |
| 175 | activation_input, |
| 176 | topk_ids=topk_ids, |
| 177 | expert_map=expert_map, |
| 178 | ) |
| 179 | |
| 180 | if output is None: |
| 181 | output = intermediate_cache3 |
| 182 | |
| 183 | a_scales2 = None |
| 184 | if input_dtype == torch.int8: |
| 185 | intermediate_cache2, a_scales2 = marlin_quant_input( |
| 186 | intermediate_cache2, input_dtype |
| 187 | ) |
| 188 | if input_global_scale2 is not None: |
| 189 | a_scales2 = a_scales2 * input_global_scale2 |
| 190 | elif input_dtype == torch.float8_e4m3fn: |
| 191 | intermediate_cache2, a_scales2 = marlin_quant_input( |
| 192 | intermediate_cache2, input_dtype |
| 193 | ) |
| 194 | |
| 195 | output = ops.moe_wna16_marlin_gemm( |
| 196 | intermediate_cache2, |
| 197 | output, |
| 198 | w2, |
| 199 | bias2, |
| 200 | w2_scale, |
| 201 | a_scales2, |
| 202 | global_scale2, |
| 203 | w2_zeros, |
| 204 | workspace, |
| 205 | sorted_token_ids, |
| 206 | expert_ids, |
| 207 | num_tokens_post_padded, |
| 208 | topk_weights, |
| 209 | moe_block_size=block_size_m, |
| 210 | top_k=1, |
| 211 | mul_topk_weights=not apply_router_weight_on_input, |
| 212 | b_q_type=quant_type, |
| 213 | size_m=M * num_topk, |
| 214 | size_n=K, |
| 215 | size_k=N, |
| 216 | use_atomic_add=False, |
| 217 | use_fp32_reduce=True, |
| 218 | is_zp_float=False, |
| 219 | ) |
| 220 | |
| 221 | return output |
| 222 | |
| 223 | |
| 224 | def fused_marlin_moe( |
| 225 | hidden_states: torch.Tensor, |
| 226 | w1: torch.Tensor, |
| 227 | w2: torch.Tensor, |
| 228 | bias1: torch.Tensor | None, |
| 229 | bias2: torch.Tensor | None, |
| 230 | w1_scale: torch.Tensor, |
| 231 | w2_scale: torch.Tensor, |
| 232 | topk_weights: torch.Tensor, |
| 233 | topk_ids: torch.Tensor, |
| 234 | quant_type_id: int, |
| 235 | apply_router_weight_on_input: bool = False, |
| 236 | global_num_experts: int = -1, |
| 237 | activation: MoEActivation = MoEActivation.SILU, |
| 238 | activation_func: Callable[..., None] | None = None, |
| 239 | moe_sum: Callable[..., torch.Tensor | None] | None = None, |
| 240 | expert_map: torch.Tensor | None = None, |
| 241 | input_global_scale1: torch.Tensor | None = None, |
| 242 | input_global_scale2: torch.Tensor | None = None, |
| 243 | global_scale1: torch.Tensor | None = None, |
| 244 | global_scale2: torch.Tensor | None = None, |
| 245 | w1_zeros: torch.Tensor | None = None, |
| 246 | w2_zeros: torch.Tensor | None = None, |
| 247 | workspace: torch.Tensor | None = None, |
| 248 | intermediate_cache13: torch.Tensor | None = None, |
| 249 | intermediate_cache2: torch.Tensor | None = None, |
| 250 | output: torch.Tensor | None = None, |
| 251 | input_dtype: torch.dtype | None = None, |
| 252 | activation_config: ApplyMoEActivationConfig | None = None, |
| 253 | ) -> torch.Tensor: |
| 254 | """This function computes a Mixture of Experts (MoE) layer using two sets of |
| 255 | weights, w1 and w2, and top-k gating mechanism. |
| 256 | |
| 257 | Args: |
| 258 | hidden_states (torch.Tensor): The input tensor to the MoE layer. |
| 259 | w1 (torch.Tensor): The first set of expert weights. |
| 260 | w2 (torch.Tensor): The second set of expert weights. |
| 261 | bias1 (torch.Tensor|None): Optional bias for the first GEMM. |
| 262 | bias2 (torch.Tensor|None): Optional bias for the second GEMM. |
| 263 | w1_scale (torch.Tensor): Scale to be used for w1. |
| 264 | w2_scale (torch.Tensor): Scale to be used for w2. |
| 265 | topk_weights (torch.Tensor): Top-k weights. |
| 266 | topk_ids (torch.Tensor): Indices of topk-k elements. |
| 267 | quant_type_id (int): Id of the `ScalarType` the expert weights are |
| 268 | quantized to. Determines the 4- or 8-bit Marlin kernel path. |
| 269 | apply_router_weight_on_input (bool): Apply the topk weights to the |
| 270 | input rather than the output. Only valid when topk is 1. |
| 271 | global_num_experts (int): Total number of experts across all expert |
| 272 | parallel shards. -1 means the local expert count. |
| 273 | activation (MoEActivation): Activation applied between the two GEMMs. |
| 274 | activation_func (Callable|None): Override for the activation |
| 275 | implementation. |
| 276 | moe_sum (Callable|None): Override for the final reduction over the |
| 277 | topk dimension. Defaults to `ops.moe_sum` / `torch.sum`. |
| 278 | expert_map (torch.Tensor|None): Maps global expert indices to the |
| 279 | local expert space of this expert parallel shard. |
| 280 | input_global_scale1 (torch.Tensor|None): NVFP4 global input scale for |
| 281 | the first GEMM. |
| 282 | input_global_scale2 (torch.Tensor|None): NVFP4 global input scale for |
| 283 | the second GEMM. |
| 284 | global_scale1 (torch.Tensor|None): NVFP4 global weight scale for w1. |
| 285 | global_scale2 (torch.Tensor|None): NVFP4 global weight scale for w2. |
| 286 | w1_zeros (torch.Tensor|None): Optional zero points to be used for w1. |
| 287 | w2_zeros (torch.Tensor|None): Optional zero points to be used for w2. |
| 288 | workspace (torch.Tensor|None): Scratch buffer for the Marlin kernels. |
| 289 | intermediate_cache13 (torch.Tensor|None): Scratch buffer for the |
| 290 | first and third intermediate activations. |
| 291 | intermediate_cache2 (torch.Tensor|None): Scratch buffer for the |
| 292 | second intermediate activation. |
| 293 | output (torch.Tensor|None): Tensor to write the result into. A new |
| 294 | tensor is allocated when omitted. |
| 295 | input_dtype (torch.dtype|None): dtype the activations are quantized |
| 296 | to. A 1-byte dtype raises the M block size to at least 16. |
| 297 | activation_config (ApplyMoEActivationConfig|None): Extra |
| 298 | configuration forwarded to the activation. |
| 299 | |
| 300 | Returns: |
| 301 | torch.Tensor: The output tensor after applying the MoE layer. |
| 302 | |
| 303 | """ |
| 304 | quant_type = ScalarType.from_id(quant_type_id) |
| 305 | assert quant_type in [ |
| 306 | scalar_types.uint4, |
| 307 | scalar_types.uint8b128, |
| 308 | scalar_types.uint4b8, |
| 309 | scalar_types.float8_e4m3fn, |
| 310 | scalar_types.float4_e2m1f, |
| 311 | ] |
| 312 | |
| 313 | bit4_scalar_types = [ |
| 314 | scalar_types.uint4, |
| 315 | scalar_types.uint4b8, |
| 316 | scalar_types.float4_e2m1f, |
| 317 | ] |
| 318 | num_bits = 4 if quant_type in bit4_scalar_types else 8 |
| 319 | |
| 320 | M, K = hidden_states.size() |
| 321 | E = w1.size(0) |
| 322 | topk = topk_ids.size(1) |
| 323 | |
| 324 | # Check constraints. |
| 325 | assert w1.size(1) * 16 == K, "Hidden size mismatch w1" |
| 326 | assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2" |
| 327 | assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" |
| 328 | assert w1.is_contiguous(), "Expert weights1 must be contiguous" |
| 329 | assert w2.is_contiguous(), "Expert weights2 must be contiguous" |
| 330 | assert hidden_states.dtype in [torch.float16, torch.bfloat16] |
| 331 | assert num_bits in [4, 8] |
| 332 | assert topk_weights.dtype == torch.float32 |
| 333 | |
| 334 | if global_num_experts == -1: |
| 335 | global_num_experts = E |
| 336 | else: |
| 337 | # Set M to estimated valid tokens per rank |
| 338 | M = math.ceil(M * E / global_num_experts) |
| 339 | |
| 340 | # M block size selection logic |
| 341 | # TODO: tune this further for specific models |
| 342 | for block_size_m in [8, 16, 32, 48, 64]: |
| 343 | if M * topk / E / block_size_m < 0.9: |
| 344 | break |
| 345 | |
| 346 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 347 | block_size_m = max(block_size_m, 16) |
| 348 | |
| 349 | sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( |
| 350 | topk_ids, |
| 351 | block_size_m, |
| 352 | global_num_experts, |
| 353 | expert_map, |
| 354 | ignore_invalid_experts=True, |
| 355 | ) |
| 356 | |
| 357 | assert activation is not None |
| 358 | moe_output = _fused_marlin_moe( |
| 359 | hidden_states=hidden_states, |
| 360 | w1=w1, |
| 361 | w2=w2, |
| 362 | bias1=bias1, |
| 363 | bias2=bias2, |
| 364 | w1_scale=w1_scale, |
| 365 | w2_scale=w2_scale, |
| 366 | topk_weights=topk_weights, |
| 367 | topk_ids=topk_ids, |
| 368 | num_topk=topk, |
| 369 | quant_type=quant_type, |
| 370 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 371 | expert_map=expert_map, |
| 372 | block_size_m=block_size_m, |
| 373 | sorted_token_ids=sorted_token_ids, |
| 374 | expert_ids=expert_ids, |
| 375 | num_tokens_post_padded=num_tokens_post_padded, |
| 376 | activation=activation, |
| 377 | activation_func=activation_func, |
| 378 | activation_config=activation_config, |
| 379 | input_global_scale1=input_global_scale1, |
| 380 | input_global_scale2=input_global_scale2, |
| 381 | global_scale1=global_scale1, |
| 382 | global_scale2=global_scale2, |
| 383 | w1_zeros=w1_zeros, |
| 384 | w2_zeros=w2_zeros, |
| 385 | workspace=workspace, |
| 386 | intermediate_cache13=intermediate_cache13, |
| 387 | intermediate_cache2=intermediate_cache2, |
| 388 | output=None, |
| 389 | input_dtype=input_dtype, |
| 390 | ).view(-1, topk, K) |
| 391 | |
| 392 | if output is None: |
| 393 | output = torch.empty_like(hidden_states) |
| 394 | |
| 395 | if moe_sum is None: |
| 396 | if expert_map is not None: |
| 397 | ops.moe_sum(moe_output, output, topk_ids, expert_map) |
| 398 | return output |
| 399 | return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output) |
| 400 | else: |
| 401 | return moe_sum(moe_output, output, topk_ids, expert_map) |
| 402 | |
| 403 | |
| 404 | def batched_fused_marlin_moe( |
| 405 | hidden_states: torch.Tensor, |
| 406 | expert_num_tokens: torch.Tensor, |
| 407 | w1: torch.Tensor, |
| 408 | w2: torch.Tensor, |
| 409 | bias1: torch.Tensor | None, |
| 410 | bias2: torch.Tensor | None, |
| 411 | w1_scale: torch.Tensor, |
| 412 | w2_scale: torch.Tensor, |
| 413 | quant_type_id: int, |
| 414 | apply_router_weight_on_input: bool = False, |
| 415 | global_num_experts: int = -1, |
| 416 | activation: MoEActivation = MoEActivation.SILU, |
| 417 | expert_map: torch.Tensor | None = None, |
| 418 | input_global_scale1: torch.Tensor | None = None, |
| 419 | input_global_scale2: torch.Tensor | None = None, |
| 420 | global_scale1: torch.Tensor | None = None, |
| 421 | global_scale2: torch.Tensor | None = None, |
| 422 | w1_zeros: torch.Tensor | None = None, |
| 423 | w2_zeros: torch.Tensor | None = None, |
| 424 | workspace: torch.Tensor | None = None, |
| 425 | intermediate_cache13: torch.Tensor | None = None, |
| 426 | intermediate_cache2: torch.Tensor | None = None, |
| 427 | output: torch.Tensor | None = None, |
| 428 | input_dtype: torch.dtype | None = None, |
| 429 | activation_func: Callable[..., None] | None = None, |
| 430 | activation_config: ApplyMoEActivationConfig | None = None, |
| 431 | ) -> torch.Tensor: |
| 432 | """This function massages the inputs so the batched hidden_states can be |
| 433 | presented as a 2D contiguous tensor that could be used with |
| 434 | _fused_marlin_moe. |
| 435 | |
| 436 | Note that both batched_fused_marlin_moe and fused_marlin_moe ultimately |
| 437 | use `ops.moe_wna16_marlin_gemm` for the gemm operation and |
| 438 | `ops.moe_mna16_marlin_gemm` supports only 2D contiguous hidden_states. |
| 439 | Note that the moe_align_block_size function indicates, |
| 440 | - What rows of the A matrix (hidden_states) to access during the |
| 441 | matmul, via sorted_ids output. |
| 442 | - What expert_id to use for each block matmul, via expert_ids output. |
| 443 | |
| 444 | In the batched version, the tokens are already grouped/batched by experts |
| 445 | they subscribe to. Due to this, we can represent the batched hidden_states |
| 446 | tensor of shape [B, MAX_TOKENS_PER_BATCH, K] as a 2D tensor of shape, |
| 447 | [B * MAX_TOKENS_PER_BATCH, K]. We may treat this a 2D contiguous tensor |
| 448 | with topk=1 as each token (row in the tensor) subscribes to exactly one |
| 449 | expert_id (which is the batch_id). With the expert_num_tokens tensor, that |
| 450 | indicates how many tokens are actually valid in each batch, the |
| 451 | batched_moe_align_block_size function constructs the sorted_ids and |
| 452 | expert_ids tensors, so only relevant/valid rows of A (hidden_states) |
| 453 | are accessed and are processed with the correct expert_ids. |
| 454 | """ |
| 455 | assert hidden_states.ndim == 3, ( |
| 456 | f"hidden states must be batched. e.g. [B, MAX_TOKENS, K]." |
| 457 | f"But got {hidden_states.size()}" |
| 458 | ) |
| 459 | |
| 460 | quant_type = ScalarType.from_id(quant_type_id) |
| 461 | assert quant_type in [ |
| 462 | scalar_types.uint4, |
| 463 | scalar_types.uint8b128, |
| 464 | scalar_types.uint4b8, |
| 465 | scalar_types.float8_e4m3fn, |
| 466 | scalar_types.float4_e2m1f, |
| 467 | ] |
| 468 | |
| 469 | bit4_scalar_types = [ |
| 470 | scalar_types.uint4, |
| 471 | scalar_types.uint4b8, |
| 472 | scalar_types.float4_e2m1f, |
| 473 | ] |
| 474 | num_bits = 4 if quant_type in bit4_scalar_types else 8 |
| 475 | |
| 476 | B, BATCH_TOKENS_MAX, K = hidden_states.size() |
| 477 | M = hidden_states.view(-1, K).size(0) |
| 478 | E = w1.size(0) |
| 479 | |
| 480 | # Check constraints. |
| 481 | assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" |
| 482 | assert hidden_states.dtype in [torch.float16, torch.bfloat16] |
| 483 | assert expert_num_tokens.size(0) == E |
| 484 | assert B == E, ( |
| 485 | "Batch must be as big as number of experts as the tokens" |
| 486 | "are sorted into the batch/expert they belong to" |
| 487 | ) |
| 488 | assert w1.size(1) * 16 == K, "Hidden size mismatch w1" |
| 489 | assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2" |
| 490 | assert w1.is_contiguous(), "Expert weights1 must be contiguous" |
| 491 | assert w2.is_contiguous(), "Expert weights2 must be contiguous" |
| 492 | assert num_bits in [4, 8] |
| 493 | |
| 494 | # Technically, the tokens are already separated by their expert ids. |
| 495 | # Hidden-States can just be squeezed to have just 2 dimensions, |
| 496 | # [B * MAX_TOKENS, K] and top_k can be interpreted as just 1. |
| 497 | topk = 1 |
| 498 | |
| 499 | # TODO(varun) : Choose a decent block size like in fused_marlin_moe |
| 500 | # Tune block_size_m based on expert capacity to reduce padding overhead. |
| 501 | block_size_m = 64 |
| 502 | for b_m in [8, 16, 32, 48, 64]: |
| 503 | if BATCH_TOKENS_MAX / b_m < 0.9: |
| 504 | block_size_m = b_m |
| 505 | break |
| 506 | |
| 507 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 508 | block_size_m = max(block_size_m, 16) |
| 509 | |
| 510 | sorted_token_ids, expert_ids, num_tokens_post_padded = batched_moe_align_block_size( |
| 511 | max_tokens_per_batch=BATCH_TOKENS_MAX, |
| 512 | block_size=block_size_m, |
| 513 | expert_num_tokens=expert_num_tokens, |
| 514 | ) |
| 515 | |
| 516 | # TODO (varun): This can be avoided by plumbing the marlin kernel to |
| 517 | # ignore topk_weights when topk_weights_ptr is a nullptr. |
| 518 | topk_weights = torch.ones( |
| 519 | (M, topk), device=hidden_states.device, dtype=torch.float32 |
| 520 | ) |
| 521 | |
| 522 | assert activation is not None |
| 523 | output = _fused_marlin_moe( |
| 524 | hidden_states=hidden_states.view(-1, K), |
| 525 | w1=w1, |
| 526 | w2=w2, |
| 527 | bias1=bias1, |
| 528 | bias2=bias2, |
| 529 | w1_scale=w1_scale, |
| 530 | w2_scale=w2_scale, |
| 531 | topk_weights=topk_weights, |
| 532 | num_topk=topk, |
| 533 | quant_type=quant_type, |
| 534 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 535 | activation=activation, |
| 536 | activation_func=activation_func, |
| 537 | activation_config=activation_config, |
| 538 | expert_map=expert_map, |
| 539 | block_size_m=block_size_m, |
| 540 | sorted_token_ids=sorted_token_ids, |
| 541 | expert_ids=expert_ids, |
| 542 | num_tokens_post_padded=num_tokens_post_padded, |
| 543 | input_global_scale1=input_global_scale1, |
| 544 | input_global_scale2=input_global_scale2, |
| 545 | global_scale1=global_scale1, |
| 546 | global_scale2=global_scale2, |
| 547 | w1_zeros=w1_zeros, |
| 548 | w2_zeros=w2_zeros, |
| 549 | workspace=workspace, |
| 550 | intermediate_cache13=intermediate_cache13, |
| 551 | intermediate_cache2=intermediate_cache2, |
| 552 | output=output.view(-1, K) if output is not None else output, |
| 553 | input_dtype=input_dtype, |
| 554 | ) |
| 555 | |
| 556 | output = output.view(B, BATCH_TOKENS_MAX, K) |
| 557 | |
| 558 | return output |
| 559 | |
| 560 | |
| 561 | class MarlinExpertsBase(mk.FusedMoEExpertsModular): |
| 562 | def __init__( |
| 563 | self, |
| 564 | moe_config: FusedMoEConfig, |
| 565 | quant_config: FusedMoEQuantConfig, |
| 566 | max_num_tokens: int | None = None, |
| 567 | num_dispatchers: int | None = None, |
| 568 | ): |
| 569 | # TODO (varun) : Enable activation quantization |
| 570 | assert ( |
| 571 | quant_config.use_mxfp4_w4a16 |
| 572 | or quant_config.use_nvfp4_w4a16 |
| 573 | or quant_config.use_int4_w4a16 |
| 574 | or quant_config.use_int8_w8a16 |
| 575 | or quant_config.use_fp8_w8a16 |
| 576 | ), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16" |
| 577 | self.input_dtype = get_marlin_input_dtype() |
| 578 | |
| 579 | super().__init__( |
| 580 | moe_config=moe_config, |
| 581 | quant_config=quant_config, |
| 582 | max_num_tokens=max_num_tokens, |
| 583 | num_dispatchers=num_dispatchers, |
| 584 | ) |
| 585 | |
| 586 | @staticmethod |
| 587 | def _supports_current_device() -> bool: |
| 588 | p = current_platform |
| 589 | return p.is_cuda() and p.has_device_capability((7, 5)) |
| 590 | |
| 591 | @staticmethod |
| 592 | def _supports_no_act_and_mul() -> bool: |
| 593 | return True |
| 594 | |
| 595 | @staticmethod |
| 596 | def _supports_quant_scheme( |
| 597 | weight_key: QuantKey | None, |
| 598 | activation_key: QuantKey | None, |
| 599 | ) -> bool: |
| 600 | # TODO(rob): add int4, int8 as integrations |
| 601 | # are migrated to use the oracle one-by-one. |
| 602 | SUPPORTED_W = [ |
| 603 | kFp8Static128BlockSym, |
| 604 | kFp8StaticChannelSym, |
| 605 | kFp8StaticTensorSym, |
| 606 | kMxfp4Static, |
| 607 | kMxfp8Static, |
| 608 | kNvfp4Static, |
| 609 | kInt4Static, |
| 610 | kInt8Static, |
| 611 | kInt4Static32, |
| 612 | kInt4StaticAsym, |
| 613 | kInt4Static32Asym, |
| 614 | ] |
| 615 | return weight_key in SUPPORTED_W |
| 616 | |
| 617 | @staticmethod |
| 618 | def _supports_activation(activation: MoEActivation) -> bool: |
| 619 | # Marlin uses apply_moe_activation() callback for activation, |
| 620 | # so any activation supported there can be used here. |
| 621 | return apply_moe_activation_supported(activation) |
| 622 | |
| 623 | @staticmethod |
| 624 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 625 | # One-sided FI-NVL all2all pairs with MarlinExperts fine (the |
| 626 | # compressed-tensors MXFP4 path runs this exact combo); only the |
| 627 | # two-sided kernels are unsupported here. |
| 628 | return not moe_parallel_config.use_fi_nvl_two_sided_kernels |
| 629 | |
| 630 | @property |
| 631 | def quant_type_id(self) -> int: |
| 632 | if self.quant_config.use_int4_w4a16: |
| 633 | if self.w1_zp is not None or self.w2_zp is not None: |
| 634 | return scalar_types.uint4.id |
| 635 | return scalar_types.uint4b8.id |
| 636 | elif self.quant_config.use_int8_w8a16: |
| 637 | return scalar_types.uint8b128.id |
| 638 | elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16: |
| 639 | return scalar_types.float4_e2m1f.id |
| 640 | elif ( |
| 641 | self.quant_config.use_fp8_w8a16 |
| 642 | and current_platform.fp8_dtype() == torch.float8_e4m3fn |
| 643 | ): |
| 644 | return scalar_types.float8_e4m3fn.id |
| 645 | else: |
| 646 | raise NotImplementedError("Unsupported quantization type.") |
| 647 | |
| 648 | def moe_problem_size( |
| 649 | self, |
| 650 | a1: torch.Tensor, |
| 651 | w1: torch.Tensor, |
| 652 | w2: torch.Tensor, |
| 653 | topk_ids: torch.Tensor, |
| 654 | ) -> tuple[int, int, int, int, int]: |
| 655 | assert w1.dim() == 3 and w2.dim() == 3 |
| 656 | |
| 657 | E = w1.size(0) |
| 658 | K = a1.size(-1) |
| 659 | N = marlin_moe_intermediate_size(w1, w2) |
| 660 | |
| 661 | if a1.dim() == 2: |
| 662 | # Make sure we are using the correct a1 (pre-permute). |
| 663 | assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" |
| 664 | M = a1.size(0) |
| 665 | else: |
| 666 | assert a1.dim() == 3 |
| 667 | assert a1.size(0) == E, f"{a1.size(0)} == {E}" |
| 668 | M = a1.size(1) # This is max_num_tokens |
| 669 | |
| 670 | assert topk_ids.dim() == 2 |
| 671 | topk = topk_ids.size(1) |
| 672 | |
| 673 | return E, M, N, K, topk |
| 674 | |
| 675 | |
| 676 | class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): |
| 677 | """Marlin-based fused MoE expert implementation.""" |
| 678 | |
| 679 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 680 | return TopKWeightAndReduceNoOP() |
| 681 | |
| 682 | @staticmethod |
| 683 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 684 | return mk.FusedMoEActivationFormat.Standard |
| 685 | |
| 686 | def workspace_shapes( |
| 687 | self, |
| 688 | M: int, |
| 689 | N: int, |
| 690 | K: int, |
| 691 | topk: int, |
| 692 | global_num_experts: int, |
| 693 | local_num_experts: int, |
| 694 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 695 | activation: MoEActivation, |
| 696 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 697 | # Modular Kernel provisions output buffer from workspace1. However in |
| 698 | # the fused_marlin_moe() function, the final torch.sum(), is defined |
| 699 | # essentially as, |
| 700 | # `torch.sum(workspace1, dim=1, out=output)` |
| 701 | # Having overlapping input and output tensors for torch.sum seems |
| 702 | # error prone and depends on how the torch.sum is implemented. |
| 703 | # For this reason we swap let the output buffer provision from |
| 704 | # workspace2. |
| 705 | |
| 706 | # Workspace/IntermediateCache allocation matching fused_marlin_moe() |
| 707 | # workspace1 = (M * topk * max(2 * N, K),) |
| 708 | # workspace2 = (M * topk, N) |
| 709 | |
| 710 | # Workspace/IntermediateCache allocation accounting for output buffer |
| 711 | # provisioning |
| 712 | workspace1 = (M * topk, max(N, K)) |
| 713 | workspace2 = (M * topk * max(2 * N, K),) |
| 714 | output = (M, K) |
| 715 | |
| 716 | return (workspace1, workspace2, output) |
| 717 | |
| 718 | def apply( |
| 719 | self, |
| 720 | output: torch.Tensor, |
| 721 | hidden_states: torch.Tensor, |
| 722 | w1: torch.Tensor, |
| 723 | w2: torch.Tensor, |
| 724 | topk_weights: torch.Tensor, |
| 725 | topk_ids: torch.Tensor, |
| 726 | activation: MoEActivation, |
| 727 | global_num_experts: int, |
| 728 | expert_map: torch.Tensor | None, |
| 729 | a1q_scale: torch.Tensor | None, |
| 730 | a2_scale: torch.Tensor | None, |
| 731 | workspace13: torch.Tensor, |
| 732 | workspace2: torch.Tensor, |
| 733 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 734 | apply_router_weight_on_input: bool, |
| 735 | ): |
| 736 | assert self.w1_scale is not None |
| 737 | assert self.w2_scale is not None |
| 738 | |
| 739 | ctx = self._lora_context |
| 740 | if ctx is None: |
| 741 | fused_marlin_moe( |
| 742 | hidden_states=hidden_states, |
| 743 | w1=w1, |
| 744 | w2=w2, |
| 745 | bias1=self.w1_bias, |
| 746 | bias2=self.w2_bias, |
| 747 | w1_scale=self.w1_scale, |
| 748 | w2_scale=self.w2_scale, |
| 749 | topk_weights=topk_weights, |
| 750 | topk_ids=topk_ids, |
| 751 | global_scale1=self.g1_alphas, |
| 752 | global_scale2=self.g2_alphas, |
| 753 | input_global_scale1=self.a1_gscale, |
| 754 | input_global_scale2=self.a2_gscale, |
| 755 | w1_zeros=self.w1_zp, |
| 756 | w2_zeros=self.w2_zp, |
| 757 | quant_type_id=self.quant_type_id, |
| 758 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 759 | global_num_experts=global_num_experts, |
| 760 | activation=activation, |
| 761 | activation_func=self.activation, |
| 762 | activation_config=self.activation_config, |
| 763 | moe_sum=self.moe_sum, |
| 764 | expert_map=expert_map, |
| 765 | output=output, |
| 766 | # Workspaces are swapped in workspace_shapes() to account for proper |
| 767 | # output buffer allocation. Please refer to workspace_shapes(). |
| 768 | intermediate_cache13=workspace2, |
| 769 | intermediate_cache2=workspace13, |
| 770 | input_dtype=self.input_dtype, |
| 771 | ) |
| 772 | return |
| 773 | |
| 774 | # LoRA path: wrap activation_func and moe_sum to inject LoRA at the |
| 775 | # two natural injection points. |
| 776 | # |
| 777 | # Marlin uses moe_align_block_size (same as TritonExperts) so |
| 778 | # intermediate_cache1 is indexed by flat (token, expert) pair index, |
| 779 | # which is compatible with add_lora_fused_moe's scatter mechanism. |
| 780 | |
| 781 | M = hidden_states.size(0) |
| 782 | top_k_num = topk_ids.size(1) |
| 783 | lora_state: dict = {} |
| 784 | |
| 785 | def activation_with_lora( |
| 786 | act_enum: MoEActivation, |
| 787 | act_output: torch.Tensor, |
| 788 | act_input: torch.Tensor, |
| 789 | *, |
| 790 | topk_ids: torch.Tensor | None = None, |
| 791 | expert_map: torch.Tensor | None = None, |
| 792 | ) -> None: |
| 793 | # act_input = intermediate_cache1 (M*topk, 2N for gated) |
| 794 | # act_output = intermediate_cache2 (M*topk, N) |
| 795 | |
| 796 | ( |
| 797 | sorted_token_ids_lora, |
| 798 | expert_ids_lora, |
| 799 | num_tokens_post_padded_lora, |
| 800 | token_lora_mapping, |
| 801 | ) = self.apply_w13_lora( |
| 802 | ctx, |
| 803 | y=act_input, |
| 804 | x=hidden_states, |
| 805 | topk_ids=topk_ids, |
| 806 | topk_weights=topk_weights, |
| 807 | expert_map=expert_map, |
| 808 | w1=w1, |
| 809 | w2=w2, |
| 810 | num_tokens=M, |
| 811 | top_k_num=top_k_num, |
| 812 | ) |
| 813 | lora_state.update( |
| 814 | { |
| 815 | "sorted": sorted_token_ids_lora, |
| 816 | "eids": expert_ids_lora, |
| 817 | "npad": num_tokens_post_padded_lora, |
| 818 | "tlm": token_lora_mapping, |
| 819 | } |
| 820 | ) |
| 821 | self.activation( |
| 822 | act_enum, |
| 823 | act_output, |
| 824 | act_input, |
| 825 | topk_ids=topk_ids, |
| 826 | expert_map=expert_map, |
| 827 | ) |
| 828 | lora_state["cache2"] = act_output |
| 829 | |
| 830 | def moe_sum_with_lora( |
| 831 | moe_out: torch.Tensor, |
| 832 | out: torch.Tensor, |
| 833 | topk_ids: torch.Tensor, |
| 834 | expert_map: torch.Tensor | None, |
| 835 | ) -> None: |
| 836 | # moe_out shape: (M, topk, K) |
| 837 | self.apply_w2_lora( |
| 838 | ctx, |
| 839 | y=moe_out, |
| 840 | x=lora_state["cache2"], |
| 841 | topk_weights=topk_weights, |
| 842 | sorted_token_ids_lora=lora_state["sorted"], |
| 843 | expert_ids_lora=lora_state["eids"], |
| 844 | num_tokens_post_padded_lora=lora_state["npad"], |
| 845 | token_lora_mapping=lora_state["tlm"], |
| 846 | num_tokens=M, |
| 847 | w1=w1, |
| 848 | w2=w2, |
| 849 | top_k_num=top_k_num, |
| 850 | ) |
| 851 | self.moe_sum(moe_out, out, topk_ids, expert_map) |
| 852 | |
| 853 | return fused_marlin_moe( |
| 854 | hidden_states=hidden_states, |
| 855 | w1=w1, |
| 856 | w2=w2, |
| 857 | bias1=self.w1_bias, |
| 858 | bias2=self.w2_bias, |
| 859 | w1_scale=self.w1_scale, |
| 860 | w2_scale=self.w2_scale, |
| 861 | topk_weights=topk_weights, |
| 862 | topk_ids=topk_ids, |
| 863 | global_scale1=self.g1_alphas, |
| 864 | global_scale2=self.g2_alphas, |
| 865 | input_global_scale1=self.a1_gscale, |
| 866 | input_global_scale2=self.a2_gscale, |
| 867 | w1_zeros=self.w1_zp, |
| 868 | w2_zeros=self.w2_zp, |
| 869 | quant_type_id=self.quant_type_id, |
| 870 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 871 | global_num_experts=global_num_experts, |
| 872 | activation=activation, |
| 873 | activation_func=activation_with_lora, |
| 874 | activation_config=self.activation_config, |
| 875 | moe_sum=moe_sum_with_lora, |
| 876 | expert_map=expert_map, |
| 877 | output=output, |
| 878 | intermediate_cache13=workspace2, |
| 879 | intermediate_cache2=workspace13, |
| 880 | input_dtype=self.input_dtype, |
| 881 | ) |
| 882 | |
| 883 | def moe_sum( |
| 884 | self, |
| 885 | input: torch.Tensor, |
| 886 | output: torch.Tensor, |
| 887 | topk_ids: torch.Tensor, |
| 888 | expert_map: torch.Tensor | None, |
| 889 | ) -> None: |
| 890 | if expert_map is not None: |
| 891 | ops.moe_sum(input, output, topk_ids, expert_map) |
| 892 | else: |
| 893 | ops.moe_sum(input, output) |
| 894 | |
| 895 | |
| 896 | class BatchedMarlinExperts(MarlinExpertsBase): |
| 897 | """Batched Marlin-based fused MoE expert implementation.""" |
| 898 | |
| 899 | def __init__( |
| 900 | self, |
| 901 | moe_config: FusedMoEConfig, |
| 902 | quant_config: FusedMoEQuantConfig, |
| 903 | max_num_tokens: int, |
| 904 | num_dispatchers: int, |
| 905 | ): |
| 906 | super().__init__( |
| 907 | moe_config=moe_config, |
| 908 | quant_config=quant_config, |
| 909 | max_num_tokens=max_num_tokens, |
| 910 | num_dispatchers=num_dispatchers, |
| 911 | ) |
| 912 | |
| 913 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 914 | return TopKWeightAndReduceDelegate() |
| 915 | |
| 916 | @staticmethod |
| 917 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 918 | return mk.FusedMoEActivationFormat.BatchedExperts |
| 919 | |
| 920 | @staticmethod |
| 921 | def _supports_activation(activation: MoEActivation) -> bool: |
| 922 | return apply_moe_activation_masked_supported(activation) |
| 923 | |
| 924 | def workspace_shapes( |
| 925 | self, |
| 926 | M: int, |
| 927 | N: int, |
| 928 | K: int, |
| 929 | topk: int, |
| 930 | global_num_experts: int, |
| 931 | local_num_experts: int, |
| 932 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 933 | activation: MoEActivation, |
| 934 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 935 | assert self.num_dispatchers is not None |
| 936 | assert self.max_num_tokens is not None |
| 937 | num_dispatchers = self.num_dispatchers |
| 938 | num_experts = local_num_experts |
| 939 | max_num_tokens = self.max_num_tokens |
| 940 | workspace13 = (num_experts * max_num_tokens * num_dispatchers, max(K, N * 2)) |
| 941 | workspace2 = (num_experts * max_num_tokens * num_dispatchers, N) |
| 942 | output = (num_experts, max_num_tokens * num_dispatchers, K) |
| 943 | return (workspace13, workspace2, output) |
| 944 | |
| 945 | def apply( |
| 946 | self, |
| 947 | output: torch.Tensor, |
| 948 | hidden_states: torch.Tensor, |
| 949 | w1: torch.Tensor, |
| 950 | w2: torch.Tensor, |
| 951 | topk_weights: torch.Tensor, |
| 952 | topk_ids: torch.Tensor, |
| 953 | activation: MoEActivation, |
| 954 | global_num_experts: int, |
| 955 | expert_map: torch.Tensor | None, |
| 956 | a1q_scale: torch.Tensor | None, |
| 957 | a2_scale: torch.Tensor | None, |
| 958 | workspace13: torch.Tensor, |
| 959 | workspace2: torch.Tensor, |
| 960 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 961 | apply_router_weight_on_input: bool, |
| 962 | ): |
| 963 | assert expert_tokens_meta is not None, "Num valid tokens per batch is required" |
| 964 | |
| 965 | def activation_func( |
| 966 | act: MoEActivation, |
| 967 | act_output: torch.Tensor, |
| 968 | act_input: torch.Tensor, |
| 969 | *, |
| 970 | topk_ids: torch.Tensor | None = None, |
| 971 | expert_map: torch.Tensor | None = None, |
| 972 | ) -> None: |
| 973 | num_experts, max_num_tokens = hidden_states.shape[:2] |
| 974 | self.activation( |
| 975 | act, |
| 976 | act_output.view(num_experts, max_num_tokens, -1), |
| 977 | act_input.view(num_experts, max_num_tokens, -1), |
| 978 | topk_ids=topk_ids, |
| 979 | expert_map=expert_map, |
| 980 | valid_token_counts=expert_tokens_meta.expert_num_tokens, |
| 981 | ) |
| 982 | |
| 983 | return batched_fused_marlin_moe( |
| 984 | hidden_states=hidden_states, |
| 985 | expert_num_tokens=expert_tokens_meta.expert_num_tokens, |
| 986 | w1=w1, |
| 987 | w2=w2, |
| 988 | bias1=self.w1_bias, |
| 989 | bias2=self.w2_bias, |
| 990 | w1_scale=self.w1_scale, |
| 991 | w2_scale=self.w2_scale, |
| 992 | quant_type_id=self.quant_type_id, |
| 993 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 994 | global_num_experts=global_num_experts, |
| 995 | activation=activation, |
| 996 | expert_map=expert_map, |
| 997 | input_global_scale1=self.a1_gscale, |
| 998 | input_global_scale2=self.a2_gscale, |
| 999 | output=output, |
| 1000 | intermediate_cache13=workspace13, |
| 1001 | intermediate_cache2=workspace2, |
| 1002 | w1_zeros=self.w1_zp, |
| 1003 | w2_zeros=self.w2_zp, |
| 1004 | input_dtype=self.input_dtype, |
| 1005 | activation_func=activation_func, |
| 1006 | activation_config=self.activation_config, |
| 1007 | ) |