Profile exact vLLM MoE serving entrypoint

runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_entrypoint.patch

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