Unified Llama 4 Scout/Maverick BI-V150 deployment bundle

patches/corexx_fp8_moe.py

@@ -0,0 +1,125 @@
+"""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