Address perf issue for allgather_fullmesh_2

[Binyang Li (Binyang2014)]

Summary

• Remove redundant local-output copies in default_allgather_fullmesh2.
• Interleave unique local-copy chunks with remote peer transfers.
• Add BF16 benchmark encoding, decoding, and correctness support.
• Support dtype/accum-qualified tuning entries with generic fallback.
• Consolidate benchmark tests into test_benchmark.py.

Motivation

In out-of-place allgather, every peer channel references the same local output buffer. The previous kernel copied each local chunk once per peer, resulting in seven redundant local writes on an 8-GPU system.

The motivating configuration showed a large latency gap between 49,504 B and 51,200 B messages.

Implementation

Remote transfers retain the existing full-mesh mapping:

peerIdx = gWid % nPeer;
chunkIdx = gWid / nPeer;

Local copies use an independent offset so each local output region is written only once:

localOffset = gWid * localCopyBytesPerWarp;

Tail copies are bounded by the remaining message size.

The benchmark updates also:

• Encode/decode BF16 raw uint16 storage correctly.
• Require exact correctness for non-reduction collectives.
• Select tuning entries by dtype and accumulation type when specified.
• Preserve generic allgather tuning entries.

Performance

Measured on 8x MI300X using BF16 out-of-place allgather:

Message size	Latency
49,504 B	7.54 us
51,200 B	7.51 us
61,440 B	7.62 us
71,680 B	7.75 us
81,920 B	8.77 us
99,008 B	8.99 us
102,400 B	8.99 us

The 49,504/51,200 B gap is reduced to approximately 0.4%.

Validation

• Built mscclpp_collectives successfully.
• Ran ./tools/lint.sh.
• Benchmark unit tests: 7 passed.
• Checked all 109 configured MI300X allgather sizes on 8 GPUs.
• All correctness checks passed with zero mismatches.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Optimizes out-of-place full-mesh allgather and makes tuning configurations dtype-aware.

Changes:

• Distributes local allgather copies across warps.
• Keys tuning entries by data and accumulation types.
• Adds BF16 benchmark selection.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
src/ext/collectives/allgather/allgather_fullmesh_2.cu	Rebalances local-copy work.
python/mscclpp_benchmark/tuning_config.py	Adds dtype-aware configuration storage.
python/mscclpp_benchmark/comm.py	Selects configurations by dtype.
python/mscclpp_benchmark/bench_collective.py	Records typed tuning results and adds BF16.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

python/mscclpp_benchmark/tuning_config.py:226

• Entries loaded from JSON do not apply the same accum = dtype default as select() and upsert(). A hand-authored entry such as { "dtype": "float16" } is stored as (float16, None), while every normal lookup defaults to (float16, float16), so the entry can never be selected. Default a missing accum to the parsed dtype during deserialization as well.

                    dtype=None if raw_entry.get("dtype") is None else str(raw_entry["dtype"]),
                    accum=None if raw_entry.get("accum") is None else str(raw_entry["accum"]),

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.