Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions collectivex/bench/ep_nccl.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ class NCCLEPBackend(EPBackend):
# pool with post-upgrade rows.
# "-static" marks the combine input bound to the full static receive plane (see
# `_bind_ht_recv_count`); "-zc" rows before it sliced that input to the received count.
# Graphed HT decode appends "-gum1": its communicator runs graph usage mode 1 (`_comm_config`).
kernel_generation = "nccl-ep-v02-ht-routed-zc-static"
SUPPORTED_MODES = ("normal", "low-latency")
CUDA_GRAPH_MODES = ("normal", "low-latency")
Expand All @@ -99,7 +100,7 @@ class NCCLEPBackend(EPBackend):
@property
def cuda_graph_supported(self) -> bool:
# HT replays at decode only, as a captured decode step runs it (engines run prefill
# uncaptured). It is slower there at EP16: see methodology, CUDA Graph Replay.
# uncaptured); see methodology, CUDA Graph Replay.
if self.mode == "normal" and getattr(self.args, "phase", None) != "decode":
return False
return super().cuda_graph_supported
Expand Down Expand Up @@ -143,6 +144,10 @@ def __init__(self, args, rank, world_size, local_rank, device):
# NCCL LL rank-major suppresses duplicate ranks in top-k order, accumulates the
# returned BF16 rank rows in FP32, and narrows only at the final output.
self.combine_reduction = "rank-fp32"
elif self.cuda_graph_supported:
# Only graphed HT captures a host NCCL collective (the routing ncclAllGather), so
# only its rows change with the communicator's graph usage mode.
self.kernel_generation = f"{type(self).kernel_generation}-gum1"
Comment on lines +147 to +150

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 (optional) When COLLX_CUDA_GRAPH=0 forces HT decode eager (methodology.md: "restores the eager pipeline"), rows now get the -gum1 kernel_generation suffix even though no graph is ever captured. self.kernel_generation is fixed at construction from the static cuda_graph_supported property (line 159), not from cuda_graph_enabled, which is what ep_harness.kernel_generation() (ep_harness.py:195-197) checks at record time to append -cudagraph. So an eager-forced decode row is stored as nccl-ep-v02-ht-routed-zc-static-gum1 with no -cudagraph suffix, a new family that no longer pools with the pre-PR eager rows (nccl-ep-v02-ht-routed-zc-static) even though its measured behavior is unchanged (mode 0/1 only diverges from mode 2 when a collective is actually captured). …

Why this was flagged

…Fix: gate the -gum1 suffix on the same dynamic cuda_graph_enabled condition the harness uses for -cudagraph, or apply it in kernel_generation() alongside that suffix, so eager-forced rows keep the pre-existing family name.

Trigger: run nccl-ep HT decode with COLLX_CUDA_GRAPH=0 (documented in methodology.md as the eager override). NCCLEPBackend.init (ep_nccl.py:159-162) sets kernel_generation using self.cuda_graph_supported, a static capability check, before COLLX_CUDA_GRAPH is consulted. ep_harness.kernel_generation() (ep_harness.py:195-197) only appends -cudagraph when backend.cuda_graph_enabled is true, so the stored family becomes ...-gum1 with no cudagraph marker, distinct from the base branch's plain nccl-ep-v02-ht-routed-zc-static for the same eager run. This fragments the durable store's history/trend for that eager configuration with no corresponding behavior change, defeating the suffix mechanism methodology.md describes for preventing incomparable rows from being silently pooled or split.

Verification: nit. The mechanism is real and reachable. ep_nccl.py:159-162 sets self.kernel_generation = f"{type(self).kernel_generation}-gum1" under elif self.cuda_graph_supported:, keyed on the static capability property. ep_nccl.py:112-118 shows cuda_graph_supported only checks mode=="normal" and phase=="decode"; it never reads COLLX_CUDA_GRAPH. The eager override is instead handled by… | nit.…

# NCCL EP's handle is explicitly reusable across dispatch/combine cycles (ep_test.py
# cached mode redispatches and recombines on one handle), so — unlike DeepEP's legacy
# low-latency Buffer — no timed component needs a fresh dispatch or a draining combine;
Expand Down Expand Up @@ -220,9 +225,21 @@ def _bootstrap_comm(self):
n = int(length.item())
uid = nccl_core.UniqueId.from_bytes(bytes(payload[:n].cpu().numpy().tobytes()))
self._comm = nccl_core.Communicator.init(
nranks=self.world_size, rank=self.rank, unique_id=uid
nranks=self.world_size, rank=self.rank, unique_id=uid, config=self._comm_config()
)

