[ROCm][MoE] Add roofline profiling scope to the Triton W4A16 MoE GEMM

[roberteg16]

Summary

This adds an apply()-level record_function profiling scope around the Triton W4A16 MoE prefill GEMMs (fused_moe_kernel_gptq_awq, via invoke_fused_moe_kernel_hybrid_triton) in HybridW4A16MoEExperts.apply(), so a profiler trace carries the dimensions the roofline tooling needs to compute real FLOPs/bytes for these kernels. It is pure profiling instrumentation: it changes no kernel behavior or numerics, and it is a no-op (a nullcontext) unless profiling scopes are explicitly enabled, so production pays nothing.

Motivation

The raw kernel name in the profiler (fused_moe_kernel_gptq_awq) exposes none of the GEMM dimensions, so the roofline tool cannot attribute FLOPs/bytes to the MoE GEMMs and they fall out of the roofline summary. Several quantities that matter are not derivable from any tensor shape in the trace: the per-call problem size (M, N, K, E, top_k); the quant group size g, since the per-group scale bytes and dequant FLOPs scale with K/g and this MoE path supports configurable group sizes (block_shape=[0, group_size] in compressed-tensors), so a non-128 g would otherwise be mis-attributed; and valid_blocks — the number of non-padding alignment blocks the kernel actually runs (num_tokens_post_padded // block_m), so that valid_blocks * block_m is the padded row count the kernel really computes (the routing-dependent padding is what drives the gap between useful and computed FLOPs).

What this adds

• vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py: a _moe_gemm_triton_scope(M, N, K, E, top_k, group_size, block_m, num_tokens_post_padded) helper that emits moe_gemm_triton {M}x{N}x{K} E={E} top_k={top_k} g={group_size} block_m={block_m} valid_blocks={valid_blocks}, and wraps both Triton prefill GEMMs (gemm1 up/gate proj and gemm2 down proj) in apply() with it. The g field is the real self._group_size, not a hardcoded 128.

Design notes

The scope is emitted at the apply() call site rather than deep inside invoke_fused_moe_kernel_hybrid_triton on purpose: an apply()-level scope survives torch.compile replay and stays attached to the compiled kernel, whereas a scope inside the launch helper is orphaned from the compiled kernel and leaves the MoE GEMM with no roofline. valid_blocks requires reading num_tokens_post_padded.item() (a device→host sync), so the scope is gated on VLLM_CUSTOM_SCOPES_FOR_PROFILING / VLLM_NVTX_SCOPES_FOR_PROFILING and returns a nullcontext when neither is set — the sync (and the scope) only happen during profiling runs. vllm bench serve --profile sets VLLM_CUSTOM_SCOPES_FOR_PROFILING=1 automatically.

[Copilot]

Pull request overview

This PR adds apply()-level profiling scopes around the Triton W4A16 MoE GEMMs in HybridW4A16MoEExperts.apply() so profiler traces include MoE GEMM problem parameters (including valid_blocks) needed for downstream roofline attribution.

Changes:

• Introduces _moe_gemm_triton_scope(...) to build a record_function name that includes GEMM dims and valid_blocks (derived from num_tokens_post_padded).
• Wraps both Triton prefill GEMM invocations (GEMM1 and GEMM2) with the new scope, while remaining a no-op unless profiling scopes are enabled.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[roberteg16]

@mgehre-amd ping

[mgehre-amd]

Please rebase to fix CI