xiejiahao/llama-4-bi150
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Unified Llama 4 Scout/Maverick BI-V150 deployment bundle

Browse Source
This commit is contained in:
xiejiahao
2026-10-26 16:24:16 +08:00
commit caddeadd15
16 changed files with 1201 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
patches/01_fp8_quant.patch Normal file
View File

@@ -0,0 +1,46 @@
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -1933,14 +1933,28 @@
if scale is None:
if use_per_token_if_dynamic:
scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32)
- torch.ops._C.dynamic_per_token_scaled_fp8_quant(
- output, input, scale, scale_ub
- )
+ # COREXX_FP8_QUANT_PATCH: native fallback (vLLM _C op missing on CoreX)
+ m = input.shape[0]
+ s = input.abs().amax(dim=1, keepdim=True).float().clamp_min(1e-12) / 448.0
+ if scale_ub is not None:
+ ub = torch.as_tensor(scale_ub, device=s.device, dtype=s.dtype)
+ s = torch.minimum(s, ub)
+ scale[:m].copy_(s)
+ if shape[0] > m:
+ scale[m:].fill_(1.0)
+ output[m:].zero_()
+ output[:m].copy_((input.float() / s).to(out_dtype))
else:
scale = torch.empty(1, device=input.device, dtype=torch.float32)
- torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
+ # COREXX_FP8_QUANT_PATCH: native fallback
+ amax = input.abs().amax().float().clamp_min(1e-12)
+ scale.fill_(amax / 448.0)
+ output.copy_((input.float() / scale).to(out_dtype))
else:
- torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape)
+ # COREXX_FP8_QUANT_PATCH: native fallback
+ if group_shape is not None and tuple(group_shape) not in ((-1, -1), (-1, 1), (1, -1)):
+ raise NotImplementedError(f"native fp8 static quant: unsupported group_shape {group_shape}")
+ output.copy_((input.float() / scale).to(out_dtype))
return output, scale
@@ -3824,6 +3838,9 @@
mul_routed_weight,
top_k,
config['BLOCK_SIZE_M'],
+ # COREXX_FP8_MOE_PATCH: pass quant scales to ixformer kernel
+ w_scale=B_scale,
+ a_scale=A_scale,
bias=bias
)

183
patches/02_fp8_moe.patch Normal file
View File

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

16
patches/03_llama4_chunked_attn.patch Normal file
View File

@@ -0,0 +1,16 @@
--- a/vllm/v1/attention/backends/utils.py
+++ b/vllm/v1/attention/backends/utils.py
@@ -346,7 +346,13 @@
return CommonAttentionMetadata(
query_start_loc_cpu=query_start_loc_cpu,
query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
+ key_start_loc=torch.from_numpy(
+ np.concatenate(
+ (np.zeros(1, dtype=np.int32), np.cumsum(seqlens_k_local, dtype=np.int32))
+ )
+ ).to(device=device, non_blocking=True),
seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
+ seq_lens_np=seqlens_k_local,
num_reqs=len(seq_lens_cpu),
num_actual_tokens=common_attn_metadata.num_actual_tokens,
max_query_len=seqlens_q_local.max(),

125
patches/corexx_fp8_moe.py Normal file
View File

@@ -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