mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Scripts and a skill to do per-layer benchmark using flashinfer (#1980)
add scripts and a skill to use flashinfer to do layerwise benchmark with
different backends.
### What does this PR do?
Type of change: ? new skill and scripts
### Usage
```bash
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \
openai/gpt-oss-120b \
--tp 4 --ep 4 --ms 8 16 \
--flashinfer_repo $HOME/flashinfer \
--workdir /tmp/gpt-oss-benchmark
```
or
```bash
python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \
--flashinfer_repo $HOME/flashinfer \
--ms 8 16 \
--nks 1280,2880 2880,1024 \
--moe_hidden_size 2880 \
--moe_intermediate_size 2880 \
--moe_num_experts 32 \
--moe_top_k 4 \
--workdir /tmp/gpt-oss-via-builtin
```
Or example prompt to trigger the skill:
```
Benchmark the kernel shapes for NVIDIA-Nemotron-3-Nano-30B-A3B
```
example output (TP1, EP4):
```
flashinfer 0.6.13; checkout ~/flashinfer3 @ 453bc6e1c6c852a4dc8a35c672f271e1e0d216a5; NVIDIA Graphics Device (sm_100 / 148 SMs / 178 GiB); 600 W power limit
GEMM
module_name,M,N,K,backend,with_quant,runtime
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,bf16,False,8.384
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,bf16,False,8.896
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cudnn,False,6.112
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cudnn,True,8.832
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cudnn,False,6.176
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cudnn,True,8.944
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cutlass,False,7.280
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cutlass,True,10.000
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cutlass,False,7.216
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cutlass,True,9.984
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_trtllm,False,7.583
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_trtllm,True,10.303
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_trtllm,False,7.520
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_trtllm,True,10.288
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cudnn,False,7.456
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cudnn,True,10.016
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cudnn,False,7.472
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cudnn,True,10.192
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutedsl,False,6.928
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutedsl,True,9.488
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutedsl,False,6.160
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutedsl,True,8.880
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutlass,False,7.232
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutlass,True,9.792
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutlass,False,8.320
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutlass,True,11.040
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_trtllm,False,7.024
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_trtllm,True,9.488
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_trtllm,False,6.976
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_trtllm,True,9.584
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,bf16,False,9.504
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,bf16,False,9.791
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cudnn,False,6.560
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cudnn,True,9.344
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cudnn,False,6.688
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cudnn,True,9.488
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cutlass,False,7.712
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cutlass,True,10.496
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cutlass,False,7.680
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cutlass,True,10.480
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_trtllm,False,8.560
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_trtllm,True,11.344
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_trtllm,False,8.623
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_trtllm,True,11.424
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cudnn,False,8.464
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cudnn,True,11.216
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cudnn,False,8.512
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cudnn,True,11.360
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutedsl,False,7.584
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutedsl,True,10.336
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutedsl,False,6.416
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutedsl,True,9.264
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutlass,False,8.992
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutlass,True,11.744
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutlass,False,8.800
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutlass,True,11.648
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_trtllm,False,7.648
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_trtllm,True,10.208
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_trtllm,False,7.104
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_trtllm,True,9.887
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,bf16,False,8.704
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,bf16,False,8.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cudnn,False,6.512
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cudnn,True,9.232
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cudnn,False,6.464
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cudnn,True,9.232
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cutlass,False,8.016
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cutlass,True,10.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cutlass,False,7.648
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cutlass,True,10.416
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_trtllm,False,7.776
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_trtllm,True,10.496
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_trtllm,False,7.904
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_trtllm,True,10.672
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cudnn,False,7.808
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cudnn,True,10.367
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cudnn,False,7.920
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cudnn,True,10.640
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutedsl,False,7.104
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutedsl,True,9.664
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutedsl,False,6.304
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutedsl,True,9.024
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutlass,False,8.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutlass,True,11.296
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutlass,False,7.583
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutlass,True,10.303
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_trtllm,False,6.784
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_trtllm,True,9.248
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_trtllm,False,7.040
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_trtllm,True,9.647
model.layers.*.mixer.o_proj,1,2688,4096,bf16,False,10.368
model.layers.*.mixer.o_proj,16,2688,4096,bf16,False,10.880
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cudnn,False,6.912
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cudnn,True,9.663
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cudnn,False,7.008
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cudnn,True,9.888
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cutlass,False,8.160
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cutlass,True,10.912
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cutlass,False,8.064
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cutlass,True,10.944
model.layers.*.mixer.o_proj,1,2688,4096,fp8_trtllm,False,8.480
model.layers.*.mixer.o_proj,1,2688,4096,fp8_trtllm,True,11.232
model.layers.*.mixer.o_proj,16,2688,4096,fp8_trtllm,False,8.752
model.layers.*.mixer.o_proj,16,2688,4096,fp8_trtllm,True,11.632
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cudnn,False,8.768
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cudnn,True,11.520
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cudnn,False,8.864
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cudnn,True,11.680
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutedsl,False,7.904
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutedsl,True,10.656
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutedsl,False,6.784
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutedsl,True,9.600
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutlass,False,9.248
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutlass,True,12.000
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutlass,False,9.264
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutlass,True,12.080
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_trtllm,False,7.184
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_trtllm,True,9.744
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_trtllm,False,7.744
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_trtllm,True,10.432
model.layers.*.mixer.out_proj,1,2688,4096,bf16,False,10.368
model.layers.*.mixer.out_proj,16,2688,4096,bf16,False,10.880
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cudnn,False,6.912
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cudnn,True,9.663
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cudnn,False,7.008
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cudnn,True,9.888
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cutlass,False,8.160
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cutlass,True,10.912
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cutlass,False,8.064
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cutlass,True,10.944
model.layers.*.mixer.out_proj,1,2688,4096,fp8_trtllm,False,8.480
model.layers.*.mixer.out_proj,1,2688,4096,fp8_trtllm,True,11.232
model.layers.*.mixer.out_proj,16,2688,4096,fp8_trtllm,False,8.752
model.layers.*.mixer.out_proj,16,2688,4096,fp8_trtllm,True,11.632
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cudnn,False,8.768
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cudnn,True,11.520
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cudnn,False,8.864
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cudnn,True,11.680
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutedsl,False,7.904
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutedsl,True,10.656
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutedsl,False,6.784
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutedsl,True,9.600
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutlass,False,9.248
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutlass,True,12.000
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutlass,False,9.264
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutlass,True,12.080
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_trtllm,False,7.184
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_trtllm,True,9.744
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_trtllm,False,7.744
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_trtllm,True,10.432
model.layers.*.mixer.in_proj,1,10304,2688,bf16,False,17.520
model.layers.*.mixer.in_proj,16,10304,2688,bf16,False,16.032
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cudnn,False,8.912
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cudnn,True,11.632
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cudnn,False,9.375
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cudnn,True,12.143
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cutlass,False,10.527
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cutlass,True,13.248
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cutlass,False,10.415
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cutlass,True,13.183
model.layers.*.mixer.in_proj,1,10304,2688,fp8_trtllm,False,9.328
model.layers.*.mixer.in_proj,1,10304,2688,fp8_trtllm,True,12.048
model.layers.*.mixer.in_proj,16,10304,2688,fp8_trtllm,False,9.392
model.layers.*.mixer.in_proj,16,10304,2688,fp8_trtllm,True,12.160
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cudnn,False,9.088
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cudnn,True,11.648
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cudnn,False,9.024
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cudnn,True,11.744
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutedsl,False,9.424
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutedsl,True,11.984
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutedsl,False,7.616
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutedsl,True,10.336
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutlass,False,9.728
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutlass,True,12.288
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutlass,False,10.144
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutlass,True,12.864
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_trtllm,False,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_trtllm,True,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_trtllm,False,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_trtllm,True,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
MoE
H=2688 F=1856 E=32 top_k=6 activation=Relu2
module_name,M,N,K,backend,with_quant,runtime
model.layers.*.mixer.experts,1,,,bf16_cutlass,False,54.271
model.layers.*.mixer.experts,16,,,bf16_cutlass,False,175.552
model.layers.*.mixer.experts,1,,,fp8_cutlass,False,41.312
model.layers.*.mixer.experts,1,,,fp8_cutlass,True,44.032
model.layers.*.mixer.experts,16,,,fp8_cutlass,False,94.799
model.layers.*.mixer.experts,16,,,fp8_cutlass,True,97.567
model.layers.*.mixer.experts,1,,,fp8_trtllm,False,26.480
model.layers.*.mixer.experts,1,,,fp8_trtllm,True,29.200
model.layers.*.mixer.experts,16,,,fp8_trtllm,False,97.680
model.layers.*.mixer.experts,16,,,fp8_trtllm,True,100.448
model.layers.*.mixer.experts,1,,,nvfp4_cutlass,False,40.944
model.layers.*.mixer.experts,1,,,nvfp4_cutlass,True,40.623
model.layers.*.mixer.experts,16,,,nvfp4_cutlass,False,70.432
model.layers.*.mixer.experts,16,,,nvfp4_cutlass,True,71.072
model.layers.*.mixer.experts,1,,,nvfp4_trtllm,False,22.784
model.layers.*.mixer.experts,1,,,nvfp4_trtllm,True,25.119
model.layers.*.mixer.experts,16,,,nvfp4_trtllm,False,57.407
model.layers.*.mixer.experts,16,,,nvfp4_trtllm,True,59.999
```
### Testing
run locally
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added a guided workflow for planning and running per-rank GEMM and
fused-MoE kernel benchmarks.
- Added automatic model inspection to derive supported kernel shapes,
tensor-parallel and expert-parallel settings, routing details, and
required benchmark parameters.
- Added FlashInfer built-in benchmark execution with BF16, FP8, NVFP4,
and activation-quantization timing.
- Added structured previews, logs, CSV reports, error details, and
dry-run validation.
- **Bug Fixes**
- Improved handling of unsupported layouts, padding, missing results,
invalid settings, and benchmark failures.
- **Tests**
- Added extensive coverage for model inspection, argument validation,
benchmark generation, quantization, routing, reporting, and error
handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
This commit is contained in:
sychen52
@@ -0,0 +1,145 @@
|
||||
---
|
||||
name: benchmark-model-kernels
|
||||
description: >-
|
||||
Inspect Hugging Face decoder layers on meta tensors and plan or run per-rank
|
||||
BF16, FP8, and NVFP4 GEMM or fused-MoE microbenchmarks with the bundled
|
||||
scripts and a local FlashInfer checkout. Use when choosing a model, GPU, TP,
|
||||
EP, or M/token-concurrency sweep; deriving common fused QKV and gate/up
|
||||
shapes without loading checkpoint weights; or using FlashInfer benchmark
|
||||
utilities. Do not use for end-to-end server throughput or request latency.
|
||||
---
|
||||
|
||||
# Benchmark Model Kernels
|
||||
|
||||
Plan a single-GPU microbenchmark for the per-rank shapes of an intended
|
||||
deployment. The scripts never load checkpoint weights, launch distributed
|
||||
workers, or measure collectives, serving throughput, or request latency.
|
||||
|
||||
## Interview, one decision per message
|
||||
|
||||
Ask exactly one unresolved decision per message with two or three concrete
|
||||
options, state the default, and wait for the answer. Recommend an option only
|
||||
when it is substantively better. Skip decisions already answered. Follow this
|
||||
order:
|
||||
|
||||
1. **Model** — a local directory, a `config.json`, or a Hub ID (equivalent; no
|
||||
default). The script builds the model on meta tensors from configuration
|
||||
only. Do not enable `--trust-remote-code` without explicit approval; pin
|
||||
`--revision <commit>` when it is needed.
|
||||
2. **TP and EP** — accept both directly, or derive them from the intended
|
||||
deployment's GPU model and count and confirm. Defaults: `TP=1`, `EP=1`.
|
||||
Derive parallelism from the intended deployment, never from the GPU that
|
||||
runs the microbenchmark. The script validates every sharding rule
|
||||
(divisibility, GQA replication, expert partitioning) and errors loudly.
|
||||
3. **M sweep** — balanced default (recommended): `1 8 64 512`; decode-focused:
|
||||
`1 4 16 32`; throughput-focused: `64 256 1024 4096`. M is roughly the
|
||||
tokens scheduled per step (decode: active sequences), not endpoint
|
||||
concurrency. With EP, an MoE row models one rank's share of the global
|
||||
batch: a global batch of B tokens corresponds to the column M = B/EP, so
|
||||
do not compare different EP values at the same M.
|
||||
4. **Shape preview** — always run this before anything else; no FlashInfer or
|
||||
GPU needed:
|
||||
|
||||
```bash
|
||||
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py <model> \
|
||||
--tp <tp> --ep <ep> --ms <m1> <m2> ... --print_only
|
||||
```
|
||||
|
||||
Review the printed shapes and the MoE tuple with the user.
|
||||
`# unsupported:` lines mean the list is partial and the script exits
|
||||
nonzero — handle those via **Manual supplements** below.
|
||||
5. **FlashInfer checkout and GPU** — the full benchmark needs a FlashInfer
|
||||
*source checkout* containing `benchmarks/flashinfer_benchmark.py`; the
|
||||
installed wheel alone is not enough. Prefer a clean checkout matching the
|
||||
installed `flashinfer` version; ask before cloning or installing anything.
|
||||
On the benchmark machine, check `nvidia-smi` and package versions, verify
|
||||
the target GPU is idle (concurrent work on the same GPU skews timings),
|
||||
and verify CUPTI timing with a tiny `bench_gpu_time(..., enable_cupti=True)`
|
||||
probe — a warning that falls back to CUDA events is a failure (the
|
||||
`cupti-python`/`nvidia-cuda-cupti` packages must match PyTorch's CUDA
|
||||
major). Ask for a GPU index only when several are visible. Pick a fresh
|
||||
workdir and state it.
|
||||
6. **Full benchmark**:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=<gpu-index> \
|
||||
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py <model> \
|
||||
--tp <tp> --ep <ep> --ms <m1> <m2> ... \
|
||||
--flashinfer_repo <flashinfer-repo> --workdir <workdir>
|
||||
```
|
||||
|
||||
Run a short plumbing check first (append
|
||||
`--dry_run_iters 1 --num_iters 3 --no_autotune`, throwaway workdir),
|
||||
especially after changing FlashInfer versions or shape logic; never present
|
||||
it as a performance result. Then run with defaults. The first fused-MoE
|
||||
build or autotune can take several minutes; its cache is reused.
|
||||
|
||||
## Rules the scripts own
|
||||
|
||||
Do not restate, re-derive, or override these — the code enforces them:
|
||||
|
||||
- `benchmark_model.py` derives `--nks` (as `N,K,module-name` triples) and
|
||||
all `--moe_*` arguments; overrides are rejected.
|
||||
- Row labels are logical shapes; the runner applies vLLM's physical padding.
|
||||
Report both shapes whenever they differ.
|
||||
- A failed case writes FlashInfer's error message (or a pointer to
|
||||
`driver.log`) into its `combined_results.csv` cell, and the command exits
|
||||
nonzero after the table is written. Never present a partial table as a
|
||||
successful benchmark; read `driver.log` before rerunning anything.
|
||||
|
||||
## Known backend and driver limits
|
||||
|
||||
- `mm_fp4` TensorRT-LLM needs `N % 128 == 0` (shuffled weight layout): its
|
||||
cell reports no result row, or on stock 0.6.x drivers an empty assertion
|
||||
that fails every `mm_fp4` backend for that shape.
|
||||
- `mm_fp8` (trtllm_low_latency) needs `K % 128 == 0`.
|
||||
- MoE rows cover FlashInfer's CUTLASS fused MoE, the trtllm-gen NVFP4 and
|
||||
per-tensor FP8 MoE, and the CuteDSL NVFP4 MoE. The CuteDSL MoE row appears
|
||||
only for Swiglu models (the kernel supports nothing else); the `cutedsl` and
|
||||
`trtllm` backends require recent GPUs and report per-case errors elsewhere.
|
||||
- Gated NVFP4 CUTLASS MoE with `2F % 128 != 0` per rank fails — vLLM raises
|
||||
instead of padding — often as a `CUDA error: misaligned address`. Prefer EP
|
||||
over TP for the experts to keep the per-rank width legal.
|
||||
- Do not benchmark a padded shape to dodge these limits and call it
|
||||
vLLM-equivalent; report the limit instead.
|
||||
|
||||
## Manual supplements
|
||||
|
||||
For each `# unsupported:` layout from the preview: inspect that module's
|
||||
forward path and the intended runtime's TP/EP sharding, derive the per-rank
|
||||
shape (never guess sharding from a weight shape alone), and benchmark only the
|
||||
missing shape:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=<gpu-index> \
|
||||
python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \
|
||||
--flashinfer_repo <flashinfer-repo> --ms <m1> <m2> ... \
|
||||
--nks <n>,<k>,<name> --workdir <workdir>
|
||||
```
|
||||
|
||||
Label these rows as manual supplements; this is not a substitute for running
|
||||
`benchmark_model.py` first. Shell-quote user-supplied paths.
|
||||
|
||||
## Report
|
||||
|
||||
Report the command, GPU, versions, TP/EP, M values, shapes (logical and
|
||||
physical where they differ), warnings, and the artifacts: `testlist.txt`,
|
||||
`driver.log` (full driver output), `builtin_results.csv` (milliseconds,
|
||||
success-only), and `combined_results.csv` (long form, microseconds: columns `module_name,
|
||||
M, N, K, backend, with_quant, runtime` in `GEMM` and `MoE` sections; modules
|
||||
fused into one GEMM are joined with `|` in `module_name`, while distinct
|
||||
same-shape modules appear as duplicated rows sharing one measurement; MoE
|
||||
rows keep the `H= F= E= top_k=` parameter line and leave N/K empty). In quantization-recipe terms:
|
||||
`bf16` rows are the unquantized W16A16 baseline, `fp8` rows are per-tensor
|
||||
W8A8, and `nvfp4` rows are W4A4. Plain quantized rows time the kernel with
|
||||
pre-quantized activations; `*_with_quant` rows add a separately measured
|
||||
activation-quantization time in the scale-factor layout that backend
|
||||
consumes, except the NVFP4 CUTLASS MoE row, which is a single fused
|
||||
measurement. MoE routing is synthetic: uniform expert
|
||||
distribution everywhere (real skewed routing is slower), and the trtllm-gen
|
||||
rows, which route in-kernel, use a fixed `renormalize` method regardless of
|
||||
the model's routing scheme. CUTLASS and CuteDSL MoE rows receive precomputed
|
||||
expert indices and exclude routing-selection cost entirely. No row includes
|
||||
the router GEMM itself. These are kernel times
|
||||
only — never describe them as end-to-end latency or throughput; they omit
|
||||
weights, layer frequency, communication, KV cache, and scheduling.
|
||||
@@ -0,0 +1,843 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Derive per-rank benchmark shapes from a Transformers model on meta tensors.
|
||||
|
||||
The script walks the instantiated model's Linear modules, fuses Q/K/V and
|
||||
gate/up projections, recognizes Mamba 2, GatedDeltaNet, and common
|
||||
routed-expert layouts, and applies the common serving/export TP layout. It
|
||||
never calls a checkpoint weight loader.
|
||||
When a decoder layout is unsupported, the derived shapes are still printed and
|
||||
the script exits nonzero; benchmark the missing shapes directly with
|
||||
benchmark_via_builtin.py.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import re
|
||||
import shlex
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ShapeError(ValueError):
|
||||
"""A model layer cannot be represented by the supported benchmark layout."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _MoeShape:
|
||||
hidden: int
|
||||
intermediate: int
|
||||
experts: int
|
||||
top_k: int
|
||||
activation: str | None = None
|
||||
# The expert container's module path. Metadata, not identity: two layouts
|
||||
# that differ only by path are still one shape for dedup purposes.
|
||||
name: str = field(default="experts", compare=False)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ExpertShape:
|
||||
hidden: int
|
||||
intermediate: int
|
||||
gated: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Kernel:
|
||||
"""One derived per-rank GEMM: N x K, labeled by its source module path(s)."""
|
||||
|
||||
n: int
|
||||
k: int
|
||||
label: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _MambaLayout:
|
||||
in_shape: tuple[int, int]
|
||||
out_shape: tuple[int, int]
|
||||
intermediate: int
|
||||
heads: int
|
||||
groups: int
|
||||
state: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _GdnLayout:
|
||||
out_shape: tuple[int, int]
|
||||
num_k_heads: int
|
||||
num_v_heads: int
|
||||
key_dim: int
|
||||
value_dim: int
|
||||
|
||||
|
||||
_PROJECTIONS = {
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"gate_up_proj",
|
||||
"down_proj",
|
||||
}
|
||||
_RESERVED = {
|
||||
"--nks",
|
||||
"--moe_name",
|
||||
"--moe_hidden_size",
|
||||
"--moe_intermediate_size",
|
||||
"--moe_num_experts",
|
||||
"--moe_top_k",
|
||||
"--moe_activation_type",
|
||||
}
|
||||
|
||||
# Modules deliberately outside the benchmarked GEMM list fall into two
|
||||
# exclusion mechanisms:
|
||||
# 1. Positional: embeddings, the LM head, and anything else outside the
|
||||
# decoder blocks never enter the audit (the `layers.<i>` position filter
|
||||
# in _unsupported_decoder_linears).
|
||||
# 2. Name-based: routing and gating projections that live inside decoder
|
||||
# blocks are excluded by these names. They are never quantized in
|
||||
# deployment recipes and vLLM dispatches them through specialized (often
|
||||
# FP32-output) paths, so a standard GEMM row would not model them anyway.
|
||||
_ROUTER_PATH_PARTS = {"router", "routers"}
|
||||
_GATING_LEAF_NAMES = {"gate", "router", "router_proj", "shared_expert_gate"}
|
||||
|
||||
|
||||
def _positive_int(value: str) -> int:
|
||||
try:
|
||||
parsed = int(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"expected a positive integer, got {value!r}") from exc
|
||||
if parsed <= 0:
|
||||
raise argparse.ArgumentTypeError(f"expected a positive integer, got {value!r}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _load_meta_model(model_ref: str, trust_remote_code: bool, revision: str | None):
|
||||
# Transformers and Accelerate are optional, heavy ModelOpt dependencies.
|
||||
try:
|
||||
from accelerate import init_empty_weights
|
||||
from transformers import AutoConfig, AutoModelForCausalLM
|
||||
except ImportError as exc:
|
||||
raise ShapeError("install ModelOpt with the 'hf' extra") from exc
|
||||
|
||||
path = Path(model_ref).expanduser()
|
||||
ref = str(path) if path.exists() else model_ref
|
||||
try:
|
||||
config = AutoConfig.from_pretrained(
|
||||
ref, trust_remote_code=trust_remote_code, revision=revision
|
||||
)
|
||||
model_kwargs = {"trust_remote_code": trust_remote_code}
|
||||
auto_map = getattr(config, "auto_map", {}) or {}
|
||||
if revision and trust_remote_code and "AutoModelForCausalLM" in auto_map:
|
||||
model_kwargs["code_revision"] = revision
|
||||
with init_empty_weights(include_buffers=True):
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_config(config, **model_kwargs)
|
||||
except Exception:
|
||||
text_config = getattr(config, "text_config", None)
|
||||
if text_config is None:
|
||||
raise
|
||||
# Multimodal wrapper configs build their text decoder
|
||||
# directly; vision towers are outside the benchmark scope.
|
||||
model = AutoModelForCausalLM.from_config(text_config, **model_kwargs)
|
||||
except Exception as exc:
|
||||
raise ShapeError(f"could not construct {model_ref!r} on meta tensors: {exc}") from exc
|
||||
|
||||
tensors = list(model.named_parameters()) + list(model.named_buffers())
|
||||
materialized = [name for name, tensor in tensors if not tensor.is_meta]
|
||||
if materialized:
|
||||
raise ShapeError(f"model construction allocated tensors: {', '.join(materialized[:3])}")
|
||||
return config, model
|
||||
|
||||
|
||||
def _linear_shape(module: Any) -> tuple[int, int] | None:
|
||||
if hasattr(module, "out_features") and hasattr(module, "in_features"):
|
||||
return int(module.out_features), int(module.in_features)
|
||||
return None
|
||||
|
||||
|
||||
def _divide(value: int, size: int, label: str) -> int:
|
||||
if value % size:
|
||||
raise ShapeError(f"{label}={value} is not divisible by {size}")
|
||||
return value // size
|
||||
|
||||
|
||||
def _path_label(parent: str, leaf: str) -> str:
|
||||
"""Label a GEMM by its module path with layer indices normalized to ``*``."""
|
||||
path = f"{parent}.{leaf}" if parent else leaf
|
||||
return re.sub(r"(?<=\.)\d+(?=\.|$)", "*", path)
|
||||
|
||||
|
||||
def _fused_qkv(
|
||||
q: tuple[int, int],
|
||||
k: tuple[int, int],
|
||||
v: tuple[int, int],
|
||||
head_dim: int,
|
||||
tp: int,
|
||||
parent: str,
|
||||
) -> _Kernel:
|
||||
if q[1] != k[1] or q[1] != v[1] or k[0] != v[0]:
|
||||
raise ShapeError(f"unsupported Q/K/V shapes under {parent}")
|
||||
q_heads = _divide(q[0], head_dim, f"{parent}.q_proj output")
|
||||
kv_heads = _divide(k[0], head_dim, f"{parent}.k_proj output")
|
||||
if kv_heads >= tp:
|
||||
_divide(q_heads, tp, f"{parent}.q_proj heads")
|
||||
_divide(kv_heads, tp, f"{parent}.k_proj heads")
|
||||
local_n = _divide(q[0] + k[0] + v[0], tp, f"{parent}.qkv")
|
||||
else:
|
||||
local_q = _divide(q_heads, tp, f"{parent}.q_proj heads") * head_dim
|
||||
_divide(tp, kv_heads, "TP/KV replication ratio")
|
||||
local_n = local_q + 2 * head_dim
|
||||
label = "|".join(_path_label(parent, leaf) for leaf in ("q_proj", "k_proj", "v_proj"))
|
||||
return _Kernel(local_n, q[1], label)
|
||||
|
||||
|
||||
def _dense_kernels(model: Any, config: Any, tp: int) -> list[_Kernel]:
|
||||
groups: dict[str, dict[str, tuple[int, int]]] = {}
|
||||
for name, module in model.named_modules():
|
||||
leaf = name.rsplit(".", 1)[-1]
|
||||
parts = name.split(".")
|
||||
# Routed experts are excluded here and handled by _moe_shapes. Shared
|
||||
# experts (for example `.shared_experts.up_proj`) intentionally do not
|
||||
# match the filter: they run densely for every token, so their
|
||||
# projections belong in the dense GEMM list.
|
||||
if (
|
||||
leaf not in _PROJECTIONS
|
||||
or ".experts." in name
|
||||
or ".local_experts." in name
|
||||
or any(part in _ROUTER_PATH_PARTS for part in parts[:-1])
|
||||
):
|
||||
continue
|
||||
shape = _linear_shape(module)
|
||||
if shape:
|
||||
groups.setdefault(name.rpartition(".")[0], {})[leaf] = shape
|
||||
|
||||
head_dim = getattr(config, "head_dim", None)
|
||||
if head_dim is None:
|
||||
head_dim = config.hidden_size // config.num_attention_heads
|
||||
kernels = []
|
||||
for parent, layers in groups.items():
|
||||
kernels += _parent_kernels(parent, layers, int(head_dim), tp)
|
||||
return list(dict.fromkeys(kernels))
|
||||
|
||||
|
||||
def _parent_kernels(
|
||||
parent: str, layers: dict[str, tuple[int, int]], head_dim: int, tp: int
|
||||
) -> list[_Kernel]:
|
||||
"""Derive one parent module's GEMMs from its recognized projections."""
|
||||
kernels = []
|
||||
qkv = {"q_proj", "k_proj", "v_proj"}
|
||||
present_qkv = qkv.intersection(layers)
|
||||
if present_qkv:
|
||||
if present_qkv != qkv:
|
||||
raise ShapeError(f"incomplete Q/K/V projections under {parent}")
|
||||
kernels.append(
|
||||
_fused_qkv(
|
||||
layers["q_proj"],
|
||||
layers["k_proj"],
|
||||
layers["v_proj"],
|
||||
head_dim,
|
||||
tp,
|
||||
parent,
|
||||
)
|
||||
)
|
||||
if "o_proj" in layers:
|
||||
# Row-parallel: TP shards K.
|
||||
n, k = layers["o_proj"]
|
||||
kernels.append(
|
||||
_Kernel(n, _divide(k, tp, f"{parent}.o_proj"), _path_label(parent, "o_proj"))
|
||||
)
|
||||
|
||||
if "gate_proj" in layers and "up_proj" in layers:
|
||||
gate_n, gate_k = layers["gate_proj"]
|
||||
up_n, up_k = layers["up_proj"]
|
||||
if gate_k != up_k:
|
||||
raise ShapeError(f"gate/up inputs differ under {parent}")
|
||||
# vLLM shards gate and up individually before merging them, so
|
||||
# each output must divide by TP, not just their sum.
|
||||
n = _divide(gate_n, tp, f"{parent}.gate_proj") + _divide(up_n, tp, f"{parent}.up_proj")
|
||||
label = "|".join(_path_label(parent, leaf) for leaf in ("gate_proj", "up_proj"))
|
||||
kernels.append(_Kernel(n, gate_k, label))
|
||||
elif "gate_proj" in layers:
|
||||
raise ShapeError(f"gate projection has no matching up projection under {parent}")
|
||||
elif "up_proj" in layers:
|
||||
# Column-parallel: TP shards N.
|
||||
n, k = layers["up_proj"]
|
||||
kernels.append(
|
||||
_Kernel(_divide(n, tp, f"{parent}.up_proj"), k, _path_label(parent, "up_proj"))
|
||||
)
|
||||
if "gate_up_proj" in layers:
|
||||
# Column-parallel: TP shards N (the checkpoint pre-fused gate and up).
|
||||
n, k = layers["gate_up_proj"]
|
||||
kernels.append(
|
||||
_Kernel(
|
||||
_divide(n, tp, f"{parent}.gate_up_proj"), k, _path_label(parent, "gate_up_proj")
|
||||
)
|
||||
)
|
||||
if "down_proj" in layers:
|
||||
# Row-parallel: TP shards K.
|
||||
n, k = layers["down_proj"]
|
||||
kernels.append(
|
||||
_Kernel(n, _divide(k, tp, f"{parent}.down_proj"), _path_label(parent, "down_proj"))
|
||||
)
|
||||
return kernels
|
||||
|
||||
|
||||
def _mamba_layout(module: Any) -> _MambaLayout | None:
|
||||
in_shape = _linear_shape(getattr(module, "in_proj", None))
|
||||
out_shape = _linear_shape(getattr(module, "out_proj", None))
|
||||
intermediate = getattr(module, "intermediate_size", None)
|
||||
heads = getattr(module, "num_heads", None)
|
||||
groups = getattr(module, "n_groups", None)
|
||||
state = getattr(module, "ssm_state_size", None)
|
||||
if (
|
||||
in_shape is None
|
||||
or out_shape is None
|
||||
or intermediate is None
|
||||
or heads is None
|
||||
or groups is None
|
||||
or state is None
|
||||
):
|
||||
return None
|
||||
return _MambaLayout(in_shape, out_shape, int(intermediate), int(heads), int(groups), int(state))
|
||||
|
||||
|
||||
def _mamba_kernels(model: Any, tp: int) -> list[_Kernel]:
|
||||
kernels = []
|
||||
for name, module in model.named_modules():
|
||||
layout = _mamba_layout(module)
|
||||
if layout is None:
|
||||
continue
|
||||
hidden = layout.in_shape[1]
|
||||
expected_in = 2 * layout.intermediate + 2 * layout.groups * layout.state + layout.heads
|
||||
if layout.in_shape[0] != expected_in or layout.out_shape != (hidden, layout.intermediate):
|
||||
raise ShapeError(f"unsupported Mamba projection shapes under {name}")
|
||||
|
||||
local_intermediate = _divide(layout.intermediate, tp, f"{name}.intermediate_size")
|
||||
local_heads = _divide(layout.heads, tp, f"{name}.num_heads")
|
||||
if layout.groups % tp == 0:
|
||||
local_groups = layout.groups // tp
|
||||
elif layout.groups == 1:
|
||||
local_groups = 1
|
||||
else:
|
||||
raise ShapeError(f"{name}.n_groups={layout.groups} is not divisible by TP={tp}")
|
||||
local_in = 2 * local_intermediate + 2 * local_groups * layout.state + local_heads
|
||||
kernels.extend(
|
||||
[
|
||||
_Kernel(local_in, hidden, _path_label(name, "in_proj")),
|
||||
_Kernel(hidden, local_intermediate, _path_label(name, "out_proj")),
|
||||
]
|
||||
)
|
||||
return list(dict.fromkeys(kernels))
|
||||
|
||||
|
||||
def _gdn_layout(module: Any) -> _GdnLayout | None:
|
||||
out_shape = _linear_shape(getattr(module, "out_proj", None))
|
||||
num_k_heads = getattr(module, "num_k_heads", None)
|
||||
num_v_heads = getattr(module, "num_v_heads", None)
|
||||
key_dim = getattr(module, "key_dim", None)
|
||||
value_dim = getattr(module, "value_dim", None)
|
||||
has_input = (
|
||||
getattr(module, "in_proj_qkvz", None) is not None
|
||||
or getattr(module, "in_proj_qkv", None) is not None
|
||||
)
|
||||
if (
|
||||
out_shape is None
|
||||
or not has_input
|
||||
or num_k_heads is None
|
||||
or num_v_heads is None
|
||||
or key_dim is None
|
||||
or value_dim is None
|
||||
):
|
||||
return None
|
||||
return _GdnLayout(out_shape, int(num_k_heads), int(num_v_heads), int(key_dim), int(value_dim))
|
||||
|
||||
|
||||
def _gdn_kernels(model: Any, tp: int) -> list[_Kernel]:
|
||||
kernels = []
|
||||
for name, module in model.named_modules():
|
||||
layout = _gdn_layout(module)
|
||||
if layout is None:
|
||||
continue
|
||||
key_dim, value_dim = layout.key_dim, layout.value_dim
|
||||
num_v_heads = layout.num_v_heads
|
||||
hidden = layout.out_shape[0]
|
||||
fused_qkvz = _linear_shape(getattr(module, "in_proj_qkvz", None))
|
||||
qkvz_label = _path_label(name, "in_proj_qkvz")
|
||||
ba_label = _path_label(name, "in_proj_ba")
|
||||
if fused_qkvz is not None:
|
||||
# Qwen3-Next stores qkvz and ba pre-fused.
|
||||
ba = _linear_shape(getattr(module, "in_proj_ba", None))
|
||||
valid = fused_qkvz == (2 * key_dim + 2 * value_dim, hidden) and ba == (
|
||||
2 * num_v_heads,
|
||||
hidden,
|
||||
)
|
||||
else:
|
||||
# Qwen3.5 stores them split, but vLLM's shared GDN mixer merges
|
||||
# qkv+z and b+a into the same two per-rank GEMMs either way.
|
||||
qkvz_label = "|".join(_path_label(name, leaf) for leaf in ("in_proj_qkv", "in_proj_z"))
|
||||
ba_label = "|".join(_path_label(name, leaf) for leaf in ("in_proj_b", "in_proj_a"))
|
||||
shapes = {
|
||||
leaf: _linear_shape(getattr(module, leaf, None))
|
||||
for leaf in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
|
||||
}
|
||||
valid = (
|
||||
shapes["in_proj_qkv"] == (2 * key_dim + value_dim, hidden)
|
||||
and shapes["in_proj_z"] == (value_dim, hidden)
|
||||
and shapes["in_proj_b"] == (num_v_heads, hidden)
|
||||
and shapes["in_proj_a"] == (num_v_heads, hidden)
|
||||
)
|
||||
if not valid or layout.out_shape != (hidden, value_dim):
|
||||
raise ShapeError(f"unsupported GatedDeltaNet projection shapes under {name}")
|
||||
_divide(layout.num_k_heads, tp, f"{name}.linear_num_key_heads")
|
||||
local_v_heads = _divide(num_v_heads, tp, f"{name}.linear_num_value_heads")
|
||||
kernels.extend(
|
||||
[
|
||||
_Kernel(
|
||||
_divide(2 * key_dim + 2 * value_dim, tp, f"{name}.qkvz"),
|
||||
hidden,
|
||||
qkvz_label,
|
||||
),
|
||||
_Kernel(2 * local_v_heads, hidden, ba_label),
|
||||
_Kernel(
|
||||
hidden,
|
||||
_divide(value_dim, tp, f"{name}.value_dim"),
|
||||
_path_label(name, "out_proj"),
|
||||
),
|
||||
]
|
||||
)
|
||||
return list(dict.fromkeys(kernels))
|
||||
|
||||
|
||||
def _expert_shape(module: Any) -> _ExpertShape | None:
|
||||
gate = getattr(module, "gate_proj", None)
|
||||
up = getattr(module, "up_proj", None)
|
||||
down = getattr(module, "down_proj", None)
|
||||
if gate is None and getattr(module, "w1", None) is not None:
|
||||
gate = getattr(module, "w1", None)
|
||||
up = getattr(module, "w3", None)
|
||||
down = getattr(module, "w2", None)
|
||||
if gate is not None:
|
||||
gate_shape, up_shape, down_shape = map(_linear_shape, (gate, up, down))
|
||||
if gate_shape is None or up_shape is None or down_shape is None:
|
||||
raise ShapeError("incomplete gated expert Linear layout")
|
||||
if gate_shape != up_shape or down_shape != (gate_shape[1], gate_shape[0]):
|
||||
raise ShapeError("unsupported gated expert Linear shapes")
|
||||
return _ExpertShape(gate_shape[1], gate_shape[0], True)
|
||||
|
||||
up_shape, down_shape = map(_linear_shape, (up, down))
|
||||
if up_shape is None and down_shape is None:
|
||||
return None
|
||||
if up_shape is None or down_shape is None or down_shape != (up_shape[1], up_shape[0]):
|
||||
raise ShapeError("unsupported non-gated expert Linear shapes")
|
||||
return _ExpertShape(up_shape[1], up_shape[0], False)
|
||||
|
||||
|
||||
def _stacked_expert_shape(
|
||||
first: Any, down: Any, factor: int, name: str, expected_hidden: int | None
|
||||
) -> _ExpertShape:
|
||||
if first.ndim != 3 or down.ndim != 3 or first.shape[0] != down.shape[0]:
|
||||
raise ShapeError(f"unsupported stacked experts at {name}")
|
||||
|
||||
first_shape = tuple(int(value) for value in first.shape)
|
||||
down_shape = tuple(int(value) for value in down.shape)
|
||||
candidates = []
|
||||
if first_shape[1] % factor == 0:
|
||||
hidden, intermediate = first_shape[2], first_shape[1] // factor
|
||||
if down_shape[1:] == (hidden, intermediate):
|
||||
candidates.append((hidden, intermediate))
|
||||
if first_shape[2] % factor == 0:
|
||||
hidden, intermediate = first_shape[1], first_shape[2] // factor
|
||||
if down_shape[1:] == (intermediate, hidden):
|
||||
candidates.append((hidden, intermediate))
|
||||
if expected_hidden is not None:
|
||||
candidates = [candidate for candidate in candidates if candidate[0] == expected_hidden]
|
||||
if len(set(candidates)) != 1:
|
||||
raise ShapeError(f"unsupported stacked expert projection shapes at {name}")
|
||||
hidden, intermediate = candidates[0]
|
||||
return _ExpertShape(hidden, intermediate, factor == 2)
|
||||
|
||||
|
||||
def _moe_activation(config: Any, gated: bool) -> str | None:
|
||||
configured = getattr(config, "mlp_hidden_act", None) or getattr(config, "hidden_act", None)
|
||||
normalized = str(configured).lower().replace("-", "_")
|
||||
if gated:
|
||||
if normalized in {"silu", "swiglu", "swish"}:
|
||||
return "Swiglu"
|
||||
# vLLM's FlashInfer MoE path serves only the tanh-approximation GELU
|
||||
# ("gelu_tanh", "gelu_pytorch_tanh", and "gelu_new" share that
|
||||
# formula). Exact gelu and quick_gelu take non-FlashInfer backends in
|
||||
# vLLM, so they are rejected here rather than timed as a proxy.
|
||||
if normalized in {"gelu_tanh", "gelu_pytorch_tanh", "gelu_new"}:
|
||||
return "Geglu"
|
||||
raise ShapeError(
|
||||
f"unsupported gated MoE activation {configured!r}; vLLM's FlashInfer MoE "
|
||||
"serves only SiLU/SwiGLU and tanh-GELU"
|
||||
)
|
||||
activations = {
|
||||
"gelu": "Gelu",
|
||||
"identity": "Identity",
|
||||
"relu": "Relu",
|
||||
"relu2": "Relu2",
|
||||
"relu_squared": "Relu2",
|
||||
"silu": "Silu",
|
||||
}
|
||||
if normalized not in activations:
|
||||
raise ShapeError(f"unsupported non-gated MoE activation {configured!r}")
|
||||
return activations[normalized]
|
||||
|
||||
|
||||
# Fallback copy of ModelOpt's _ACTIVE_MOE_TOP_K_ATTRS for environments without
|
||||
# ModelOpt installed; a test asserts it stays in sync with the canonical list.
|
||||
_MOE_TOP_K_ATTRS_FALLBACK = (
|
||||
"num_experts_per_tok",
|
||||
"num_experts_per_token",
|
||||
"moe_top_k",
|
||||
"top_k",
|
||||
"num_selected_experts",
|
||||
)
|
||||
|
||||
|
||||
def _top_k(config: Any) -> int | None:
|
||||
# ModelOpt's AutoQuantize cost model owns the canonical attribute list, so
|
||||
# benchmark rows and AutoQuantize agree on how a config declares top_k.
|
||||
try:
|
||||
from modelopt.torch.quantization._auto_quantize_cost import _ACTIVE_MOE_TOP_K_ATTRS
|
||||
|
||||
attrs = _ACTIVE_MOE_TOP_K_ATTRS
|
||||
except ImportError:
|
||||
attrs = _MOE_TOP_K_ATTRS_FALLBACK
|
||||
|
||||
for attr in attrs:
|
||||
value = getattr(config, attr, None)
|
||||
if value is not None:
|
||||
return int(value)
|
||||
return None
|
||||
|
||||
|
||||
def _moe_shapes(model: Any, config: Any) -> set[_MoeShape]:
|
||||
shapes = set()
|
||||
top_k = _top_k(config)
|
||||
for name, module in model.named_modules():
|
||||
expert_container = name.rsplit(".", 1)[-1] in {"experts", "local_experts"}
|
||||
if expert_container:
|
||||
expert_modules = list(module.children())
|
||||
shape = _expert_shape(expert_modules[0]) if expert_modules else None
|
||||
if shape:
|
||||
if top_k is None:
|
||||
raise ShapeError("could not determine MoE top_k")
|
||||
if any(_expert_shape(expert) != shape for expert in expert_modules[1:]):
|
||||
raise ShapeError(f"experts under {name} do not share one Linear layout")
|
||||
shapes.add(
|
||||
_MoeShape(
|
||||
shape.hidden,
|
||||
shape.intermediate,
|
||||
len(expert_modules),
|
||||
top_k,
|
||||
_moe_activation(config, shape.gated),
|
||||
_path_label(*name.rpartition(".")[::2]),
|
||||
)
|
||||
)
|
||||
|
||||
params = dict(module.named_parameters(recurse=False))
|
||||
down = params.get("down_proj")
|
||||
if params.get("gate_up_proj") is not None:
|
||||
first, factor = params["gate_up_proj"], 2
|
||||
elif params.get("up_proj") is not None:
|
||||
first, factor = params["up_proj"], 1
|
||||
else:
|
||||
expert_params = [
|
||||
f"{param_name}{tuple(param.shape)}"
|
||||
for param_name, param in params.items()
|
||||
if param.ndim >= 2
|
||||
]
|
||||
if expert_container and expert_params:
|
||||
raise ShapeError(
|
||||
f"unsupported stacked expert parameters at {name}: " + ", ".join(expert_params)
|
||||
)
|
||||
continue
|
||||
if down is None:
|
||||
raise ShapeError(f"stacked experts at {name} have no down projection")
|
||||
if top_k is None:
|
||||
raise ShapeError("could not determine MoE top_k")
|
||||
expected_hidden = getattr(config, "moe_latent_size", None) or getattr(
|
||||
config, "hidden_size", None
|
||||
)
|
||||
shape = _stacked_expert_shape(
|
||||
first,
|
||||
down,
|
||||
factor,
|
||||
name,
|
||||
int(expected_hidden) if expected_hidden is not None else None,
|
||||
)
|
||||
shapes.add(
|
||||
_MoeShape(
|
||||
shape.hidden,
|
||||
shape.intermediate,
|
||||
int(first.shape[0]),
|
||||
top_k,
|
||||
_moe_activation(config, shape.gated),
|
||||
_path_label(*name.rpartition(".")[::2]),
|
||||
)
|
||||
)
|
||||
return shapes
|
||||
|
||||
|
||||
def _declared_expert_count(config: Any) -> int | None:
|
||||
if _top_k(config) is None:
|
||||
return None
|
||||
for attr in ("n_routed_experts", "num_local_experts", "num_experts"):
|
||||
value = getattr(config, attr, None)
|
||||
if value is not None and int(value) > 0:
|
||||
return int(value)
|
||||
return None
|
||||
|
||||
|
||||
def _mixer_claimed_projections(model: Any) -> set[str]:
|
||||
"""Full paths of the Linears the Mamba and GDN recognizers already derive."""
|
||||
claimed: set[str] = set()
|
||||
for parent, module in model.named_modules():
|
||||
if _mamba_layout(module) is not None:
|
||||
claimed.update({f"{parent}.in_proj", f"{parent}.out_proj"})
|
||||
if _gdn_layout(module) is not None:
|
||||
claimed.update(
|
||||
f"{parent}.{leaf}"
|
||||
for leaf in (
|
||||
"in_proj_qkvz",
|
||||
"in_proj_ba",
|
||||
"in_proj_qkv",
|
||||
"in_proj_z",
|
||||
"in_proj_b",
|
||||
"in_proj_a",
|
||||
"out_proj",
|
||||
)
|
||||
)
|
||||
return claimed
|
||||
|
||||
|
||||
def _unsupported_decoder_linears(
|
||||
model: Any, routed_experts_handled: bool = False
|
||||
) -> list[tuple[str, int, int]]:
|
||||
claimed = _mixer_claimed_projections(model)
|
||||
layouts: dict[tuple[str, int, int], str] = {}
|
||||
for name, module in model.named_modules():
|
||||
shape = _linear_shape(module)
|
||||
if shape is None:
|
||||
continue
|
||||
parts = name.split(".")
|
||||
in_decoder = any(
|
||||
part in {"block", "blocks", "h", "layer", "layers"}
|
||||
and index + 1 < len(parts)
|
||||
and parts[index + 1].isdigit()
|
||||
for index, part in enumerate(parts)
|
||||
)
|
||||
leaf = parts[-1]
|
||||
if not in_decoder:
|
||||
continue
|
||||
if any(part in _ROUTER_PATH_PARTS for part in parts):
|
||||
continue
|
||||
if routed_experts_handled and any(part in {"experts", "local_experts"} for part in parts):
|
||||
continue
|
||||
if leaf in _PROJECTIONS or leaf in _GATING_LEAF_NAMES or name in claimed:
|
||||
continue
|
||||
layouts.setdefault((leaf, *shape), name)
|
||||
return [(name, n, k) for (leaf, n, k), name in layouts.items()]
|
||||
|
||||
|
||||
def _audited_moe_shape(model: Any, config: Any) -> tuple[_MoeShape | None, bool, list[str]]:
|
||||
"""Derive the model's global routed-expert shape, or its audit problems.
|
||||
|
||||
Returns ``(shape, experts_recognized, problems)``; the shape is ``None``
|
||||
whenever any problem is found, so the audit findings are reported instead
|
||||
of a masking per-rank ShapeError.
|
||||
"""
|
||||
problems: list[str] = []
|
||||
try:
|
||||
moe_shapes = _moe_shapes(model, config)
|
||||
except ShapeError as exc:
|
||||
moe_shapes = set()
|
||||
problems.append(str(exc))
|
||||
declared_experts = _declared_expert_count(config)
|
||||
if not problems and declared_experts is not None:
|
||||
if not moe_shapes:
|
||||
problems.append(
|
||||
f"model declares {declared_experts} routed experts but no supported expert "
|
||||
"GEMM layout was found"
|
||||
)
|
||||
elif any(shape.experts != declared_experts for shape in moe_shapes):
|
||||
found = sorted({shape.experts for shape in moe_shapes})
|
||||
problems.append(
|
||||
f"model declares {declared_experts} routed experts but instantiated layouts "
|
||||
f"have expert counts {found}"
|
||||
)
|
||||
experts_recognized = bool(moe_shapes)
|
||||
if len(moe_shapes) > 1:
|
||||
problems.append("model contains multiple routed-expert layouts")
|
||||
if problems:
|
||||
moe_shapes = set()
|
||||
return next(iter(moe_shapes), None), experts_recognized, problems
|
||||
|
||||
|
||||
def _shard_moe(moe: _MoeShape, tp: int, ep: int) -> _MoeShape:
|
||||
"""Apply vLLM's EP/TP partitioning to the global MoE shape."""
|
||||
if ep != 1 and ep % tp:
|
||||
raise ShapeError(
|
||||
f"EP={ep} is not a multiple of TP={tp}; vLLM expert parallelism spans TP x DP, "
|
||||
"so no modeled serving layout matches this combination — if it is intentional "
|
||||
"(e.g. Megatron-style EP), benchmark the per-rank expert shape directly with "
|
||||
"benchmark_via_builtin.py"
|
||||
)
|
||||
local_experts = _divide(moe.experts, ep, "expert count")
|
||||
intermediate = moe.intermediate
|
||||
if ep == 1:
|
||||
intermediate = _divide(intermediate, tp, "expert intermediate size")
|
||||
if moe.top_k > local_experts:
|
||||
raise ShapeError("top_k exceeds the per-rank expert count")
|
||||
return _MoeShape(moe.hidden, intermediate, local_experts, moe.top_k, moe.activation, moe.name)
|
||||
|
||||
|
||||
def _inspect_model(
|
||||
model: Any, config: Any, tp: int, ep: int
|
||||
) -> tuple[list[_Kernel], _MoeShape | None, list[str]]:
|
||||
config = getattr(config, "text_config", None) or config
|
||||
kernels = (
|
||||
_dense_kernels(model, config, tp) + _mamba_kernels(model, tp) + _gdn_kernels(model, tp)
|
||||
)
|
||||
moe, experts_recognized, problems = _audited_moe_shape(model, config)
|
||||
if moe is None:
|
||||
if ep != 1 and not problems:
|
||||
raise ShapeError("EP requires routed experts")
|
||||
else:
|
||||
moe = _shard_moe(moe, tp, ep)
|
||||
unsupported = _unsupported_decoder_linears(model, routed_experts_handled=experts_recognized)
|
||||
if unsupported:
|
||||
details = ", ".join(f"{name} ({n}x{k})" for name, n, k in unsupported)
|
||||
problems.append(f"unsupported decoder Linear GEMM layout(s): {details}")
|
||||
if not kernels and moe is None and not problems:
|
||||
raise ShapeError("no dense benchmark shapes found")
|
||||
return kernels, moe, problems
|
||||
|
||||
|
||||
def _command(
|
||||
kernels: list[_Kernel],
|
||||
moe: _MoeShape | None,
|
||||
passthrough: list[str],
|
||||
) -> list[str]:
|
||||
command: list[str] = []
|
||||
if kernels:
|
||||
# One N,K,NAME argument per derived kernel: same-shape kernels from
|
||||
# different modules keep separate names and become duplicated rows.
|
||||
command += ["--nks", *(f"{kernel.n},{kernel.k},{kernel.label}" for kernel in kernels)]
|
||||
if moe:
|
||||
command += [
|
||||
"--moe_hidden_size",
|
||||
str(moe.hidden),
|
||||
"--moe_intermediate_size",
|
||||
str(moe.intermediate),
|
||||
"--moe_num_experts",
|
||||
str(moe.experts),
|
||||
"--moe_top_k",
|
||||
str(moe.top_k),
|
||||
"--moe_name",
|
||||
moe.name,
|
||||
]
|
||||
if moe.activation:
|
||||
command += ["--moe_activation_type", moe.activation]
|
||||
return command + passthrough
|
||||
|
||||
|
||||
_RUNNER_PATH = Path(__file__).with_name("benchmark_via_builtin.py")
|
||||
|
||||
|
||||
def _load_runner() -> Any:
|
||||
spec = importlib.util.spec_from_file_location("benchmark_via_builtin", _RUNNER_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _print_preview(
|
||||
model: Any,
|
||||
config: Any,
|
||||
tp: int,
|
||||
ep: int,
|
||||
kernels: list[_Kernel],
|
||||
moe: _MoeShape | None,
|
||||
problems: list[str],
|
||||
) -> None:
|
||||
print(f"# {type(model).__name__} ({getattr(config, 'model_type', '?')}), TP={tp}, EP={ep}")
|
||||
print(
|
||||
"# layout: Transformers meta model; fused QKV and gate/up; "
|
||||
"Mamba 2, GatedDeltaNet, and routed experts"
|
||||
)
|
||||
for kernel in dict.fromkeys(kernels):
|
||||
print(f"# {kernel.n}x{kernel.k} <- {kernel.label}")
|
||||
if moe:
|
||||
activation = f" activation={moe.activation}" if moe.activation else ""
|
||||
print(
|
||||
f"# MoE: H={moe.hidden} F={moe.intermediate} E={moe.experts} "
|
||||
f"top_k={moe.top_k}{activation}"
|
||||
)
|
||||
if ep > 1:
|
||||
print(
|
||||
f"# MoE sharding: EP={ep} partitions whole experts; "
|
||||
"expert width stays intact (expert-TP=1)"
|
||||
)
|
||||
elif tp > 1:
|
||||
print(f"# MoE sharding: TP={tp} shards the expert intermediate width (EP=1)")
|
||||
for problem in problems:
|
||||
print(f"# unsupported: {problem}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Parse arguments, derive shapes, and optionally run the benchmark."""
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("model", help="Hub ID, local model directory, or config.json")
|
||||
parser.add_argument("--tp", type=_positive_int, default=1, help="tensor parallel size, e.g. 8")
|
||||
parser.add_argument("--ep", type=_positive_int, default=1, help="expert parallel size, e.g. 8")
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
parser.add_argument("--revision", help="Hugging Face branch, tag, or commit")
|
||||
parser.add_argument("--print_only", action="store_true")
|
||||
args, passthrough = parser.parse_known_args()
|
||||
for token in passthrough:
|
||||
if token.split("=", 1)[0] in _RESERVED:
|
||||
parser.error("derived --nks/--nk_names/--moe_* shapes cannot be overridden")
|
||||
|
||||
try:
|
||||
config, model = _load_meta_model(args.model, args.trust_remote_code, args.revision)
|
||||
kernels, moe, problems = _inspect_model(model, config, args.tp, args.ep)
|
||||
except ShapeError as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
_print_preview(model, config, args.tp, args.ep, kernels, moe, problems)
|
||||
if problems:
|
||||
parser.error(
|
||||
"the derived shapes above are incomplete; validate each unsupported layout's "
|
||||
"TP/EP sharding and benchmark it directly with benchmark_via_builtin.py"
|
||||
)
|
||||
command = _command(kernels, moe, passthrough)
|
||||
print(">>> " + shlex.join([sys.executable, str(_RUNNER_PATH), *command]))
|
||||
if not args.print_only:
|
||||
_load_runner().main(command)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,867 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Run FlashInfer's built-in GEMM and fused-MoE microbenchmarks.
|
||||
|
||||
Plain rows contain kernel time. Most ``*_with_quant`` rows add a separately
|
||||
measured activation-quantization time in the scale-factor layout the backend
|
||||
consumes; the NVFP4 CUTLASS MoE row is instead a single fused measurement.
|
||||
Logical shapes label each case while backend-specific physical padding follows
|
||||
vLLM. A local FlashInfer source checkout is required for its benchmark driver
|
||||
and utilities.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import os
|
||||
import shlex
|
||||
import subprocess # nosec B404
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import TextIO
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from vllm import _custom_ops as vllm_ops
|
||||
except ImportError:
|
||||
vllm_ops = None
|
||||
|
||||
_ERROR_CASE_PREFIX = "[ERROR] Error running test:"
|
||||
_ERROR_MESSAGE_PREFIX = "[ERROR] Error:"
|
||||
_FP8_QUANT_UNAVAILABLE = "ERROR: vLLM is unavailable for FP8 activation quantization"
|
||||
_MOE_ACTIVATIONS = (
|
||||
"Gelu",
|
||||
"Relu",
|
||||
"Silu",
|
||||
"Swiglu",
|
||||
"Geglu",
|
||||
"SwigluBias",
|
||||
"Relu2",
|
||||
"SwigluStep",
|
||||
"Identity",
|
||||
)
|
||||
_ResultValue = float | str
|
||||
|
||||
|
||||
@dataclass(frozen=True, order=True)
|
||||
class _QuantSpec:
|
||||
"""One shared activation-quantization measurement of an m-by-k BF16 tile.
|
||||
|
||||
Attributes:
|
||||
dtype: Quantized element type, ``"nvfp4"`` or ``"fp8"``.
|
||||
layout: Scale-factor layout the consuming kernel expects: ``"128x4"``,
|
||||
``"8x4"``, or ``"linear"`` for NVFP4; ``"static"`` for FP8.
|
||||
m: Token count of the activation tile.
|
||||
k: Inner dimension of the activation tile.
|
||||
"""
|
||||
|
||||
dtype: str
|
||||
layout: str
|
||||
m: int
|
||||
k: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Case:
|
||||
"""One driver invocation and, once it has run, its outcome.
|
||||
|
||||
Attributes:
|
||||
section: ``"gemm"`` or ``"moe"``.
|
||||
tag: Case label passed to the driver as ``--case_tag``; validates
|
||||
returned rows and labels the artifacts. Never an internal key.
|
||||
backend: Value of the output ``backend`` column.
|
||||
m: Token count.
|
||||
n: Logical output size; the driver may run a padded physical shape
|
||||
from ``argv``. ``None`` for MoE cases.
|
||||
k: Logical reduction size, as ``n``.
|
||||
argv: Driver arguments, before ``--case_tag``/``--output_path``.
|
||||
with_quant: ``with_quant`` column of the measured row itself; ``True``
|
||||
only for the fused NVFP4 CUTLASS MoE measurement.
|
||||
quant: Spec of the activation-quantization time a derived
|
||||
``with_quant`` row adds to ``result``.
|
||||
result: Median kernel time in microseconds, an ``ERROR: ...`` message,
|
||||
or ``None`` until the case has run.
|
||||
quant_result: Measured time (or error) of ``quant``, recorded by
|
||||
``_attach_quant_times``.
|
||||
"""
|
||||
|
||||
section: str
|
||||
tag: str
|
||||
backend: str
|
||||
m: int
|
||||
n: int | None
|
||||
k: int | None
|
||||
argv: list[str]
|
||||
with_quant: bool = False
|
||||
quant: _QuantSpec | None = None
|
||||
result: _ResultValue | None = None
|
||||
quant_result: _ResultValue | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _MoeShape:
|
||||
"""The per-rank fused-MoE problem all MoE cases share.
|
||||
|
||||
Attributes:
|
||||
hidden: Model hidden size.
|
||||
intermediate: Per-rank expert width.
|
||||
experts: Per-rank expert count.
|
||||
top_k: Experts activated per token.
|
||||
activation: FlashInfer activation name; ``None`` means the driver
|
||||
default (gated SwiGLU).
|
||||
name: Module path of the expert container, used as the MoE rows'
|
||||
``module_name``.
|
||||
"""
|
||||
|
||||
hidden: int
|
||||
intermediate: int
|
||||
experts: int
|
||||
top_k: int
|
||||
activation: str | None = None
|
||||
name: str = "experts"
|
||||
|
||||
def label(self) -> str:
|
||||
label = f"H={self.hidden} F={self.intermediate} E={self.experts} top_k={self.top_k}"
|
||||
if self.activation:
|
||||
label += f" activation={self.activation}"
|
||||
return label
|
||||
|
||||
|
||||
def _positive_int(value: str) -> int:
|
||||
try:
|
||||
parsed = int(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"{value!r} is not an integer") from exc
|
||||
if parsed <= 0:
|
||||
raise argparse.ArgumentTypeError(f"{value!r} is not a positive integer")
|
||||
return parsed
|
||||
|
||||
|
||||
def _nk_arg(value: str) -> tuple[int, int, str | None]:
|
||||
"""Parse one GEMM shape argument: ``N,K`` or ``N,K,NAME``.
|
||||
|
||||
Module names never contain commas, so a fourth field is a typo (two
|
||||
shapes missing their separating space) and is rejected loudly.
|
||||
"""
|
||||
fields = value.split(",")
|
||||
if len(fields) not in (2, 3) or (len(fields) == 3 and not fields[2]):
|
||||
raise argparse.ArgumentTypeError(f"expected N,K or N,K,NAME, got {value!r}")
|
||||
try:
|
||||
n, k = _positive_int(fields[0]), _positive_int(fields[1])
|
||||
except argparse.ArgumentTypeError as exc:
|
||||
raise argparse.ArgumentTypeError(f"expected positive N and K in {value!r}: {exc}") from exc
|
||||
return n, k, fields[2] if len(fields) == 3 else None
|
||||
|
||||
|
||||
def _labels_by_nk(nk_args: list[tuple[int, int, str | None]]) -> dict[tuple[int, int], list[str]]:
|
||||
"""Map each unique N,K, in first-seen order, to its module labels.
|
||||
|
||||
Unnamed shapes label themselves ``"NxK"``.
|
||||
"""
|
||||
labels_by_nk: dict[tuple[int, int], list[str]] = {}
|
||||
for n, k, name in nk_args:
|
||||
labels = labels_by_nk.setdefault((n, k), [])
|
||||
label = name if name is not None else f"{n}x{k}"
|
||||
if label not in labels:
|
||||
labels.append(label)
|
||||
return labels_by_nk
|
||||
|
||||
|
||||
def _round_up(value: int, alignment: int) -> int:
|
||||
return (value + alignment - 1) // alignment * alignment
|
||||
|
||||
|
||||
def _parse_driver_error(lines: list[str]) -> str | None:
|
||||
"""Extract the driver's error message from one case's output.
|
||||
|
||||
Returns the message, ``""`` when the driver reported an error without a
|
||||
message, or ``None`` when no error was reported. Each case runs in its own
|
||||
driver process, so any reported error belongs to that case.
|
||||
"""
|
||||
error = None
|
||||
pending = False
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped.startswith(_ERROR_CASE_PREFIX):
|
||||
pending = True
|
||||
elif stripped.startswith(_ERROR_MESSAGE_PREFIX) and pending:
|
||||
error = stripped.removeprefix(_ERROR_MESSAGE_PREFIX).strip().replace(",", ";")
|
||||
pending = False
|
||||
return error
|
||||
|
||||
|
||||
def _failure_message(case_output: list[str], returncode: int, driver_log: Path) -> str:
|
||||
"""Classify a case that produced no result row into an ``ERROR: ...`` cell."""
|
||||
message = _parse_driver_error(case_output)
|
||||
if message:
|
||||
return f"ERROR: {message}"
|
||||
if message is not None:
|
||||
return (
|
||||
"ERROR: FlashInfer reported an error without a message (empty exception); "
|
||||
f"see {driver_log}"
|
||||
)
|
||||
if returncode:
|
||||
return (
|
||||
f"ERROR: FlashInfer driver exited with status {returncode} for this case; "
|
||||
f"see {driver_log}"
|
||||
)
|
||||
return f"ERROR: FlashInfer produced no result row and no error message; see {driver_log}"
|
||||
|
||||
|
||||
def _run_case(benchmarks_dir: Path, argv: list[str], log: TextIO) -> tuple[int, list[str]]:
|
||||
# Each case gets its own driver process: a fatal CUDA fault (for example a
|
||||
# misaligned address) permanently poisons the CUDA context, so sharing one
|
||||
# process would fail every later case (verified empirically). This invokes
|
||||
# the explicitly selected FlashInfer checkout without a shell.
|
||||
process = subprocess.Popen( # nosec B603
|
||||
[sys.executable, "flashinfer_benchmark.py", *argv],
|
||||
cwd=benchmarks_dir,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
assert process.stdout is not None
|
||||
lines = []
|
||||
for line in process.stdout:
|
||||
print(line, end="", flush=True)
|
||||
log.write(line)
|
||||
lines.append(line)
|
||||
return process.wait(), lines
|
||||
|
||||
|
||||
def _gpu_description() -> str:
|
||||
try:
|
||||
import torch
|
||||
|
||||
# The driver name can be a placeholder on pre-release GPUs, so record
|
||||
# compute capability, SM count, and memory to pin down the exact part.
|
||||
properties = torch.cuda.get_device_properties(0)
|
||||
name = (
|
||||
f"{properties.name} (sm_{properties.major}{properties.minor} / "
|
||||
f"{properties.multi_processor_count} SMs / "
|
||||
f"{properties.total_memory / (1 << 30):.0f} GiB)"
|
||||
)
|
||||
except Exception:
|
||||
return "unknown GPU"
|
||||
watts = "unknown power limit"
|
||||
try:
|
||||
import pynvml
|
||||
|
||||
pynvml.nvmlInit()
|
||||
try:
|
||||
# NVML does not honor CUDA_VISIBLE_DEVICES, so map the first
|
||||
# visible device back to its physical NVML handle.
|
||||
visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")[0].strip()
|
||||
if visible.startswith(("GPU-", "MIG-")):
|
||||
handle = pynvml.nvmlDeviceGetHandleByUUID(visible)
|
||||
else:
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(int(visible) if visible else 0)
|
||||
limit = pynvml.nvmlDeviceGetPowerManagementLimit(handle)
|
||||
watts = f"{limit / 1000:.0f} W power limit"
|
||||
finally:
|
||||
pynvml.nvmlShutdown()
|
||||
except Exception:
|
||||
pass
|
||||
return f"{name}; {watts}"
|
||||
|
||||
|
||||
def _environment_header(flashinfer_repo: Path) -> str:
|
||||
try:
|
||||
import flashinfer
|
||||
|
||||
version = flashinfer.__version__
|
||||
except Exception:
|
||||
version = "unknown"
|
||||
try:
|
||||
# Reads the revision of the explicitly selected checkout, no shell.
|
||||
result = subprocess.run( # nosec B603 B607
|
||||
["git", "-C", str(flashinfer_repo), "rev-parse", "HEAD"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
check=False,
|
||||
)
|
||||
revision = result.stdout.strip() or "unknown"
|
||||
except OSError:
|
||||
revision = "unknown"
|
||||
return (
|
||||
f"flashinfer {version}; checkout {flashinfer_repo.resolve()} @ {revision}; "
|
||||
f"{_gpu_description()}"
|
||||
)
|
||||
|
||||
|
||||
def _write_builtin(path: Path, rows: list[dict[str, str]]) -> None:
|
||||
fieldnames: dict[str, None] = {}
|
||||
for row in rows:
|
||||
for key in row:
|
||||
fieldnames.setdefault(key, None)
|
||||
with path.open("w", newline="") as stream:
|
||||
writer = csv.DictWriter(
|
||||
stream, fieldnames=list(fieldnames), restval="", lineterminator="\n"
|
||||
)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def _gemm_cases(
|
||||
ms: list[int],
|
||||
nks: list[tuple[int, int]],
|
||||
common: list[str],
|
||||
) -> list[_Case]:
|
||||
cases: list[_Case] = []
|
||||
|
||||
def add(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
backend: str,
|
||||
routine: str,
|
||||
driver_backend: str,
|
||||
run_n: int | None = None,
|
||||
run_k: int | None = None,
|
||||
extra: list[str] | None = None,
|
||||
quant: _QuantSpec | None = None,
|
||||
) -> None:
|
||||
cases.append(
|
||||
_Case(
|
||||
section="gemm",
|
||||
tag=f"gemm_{backend}_MxNxK={m}x{n}x{k}",
|
||||
backend=backend,
|
||||
m=m,
|
||||
n=n,
|
||||
k=k,
|
||||
argv=[
|
||||
"--routine",
|
||||
routine,
|
||||
"--backends",
|
||||
driver_backend,
|
||||
*(extra or []),
|
||||
"--m",
|
||||
str(m),
|
||||
"--n",
|
||||
str(run_n if run_n is not None else n),
|
||||
"--k",
|
||||
str(run_k if run_k is not None else k),
|
||||
*common,
|
||||
],
|
||||
quant=quant,
|
||||
)
|
||||
)
|
||||
|
||||
for m in ms:
|
||||
for n, k in nks:
|
||||
# Physical padding follows vLLM: dense NVFP4 on cuDNN, CUTLASS,
|
||||
# and CuteDSL pads N and K to multiples of 32; trtllm keeps the
|
||||
# exact shape with the shuffled layout; BF16 and FP8 stay exact.
|
||||
add(m, n, k, "bf16", "mm_bf16", "cudnn")
|
||||
for row_suffix, driver_backend in (
|
||||
("cudnn", "cudnn"),
|
||||
("cutlass", "cutlass"),
|
||||
("cutedsl", "cute-dsl"),
|
||||
("trtllm", "trtllm"),
|
||||
):
|
||||
layout = "128x4" if driver_backend != "trtllm" or m > 32 else "8x4"
|
||||
extra = ["--use_nvfp4"]
|
||||
if layout == "128x4":
|
||||
extra.append("--use_128x4_sf_layout")
|
||||
run_n, run_k = n, k
|
||||
if driver_backend != "trtllm":
|
||||
run_n, run_k = _round_up(n, 32), _round_up(k, 32)
|
||||
add(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
f"nvfp4_{row_suffix}",
|
||||
"mm_fp4",
|
||||
driver_backend,
|
||||
run_n,
|
||||
run_k,
|
||||
extra,
|
||||
_QuantSpec("nvfp4", layout, m, run_k),
|
||||
)
|
||||
for driver_backend in ("cudnn", "cutlass"):
|
||||
add(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
f"fp8_{driver_backend}",
|
||||
"bmm_fp8",
|
||||
driver_backend,
|
||||
extra=["--batch_size", "1"],
|
||||
quant=_QuantSpec("fp8", "static", m, k),
|
||||
)
|
||||
add(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
"fp8_trtllm",
|
||||
"mm_fp8",
|
||||
"trtllm_low_latency",
|
||||
quant=_QuantSpec("fp8", "static", m, k),
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
def _moe_cases(ms: list[int], shape: _MoeShape, common: list[str]) -> list[_Case]:
|
||||
cases: list[_Case] = []
|
||||
|
||||
def add(
|
||||
backend: str,
|
||||
routine: str,
|
||||
intermediate: int,
|
||||
hidden: int = shape.hidden,
|
||||
extra: list[str] | None = None,
|
||||
quant: tuple[str, str] | None = None,
|
||||
with_quant: bool = False,
|
||||
) -> None:
|
||||
suffix = "_with_quant" if with_quant else ""
|
||||
cases.extend(
|
||||
_Case(
|
||||
section="moe",
|
||||
tag=f"moe_{backend}_moe{suffix}_M={m}",
|
||||
backend=backend,
|
||||
m=m,
|
||||
n=None,
|
||||
k=None,
|
||||
with_quant=with_quant,
|
||||
quant=_QuantSpec(*quant, m, hidden) if quant else None,
|
||||
argv=[
|
||||
"--routine",
|
||||
routine,
|
||||
"--num_tokens",
|
||||
str(m),
|
||||
"--hidden_size",
|
||||
str(hidden),
|
||||
"--num_experts",
|
||||
str(shape.experts),
|
||||
"--top_k",
|
||||
str(shape.top_k),
|
||||
*(["--activation-type", shape.activation] if shape.activation else []),
|
||||
"--intermediate_size",
|
||||
str(intermediate),
|
||||
*(extra or []),
|
||||
*common,
|
||||
],
|
||||
)
|
||||
for m in ms
|
||||
)
|
||||
|
||||
gated = shape.activation is None or shape.activation in {
|
||||
"Swiglu",
|
||||
"Geglu",
|
||||
"SwigluBias",
|
||||
"SwigluStep",
|
||||
}
|
||||
# Pad a dimension only when vLLM pads it. FP8 per-tensor (CUTLASS and
|
||||
# trtllm-gen) pads the intermediate to 16 gated / 128 non-gated. NVFP4
|
||||
# CUTLASS pads non-gated intermediate up to the 128-aligned swizzled scale
|
||||
# rows but raises instead of padding gated, so gated stays exact and may
|
||||
# fail like vLLM. NVFP4 trtllm-gen additionally pads hidden to 256.
|
||||
fp8_intermediate = _round_up(shape.intermediate, 16 if gated else 128)
|
||||
nvfp4_intermediate = shape.intermediate if gated else _round_up(shape.intermediate, 128)
|
||||
|
||||
add("bf16_cutlass", "cutlass_fused_moe", shape.intermediate)
|
||||
add(
|
||||
"fp8_cutlass",
|
||||
"cutlass_fused_moe",
|
||||
fp8_intermediate,
|
||||
extra=["--cutlass_variant", "fp8"],
|
||||
quant=("fp8", "static"),
|
||||
)
|
||||
add(
|
||||
"nvfp4_cutlass",
|
||||
"cutlass_fused_moe",
|
||||
nvfp4_intermediate,
|
||||
extra=["--cutlass_variant", "nvfp4", "--quantized_input"],
|
||||
)
|
||||
# The NVFP4 CUTLASS with_quant row is its own fused driver measurement
|
||||
# (unquantized input), not a derived base-plus-quant-time row.
|
||||
add(
|
||||
"nvfp4_cutlass",
|
||||
"cutlass_fused_moe",
|
||||
nvfp4_intermediate,
|
||||
extra=["--cutlass_variant", "nvfp4"],
|
||||
with_quant=True,
|
||||
)
|
||||
# Routing is synthetic in this benchmark (uniform random logits), so the
|
||||
# trtllm-gen rows, which route in-kernel, use a fixed renormalize method
|
||||
# to stay comparable across models; the model's real routing scheme is
|
||||
# not derivable from its config alone. CUTLASS and CuteDSL rows receive
|
||||
# precomputed indices and have no routing stage to time.
|
||||
add(
|
||||
"fp8_trtllm",
|
||||
"trtllm_fp8_per_tensor_scale_moe",
|
||||
fp8_intermediate,
|
||||
extra=["--routing_method", "renormalize"],
|
||||
quant=("fp8", "static"),
|
||||
)
|
||||
add(
|
||||
"nvfp4_trtllm",
|
||||
"trtllm_fp4_block_scale_moe",
|
||||
fp8_intermediate,
|
||||
hidden=_round_up(shape.hidden, 256),
|
||||
extra=["--routing_method", "renormalize"],
|
||||
quant=("nvfp4", "linear"),
|
||||
)
|
||||
if shape.activation in (None, "Swiglu"):
|
||||
# FlashInfer's CuteDSL fused MoE supports only gated Swiglu.
|
||||
add(
|
||||
"nvfp4_cutedsl",
|
||||
"cute_dsl_fp4_block_scale_moe",
|
||||
shape.intermediate,
|
||||
quant=("nvfp4", "linear"),
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
def _nvfp4_runner(tensor: torch.Tensor, layout: str):
|
||||
import flashinfer
|
||||
|
||||
global_scale = (448 * 6) / tensor.float().abs().nan_to_num().max()
|
||||
if layout == "linear":
|
||||
# The trtllm-gen and CuteDSL fused-MoE kernels consume activation
|
||||
# scale factors in linear (unswizzled) layout.
|
||||
def linear_kernel(value, scale):
|
||||
return flashinfer.fp4_quantize(value, scale, is_sf_swizzled_layout=False)
|
||||
|
||||
return linear_kernel, (tensor, global_scale)
|
||||
sf_layout = (
|
||||
flashinfer.SfLayout.layout_128x4 if layout == "128x4" else flashinfer.SfLayout.layout_8x4
|
||||
)
|
||||
|
||||
def kernel(value, scale):
|
||||
return flashinfer.nvfp4_quantize(value, scale, sfLayout=sf_layout, do_shuffle=False)
|
||||
|
||||
return kernel, (tensor, global_scale)
|
||||
|
||||
|
||||
def _fp8_runner(tensor: torch.Tensor):
|
||||
import torch
|
||||
|
||||
scale = tensor.abs().max().float() / torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
def kernel(value, value_scale):
|
||||
quantized, _ = vllm_ops.scaled_fp8_quant(value.contiguous(), value_scale)
|
||||
return quantized
|
||||
|
||||
return kernel, (tensor, scale)
|
||||
|
||||
|
||||
def _attach_quant_times(
|
||||
cases: list[_Case], dry_runs: int, iterations: int, cuda_graph: bool
|
||||
) -> None:
|
||||
"""Measure each distinct quantization spec once and attach shared times."""
|
||||
specs = sorted(
|
||||
{case.quant for case in cases if case.quant is not None and isinstance(case.result, float)}
|
||||
)
|
||||
results: dict[_QuantSpec, _ResultValue] = {}
|
||||
if vllm_ops is None and any(spec.dtype == "fp8" for spec in specs):
|
||||
print(f"[WARN] {_FP8_QUANT_UNAVAILABLE.removeprefix('ERROR: ')}")
|
||||
results = {spec: _FP8_QUANT_UNAVAILABLE for spec in specs if spec.dtype == "fp8"}
|
||||
specs = [spec for spec in specs if spec.dtype != "fp8"]
|
||||
if specs:
|
||||
# The GPU stack is imported lazily so shape planning, result parsing,
|
||||
# and their tests work without FlashInfer or torch installed.
|
||||
import numpy as np
|
||||
import torch
|
||||
from flashinfer.testing import bench_gpu_time
|
||||
|
||||
for spec in specs:
|
||||
tensor = torch.randn(spec.m, spec.k, device="cuda", dtype=torch.bfloat16)
|
||||
kernel, inputs = (
|
||||
_nvfp4_runner(tensor, spec.layout) if spec.dtype == "nvfp4" else _fp8_runner(tensor)
|
||||
)
|
||||
times = bench_gpu_time(
|
||||
fn=kernel,
|
||||
input_args=inputs,
|
||||
dry_run_iters=dry_runs,
|
||||
repeat_iters=iterations,
|
||||
enable_cupti=True,
|
||||
use_cuda_graph=cuda_graph,
|
||||
cold_l2_cache=True,
|
||||
sleep_after_run=True,
|
||||
)
|
||||
results[spec] = float(np.median(times)) * 1000
|
||||
for case in cases:
|
||||
if case.quant is not None and isinstance(case.result, float):
|
||||
case.quant_result = results[case.quant]
|
||||
|
||||
|
||||
def _format_result(value: _ResultValue) -> str:
|
||||
if isinstance(value, float):
|
||||
return f"{value:.3f}"
|
||||
return value
|
||||
|
||||
|
||||
def _output_rows(case: _Case) -> list[tuple[bool, _ResultValue]]:
|
||||
"""Expand a case into its (with_quant, runtime) output rows.
|
||||
|
||||
A case with a quant spec gets a derived ``with_quant`` row adding the
|
||||
shared activation-quantization time; an error on either measurement
|
||||
propagates into the derived row.
|
||||
"""
|
||||
assert case.result is not None
|
||||
rows = [(case.with_quant, case.result)]
|
||||
if case.quant is None:
|
||||
return rows
|
||||
if isinstance(case.result, str):
|
||||
rows.append((True, case.result))
|
||||
return rows
|
||||
quant_result = case.quant_result
|
||||
assert quant_result is not None
|
||||
rows.append(
|
||||
(True, case.result + quant_result if isinstance(quant_result, float) else quant_result)
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def _write_results(
|
||||
path: Path,
|
||||
cases: list[_Case],
|
||||
labels_by_nk: dict[tuple[int, int], list[str]],
|
||||
header: str | None = None,
|
||||
moe_shape: _MoeShape | None = None,
|
||||
) -> None:
|
||||
columns = ["module_name", "M", "N", "K", "backend", "with_quant", "runtime"]
|
||||
gemm = [case for case in cases if case.section == "gemm" and case.result is not None]
|
||||
moe = [case for case in cases if case.section == "moe" and case.result is not None]
|
||||
with path.open("w", newline="") as stream:
|
||||
writer = csv.writer(stream, lineterminator="\n")
|
||||
if header:
|
||||
writer.writerow([header])
|
||||
if gemm:
|
||||
writer.writerow(["GEMM"])
|
||||
writer.writerow(columns)
|
||||
for (n, k), labels in labels_by_nk.items():
|
||||
group = sorted(
|
||||
(case for case in gemm if (case.n, case.k) == (n, k)),
|
||||
key=lambda case: (case.backend, case.m),
|
||||
)
|
||||
# Modules fused into one GEMM are joined with "|" inside one
|
||||
# name; distinct same-shape modules each get their own rows,
|
||||
# duplicating the shared measurement.
|
||||
for label in labels:
|
||||
for case in group:
|
||||
for with_quant, value in _output_rows(case):
|
||||
writer.writerow(
|
||||
[
|
||||
label,
|
||||
case.m,
|
||||
n,
|
||||
k,
|
||||
case.backend,
|
||||
with_quant,
|
||||
_format_result(value),
|
||||
]
|
||||
)
|
||||
if moe:
|
||||
if gemm:
|
||||
writer.writerow([])
|
||||
writer.writerow(["MoE"])
|
||||
if moe_shape is not None:
|
||||
writer.writerow([moe_shape.label()])
|
||||
writer.writerow(columns)
|
||||
for case in sorted(moe, key=lambda case: (case.backend, case.m, case.with_quant)):
|
||||
for with_quant, value in _output_rows(case):
|
||||
writer.writerow(
|
||||
[
|
||||
moe_shape.name if moe_shape is not None else "experts",
|
||||
case.m,
|
||||
"",
|
||||
"",
|
||||
case.backend,
|
||||
with_quant,
|
||||
_format_result(value),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--flashinfer_repo",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="checkout containing benchmarks/flashinfer_benchmark.py",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ms",
|
||||
type=_positive_int,
|
||||
nargs="+",
|
||||
default=[1, 8, 64, 512],
|
||||
help="token counts, for example: 1 8 64 512",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--nks",
|
||||
type=_nk_arg,
|
||||
nargs="+",
|
||||
help="GEMM shapes as N,K or N,K,NAME, e.g. 4096,4096 2688,4096,mixer.o_proj",
|
||||
)
|
||||
parser.add_argument("--dry_run_iters", type=_positive_int, help="warmup iterations, e.g. 5")
|
||||
parser.add_argument("--num_iters", type=_positive_int, help="timed iterations, e.g. 30")
|
||||
parser.add_argument("--no_cuda_graph", action="store_true")
|
||||
parser.add_argument("--no_autotune", action="store_true")
|
||||
parser.add_argument(
|
||||
"--moe_hidden_size", type=_positive_int, help="model hidden size, e.g. 4096"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--moe_intermediate_size", type=_positive_int, help="expert width, e.g. 14336"
|
||||
)
|
||||
parser.add_argument("--moe_num_experts", type=_positive_int, help="local expert count, e.g. 8")
|
||||
parser.add_argument("--moe_top_k", type=_positive_int, help="experts per token, e.g. 2")
|
||||
parser.add_argument(
|
||||
"--moe_name",
|
||||
default="experts",
|
||||
help="expert container module path, e.g. model.layers.*.mlp.experts",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--moe_activation_type",
|
||||
choices=_MOE_ACTIVATIONS,
|
||||
help="FlashInfer activation, e.g. Swiglu",
|
||||
)
|
||||
parser.add_argument("--workdir", type=Path, default=Path("benchmark_via_builtin_out"))
|
||||
return parser
|
||||
|
||||
|
||||
def _execute_cases(
|
||||
cases: list[_Case], benchmarks_dir: Path, workdir: Path, driver_log: Path, header: str
|
||||
) -> list[dict[str, str]]:
|
||||
"""Run each case in its own driver process and record outcomes on it.
|
||||
|
||||
Returns the raw driver CSV rows for ``builtin_results.csv``.
|
||||
"""
|
||||
case_csv = workdir / "case_result.csv"
|
||||
rows: list[dict[str, str]] = []
|
||||
with driver_log.open("w") as log:
|
||||
print(header, flush=True)
|
||||
log.write(header + "\n")
|
||||
for case in cases:
|
||||
marker = f"[CASE] {case.tag}\n"
|
||||
print(marker, end="", flush=True)
|
||||
log.write(marker)
|
||||
case_csv.unlink(missing_ok=True)
|
||||
returncode, case_output = _run_case(
|
||||
benchmarks_dir,
|
||||
[*case.argv, "--case_tag", case.tag, "--output_path", str(case_csv.resolve())],
|
||||
log,
|
||||
)
|
||||
case_rows = []
|
||||
if case_csv.is_file():
|
||||
with case_csv.open(newline="") as stream:
|
||||
# A row that does not carry this case's tag cannot be
|
||||
# trusted as this case's measurement; fail the case.
|
||||
case_rows = [
|
||||
row for row in csv.DictReader(stream) if row.get("case_tag") == case.tag
|
||||
]
|
||||
if case_rows:
|
||||
rows.extend(case_rows)
|
||||
case.result = float(case_rows[-1]["median_time"]) * 1000
|
||||
else:
|
||||
case.result = _failure_message(case_output, returncode, driver_log)
|
||||
case_csv.unlink(missing_ok=True)
|
||||
return rows
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
"""Validate inputs, run the FlashInfer driver, and combine its results."""
|
||||
parser = _parser()
|
||||
args = parser.parse_args(argv)
|
||||
ms = list(dict.fromkeys(args.ms))
|
||||
labels_by_nk = _labels_by_nk(args.nks or [])
|
||||
moe_values = (
|
||||
args.moe_hidden_size,
|
||||
args.moe_intermediate_size,
|
||||
args.moe_num_experts,
|
||||
args.moe_top_k,
|
||||
)
|
||||
if any(moe_values) and not all(moe_values):
|
||||
parser.error("all four --moe_* shape arguments are required together")
|
||||
moe_shape = None
|
||||
if all(moe_values):
|
||||
moe_shape = _MoeShape(
|
||||
args.moe_hidden_size,
|
||||
args.moe_intermediate_size,
|
||||
args.moe_num_experts,
|
||||
args.moe_top_k,
|
||||
args.moe_activation_type,
|
||||
args.moe_name,
|
||||
)
|
||||
if not labels_by_nk and moe_shape is None:
|
||||
parser.error("pass --nks and/or all four --moe_* shape arguments")
|
||||
if moe_shape is not None and moe_shape.top_k > moe_shape.experts:
|
||||
parser.error("--moe_top_k cannot exceed --moe_num_experts")
|
||||
|
||||
benchmarks_dir = args.flashinfer_repo / "benchmarks"
|
||||
driver = benchmarks_dir / "flashinfer_benchmark.py"
|
||||
if not driver.is_file():
|
||||
parser.error(f"{driver} does not exist")
|
||||
|
||||
common = []
|
||||
if args.dry_run_iters is not None:
|
||||
common += ["--dry_run_iters", str(args.dry_run_iters)]
|
||||
if args.num_iters is not None:
|
||||
common += ["--num_iters", str(args.num_iters)]
|
||||
if not args.no_autotune:
|
||||
common.append("--autotune")
|
||||
if args.no_cuda_graph:
|
||||
common.append("--no_cuda_graph")
|
||||
|
||||
cases = _gemm_cases(ms, list(labels_by_nk), common)
|
||||
if moe_shape is not None:
|
||||
cases += _moe_cases(ms, moe_shape, common)
|
||||
|
||||
args.workdir.mkdir(parents=True, exist_ok=True)
|
||||
testlist = args.workdir / "testlist.txt"
|
||||
builtin_csv = args.workdir / "builtin_results.csv"
|
||||
combined_csv = args.workdir / "combined_results.csv"
|
||||
driver_log = args.workdir / "driver.log"
|
||||
if builtin_csv.exists() or combined_csv.exists():
|
||||
parser.error(f"{args.workdir} already contains results; choose a fresh --workdir")
|
||||
testlist.write_text(
|
||||
"\n".join(shlex.join([*case.argv, "--case_tag", case.tag]) for case in cases) + "\n"
|
||||
)
|
||||
|
||||
header = _environment_header(args.flashinfer_repo)
|
||||
rows = _execute_cases(cases, benchmarks_dir, args.workdir, driver_log, header)
|
||||
if rows:
|
||||
_write_builtin(builtin_csv, rows)
|
||||
|
||||
_attach_quant_times(
|
||||
cases,
|
||||
args.dry_run_iters if args.dry_run_iters is not None else 5,
|
||||
args.num_iters if args.num_iters is not None else 30,
|
||||
not args.no_cuda_graph,
|
||||
)
|
||||
_write_results(combined_csv, cases, labels_by_nk, header, moe_shape)
|
||||
print(f"Wrote {combined_csv}")
|
||||
failed = [case.tag for case in cases if isinstance(case.result, str)]
|
||||
if failed:
|
||||
raise RuntimeError(
|
||||
"FlashInfer failed benchmark cases: "
|
||||
+ ", ".join(failed)
|
||||
+ f"; wrote failure details to {combined_csv}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,640 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch")
|
||||
pytest.importorskip("accelerate")
|
||||
transformers = pytest.importorskip("transformers")
|
||||
|
||||
from torch import nn
|
||||
|
||||
SCRIPT = Path(__file__).parents[1] / "scripts" / "benchmark_model.py"
|
||||
SPEC = importlib.util.spec_from_file_location("benchmark_model", SCRIPT)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
benchmark_model = importlib.util.module_from_spec(SPEC)
|
||||
sys.modules[SPEC.name] = benchmark_model
|
||||
SPEC.loader.exec_module(benchmark_model)
|
||||
|
||||
|
||||
def _save(tmp_path, config):
|
||||
config.save_pretrained(tmp_path)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _preview(model_ref, monkeypatch, capsys, *, tp=1, ep=1):
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[str(SCRIPT), str(model_ref), "--tp", str(tp), "--ep", str(ep), "--print_only"],
|
||||
)
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
benchmark_model.main()
|
||||
return capsys.readouterr().out
|
||||
|
||||
|
||||
def _llama_config():
|
||||
return transformers.LlamaConfig(
|
||||
vocab_size=128,
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=8,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
|
||||
|
||||
def _nemotron_h_config(*, n_groups=2):
|
||||
return transformers.NemotronHConfig(
|
||||
vocab_size=128,
|
||||
hidden_size=32,
|
||||
layers_block_type=["mamba", "moe"],
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=8,
|
||||
intermediate_size=40,
|
||||
use_mamba_kernels=False,
|
||||
ssm_state_size=4,
|
||||
mamba_num_heads=4,
|
||||
mamba_head_dim=8,
|
||||
n_groups=n_groups,
|
||||
conv_kernel=4,
|
||||
expand=1,
|
||||
n_routed_experts=4,
|
||||
n_shared_experts=1,
|
||||
moe_intermediate_size=48,
|
||||
moe_shared_expert_intermediate_size=40,
|
||||
num_experts_per_tok=2,
|
||||
)
|
||||
|
||||
|
||||
def test_llama_meta_walk_fuses_common_projections(tmp_path, monkeypatch, capsys):
|
||||
model_dir = _save(tmp_path, _llama_config())
|
||||
|
||||
output = _preview(model_dir, monkeypatch, capsys, tp=2)
|
||||
|
||||
assert "layout: Transformers meta model; fused QKV and gate/up" in output
|
||||
assert (
|
||||
"32x32 <- model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj|model.layers.*.self_attn.v_proj"
|
||||
in output
|
||||
)
|
||||
assert "32x16 <- model.layers.*.self_attn.o_proj" in output
|
||||
assert "64x32 <- model.layers.*.mlp.gate_proj|model.layers.*.mlp.up_proj" in output
|
||||
# Same-shape kernels keep separate N,K,NAME arguments.
|
||||
assert "--nks " in output
|
||||
assert output.count("32,32,model.layers.") == 2
|
||||
assert "128x32" not in output # The output head is outside this benchmark.
|
||||
|
||||
|
||||
def test_gqa_kv_heads_are_replicated_when_tp_exceeds_kv_heads(tmp_path):
|
||||
config = _llama_config()
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
|
||||
kernels, _, problems = benchmark_model._inspect_model(model, config, tp=4, ep=1)
|
||||
|
||||
assert (
|
||||
benchmark_model._Kernel(
|
||||
24,
|
||||
32,
|
||||
"model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj|model.layers.*.self_attn.v_proj",
|
||||
)
|
||||
in kernels
|
||||
)
|
||||
assert problems == []
|
||||
|
||||
|
||||
def test_meta_loader_never_materializes_model_tensors(tmp_path):
|
||||
model_dir = _save(tmp_path, _llama_config())
|
||||
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir / "config.json"), False, None)
|
||||
|
||||
tensors = list(model.named_parameters()) + list(model.named_buffers())
|
||||
assert tensors and all(tensor.is_meta for _, tensor in tensors)
|
||||
|
||||
|
||||
def test_revision_does_not_reach_registered_model_constructor(tmp_path):
|
||||
model_dir = _save(tmp_path, _llama_config())
|
||||
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, "main")
|
||||
|
||||
assert type(model).__name__ == "LlamaForCausalLM"
|
||||
|
||||
|
||||
def test_mixtral_modulelist_experts_use_ep(tmp_path):
|
||||
config = transformers.MixtralConfig(
|
||||
vocab_size=128,
|
||||
hidden_size=32,
|
||||
intermediate_size=48,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
num_local_experts=4,
|
||||
num_experts_per_tok=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
|
||||
_, moe, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2)
|
||||
|
||||
assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu")
|
||||
|
||||
|
||||
def test_gpt_oss_direct_expert_tensors_are_inspected(tmp_path):
|
||||
config = transformers.GptOssConfig(
|
||||
vocab_size=128,
|
||||
hidden_size=32,
|
||||
intermediate_size=48,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=8,
|
||||
num_local_experts=4,
|
||||
num_experts_per_tok=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
|
||||
_, moe, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2)
|
||||
|
||||
assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu")
|
||||
|
||||
|
||||
def test_nemotron_h_mamba_and_stacked_experts_are_inspected(tmp_path):
|
||||
config = _nemotron_h_config()
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
|
||||
experts = next(module for name, module in model.named_modules() if name.endswith(".experts"))
|
||||
kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1)
|
||||
|
||||
assert experts.up_proj.ndim == experts.down_proj.ndim == 3
|
||||
assert benchmark_model._Kernel(42, 32, "model.layers.*.mixer.in_proj") in kernels
|
||||
assert benchmark_model._Kernel(32, 16, "model.layers.*.mixer.out_proj") in kernels
|
||||
assert benchmark_model._Kernel(20, 32, "model.layers.*.mixer.shared_experts.up_proj") in kernels
|
||||
assert (
|
||||
benchmark_model._Kernel(32, 20, "model.layers.*.mixer.shared_experts.down_proj") in kernels
|
||||
)
|
||||
assert moe == benchmark_model._MoeShape(32, 24, 4, 2, "Relu2")
|
||||
assert problems == []
|
||||
command = benchmark_model._command(kernels, moe, [])
|
||||
assert command[command.index("--moe_activation_type") + 1] == "Relu2"
|
||||
|
||||
with pytest.raises(benchmark_model.ShapeError, match=r"n_groups=2.*TP=4"):
|
||||
benchmark_model._inspect_model(model, config, tp=4, ep=1)
|
||||
|
||||
config.n_routed_experts = 8
|
||||
_, _, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1)
|
||||
assert any("declares 8" in problem and "[4]" in problem for problem in problems)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("config_cls_name", "model_cls_name"),
|
||||
[
|
||||
("Qwen3NextConfig", "Qwen3NextForCausalLM"),
|
||||
("Qwen3_5MoeTextConfig", "Qwen3_5MoeForCausalLM"),
|
||||
],
|
||||
)
|
||||
def test_gated_delta_net_kernels_are_derived(config_cls_name, model_cls_name):
|
||||
config_cls = getattr(transformers, config_cls_name, None)
|
||||
model_cls = getattr(transformers, model_cls_name, None)
|
||||
if config_cls is None or model_cls is None:
|
||||
pytest.skip(f"transformers does not provide {config_cls_name}")
|
||||
assert config_cls is not None and model_cls is not None
|
||||
from accelerate import init_empty_weights
|
||||
|
||||
config = config_cls(
|
||||
vocab_size=128,
|
||||
hidden_size=32,
|
||||
num_hidden_layers=1,
|
||||
layer_types=["linear_attention"],
|
||||
linear_num_key_heads=2,
|
||||
linear_num_value_heads=4,
|
||||
linear_key_head_dim=8,
|
||||
linear_value_head_dim=8,
|
||||
linear_conv_kernel_dim=4,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=8,
|
||||
intermediate_size=32,
|
||||
num_experts=4,
|
||||
num_experts_per_tok=2,
|
||||
moe_intermediate_size=16,
|
||||
shared_expert_intermediate_size=16,
|
||||
decoder_sparse_step=1,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
with init_empty_weights(include_buffers=True):
|
||||
model = model_cls(config)
|
||||
|
||||
kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1)
|
||||
|
||||
# key_dim = 2 heads x 8 = 16, value_dim = 4 heads x 8 = 32; vLLM's shared
|
||||
# GDN mixer runs qkv+z and b+a as two fused per-rank GEMMs (both the
|
||||
# pre-fused Qwen3-Next and split Qwen3.5 checkpoint layouts).
|
||||
prefix = "model.layers.*.linear_attn"
|
||||
if config_cls_name == "Qwen3NextConfig":
|
||||
qkvz_label, ba_label = f"{prefix}.in_proj_qkvz", f"{prefix}.in_proj_ba"
|
||||
else:
|
||||
qkvz_label = f"{prefix}.in_proj_qkv|{prefix}.in_proj_z"
|
||||
ba_label = f"{prefix}.in_proj_b|{prefix}.in_proj_a"
|
||||
assert benchmark_model._Kernel(48, 32, qkvz_label) in kernels
|
||||
assert benchmark_model._Kernel(4, 32, ba_label) in kernels
|
||||
assert benchmark_model._Kernel(32, 16, f"{prefix}.out_proj") in kernels
|
||||
assert problems == []
|
||||
assert moe == benchmark_model._MoeShape(32, 8, 4, 2, "Swiglu")
|
||||
|
||||
with pytest.raises(
|
||||
benchmark_model.ShapeError, match=r"linear_num_key_heads=2 is not divisible by 4"
|
||||
):
|
||||
benchmark_model._inspect_model(model, config, tp=4, ep=1)
|
||||
|
||||
|
||||
def test_expert_audit_problem_is_not_masked_by_per_rank_validation(tmp_path):
|
||||
config = _nemotron_h_config()
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
config.n_routed_experts = 8
|
||||
|
||||
# EP=3 does not divide the instantiated expert count; the audit mismatch
|
||||
# must still be reported instead of a masking divisibility error.
|
||||
kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=3)
|
||||
|
||||
assert kernels
|
||||
assert moe is None
|
||||
assert any("declares 8" in problem for problem in problems)
|
||||
|
||||
|
||||
def test_moe_only_model_benchmarks_without_dense_kernels():
|
||||
class Expert(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.experts = nn.ModuleList([Expert() for _ in range(4)])
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(
|
||||
hidden_size=32,
|
||||
num_attention_heads=4,
|
||||
num_experts_per_tok=2,
|
||||
mlp_hidden_act="relu2",
|
||||
)
|
||||
|
||||
kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1)
|
||||
|
||||
assert kernels == []
|
||||
assert problems == []
|
||||
assert moe == benchmark_model._MoeShape(32, 48, 4, 2, "Relu2")
|
||||
command = benchmark_model._command(kernels, moe, [])
|
||||
assert "--nks" not in command
|
||||
assert command[:2] == ["--moe_hidden_size", "32"]
|
||||
|
||||
|
||||
def test_legacy_nongated_modulelist_experts_are_inspected():
|
||||
class Expert(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
|
||||
model = nn.Module()
|
||||
model.experts = nn.ModuleList([Expert() for _ in range(4)])
|
||||
config = SimpleNamespace(num_experts_per_tok=2, mlp_hidden_act="relu2")
|
||||
|
||||
assert benchmark_model._moe_shapes(model, config) == {
|
||||
benchmark_model._MoeShape(32, 48, 4, 2, "Relu2")
|
||||
}
|
||||
|
||||
|
||||
def test_gate_and_up_projections_must_shard_individually():
|
||||
class Mlp(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(32, 6, bias=False)
|
||||
self.up_proj = nn.Linear(32, 6, bias=False)
|
||||
self.down_proj = nn.Linear(6, 32, bias=False)
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mlp = Mlp()
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(hidden_size=32, num_attention_heads=4)
|
||||
|
||||
# The summed width (12) divides by TP=4, but vLLM shards gate and up
|
||||
# individually, so the per-projection width (6) must divide too.
|
||||
with pytest.raises(benchmark_model.ShapeError, match=r"gate_proj=6 is not divisible by 4"):
|
||||
benchmark_model._inspect_model(model, config, tp=4, ep=1)
|
||||
|
||||
|
||||
def test_top_k_fallback_matches_the_modelopt_list():
|
||||
auto_quantize_cost = pytest.importorskip("modelopt.torch.quantization._auto_quantize_cost")
|
||||
|
||||
assert benchmark_model._MOE_TOP_K_ATTRS_FALLBACK == auto_quantize_cost._ACTIVE_MOE_TOP_K_ATTRS
|
||||
|
||||
|
||||
def test_top_k_covers_the_modelopt_attribute_aliases():
|
||||
assert benchmark_model._top_k(SimpleNamespace(num_selected_experts=2)) == 2
|
||||
assert benchmark_model._top_k(SimpleNamespace(num_experts_per_token=4)) == 4
|
||||
assert benchmark_model._top_k(SimpleNamespace(top_k=6)) == 6
|
||||
assert benchmark_model._top_k(SimpleNamespace(num_experts_per_tok=8, top_k=50)) == 8
|
||||
assert benchmark_model._top_k(SimpleNamespace()) is None
|
||||
|
||||
|
||||
def test_gated_moe_activation_is_derived_or_rejected():
|
||||
assert (
|
||||
benchmark_model._moe_activation(SimpleNamespace(hidden_act="gelu_pytorch_tanh"), True)
|
||||
== "Geglu"
|
||||
)
|
||||
assert benchmark_model._moe_activation(SimpleNamespace(hidden_act="gelu_new"), True) == "Geglu"
|
||||
# Exact gelu and quick_gelu are not served by vLLM's FlashInfer MoE path,
|
||||
# so they must be rejected instead of timed via the tanh-GELU kernel.
|
||||
for activation in ("gelu", "quick_gelu", "relu"):
|
||||
with pytest.raises(benchmark_model.ShapeError, match="unsupported gated MoE activation"):
|
||||
benchmark_model._moe_activation(SimpleNamespace(hidden_act=activation), True)
|
||||
|
||||
|
||||
def test_mamba_single_group_is_replicated_across_tp():
|
||||
class Mixer(nn.Module):
|
||||
intermediate_size = 32
|
||||
num_heads = 4
|
||||
n_groups = 1
|
||||
ssm_state_size = 4
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.in_proj = nn.Linear(32, 76, bias=False)
|
||||
self.out_proj = nn.Linear(32, 32, bias=False)
|
||||
|
||||
model = nn.Module()
|
||||
model.mixer = Mixer()
|
||||
|
||||
assert benchmark_model._mamba_kernels(model, tp=2) == [
|
||||
benchmark_model._Kernel(42, 32, "mixer.in_proj"),
|
||||
benchmark_model._Kernel(32, 16, "mixer.out_proj"),
|
||||
]
|
||||
|
||||
|
||||
def test_unrecognized_decoder_linear_is_reported():
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
self.unknown_proj = nn.Linear(32, 48, bias=False)
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
|
||||
assert benchmark_model._unsupported_decoder_linears(model) == [
|
||||
("layers.0.unknown_proj", 48, 32)
|
||||
]
|
||||
config = SimpleNamespace(hidden_size=32, num_attention_heads=4)
|
||||
kernels, _, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1)
|
||||
|
||||
assert benchmark_model._Kernel(48, 32, "layers.*.up_proj") in kernels
|
||||
assert problems == ["unsupported decoder Linear GEMM layout(s): layers.0.unknown_proj (48x32)"]
|
||||
|
||||
|
||||
def test_partial_inventory_is_printed_when_the_audit_fails(monkeypatch, capsys):
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
self.unknown_proj = nn.Linear(32, 48, bias=False)
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(hidden_size=32, num_attention_heads=4, model_type="test")
|
||||
monkeypatch.setattr(benchmark_model, "_load_meta_model", lambda *_: (config, model))
|
||||
monkeypatch.setattr(sys, "argv", [str(SCRIPT), "unused/model", "--print_only"])
|
||||
|
||||
with pytest.raises(SystemExit, match="2"):
|
||||
benchmark_model.main()
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "# 48x32 <- layers.*.up_proj" in captured.out
|
||||
assert "# unsupported: unsupported decoder Linear GEMM layout(s)" in captured.out
|
||||
assert "unknown_proj (48x32)" in captured.out
|
||||
assert "benchmark_via_builtin.py" in captured.err
|
||||
|
||||
|
||||
def test_declared_moe_without_supported_experts_is_reported():
|
||||
class Mlp(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mlp = Mlp()
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(
|
||||
hidden_size=32,
|
||||
num_attention_heads=4,
|
||||
num_local_experts=4,
|
||||
num_experts_per_tok=2,
|
||||
)
|
||||
|
||||
_, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1)
|
||||
|
||||
assert moe is None
|
||||
assert problems == [
|
||||
"model declares 4 routed experts but no supported expert GEMM layout was found"
|
||||
]
|
||||
|
||||
|
||||
def test_command_keeps_same_shape_kernels_as_separate_named_pairs():
|
||||
kernels = [
|
||||
benchmark_model._Kernel(64, 32, "a_proj|b_proj"),
|
||||
benchmark_model._Kernel(64, 32, "c_proj"),
|
||||
benchmark_model._Kernel(32, 64, "down_proj"),
|
||||
]
|
||||
|
||||
command = benchmark_model._command(kernels, None, [])
|
||||
|
||||
assert command == [
|
||||
"--nks",
|
||||
"64,32,a_proj|b_proj",
|
||||
"64,32,c_proj",
|
||||
"32,64,down_proj",
|
||||
]
|
||||
|
||||
|
||||
def test_moe_name_carries_the_expert_container_path(tmp_path):
|
||||
config = _nemotron_h_config()
|
||||
model_dir = _save(tmp_path, config)
|
||||
_, model = benchmark_model._load_meta_model(str(model_dir), False, None)
|
||||
|
||||
_, moe, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1)
|
||||
|
||||
assert problems == []
|
||||
assert moe is not None
|
||||
# The MoE row is labeled by its real container path, not a generic name,
|
||||
# and the path survives per-rank sharding.
|
||||
assert moe.name == "model.layers.*.mixer.experts"
|
||||
command = benchmark_model._command([], moe, [])
|
||||
assert command[command.index("--moe_name") + 1] == "model.layers.*.mixer.experts"
|
||||
# The path is metadata: two layouts differing only by it stay one shape.
|
||||
assert benchmark_model._MoeShape(8, 4, 2, 1, None, "a") == benchmark_model._MoeShape(
|
||||
8, 4, 2, 1, None, "b"
|
||||
)
|
||||
|
||||
|
||||
def test_moe_sharding_interpretation_is_printed(monkeypatch, capsys):
|
||||
class Expert(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.experts = nn.ModuleList([Expert() for _ in range(6)])
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(
|
||||
hidden_size=32,
|
||||
num_attention_heads=4,
|
||||
num_experts_per_tok=2,
|
||||
mlp_hidden_act="relu2",
|
||||
model_type="test",
|
||||
)
|
||||
monkeypatch.setattr(benchmark_model, "_load_meta_model", lambda *_: (config, model))
|
||||
monkeypatch.setattr(
|
||||
sys, "argv", [str(SCRIPT), "unused/model", "--tp", "2", "--ep", "2", "--print_only"]
|
||||
)
|
||||
|
||||
benchmark_model.main()
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert "# MoE sharding: EP=2 partitions whole experts" in output
|
||||
|
||||
# EP that is not a multiple of TP matches no modeled serving layout.
|
||||
with pytest.raises(
|
||||
benchmark_model.ShapeError, match=r"EP=2 is not a multiple of TP=4.*benchmark_via_builtin"
|
||||
):
|
||||
benchmark_model._inspect_model(model, config, tp=4, ep=2)
|
||||
|
||||
|
||||
def test_runner_is_invoked_in_process(tmp_path, monkeypatch):
|
||||
model_dir = _save(tmp_path, _llama_config())
|
||||
launched = []
|
||||
runner = SimpleNamespace(main=launched.append)
|
||||
monkeypatch.setattr(benchmark_model, "_load_runner", lambda: runner)
|
||||
monkeypatch.setattr(sys, "argv", [str(SCRIPT), str(model_dir), "--ms", "1", "16"])
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
|
||||
benchmark_model.main()
|
||||
|
||||
assert launched == [
|
||||
[
|
||||
"--nks",
|
||||
"64,32,model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj"
|
||||
"|model.layers.*.self_attn.v_proj",
|
||||
"32,32,model.layers.*.self_attn.o_proj",
|
||||
"128,32,model.layers.*.mlp.gate_proj|model.layers.*.mlp.up_proj",
|
||||
"32,64,model.layers.*.mlp.down_proj",
|
||||
"--ms",
|
||||
"1",
|
||||
"16",
|
||||
]
|
||||
]
|
||||
|
||||
|
||||
def test_router_gate_projection_is_not_treated_as_an_mlp():
|
||||
class Mlp(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(32, 48, bias=False)
|
||||
self.down_proj = nn.Linear(48, 32, bias=False)
|
||||
|
||||
class Router(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(32, 4, bias=False)
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mlp = Mlp()
|
||||
self.router = Router()
|
||||
|
||||
model = nn.Module()
|
||||
model.layers = nn.ModuleList([Block()])
|
||||
config = SimpleNamespace(hidden_size=32, num_attention_heads=4)
|
||||
|
||||
kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1)
|
||||
|
||||
assert moe is None
|
||||
assert problems == []
|
||||
assert kernels == [
|
||||
benchmark_model._Kernel(48, 32, "layers.*.mlp.up_proj"),
|
||||
benchmark_model._Kernel(32, 48, "layers.*.mlp.down_proj"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("option", "value"), [("--nks", "1,1"), ("--moe_activation_type", "Relu2")]
|
||||
)
|
||||
def test_derived_shapes_cannot_be_overridden(monkeypatch, capsys, option, value):
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[str(SCRIPT), "unused/model", option, value, "--print_only"],
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit, match="2"):
|
||||
benchmark_model.main()
|
||||
assert "cannot be overridden" in capsys.readouterr().err
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("option", "value"), [("--tp", "0"), ("--ep", "-1")])
|
||||
def test_parallel_sizes_must_be_positive(monkeypatch, capsys, option, value):
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[str(SCRIPT), "unused/model", option, value, "--print_only"],
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit, match="2"):
|
||||
benchmark_model.main()
|
||||
assert "expected a positive integer" in capsys.readouterr().err
|
||||
@@ -0,0 +1,455 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
SCRIPT = Path(__file__).parents[1] / "scripts" / "benchmark_via_builtin.py"
|
||||
SPEC = importlib.util.spec_from_file_location("benchmark_via_builtin", SCRIPT)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
benchmark = importlib.util.module_from_spec(SPEC)
|
||||
sys.modules[SPEC.name] = benchmark
|
||||
SPEC.loader.exec_module(benchmark)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "a,2", "0,2", "2,-1", "1,2,", "1,2,3,4"])
|
||||
def test_nk_arg_rejects_invalid_values(value):
|
||||
with pytest.raises(argparse.ArgumentTypeError):
|
||||
benchmark._nk_arg(value)
|
||||
|
||||
|
||||
def test_nk_arg_parses_optional_names():
|
||||
assert benchmark._nk_arg("1,2") == (1, 2, None)
|
||||
assert benchmark._nk_arg("1,2,attn.q_proj|attn.k_proj") == (1, 2, "attn.q_proj|attn.k_proj")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "-1"])
|
||||
def test_positive_int_rejects_non_positive_values(value):
|
||||
with pytest.raises(argparse.ArgumentTypeError, match="positive integer"):
|
||||
benchmark._positive_int(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"option",
|
||||
[
|
||||
"--ms",
|
||||
"--dry_run_iters",
|
||||
"--num_iters",
|
||||
"--moe_hidden_size",
|
||||
"--moe_intermediate_size",
|
||||
"--moe_num_experts",
|
||||
"--moe_top_k",
|
||||
],
|
||||
)
|
||||
def test_parser_rejects_zero_for_numeric_options(option, capsys):
|
||||
with pytest.raises(SystemExit, match="2"):
|
||||
benchmark._parser().parse_args(["--flashinfer_repo", "/unused", option, "0"])
|
||||
assert "not a positive integer" in capsys.readouterr().err
|
||||
|
||||
|
||||
def test_gemm_cases_are_data_driven_and_preserve_requested_shapes():
|
||||
cases = benchmark._gemm_cases([1], [(65, 129)], [])
|
||||
|
||||
assert {case.backend for case in cases} == {
|
||||
"bf16",
|
||||
"nvfp4_cudnn",
|
||||
"nvfp4_cutlass",
|
||||
"nvfp4_cutedsl",
|
||||
"nvfp4_trtllm",
|
||||
"fp8_cudnn",
|
||||
"fp8_cutlass",
|
||||
"fp8_trtllm",
|
||||
}
|
||||
assert len({case.tag for case in cases}) == len(cases)
|
||||
assert {(case.m, case.n, case.k) for case in cases} == {(1, 65, 129)}
|
||||
physical_shapes = {
|
||||
case.backend: (
|
||||
case.argv[case.argv.index("--n") + 1],
|
||||
case.argv[case.argv.index("--k") + 1],
|
||||
)
|
||||
for case in cases
|
||||
}
|
||||
assert physical_shapes["nvfp4_cudnn"] == ("96", "160")
|
||||
assert physical_shapes["nvfp4_cutlass"] == ("96", "160")
|
||||
assert physical_shapes["nvfp4_cutedsl"] == ("96", "160")
|
||||
assert physical_shapes["nvfp4_trtllm"] == ("65", "129")
|
||||
assert all(
|
||||
shape == ("65", "129")
|
||||
for name, shape in physical_shapes.items()
|
||||
if not name.startswith("nvfp4_")
|
||||
)
|
||||
assert {case.quant for case in cases if case.quant is not None} == {
|
||||
benchmark._QuantSpec("nvfp4", "128x4", 1, 160),
|
||||
benchmark._QuantSpec("nvfp4", "8x4", 1, 129),
|
||||
benchmark._QuantSpec("fp8", "static", 1, 129),
|
||||
}
|
||||
|
||||
|
||||
def test_labels_by_nk_maps_duplicate_and_unnamed_shapes():
|
||||
labels_by_nk = benchmark._labels_by_nk(
|
||||
[(32, 64, "o_proj"), (32, 64, "out_proj"), (128, 256, "qkv_proj")]
|
||||
)
|
||||
|
||||
assert labels_by_nk == {(32, 64): ["o_proj", "out_proj"], (128, 256): ["qkv_proj"]}
|
||||
# Unnamed shapes are deduplicated in first-seen order and label themselves.
|
||||
assert benchmark._labels_by_nk([(32, 64, None), (32, 64, None), (2, 3, None)]) == {
|
||||
(32, 64): ["32x64"],
|
||||
(2, 3): ["2x3"],
|
||||
}
|
||||
|
||||
|
||||
def test_moe_cases_and_derived_with_quant_rows():
|
||||
cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2), [])
|
||||
assert [(case.backend, case.with_quant) for case in cases] == [
|
||||
("bf16_cutlass", False),
|
||||
("fp8_cutlass", False),
|
||||
("nvfp4_cutlass", False),
|
||||
("nvfp4_cutlass", True),
|
||||
("fp8_trtllm", False),
|
||||
("nvfp4_trtllm", False),
|
||||
("nvfp4_cutedsl", False),
|
||||
]
|
||||
intermediate_sizes = {
|
||||
(case.backend, case.with_quant): case.argv[case.argv.index("--intermediate_size") + 1]
|
||||
for case in cases
|
||||
}
|
||||
assert intermediate_sizes == {
|
||||
("bf16_cutlass", False): "50",
|
||||
("nvfp4_cutlass", False): "50",
|
||||
("nvfp4_cutlass", True): "50",
|
||||
("fp8_cutlass", False): "64",
|
||||
("fp8_trtllm", False): "64",
|
||||
("nvfp4_trtllm", False): "64",
|
||||
("nvfp4_cutedsl", False): "50",
|
||||
}
|
||||
hidden_sizes = {case.backend: case.argv[case.argv.index("--hidden_size") + 1] for case in cases}
|
||||
# Only the trtllm-gen NVFP4 MoE pads hidden (vLLM pads it to 256).
|
||||
assert hidden_sizes["nvfp4_trtllm"] == "256"
|
||||
assert all(value == "32" for row, value in hidden_sizes.items() if row != "nvfp4_trtllm")
|
||||
routines = {case.backend: case.argv[case.argv.index("--routine") + 1] for case in cases}
|
||||
assert routines["fp8_trtllm"] == "trtllm_fp8_per_tensor_scale_moe"
|
||||
assert routines["nvfp4_trtllm"] == "trtllm_fp4_block_scale_moe"
|
||||
assert routines["nvfp4_cutedsl"] == "cute_dsl_fp4_block_scale_moe"
|
||||
# Only the trtllm-gen rows route in-kernel; they use a fixed renormalize
|
||||
# method so rows stay comparable across models.
|
||||
for case in cases:
|
||||
if case.backend.endswith("_trtllm"):
|
||||
assert case.argv[case.argv.index("--routing_method") + 1] == "renormalize"
|
||||
else:
|
||||
assert "--routing_method" not in case.argv
|
||||
quants = {(case.backend, case.with_quant): case.quant for case in cases}
|
||||
# FP8 rows share the static activation-quant timing; the trtllm-gen and
|
||||
# CuteDSL NVFP4 rows use the linear-scale-layout quantize their kernels
|
||||
# consume (trtllm at the 256-padded hidden). The NVFP4 CUTLASS pair is a
|
||||
# direct fused measurement instead.
|
||||
assert quants[("fp8_cutlass", False)] == benchmark._QuantSpec("fp8", "static", 1, 32)
|
||||
assert quants[("fp8_trtllm", False)] == benchmark._QuantSpec("fp8", "static", 1, 32)
|
||||
assert quants[("nvfp4_trtllm", False)] == benchmark._QuantSpec("nvfp4", "linear", 1, 256)
|
||||
assert quants[("nvfp4_cutedsl", False)] == benchmark._QuantSpec("nvfp4", "linear", 1, 32)
|
||||
assert quants[("bf16_cutlass", False)] is None
|
||||
assert quants[("nvfp4_cutlass", False)] is None
|
||||
assert quants[("nvfp4_cutlass", True)] is None
|
||||
|
||||
fp8 = next(case for case in cases if case.backend == "fp8_cutlass")
|
||||
assert fp8.quant is not None
|
||||
fp8.result = 1.0
|
||||
fp8.quant_result = 2.0
|
||||
|
||||
assert benchmark._output_rows(fp8) == [(False, 1.0), (True, 3.0)]
|
||||
|
||||
|
||||
def test_swiglustep_uses_gated_fp8_alignment():
|
||||
cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "SwigluStep"), [])
|
||||
|
||||
fp8 = next(case for case in cases if case.backend == "fp8_cutlass")
|
||||
assert fp8.argv[fp8.argv.index("--intermediate_size") + 1] == "64"
|
||||
|
||||
|
||||
def test_non_gated_moe_pads_fp8_and_nvfp4_intermediate_to_128():
|
||||
cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "Relu2"), [])
|
||||
|
||||
intermediate_sizes = {
|
||||
(case.backend, case.with_quant): case.argv[case.argv.index("--intermediate_size") + 1]
|
||||
for case in cases
|
||||
}
|
||||
# The Swiglu-only CuteDSL MoE row is not emitted for non-gated activations.
|
||||
assert intermediate_sizes == {
|
||||
("bf16_cutlass", False): "50",
|
||||
("fp8_cutlass", False): "128",
|
||||
("nvfp4_cutlass", False): "128",
|
||||
("nvfp4_cutlass", True): "128",
|
||||
("fp8_trtllm", False): "128",
|
||||
("nvfp4_trtllm", False): "128",
|
||||
}
|
||||
|
||||
|
||||
def test_unavailable_fp8_quantization_is_written_as_an_error(monkeypatch, capsys, tmp_path):
|
||||
case = benchmark._Case(
|
||||
section="gemm",
|
||||
tag="gemm_fp8_cutlass_MxNxK=1x32x64",
|
||||
backend="fp8_cutlass",
|
||||
m=1,
|
||||
n=32,
|
||||
k=64,
|
||||
argv=[],
|
||||
quant=benchmark._QuantSpec("fp8", "static", 1, 64),
|
||||
result=1.0,
|
||||
)
|
||||
monkeypatch.setattr(benchmark, "vllm_ops", None)
|
||||
|
||||
benchmark._attach_quant_times([case], 1, 1, False)
|
||||
output = tmp_path / "combined_results.csv"
|
||||
benchmark._write_results(output, [case], {(32, 64): ["32x64"]})
|
||||
|
||||
assert case.quant_result == benchmark._FP8_QUANT_UNAVAILABLE
|
||||
assert "[WARN] vLLM is unavailable for FP8 activation quantization" in capsys.readouterr().out
|
||||
assert "32x64,1,32,64,fp8_cutlass,False,1.000\n" in output.read_text()
|
||||
assert f"32x64,1,32,64,fp8_cutlass,True,{benchmark._FP8_QUANT_UNAVAILABLE}\n" in (
|
||||
output.read_text()
|
||||
)
|
||||
|
||||
|
||||
def test_driver_errors_are_added_to_kernel_and_with_quant_rows(tmp_path):
|
||||
case = benchmark._Case(
|
||||
section="gemm",
|
||||
tag="gemm_fp8_trtllm_MxNxK=8x1280x2880",
|
||||
backend="fp8_trtllm",
|
||||
m=8,
|
||||
n=1280,
|
||||
k=2880,
|
||||
argv=[],
|
||||
quant=benchmark._QuantSpec("fp8", "static", 8, 2880),
|
||||
)
|
||||
output = [
|
||||
f"[ERROR] Error running test: --routine mm_fp8 --case_tag {case.tag}\n",
|
||||
"[ERROR] Error: K must be divisible by 128, got 2880\n",
|
||||
]
|
||||
|
||||
message = benchmark._parse_driver_error(output)
|
||||
assert message == "K must be divisible by 128; got 2880"
|
||||
case.result = f"ERROR: {message}"
|
||||
csv_path = tmp_path / "combined_results.csv"
|
||||
benchmark._write_results(csv_path, [case], {(1280, 2880): ["1280x2880"]})
|
||||
|
||||
expected = "ERROR: K must be divisible by 128; got 2880"
|
||||
assert f"1280x2880,8,1280,2880,fp8_trtllm,False,{expected}\n" in csv_path.read_text()
|
||||
assert f"1280x2880,8,1280,2880,fp8_trtllm,True,{expected}\n" in csv_path.read_text()
|
||||
|
||||
|
||||
def test_empty_driver_error_has_no_synthetic_reason():
|
||||
output = [
|
||||
"[ERROR] Error running test: --routine mm_fp4 --case_tag gemm_nvfp4\n",
|
||||
"[ERROR] Error:\n",
|
||||
]
|
||||
|
||||
assert benchmark._parse_driver_error(output) == ""
|
||||
# An error line without any message line is not a driver-reported error.
|
||||
assert benchmark._parse_driver_error(output[:1]) is None
|
||||
assert benchmark._parse_driver_error([]) is None
|
||||
|
||||
|
||||
def test_write_results_emits_long_form_rows(tmp_path):
|
||||
cases = [
|
||||
benchmark._Case("gemm", "a", "bf16", 1, 2, 3, [], result=1.25),
|
||||
benchmark._Case(
|
||||
"gemm",
|
||||
"b",
|
||||
"fp8_cutlass",
|
||||
8,
|
||||
4,
|
||||
5,
|
||||
[],
|
||||
quant=benchmark._QuantSpec("fp8", "static", 8, 5),
|
||||
result=3.5,
|
||||
quant_result=1.0,
|
||||
),
|
||||
benchmark._Case(
|
||||
"moe",
|
||||
"c",
|
||||
"fp8_cutlass",
|
||||
8,
|
||||
None,
|
||||
None,
|
||||
[],
|
||||
quant=benchmark._QuantSpec("fp8", "static", 8, 32),
|
||||
result=2.0,
|
||||
quant_result=0.5,
|
||||
),
|
||||
# A case that never produced a result or an error emits no row.
|
||||
benchmark._Case("gemm", "d", "fp8_trtllm", 1, 2, 3, []),
|
||||
]
|
||||
output = tmp_path / "combined_results.csv"
|
||||
benchmark._write_results(
|
||||
output,
|
||||
cases,
|
||||
{(4, 5): ["q_proj|k_proj|v_proj"], (2, 3): ["in_proj", "out_proj"]},
|
||||
header="flashinfer test-header",
|
||||
moe_shape=benchmark._MoeShape(32, 50, 4, 2, "Relu2", "model.layers.*.mlp.experts"),
|
||||
)
|
||||
|
||||
assert output.read_text() == (
|
||||
"flashinfer test-header\n"
|
||||
"GEMM\n"
|
||||
"module_name,M,N,K,backend,with_quant,runtime\n"
|
||||
"q_proj|k_proj|v_proj,8,4,5,fp8_cutlass,False,3.500\n"
|
||||
"q_proj|k_proj|v_proj,8,4,5,fp8_cutlass,True,4.500\n"
|
||||
"in_proj,1,2,3,bf16,False,1.250\n"
|
||||
"out_proj,1,2,3,bf16,False,1.250\n"
|
||||
"\n"
|
||||
"MoE\n"
|
||||
"H=32 F=50 E=4 top_k=2 activation=Relu2\n"
|
||||
"module_name,M,N,K,backend,with_quant,runtime\n"
|
||||
"model.layers.*.mlp.experts,8,,,fp8_cutlass,False,2.000\n"
|
||||
"model.layers.*.mlp.experts,8,,,fp8_cutlass,True,2.500\n"
|
||||
)
|
||||
|
||||
|
||||
def test_top_k_cannot_exceed_expert_count(monkeypatch, capsys):
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[
|
||||
str(SCRIPT),
|
||||
"--flashinfer_repo",
|
||||
"/unused",
|
||||
"--moe_hidden_size",
|
||||
"4",
|
||||
"--moe_intermediate_size",
|
||||
"8",
|
||||
"--moe_num_experts",
|
||||
"1",
|
||||
"--moe_top_k",
|
||||
"2",
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit, match="2"):
|
||||
benchmark.main()
|
||||
assert "--moe_top_k cannot exceed --moe_num_experts" in capsys.readouterr().err
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("returncode", "expected_reason"),
|
||||
[
|
||||
(0, "FlashInfer produced no result row"),
|
||||
(1, "FlashInfer driver exited with status 1"),
|
||||
],
|
||||
)
|
||||
def test_missing_builtin_results_still_writes_combined_errors(
|
||||
monkeypatch, tmp_path, returncode, expected_reason
|
||||
):
|
||||
benchmarks_dir = tmp_path / "flashinfer" / "benchmarks"
|
||||
benchmarks_dir.mkdir(parents=True)
|
||||
(benchmarks_dir / "flashinfer_benchmark.py").write_text("")
|
||||
workdir = tmp_path / "results"
|
||||
monkeypatch.setattr(benchmark, "_run_case", lambda *_: (returncode, []))
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[
|
||||
str(SCRIPT),
|
||||
"--flashinfer_repo",
|
||||
str(benchmarks_dir.parent),
|
||||
"--ms",
|
||||
"1",
|
||||
"--nks",
|
||||
"2,3",
|
||||
"--workdir",
|
||||
str(workdir),
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="FlashInfer failed benchmark cases"):
|
||||
benchmark.main()
|
||||
|
||||
assert not (workdir / "builtin_results.csv").exists()
|
||||
combined = (workdir / "combined_results.csv").read_text()
|
||||
assert f"2x3,1,2,3,bf16,False,ERROR: {expected_reason}" in combined
|
||||
assert "driver.log" in combined
|
||||
# The reproducibility header leads both the combined CSV and driver.log.
|
||||
assert combined.splitlines()[0].startswith("flashinfer ")
|
||||
assert (workdir / "driver.log").read_text().startswith("flashinfer ")
|
||||
|
||||
|
||||
def test_case_rows_with_foreign_tags_are_treated_as_failures(monkeypatch, tmp_path):
|
||||
benchmarks_dir = tmp_path / "flashinfer" / "benchmarks"
|
||||
benchmarks_dir.mkdir(parents=True)
|
||||
(benchmarks_dir / "flashinfer_benchmark.py").write_text("")
|
||||
workdir = tmp_path / "results"
|
||||
|
||||
def fake_run_case(benchmarks_dir, argv, log):
|
||||
output = Path(argv[argv.index("--output_path") + 1])
|
||||
output.write_text("case_tag,median_time\nsomeone_else,0.001\n")
|
||||
return 0, []
|
||||
|
||||
monkeypatch.setattr(benchmark, "_run_case", fake_run_case)
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[
|
||||
str(SCRIPT),
|
||||
"--flashinfer_repo",
|
||||
str(benchmarks_dir.parent),
|
||||
"--ms",
|
||||
"1",
|
||||
"--nks",
|
||||
"2,3",
|
||||
"--workdir",
|
||||
str(workdir),
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="FlashInfer failed benchmark cases"):
|
||||
benchmark.main()
|
||||
|
||||
combined = (workdir / "combined_results.csv").read_text()
|
||||
assert "no result row" in combined
|
||||
|
||||
|
||||
def test_run_case_streams_and_appends_to_the_driver_log(tmp_path, capsys):
|
||||
benchmarks_dir = tmp_path / "benchmarks"
|
||||
benchmarks_dir.mkdir()
|
||||
(benchmarks_dir / "flashinfer_benchmark.py").write_text(
|
||||
"print('line one')\nprint('line two')\n"
|
||||
)
|
||||
driver_log = tmp_path / "driver.log"
|
||||
|
||||
with driver_log.open("w") as log:
|
||||
returncode, lines = benchmark._run_case(benchmarks_dir, ["--unused"], log)
|
||||
|
||||
assert returncode == 0
|
||||
assert lines == ["line one\n", "line two\n"]
|
||||
assert driver_log.read_text() == "line one\nline two\n"
|
||||
assert "line one" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_write_builtin_merges_heterogeneous_row_columns(tmp_path):
|
||||
path = tmp_path / "builtin_results.csv"
|
||||
|
||||
benchmark._write_builtin(
|
||||
path,
|
||||
[
|
||||
{"routine": "mm_bf16", "median_time": "0.004", "case_tag": "a"},
|
||||
{"routine": "cutlass_fused_moe", "case_tag": "b", "num_experts": "8"},
|
||||
],
|
||||
)
|
||||
|
||||
assert path.read_text() == (
|
||||
"routine,median_time,case_tag,num_experts\nmm_bf16,0.004,a,\ncutlass_fused_moe,,b,8\n"
|
||||
)
|
||||
Reference in New Issue
Block a user