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

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,