Add perf model coverage for DeepEP EP communication ops

[gphuang]

Closes #518

Summary

DeepEPDispatch, DeepEPCombine, DeepEPDispatchBackward, and DeepEPCombineBackward are all-to-all token-routing ops used by Primus Turbo's Deep Expert Parallelism (EP) in MoE models such as DeepSeek V2 Lite and GPT-OSS-20B. They accounted for ~23.8% and ~7.9% of total GPU time respectively, but had zero perf model coverage and were counted as other. This PR adds an EP_Communication category and bandwidth model for all four ops.

Changes

perf_model.py — new EPComm base class and four concrete subclasses:

• EPComm — base for EP communication ops; flops() = 0, bytes() = num_tokens × hidden_dim × bpe (pure all-to-all, no matrix multiply).
• deepep_dispatch — parses DeepEPDispatch; Input Dims[0] = (num_tokens_local, hidden_dim).
• deepep_combine — parses DeepEPCombine; Input Dims[0] = (num_tokens_dispatched, hidden_dim).
• deepep_dispatch_backward / deepep_combine_backward — backward passes; same shape conventions as their forward counterparts in reverse.

torch_op_mapping.py — registers all four ops in op_to_perf_model_class_map, adds EPComm: "EP_Communication" to dict_base_class2category, and adds an EP_Communication branch in categorize_torch_op so these ops are no longer counted as other.

Tests

tests/test_deepep_ops.py — 19 tests covering op mapping, EP_Communication categorization, bytes() values for all four ops (BF16 and FP32), flops() = 0, relative magnitude (combine > dispatch), unknown dtype handling, and inheritance from EPComm.

Note

The EPComm base class currently sits in perf_model.py. The docstring explains the placement rationale:

• EPComm ops are bandwidth-bound (TB/s), not compute-bound — flops() = 0 is intentional
• The natural alternative would be alongside NcclAnalyser, but DeepEP uses intra-node shared memory, not NCCL
• There's no existing EP-specific comm analysis module to extend
• Placing it in perf_model.py gives it a named category (EP_Communication) and bandwidth metric in the report
• This is consistent with other Primus-specific ops (MoEDispatch, causal_conv1d, FusedRoPE, CrossEntropy) which also live in core as simple perf models

[Copilot]

Pull request overview

Adds performance-model coverage and categorization for DeepEP expert-parallel token-routing communication ops so they’re no longer reported as other in TraceLens’ perf reports.

Changes:

• Introduces a new EPComm perf-model base class plus four DeepEP-specific subclasses with flops() = 0 and a bytes-moved model derived from (num_tokens, hidden_dim, dtype).
• Registers DeepEPDispatch/Combine (and their backward variants) in op_to_perf_model_class_map and categorizes them under the new EP_Communication category.
• Adds a new pytest module with mapping, categorization, and bytes()/dtype coverage tests for the DeepEP ops.

Reviewed changes

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

File	Description
TraceLens/PerfModel/perf_model.py	Adds EPComm + DeepEP subclasses and implements a bytes-moved model for EP communication ops.
TraceLens/PerfModel/torch_op_mapping.py	Maps DeepEP op names to the new perf-model classes and introduces the EP_Communication category.
tests/test_deepep_ops.py	New unit tests validating mapping, categorization, and bytes()/flops() behavior for DeepEP ops.

---

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

[Copilot]

Pull request overview

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

---

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

[Copilot]

Pull request overview

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

---

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

[ajassani]

Review Notes

Looks solid overall — clean model, good test coverage. A few known limitations to document as TODOs for follow-up:

Known Limitations

1. Effective vs true interconnect BW — The current TB/s is effective throughput (bytes / total kernel time), which includes straggler sync/wait time. True interconnect BW (excluding sync overhead) requires tree perf / NCCL analyzer integration — planned for later.

2. Label as interconnect BW, not HBM — These ops move data over the interconnect, not HBM. Please add a column named Effective Interconnect TB/s in the report to distinguish from HBM-based TB/s used by other ops. "Effective" signals that the link isn't active for the full duration (straggler wait, sync overhead).

3. Bytes may be an upper bound — bytes() assumes every token crosses the interconnect, but tokens routed to a local expert stay on-rank. Need to check what routing metadata is actually available in real traces.

Please add a note in the EPComm docstring flagging these so they are tracked.

Request

Could you DM @ajassani an example DeepEP trace? Want to verify what metadata is available in the trace event args.

[gphuang]

Thanks for the review @ajassani — all three known limitations are now documented in the EPComm docstring (commit 8a5f105):

1. Effective vs true interconnect BW — noted that current TB/s includes straggler sync/wait; true interconnect BW requires tree-perf / NCCL-analyzer integration (planned follow-up).
2. Label as interconnect BW, not HBM — documented the need for an "Effective Interconnect TB/s" column to distinguish from HBM-based TB/s.
3. bytes() upper bound — noted that tokens routed to local experts stay on-rank; need to check what routing metadata is available in real trace args.

Will DM you an example DeepEP trace.