[PR #671] Make get_rdma_gbs RoCE-LAG aware #694

Open
opened 2026-10-27 21:05:20 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/671
Author: @zfan2356
Created: 6/25/2026
Status: 🔄 Open

Base: main ← Head: fix/rdma-gbs-lag-aware


📝 Commits (3)

  • a7113cd Make get_rdma_gbs RoCE-LAG aware
  • c3cb61d Add tests for get_rdma_gbs LAG-aware path
  • 2c5745a Pass explicit comp_mask=NUM_LAG_PORTS to query_mlx5_device

📊 Changes

2 files changed (+184 additions, -2 deletions)

View changed files

📝 deep_ep/utils/envs.py (+54 -2)
➕ tests/utils/test_envs.py (+130 -0)

📄 Description

Background

get_rdma_gbs reads ibstat's Rate: to feed ElasticBuffer.get_theoretical_num_sms. On Mellanox RoCE LAG (e.g., CX-7 with two 200 Gb/s NDR rails bonded into one mlx5_bond_*), ibstat reports the per-rail rate — the mlx5 driver hashes traffic across both rails, so a single QP only ever uses one, but multiple QPs reach the bond's full aggregate. The probe therefore returns half of the real per-NIC capacity and the downstream SM-count calculation picks the wrong bottleneck.

Same root cause as #614.

Measurement (H800 + CX-7, mlx5_bond_* 2-rail LAG)

Probe Result
ibstat 'mlx5_bond_1' .. Rate: 200 Gb/s (per rail)
ib_write_bw -d mlx5_bond_1 -F (1 QP) 183 Gb/s ≈ 22.9 GB/s
ib_write_bw -d mlx5_bond_1 -F -q 8 372 Gb/s ≈ 46.6 GB/s
Mlx5Context.query_mlx5_device(comp_mask=MASK_NUM_LAG_PORTS).num_lag_ports 2

Real per-NIC bandwidth is ~50 GB/s, not 25 — a 2x error.

Fix

Query num_lag_ports via pyverbs and multiply the ibstat rate by it. Same approach as Mooncake's transfer engine, see mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp:266 for the C++ equivalent (mlx5dv_query_device + MLX5DV_CONTEXT_MASK_NUM_LAG_PORTS).

ctx = Mlx5Context(attr=Mlx5DVContextAttr(), name=nic_name)
dv = ctx.query_mlx5_device(comp_mask=1 << 9)  # MLX5DV_CONTEXT_MASK_NUM_LAG_PORTS
num_lag_ports = dv.num_lag_ports

Gotcha: query_mlx5_device()'s default comp_mask=-1 ORs only the masks pyverbs knows about; at least the version I tested against didn't include NUM_LAG_PORTS, so num_lag_ports came back as 0. The bit must be passed explicitly. Hardcoded from infiniband/mlx5dv.h since pyverbs doesn't re-export the constant.

Also adds an EP_RDMA_GBS=<gbps> env override (also /8 to GB/s) so users with broken probes or non-Mellanox fabrics can short-circuit detection — addresses the request in #614.

Falls back cleanly to the original single-rail behaviour when pyverbs is missing, when the NIC isn't mlx5, or when LAG isn't enabled.

Tests

tests/utils/test_envs.py — three local checks (no distributed setup needed) that auto-skip when pyverbs / ibstat / mlx5_bond_* aren't present. On H800 + CX-7 LAG the test prints:

get_rdma_gbs('mlx5_bond_1') = 50.0 GB/s  (per-port 200 Gb/s x 2 rails / 8)
LAG aggregation detected: 400 Gb/s -> 50.0 GB/s
ALL TESTS PASSED

Compatibility

No new hard dependencies; pyverbs is optional (import is try/except). No API changes. Only the returned value changes, and only when LAG is actually in effect.

Fixes #614. Credit @michaelchen1996 env-override sketch.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/671 **Author:** [@zfan2356](https://github.com/zfan2356) **Created:** 6/25/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/rdma-gbs-lag-aware` --- ### 📝 Commits (3) - [`a7113cd`](https://github.com/deepseek-ai/DeepEP/commit/a7113cd0deff599d88a52c54df6641dff091e368) Make get_rdma_gbs RoCE-LAG aware - [`c3cb61d`](https://github.com/deepseek-ai/DeepEP/commit/c3cb61de428617c655283a845676f8446b6ce6c2) Add tests for get_rdma_gbs LAG-aware path - [`2c5745a`](https://github.com/deepseek-ai/DeepEP/commit/2c5745a23ed1f985dc166588e94f502d64ab8780) Pass explicit comp_mask=NUM_LAG_PORTS to query_mlx5_device ### 📊 Changes **2 files changed** (+184 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/utils/envs.py` (+54 -2) ➕ `tests/utils/test_envs.py` (+130 -0) </details> ### 📄 Description ## Background `get_rdma_gbs` reads `ibstat`'s `Rate:` to feed `ElasticBuffer.get_theoretical_num_sms`. On Mellanox **RoCE LAG** (e.g., CX-7 with two 200 Gb/s NDR rails bonded into one `mlx5_bond_*`), `ibstat` reports the **per-rail** rate — the mlx5 driver hashes traffic across both rails, so a single QP only ever uses one, but multiple QPs reach the bond's full aggregate. The probe therefore returns **half** of the real per-NIC capacity and the downstream SM-count calculation picks the wrong bottleneck. Same root cause as #614. ## Measurement (H800 + CX-7, `mlx5_bond_*` 2-rail LAG) | Probe | Result | |---|---| | `ibstat 'mlx5_bond_1' .. Rate:` | 200 Gb/s (per rail) | | `ib_write_bw -d mlx5_bond_1 -F` (1 QP) | 183 Gb/s ≈ 22.9 GB/s | | `ib_write_bw -d mlx5_bond_1 -F -q 8` | **372 Gb/s ≈ 46.6 GB/s** | | `Mlx5Context.query_mlx5_device(comp_mask=MASK_NUM_LAG_PORTS).num_lag_ports` | **2** | Real per-NIC bandwidth is ~50 GB/s, not 25 — a 2x error. ## Fix Query `num_lag_ports` via pyverbs and multiply the ibstat rate by it. Same approach as Mooncake's transfer engine, see [`mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp:266`](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp#L266) for the C++ equivalent (`mlx5dv_query_device` + `MLX5DV_CONTEXT_MASK_NUM_LAG_PORTS`). ```python ctx = Mlx5Context(attr=Mlx5DVContextAttr(), name=nic_name) dv = ctx.query_mlx5_device(comp_mask=1 << 9) # MLX5DV_CONTEXT_MASK_NUM_LAG_PORTS num_lag_ports = dv.num_lag_ports ``` **Gotcha**: `query_mlx5_device()`'s default `comp_mask=-1` ORs only the masks pyverbs knows about; at least the version I tested against didn't include `NUM_LAG_PORTS`, so `num_lag_ports` came back as 0. The bit must be passed explicitly. Hardcoded from `infiniband/mlx5dv.h` since pyverbs doesn't re-export the constant. Also adds an **`EP_RDMA_GBS=<gbps>` env override** (also `/8` to GB/s) so users with broken probes or non-Mellanox fabrics can short-circuit detection — addresses the request in #614. Falls back cleanly to the original single-rail behaviour when `pyverbs` is missing, when the NIC isn't mlx5, or when LAG isn't enabled. ## Tests `tests/utils/test_envs.py` — three local checks (no distributed setup needed) that auto-skip when pyverbs / ibstat / `mlx5_bond_*` aren't present. On H800 + CX-7 LAG the test prints: ``` get_rdma_gbs('mlx5_bond_1') = 50.0 GB/s (per-port 200 Gb/s x 2 rails / 8) LAG aggregation detected: 400 Gb/s -> 50.0 GB/s ALL TESTS PASSED ``` ## Compatibility No new hard dependencies; `pyverbs` is optional (import is `try/except`). No API changes. Only the returned value changes, and only when LAG is actually in effect. Fixes #614. Credit @michaelchen1996 env-override sketch. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#694
No description provided.