@staticmethod
def _comm_config():
"""Graph usage mode 1: one graph at a time on this communicator, never concurrent with
uncaptured work on it -- how the harness and a captured decode step both use it.

The default (2, "mixing") wraps every captured collective in an external event wait and
an event record so graph and uncaptured work can interleave (NCCL strongstream.cc). Traced
on h100 HT decode EP8, that left ~14us idle before each routing ncclAllGather and made the
graphed pair period 1.14x eager; mode 0/1 removes the gap (1.03x at T=1, 1.00x at T=256).
"""
return nccl_core.NCCLConfig(graph_usage_mode=1)

# ---- buffer construction -----------------------------------------------------------------

def create_buffer(self, spec):
Expand Down
9 changes: 6 additions & 3 deletions collectivex/docs/methodology.md
Original file line number Diff line number Diff line change
Expand Up @@ -313,9 +313,12 @@ under `CUDAGraph.replay()` by default: each library's best measured configuratio
every check, without changing its contract.

- **nccl-ep** low-latency and HT **decode**. HT decode replays as a captured decode step would
run it, although it is 3-11% slower than eager at h100/h200 EP16: captured, the per-step routing
`ncclAllGather` is a proxy-driven cross-node collective that NCCL fronts with a host-callback
node on every replay (about +50 µs of dispatch; nothing within one node). HT prefill stays eager.
run it. Its per-step routing `ncclAllGather` is the one host NCCL collective in any capture, so
nccl-ep's communicator runs NCCL graph usage mode 1 (one graph at a time): the default mixing
mode wraps every captured collective in an external event wait and record, which left ~14 µs
idle before each all-gather (h100 EP8 trace). Rows carry the `-gum1` generation suffix. Across
nodes the captured all-gather is still proxy-driven, and NCCL fronts it with a host-callback
node on every replay (about +50 µs of dispatch at EP16). HT prefill stays eager.
- **flashinfer-ep** decode; prefill stays eager (graphs change nothing there).
- **uccl-ep** low-latency, intranode only, except b200 FP8 (faster eager). Normal mode host-syncs.
- **deepep-v2** low-latency and normal **decode**, run as vLLM's graphed `deepep_v2` decode runs
Expand Down
18 changes: 18 additions & 0 deletions collectivex/tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,24 @@ def test_only_ll_rank_major_selects_the_nccl_fp32_reduction(self):
self.assertEqual(ll.combine_reduction, "rank-fp32")
self.assertEqual(getattr(ht, "combine_reduction", "domain-fp32"), "domain-fp32")

def test_graphed_ht_decode_runs_graph_usage_mode_one(self):
module = self._module()
module.dist.group = types.SimpleNamespace(WORLD=object())

def base_init(instance, options, rank, world_size, local_rank, device):
instance.args = options
instance.mode = options.mode

common = dict(experts=384, hidden=7168, topk=6, scale_up_domain=8, mode="normal")
with mock.patch.object(module.EPBackend, "__init__", base_init), \
mock.patch.dict(os.environ, {}, clear=True):
decode = module.NCCLEPBackend(types.SimpleNamespace(phase="decode", **common), 0, 16, 0, "cuda:0")
prefill = module.NCCLEPBackend(types.SimpleNamespace(phase="prefill", **common), 0, 16, 0, "cuda:0")
self.assertEqual(decode.kernel_generation, "nccl-ep-v02-ht-routed-zc-static-gum1")
self.assertEqual(prefill.kernel_generation, "nccl-ep-v02-ht-routed-zc-static")
with mock.patch.object(module.nccl_core, "NCCLConfig", lambda **kw: kw, create=True):
self.assertEqual(module.NCCLEPBackend._comm_config(), {"graph_usage_mode": 1})

def test_ll_layout_selector_restores_the_expert_major_contract(self):
module = self._module()
em = self._construct(module, 8, COLLX_NCCL_LL_LAYOUT="expert-major")
Expand Down
Loading