--- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1645,7 +1645,38 @@ # check constraints if use_fp8_w8a8 or use_int8_w8a8 or use_int8_w8a16 or use_int4_w4a16 or w1_scale or \ w2_scale or w1_zp or w2_zp or a1_scale or a2_scale: - raise ValueError("Quantized MoE is not supported") + # COREXX_FP8_MOE_PATCH: generic fused_experts_impl supports quantized MoE + res = fused_experts_impl( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=False, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + ocp_mx_scheme=ocp_mx_scheme, + per_channel_quant=per_channel_quant, + global_num_experts=global_num_experts, + expert_map=expert_map, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + w1_bias=w1_bias, + w2_bias=w2_bias, + ) + if output is not None: + output.copy_(res) + return output + return res attn_metadata = get_forward_context().attn_metadata use_ep = expert_map is not None @@ -1912,8 +1943,32 @@ w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: + # COREXX_FP8_MOE_PATCH: dedicated per-expert BF16 fallback for FP8 MoE + if hidden_states.dtype in ( + torch.float8_e4m3fn, torch.float8_e5m2, + torch.float8_e4m3fnuz, torch.float8_e5m2fnuz, + ): + from vllm.model_executor.layers.fused_moe.corexx_fp8_moe import corexx_fp8_moe + return corexx_fp8_moe( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + w1_bias=w1_bias, + w2_bias=w2_bias, + ) + # Convert string activation to enum for internal use - activation_enum = MoEActivation.from_str(activation) # Check constraints. if use_int4_w4a16: @@ -1937,7 +1992,12 @@ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" assert w1.stride(-1) == 1, "Stride of last dimension must be 1" assert w2.stride(-1) == 1, "Stride of last dimension must be 1" - assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] + # COREXX_FP8_MOE_PATCH: modular kernel may pass pre-quantized FP8 input + assert hidden_states.dtype in [ + torch.float32, torch.float16, torch.bfloat16, + torch.float8_e4m3fn, torch.float8_e5m2, + torch.float8_e4m3fnuz, torch.float8_e5m2fnuz, + ] num_tokens = hidden_states.size(0) E, N, _ = w1.size() @@ -1949,6 +2009,9 @@ # https://github.com/vllm-project/vllm/issues/5938 CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE M = min(num_tokens, CHUNK_SIZE) + # COREXX_FP8_MOE_PATCH: FP8 input still needs BF16 compute buffers/output + io_dtype = hidden_states.dtype if hidden_states.dtype in ( + torch.float32, torch.float16, torch.bfloat16) else torch.bfloat16 config_dtype = _get_config_dtype_str( use_fp8_w8a8=use_fp8_w8a8, @@ -1982,7 +2045,7 @@ cache13 = torch.empty( M * top_k_num * max(N, K), device=hidden_states.device, - dtype=hidden_states.dtype, + dtype=io_dtype, ) intermediate_cache1 = cache13[: M * top_k_num * N].view(M, top_k_num, N) intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K) @@ -1994,19 +2057,22 @@ intermediate_cache2 = torch.empty( (M * top_k_num, activation_out_dim), device=hidden_states.device, - dtype=hidden_states.dtype, + dtype=io_dtype, ) - if hidden_states.dtype == torch.bfloat16: + if io_dtype == torch.bfloat16: compute_type = tl.bfloat16 - elif hidden_states.dtype == torch.float16: + elif io_dtype == torch.float16: compute_type = tl.float16 - elif hidden_states.dtype == torch.float32: + elif io_dtype == torch.float32: compute_type = tl.float32 else: - raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") + raise ValueError(f"Unsupported compute_type: {io_dtype}") - out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states) + if inplace and io_dtype == hidden_states.dtype: + out_hidden_states = hidden_states + else: + out_hidden_states = torch.empty_like(hidden_states, dtype=io_dtype) if ocp_mx_scheme is not None: # TODO: On platforms for which `current_platform.supports_mx()` is True @@ -2064,14 +2130,24 @@ curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx] curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx] - qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input( - A=curr_hidden_states, - A_scale=a1_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_channel_quant, - block_shape=block_shape, - ocp_mx_scheme=ocp_mx_scheme, - ) + # COREXX_FP8_MOE_PATCH: use pre-quantized FP8 input as-is + if curr_hidden_states.dtype in ( + torch.float8_e4m3fn, torch.float8_e5m2, + torch.float8_e4m3fnuz, torch.float8_e5m2fnuz, + ): + qcurr_hidden_states = curr_hidden_states + a1q_scale = a1_scale + if a1q_scale is None: + raise ValueError("pre-quantized MoE input requires a1_scale") + else: + qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input( + A=curr_hidden_states, + A_scale=a1_scale, + quant_dtype=quant_dtype, + per_act_token_quant=per_channel_quant, + block_shape=block_shape, + ocp_mx_scheme=ocp_mx_scheme, + ) # # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k # # activates only a small fraction of total experts @@ -2305,8 +2381,9 @@ self.quant_config.w2_scale, self.quant_config.w1_zp, self.quant_config.w2_zp, - self.quant_config.a1_scale, - self.quant_config.a2_scale, + # COREXX_FP8_MOE_PATCH: prefer per-call dynamic scales + a1q_scale if a1q_scale is not None else self.quant_config.a1_scale, + a2_scale if a2_scale is not None else self.quant_config.a2_scale, self.quant_config.block_shape, self.quant_config.w1_bias, self.quant_config.w2_bias,