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

Profile exact vLLM MoE serving entrypoint

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 18:41:23 +08:00
parent d9df3003dd
commit c9a5f0a0c3
1 changed files with 127 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

127
runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_entrypoint.patch Normal file
View File

@@ -0,0 +1,127 @@
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
@@ -33,6 +33,7 @@ try:
from vllm import _custom_ops as ops
# Import vLLM 0.10.x functions
from vllm.model_executor.layers.fused_moe.fused_moe import (
+ fused_experts,
fused_moe_kernel,
invoke_fused_moe_kernel,
moe_align_block_size,
@@ -531,82 +532,58 @@ def profile_fused_moe_kernel(
block_shape=block_shape,
)
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
- topk_ids,
- config["BLOCK_SIZE_M"],
- align_num_experts,
- expert_map=expert_map,
- )
-
- output_dtype = base_dtype
- intermediate_cache1 = torch.empty(
- num_tokens,
- top_k,
- w1.shape[1],
- device=device,
- dtype=output_dtype,
- )
- intermediate_cache2 = torch.empty(
- 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(
- A=A,
+ fused_experts(
+ hidden_states=A,
w1=w1,
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,
- num_tokens_post_padded=num_tokens_post_padded,
- top_k=top_k,
- config=config,
- expert_hidden_dim_per_partition=expert_hidden_dim_per_partition,
- block_dims=block_dims,
- A_scale=None,
+ topk_ids=topk_ids,
+ inplace=True,
+ global_num_experts=align_num_experts,
+ expert_map=expert_map,
+ use_fp8_w8a8=use_fp8,
+ per_channel_quant=per_channel_quant,
w1_scale=w1_scale,
w2_scale=w2_scale,
- use_fp8=use_fp8,
- per_channel_quant=per_channel_quant,
block_shape=block_shape,
)
+ def _alignment_step() -> None:
+ moe_align_block_size(
+ topk_ids,
+ config["BLOCK_SIZE_M"],
+ align_num_experts,
+ expert_map=expert_map,
+ )
+
for _ in range(warmup_steps):
_step()
torch.cuda.synchronize()
if profile_method == "record_function":
- return _collect_record_function_stats(
- step_fn=_step,
- active_steps=active_steps,
- output_dir=output_dir,
- operation_name="moe_grouped_gemm",
+ raise ValueError(
+ "Serving-entrypoint MoE profiling requires cuda_event so the local "
+ "alignment component can be subtracted without double counting."
)
- return _collect_cuda_event_stats(
+ serving_stats = _collect_cuda_event_stats(
step_fn=_step,
active_steps=active_steps,
)
+ alignment_stats = _collect_cuda_event_stats(
+ step_fn=_alignment_step,
+ active_steps=active_steps,
+ )
+ return {
+ "min": max(0.0, serving_stats["min"] - alignment_stats["max"]),
+ "max": max(0.0, serving_stats["max"] - alignment_stats["min"]),
+ "mean": max(0.0, serving_stats["mean"] - alignment_stats["mean"]),
+ "median": max(0.0, serving_stats["median"] - alignment_stats["median"]),
+ "std": (
+ serving_stats["std"] ** 2 + alignment_stats["std"] ** 2
+ ) ** 0.5,
+ }
def generate_expert_weights(