Fix low-latency EP epoch race with host-assigned dispatch epochs

[Caio Rocha (caiocbr)]

Summary

Assign low-latency dispatch epochs once in MoELowLatencyRuntime and pass the epoch to dispatch and combine through Workload.

Previously, each CUDA block derived the epoch from device workspace state while block 0 updated that state without a grid-wide ordering guarantee. Late blocks could therefore observe a different epoch and wait indefinitely on readiness flags.

Changes

• Add a host-owned dispatch epoch counter to MoELowLatencyRuntime.
• Increment the epoch once before each dispatch.
• Pass the same epoch to dispatch and its subsequent combine operation.
• Use the workload epoch for rank-major synchronization.
• Remove device-side epoch generation, publication, and workspace storage.

This gives every block in a launch one immutable epoch value.

Validation

• Passed eager rank-major tests at capacities 1 and 2.
• Passed eager expert-major tests at capacity 2.
• Passed bit-exact 2-rank and 4-rank multi-operation CUDA Graph replays.
• Passed all 10 SGLang dispatcher tests.
• Completed six 32-rank serving rounds with 40/40 responses and no scheduler watchdog timeouts.

[azure-pipelines]

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

[Copilot]

Pull request overview

Moves low-latency dispatch epoch assignment from GPU workspace state to the host runtime.

Changes:

• Adds host-managed dispatch epochs to workloads.
• Uses workload epochs for dispatch/combine synchronization.
• Removes the workspace epoch allocation.

Reviewed changes

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

Show a summary per file

File	Description
src/ext/ep/low_latency/dispatch.cu	Consumes workload epoch.
src/ext/ep/low_latency/config.cuh	Removes workspace epoch storage.
src/ext/ep/low_latency/combine.cu	Uses workload epoch for synchronization.
src/ext/ep/ll_runtime.hpp	Adds host epoch state.
src/ext/ep/ll_runtime.cc	Assigns and propagates epochs.
src/ext/ep/include/api.cuh	Adds epoch to Workload.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.