126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
"""CoreX fallback for FP8 W8A8 MoE on Iluvatar BI-V150.
|
|
|
|
The CoreX ixformer MoE kernels are W16A16 only and torch._scaled_mm supports
|
|
only tensor-wise scales, so per-token/per-channel FP8 MoE cannot use them.
|
|
This module dequantizes one expert at a time (each expert shard is only a few
|
|
hundred MB on a TP rank), runs BF16 matmuls, and scatters the result.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
_FP8_DTYPES = (
|
|
torch.float8_e4m3fn,
|
|
torch.float8_e5m2,
|
|
torch.float8_e4m3fnuz,
|
|
torch.float8_e5m2fnuz,
|
|
)
|
|
_LOGGED = False
|
|
|
|
|
|
def _dequant_weight(w: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor:
|
|
"""w: [rows, cols] FP8, scale: per-output-row (len rows) or scalar."""
|
|
if scale is None:
|
|
return w.to(torch.bfloat16)
|
|
s = scale.reshape(-1)
|
|
if s.numel() == 1:
|
|
return (w.float() * s[0]).to(torch.bfloat16)
|
|
if s.numel() != w.shape[0]:
|
|
# try to broadcast a [1, rows] style scale
|
|
s = s.reshape(w.shape[0], -1)[:, 0]
|
|
return (w.float() * s.reshape(-1, 1)).to(torch.bfloat16)
|
|
|
|
|
|
def corexx_fp8_moe(
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
activation,
|
|
apply_router_weight_on_input: bool,
|
|
global_num_experts: int,
|
|
expert_map: torch.Tensor | None,
|
|
w1_scale: torch.Tensor | None,
|
|
w2_scale: torch.Tensor | None,
|
|
a1_scale: torch.Tensor | None,
|
|
a2_scale: torch.Tensor | None,
|
|
block_shape,
|
|
w1_bias: torch.Tensor | None,
|
|
w2_bias: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
global _LOGGED
|
|
if expert_map is not None:
|
|
raise NotImplementedError("CoreX FP8 MoE fallback does not support expert_map/EP")
|
|
if block_shape is not None:
|
|
raise NotImplementedError("CoreX FP8 MoE fallback does not support block FP8")
|
|
if not _LOGGED:
|
|
logger.info(
|
|
"CoreX FP8 MoE fallback: hidden=%s %s a1_scale=%s w1=%s w1_scale=%s w2=%s w2_scale=%s topk=%s",
|
|
tuple(hidden_states.shape), hidden_states.dtype,
|
|
None if a1_scale is None else tuple(a1_scale.shape),
|
|
tuple(w1.shape), None if w1_scale is None else tuple(w1_scale.shape),
|
|
tuple(w2.shape), None if w2_scale is None else tuple(w2_scale.shape),
|
|
tuple(topk_ids.shape),
|
|
)
|
|
_LOGGED = True
|
|
|
|
if hidden_states.dtype in _FP8_DTYPES:
|
|
if a1_scale is None:
|
|
raise ValueError("FP8 MoE input requires a1_scale")
|
|
A = (hidden_states.float() * a1_scale.reshape(-1, 1)).to(torch.bfloat16)
|
|
else:
|
|
A = hidden_states
|
|
|
|
E = w1.shape[0]
|
|
N = w1.shape[1]
|
|
K = w1.shape[2]
|
|
I = w2.shape[2]
|
|
assert N == 2 * I, f"expected w1 N=2*I, got {N} vs {I}"
|
|
|
|
act_name = getattr(activation, "value", activation)
|
|
|
|
M = A.shape[0]
|
|
out = torch.zeros((M, K), dtype=torch.bfloat16, device=A.device)
|
|
topk = topk_ids.shape[1]
|
|
|
|
# COREXX_FP8_MOE_PATCH: only touch experts that actually received tokens
|
|
active_experts = torch.unique(topk_ids).tolist()
|
|
for e in active_experts:
|
|
if e < 0 or e >= E:
|
|
continue
|
|
# rows of A routed to expert e (global id == local id when expert_map is None)
|
|
for k in range(topk):
|
|
idx = (topk_ids[:, k] == e).nonzero(as_tuple=False).flatten()
|
|
if idx.numel() == 0:
|
|
continue
|
|
x = A.index_select(0, idx)
|
|
if apply_router_weight_on_input:
|
|
x = x * topk_weights[idx, k].to(x.dtype).unsqueeze(1)
|
|
w1_e = _dequant_weight(w1[e], None if w1_scale is None else w1_scale[e])
|
|
h = x @ w1_e.t()
|
|
if w1_bias is not None:
|
|
h = h + w1_bias[e]
|
|
gate = h[:, :I]
|
|
up = h[:, I:N]
|
|
if act_name == "silu":
|
|
act = F.silu(gate) * up
|
|
elif act_name == "gelu":
|
|
act = F.gelu(gate) * up
|
|
else:
|
|
raise NotImplementedError(f"CoreX FP8 MoE activation {act_name}")
|
|
w2_e = _dequant_weight(w2[e], None if w2_scale is None else w2_scale[e])
|
|
y = act @ w2_e.t()
|
|
if w2_bias is not None:
|
|
y = y + w2_bias[e]
|
|
if not apply_router_weight_on_input:
|
|
y = y * topk_weights[idx, k].to(y.dtype).unsqueeze(1)
|
|
out.index_add_(0, idx, y)
|
|
del w1_e, w2_e, h, act, y
|
|
|
|
return out
|