184 lines
7.5 KiB
Diff
184 lines
7.5 KiB
Diff
|
--- 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,
|