Skip to content

[Misc] Minor enhancement of benchmark_moe #22068

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 1, 2025
Merged

[Misc] Minor enhancement of benchmark_moe #22068

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion benchmarks/kernels/benchmark_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,13 @@
FP8_DTYPE = current_platform.fp8_dtype()


def ensure_divisibility(numerator, denominator):
"""Ensure that numerator is divisible by the denominator."""
assert numerator % denominator == 0, (
"intermediate_size {} is not divisible by tp {}.".format(numerator, denominator)
)

Comment on lines +25 to +30
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Aug 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since this function is only used once, it can be inlined at the call site for simplicity. This also avoids introducing a new helper function with a generic name but a specific implementation, which could be misleading for future maintenance.


class BenchmarkConfig(TypedDict):
BLOCK_SIZE_M: int
BLOCK_SIZE_N: int
Expand Down Expand Up @@ -603,7 +610,7 @@ def main(args: argparse.Namespace):
topk = config.num_experts_per_tok
intermediate_size = config.intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size

ensure_divisibility(intermediate_size, args.tp_size)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Aug 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Inlining the divisibility check here makes the code more direct and removes the need for a separate helper function.

Suggested change
ensure_divisibility(intermediate_size, args.tp_size)
assert intermediate_size % args.tp_size == 0, f"intermediate_size {intermediate_size} is not divisible by tp {args.tp_size}"

hidden_size = config.hidden_size
dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
Expand Down