gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align Frontier FP8 profiling with vLLM runtime

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 18:34:13 +08:00
parent 9c8570f36b
commit 684a2de413
3 changed files with 525 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

216
runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_path.patch Normal file
View File

@@ -0,0 +1,216 @@
diff --git a/frontier/profiling/moe/moe_vllm_kernel.py b/frontier/profiling/moe/moe_vllm_kernel.py
--- a/frontier/profiling/moe/moe_vllm_kernel.py
+++ b/frontier/profiling/moe/moe_vllm_kernel.py
@@ -34,6 +34,7 @@ try:
import vllm
VLLM_VERSION = vllm.__version__
+ from vllm import _custom_ops as ops
# Import vLLM 0.10.x functions
from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_moe_kernel,
@@ -232,7 +233,7 @@ def _invoke_kernel(
"""
# Determine compute_type - for FP8, we accumulate in FP16/BF16
if use_fp8:
- compute_type = tl.float16 # FP8 accumulates in FP16
+ compute_type = tl.bfloat16
else:
dtype = A.dtype
if dtype == torch.bfloat16:
@@ -275,7 +276,9 @@ def _run_fused_moe_iteration(
w1: torch.Tensor,
w2: torch.Tensor,
intermediate_cache1: torch.Tensor,
intermediate_cache2: torch.Tensor,
+ intermediate_cache3: torch.Tensor,
+ output: torch.Tensor,
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
@@ -292,8 +295,17 @@ def _run_fused_moe_iteration(
per_channel_quant: bool = False,
block_shape: Optional[List[int]] = None,
) -> None:
+ first_input = A
+ first_A_scale = A_scale
+ if use_fp8:
+ group_size = block_dims[1] if block_dims else 128
+ first_input, first_A_scale = quantize_activations_to_fp8(
+ A,
+ group_size=group_size,
+ )
+
_invoke_kernel(
- A=A.contiguous(),
+ A=first_input.contiguous(),
B=w1.contiguous(),
C=intermediate_cache1.contiguous(),
topk_weights=topk_weights.contiguous(),
@@ -305,15 +316,18 @@ def _run_fused_moe_iteration(
mul_routed_weight=False,
top_k=top_k,
config=config,
- A_scale=A_scale,
+ A_scale=first_A_scale,
B_scale=w1_scale,
use_fp8=use_fp8,
per_channel_quant=per_channel_quant,
block_shape=block_shape,
)
- intermediate_cache1_flat = intermediate_cache1.view(-1, intermediate_cache1.shape[-1])
- intermediate_cache2_input = intermediate_cache1_flat[:, :expert_hidden_dim_per_partition].contiguous()
+ torch.ops._C.silu_and_mul(
+ intermediate_cache2,
+ intermediate_cache1.view(-1, intermediate_cache1.shape[-1]),
+ )
+ second_input = intermediate_cache2
intermediate_A_scale = None
if use_fp8:
@@ -321,13 +334,13 @@ def _run_fused_moe_iteration(
group_size = block_dims[1] if block_dims else 128
- intermediate_cache2_input, intermediate_A_scale = quantize_activations_to_fp8(
- intermediate_cache2_input,
+ second_input, intermediate_A_scale = quantize_activations_to_fp8(
+ intermediate_cache2,
group_size=group_size,
)
_invoke_kernel(
- A=intermediate_cache2_input,
+ A=second_input,
B=w2.contiguous(),
- C=intermediate_cache2.contiguous(),
+ C=intermediate_cache3.contiguous(),
topk_weights=topk_weights.contiguous(),
sorted_token_ids=sorted_token_ids.contiguous(),
expert_ids=expert_ids.contiguous(),
@@ -335,4 +350,6 @@ def _run_fused_moe_iteration(
)
+
+ ops.moe_sum(intermediate_cache3, output)
def _collect_cuda_event_stats(step_fn, active_steps: int) -> Dict:
@@ -493,6 +508,5 @@ def profile_fused_moe_kernel(
w1_scale = None
w2_scale = None
- A_scale = None
block_dims = _validate_block_shape(block_shape)
if use_fp8:
@@ -509,10 +521,8 @@ def profile_fused_moe_kernel(
per_channel=per_channel_quant,
block_shape=block_shape,
)
- group_size = block_dims[1] if block_dims else 128
- A, A_scale = quantize_activations_to_fp8(A, group_size=group_size)
- config_dtype = get_config_dtype_str(base_dtype)
+ config_dtype = get_config_dtype_str(base_dtype, use_fp8_w8a8=use_fp8)
config = try_get_optimal_moe_config(
w1_shape=w1.shape,
w2_shape=w2.shape,
@@ -535,13 +544,25 @@ def profile_fused_moe_kernel(
device=device,
dtype=output_dtype,
)
intermediate_cache2 = torch.empty(
- num_tokens,
- top_k,
- hidden_dim,
+ num_tokens * top_k,
+ expert_hidden_dim_per_partition,
+ device=device,
+ dtype=output_dtype,
+ )
+ intermediate_cache3 = torch.empty(
+ num_tokens,
+ top_k,
+ hidden_dim,
device=device,
dtype=output_dtype,
)
+ output = torch.empty(
+ num_tokens,
+ hidden_dim,
+ device=device,
+ dtype=output_dtype,
+ )
def _step() -> None:
_run_fused_moe_iteration(
@@ -552,6 +571,8 @@ def profile_fused_moe_kernel(
w2=w2,
intermediate_cache1=intermediate_cache1,
intermediate_cache2=intermediate_cache2,
+ intermediate_cache3=intermediate_cache3,
+ output=output,
topk_weights=topk_weights,
sorted_token_ids=sorted_token_ids,
expert_ids=expert_ids,
@@ -562,6 +583,6 @@ def profile_fused_moe_kernel(
expert_hidden_dim_per_partition=expert_hidden_dim_per_partition,
block_dims=block_dims,
- A_scale=A_scale,
+ A_scale=None,
w1_scale=w1_scale,
w2_scale=w2_scale,
use_fp8=use_fp8,
diff --git a/frontier/profiling/moe/moe_impl.py b/frontier/profiling/moe/moe_impl.py
--- a/frontier/profiling/moe/moe_impl.py
+++ b/frontier/profiling/moe/moe_impl.py
@@ -245,10 +245,12 @@ class MoETokenShuffler(nn.Module):
def __init__(
self,
num_experts: int,
router_topk: int,
hidden_dim: int,
expert_hidden_dim: int,
dtype: torch.dtype,
use_gated: bool,
num_local_experts: Optional[int] = None,
+ use_fp8: bool = False,
+ block_shape: Optional[list[int]] = None,
):
@@ -264,6 +266,8 @@ class MoETokenShuffler(nn.Module):
self.router_topk = router_topk
self.hidden_dim = hidden_dim
self.expert_hidden_dim = expert_hidden_dim
self.dtype = dtype
self.use_gated = use_gated
+ self.use_fp8 = use_fp8
+ self.block_shape = block_shape
self._block_size_cache = {}
@@ -325,9 +329,12 @@ class MoETokenShuffler(nn.Module):
- config_dtype = get_config_dtype_str(dtype=self.dtype)
+ config_dtype = get_config_dtype_str(
+ dtype=self.dtype,
+ use_fp8_w8a8=self.use_fp8,
+ )
config = try_get_optimal_moe_config(
w1_shape=w1_shape,
w2_shape=w2_shape,
top_k=self.router_topk,
dtype=config_dtype,
M=num_tokens,
- block_shape=None,
+ block_shape=self.block_shape,
)
diff --git a/frontier/profiling/moe/moe_wrapper.py b/frontier/profiling/moe/moe_wrapper.py
--- a/frontier/profiling/moe/moe_wrapper.py
+++ b/frontier/profiling/moe/moe_wrapper.py
@@ -149,9 +149,11 @@ class MoEWrapper:
self.shuffler = MoETokenShuffler(
num_experts=self.num_experts,
num_local_experts=self.num_experts_per_device,
router_topk=self.router_topk,
hidden_dim=self.hidden_dim,
expert_hidden_dim=self.expert_hidden_dim,
dtype=self._dtype,
use_gated=self.use_gated,
+ use_fp8=self.use_fp8,
+ block_shape=self.block_shape,
).to(dtype=self._dtype).cuda().eval()