perf(ep): add MI350X IntraNode/IntraNodeLL tuning configs for EP2/4/8

[kudomcho]

Motivation

MI350X (gfx950, 256 CUs) had no IntraNode EP tuning configs. detect_gpu_model() returned None because torch.cuda.get_device_properties(0).name is empty on MI350X, so the tuning config manager could not locate any configs at runtime.

The shipped MI355X configs (tuned for 304 CUs) are suboptimal on MI350X — the different CU count changes occupancy sweet spots for block_num and warp_per_block. In some cases, MI355X configs hit performance cliffs on MI350X where BW drops 3-6x compared to nearby configs.

Technical Details

Code changes

1. python/mori/ops/tuning_config.py — Added _ARCH_CU_TO_MODEL fallback table to detect_gpu_model(). When the device name string is empty, falls back to (gcnArchName, multi_processor_count) lookup: (gfx950, 256) → "mi350x".

2. tests/python/ops/bench_dispatch_combine.py — Added --kernel-type CLI arg (IntraNode | IntraNodeLL) to support tuning both kernel types. Threads through to EpDispatchCombineConfig and JSON config save.

Tuning configs (10 new JSON files)

• gfx950_mi350x_IntraNode_ep{2,4,8}_{dispatch,combine}.json — 6 files
• gfx950_mi350x_IntraNodeLL_ep8_{dispatch,combine}.json — 2 files (combine is new, MI355X never shipped it)

Sweep: full scope, 35 block_num × 9 warp_per_block = 315 configs per shape.

Dtype coverage:

Kernel	Phase	Dtypes tuned	Token counts
IntraNode	dispatch	fp4 (hidden=3584), fp8_e4m3 (hidden=7168)	64, 128, 256, 512, 1024, 2048, 4096
IntraNode	combine	bf16 × fp8_direct_cast quant (P2P + non-P2P)	64, 128, 256, 512, 1024, 2048, 4096
IntraNode	combine	bf16 × no quant (P2P + non-P2P)	64, 128, 256, 512, 1024, 2048, 4096
IntraNodeLL	dispatch	bf16 (hidden=7168)	1, 32, 64
IntraNodeLL	combine	bf16 × no quant	1, 32, 64

Performance Results

All numbers measured on MI350X hardware. Baseline = MI355X shipped bn/wpb configs run on MI350X.

IntraNodeLL EP8 — largest gains

Tokens	Baseline bn/wpb	Baseline BW	Tuned bn/wpb	Tuned BW	Gain
1	1/9	2.8 GB/s	144/16	3.95 GB/s	+41%
32	32/9	79.0 GB/s	32/10	95.0 GB/s	+20%
64	64/9	127.9 GB/s	72/10	148.8 GB/s	+16%

MI355X's wpb=9 is suboptimal on MI350X. Tuned configs use wpb=10 or wpb=16.

IntraNode — performance cliff fixes

MI355X's bn=256 hits severe performance cliffs on MI350X at specific shapes. Tuned configs avoid them:

EP	Shape	Baseline bn/wpb	Baseline BW	Tuned bn/wpb	Tuned BW	Gain
EP4	combine fp8dc/nP2P/128t	256/8	124.6 GB/s	224/8	186.2 GB/s	+49%
EP2	combine fp8dc/nP2P/4096t	256/16	167.2 GB/s	256/16	206.8 GB/s	+24%
EP2	dispatch fp4/256t	256/8	6.7 GB/s	208/10	43.2 GB/s	6.5x

IntraNode — general tuning improvements

EP	Phase	Shape	Baseline	Tuned	Gain
EP8	dispatch	fp4/4096t	222.5 GB/s	237.4 GB/s	+6.7%
EP8	combine	none/nP2P/64t	144.0 GB/s	156.5 GB/s	+8.7%
EP8	combine	fp8dc/nP2P/256t	409.3 GB/s	421.3 GB/s	+3.0%
EP4	combine	none/nP2P/64t	93.2 GB/s	98.7 GB/s	+5.8%
EP4	combine	fp8dc/P2P/64t	99.0 GB/s	104.4 GB/s	+5.4%
EP2	combine	none/nP2P/64t	54.1 GB/s	57.0 GB/s	+5.3%

P2P vs non-P2P decision boundary

• With fp8_direct_cast quant: non-P2P wins at all token counts for EP2/EP4; wins at ≥128 tokens for EP8. Non-P2P advantage grows as EP size decreases (EP2: up to +79%, EP8: up to +51%).
• Without quant: P2P wins at all token counts across all EP sizes (+3-19%).

Documentation

5 result docs under docs/:

• ep{2,4,8}_mi350x_tuning_results.md — recommended configs, before/after BW tables, P2P/nP2P analysis
• ep8_mi350x_IntraNodeLL_tuning_results.md — LL-specific results
• mi350x_profiler_analysis.md — static occupancy analysis, shared mem pressure

Reproducer

Hardware: 8x MI350X (gfx950, 256 CUs, 288 GB HBM3e)

# Build MORI (UMBP not required for EP tuning)
pip install -e . --config-settings=cmake.args="-DBUILD_UMBP=OFF"

# Run the full IntraNode sweep for all EP sizes (EP2/4/8 × 4 groups each)
export HSA_NO_SCRATCH_RECLAIM=1
bash tools/run_all_intranode_tuning.sh --tuning-scope full

# Or run a single group manually, e.g. EP8 fp4+fp8_direct_cast+zero-copy:
bash tools/batch_intranode_tuning.sh \
  --world-size 8 --dtype fp4 --combine-dtype bf16 \
  --quant-type fp8_direct_cast --tuning-scope full

# IntraNodeLL tuning (per token count):
python tests/python/ops/bench_dispatch_combine.py \
  --cmd tuning --world-size 8 --max-tokens 64 \
  --dtype bf16 --kernel-type IntraNodeLL --save-tuning-config auto

# Verify a specific tuned config:
python tests/python/ops/bench_dispatch_combine.py \
  --cmd bench --world-size 8 --max-tokens 512 \
  --dtype fp4 --combine-dtype bf16 --quant-type fp8_direct_cast \
  --dispatch-block-num 216 --dispatch-warp-per-block 10 \
  --combine-block-num 224 --combine-warp-per-block 16

# Configs are saved to:
ls python/mori/ops/tuning_configs/gfx950_mi350x_*

Test Plan

• Full-scope sweep completed for EP2/4/8 IntraNode (4 groups × 7 token counts × 315 configs each)
• Full-scope sweep completed for EP8 IntraNodeLL (3 token counts × 315 configs each)
• Before/after BW comparison verified against MI355X baseline configs measured on MI350X
• P2P vs non-P2P combine compared for all EP sizes with and without quant
• --kernel-type IntraNodeLL bench run verified
• Run existing IntraNode tests: pytest tests/python/ops/test_dispatch_combine_intranode.py

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

🤖 Generated with Claude Code