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(