`core.transformer.moe.ops.virtual_expert_triton`#

Triton kernels for virtual-expert planning and intra-node virtual-expert transport.

The planner kernel histograms the router’s routes, exchanges the histograms between ranks, computes the deterministic virtual-expert placement and maps every route to its runtime expert, writing HybridEP’s inputs in one cooperative launch. The transport kernels move only weights and gradients: owners push weights straight into their peers’ virtual-expert slots in symmetric memory and pull the virtual-expert gradients back into native wgrad staging. Within the reserved SM budget TMA reads feed coalesced peer stores for weight pushes; gradient returns use tensor-descriptor peer reads under Triton’s loop pipeliner. Ordering against the expert GEMMs is stream order plus the layer’s collectives: the planner’s exchange precedes every push and the reduction’s rendezvous brackets every gradient exchange. When native gradients accumulate directly into DDP buffers, virtual-expert gradient slots are cleared before the accumulating grouped GEMM; other layouts use overwrite semantics for their staging buffers.

Module Contents#

Classes#

VirtualExpertPlannerWorkspace

Planner scratch for one expert layout and EP group; every plan overwrites it. gathered_counts is this rank’s [ep_size, num_experts] NCCL symmetric-memory window: the kernel publishes the local histogram into every peer’s window and reads the peers’ rows from its own, so planning needs no collective. scratch is the int32 field arena of

Functions#

`_check_triton_version`	MoonEP style load balancing requires Triton’s context local allocator introduced in 3.5.
`_check_barrier_timeout`	Advance the poll count and sample elapsed time from the wait’s fixed start.
`_grid_sync`	Self-resetting cooperative-grid barrier.
`_argmax_lowest`	`(max, id)` over the valid entries; ties take the lowest id (`SENTINEL` if none).
`_argmax_highest`	`(max, id)` over the valid entries; ties take the highest id (`-1` if none).
`_prefix_maximum`
`_plan_virtual_expert_routes_kernel`	Plan histogram, placement, and stable route remapping in one 32-block launch.
`_scratch_layout`	Fields of the planner’s int32 scratch arena as `name -> (offset, shape)` plus its total size; mirrors the field offsets in `_plan_virtual_expert_routes_kernel`.
`launch_virtual_expert_planner`	Plan one layer’s routes in a fused 32-block cooperative launch.
`_handshake`	Flip one self-resetting signal per rank pair until every peer has flipped.
`_rendezvous`	Raise this rank’s word in every peer’s pad, then consume the word each peer raised in ours: a self-resetting release/acquire handshake over one int32 word per ordered rank pair. `peers` holds the peer ranks, `valid` the live lanes.
`_cross_rank_barrier`	Publish ordinary stores and acquire peer writes before subsequent TMA reads.
`_push_fc_layer`	Push this block’s share of one component into every virtual-expert slot.
`_virtual_expert_weight_push_kernel`	Push every owner-local expert into its virtual-expert slots and rendezvous.
`_staging_pointer`	Return one native wgrad base as a pointer whose alignment Triton knows.
`_virtual_expert_grad_reduce_kernel`	Reduce every peer’s virtual-expert gradients into native wgrad staging.
`_transport_tile`	Return the largest transport tile (in the components’ unit) dividing every component.
`_check_table`	Validate a kernel input table (pointer tables are int64 `[L]`, plans int32 `[W, L]`).
`_barrier_scratch`	The inert word that retargets peer lanes which already signalled.
`_source_scratch`	The table the reduction compacts each expert’s sources into; one per device and shape so its address is stable under CUDA-graph capture. Launches sharing this storage must be serialized, as the bridge does on its shared gradient stream.
`_launch_with_descriptor_allocator`	Scope Triton’s allocator override to this launch, including exception paths.
`launch_virtual_expert_weight_prefetch`	Push BF16 or native MXFP8 owner weights into destination virtual slots.
`launch_virtual_expert_grad_reduce`	Accumulate every peer’s virtual-expert gradients into native wgrad staging.

Data#

`MAX_VIRTUAL_EXPERT_WEIGHT_SMS`
`MAX_VIRTUAL_EXPERT_EP_SIZE`
`_ROW`
`_MAX_TILE_BYTES`
`_MAX_SCALE_TILE_BYTES`
`_NUM_STAGES`
`_PUSH_NUM_WARPS`
`_GRAD_NUM_WARPS`
`_GRID_SYNC_TAG`
`_SIGNAL_STRIDE`
`_BARRIER_TIMEOUT_NS`
`_BARRIER_CLOCK_INTERVAL`
`PLANNER_PROGRAMS`
`_PLANNER_PROGRAMS`
`_ROUTE_TILE`
`_FLAG_STRIDE`

API#

core.transformer.moe.ops.virtual_expert_triton.MAX_VIRTUAL_EXPERT_WEIGHT_SMS#: 32

core.transformer.moe.ops.virtual_expert_triton.MAX_VIRTUAL_EXPERT_EP_SIZE#: 64

core.transformer.moe.ops.virtual_expert_triton._check_triton_version() → None#: MoonEP style load balancing requires Triton’s context local allocator introduced in 3.5.

core.transformer.moe.ops.virtual_expert_triton._ROW#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._MAX_TILE_BYTES#: 32768

core.transformer.moe.ops.virtual_expert_triton._MAX_SCALE_TILE_BYTES#: 8192

core.transformer.moe.ops.virtual_expert_triton._NUM_STAGES#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._PUSH_NUM_WARPS#: 4

core.transformer.moe.ops.virtual_expert_triton._GRAD_NUM_WARPS#: 8

core.transformer.moe.ops.virtual_expert_triton._GRID_SYNC_TAG#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._SIGNAL_STRIDE#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._BARRIER_TIMEOUT_NS#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._BARRIER_CLOCK_INTERVAL#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton.PLANNER_PROGRAMS#: 32

core.transformer.moe.ops.virtual_expert_triton._PLANNER_PROGRAMS#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._ROUTE_TILE#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._FLAG_STRIDE#: ‘constexpr(…)’

core.transformer.moe.ops.virtual_expert_triton._check_barrier_timeout(start, polls, LABEL: triton.language.constexpr)#

Advance the poll count and sample elapsed time from the wait’s fixed start.

Expiry can overshoot by one sampling interval of polling. A PTX trap is enabled even with TRITON_DEBUG=0 and aborts the kernel, including blocks waiting at another grid barrier. CUDA reports an asynchronous launch failure to the caller; the process must be restarted. Other ranks still in these waits expire independently. This does not time out unrelated collectives or a GPU instruction that itself stops making progress.

core.transformer.moe.ops.virtual_expert_triton._grid_sync( grid_barrier, TAG: triton.language.constexpr, NUM_SMS: triton.language.constexpr, )#: Self-resetting cooperative-grid barrier.

core.transformer.moe.ops.virtual_expert_triton._argmax_lowest( values, ids, valid, SENTINEL: triton.language.constexpr, )#: (max, id) over the valid entries; ties take the lowest id (SENTINEL if none).

core.transformer.moe.ops.virtual_expert_triton._argmax_highest(values, ids, valid)#: (max, id) over the valid entries; ties take the highest id (-1 if none).

core.transformer.moe.ops.virtual_expert_triton._prefix_maximum(a, b)#

core.transformer.moe.ops.virtual_expert_triton._plan_virtual_expert_routes_kernel( top_indices, virtual_experts, experts_to_copy, scratch, counts_sym_mem, source_rank, peer_bases, signal_bases, NUM_TOKENS: triton.language.constexpr, ROUTER_TOPK: triton.language.constexpr, EP_SIZE: triton.language.constexpr, NUM_EXPERTS: triton.language.constexpr, ALIGNMENT: triton.language.constexpr = 1, )#

Plan histogram, placement, and stable route remapping in one 32-block launch.

Up to 32 programs exchange histograms and place experts, striding over EP ranks when the group is larger than the grid. HybridEP dispatch orders successive plans across ranks, so a new histogram cannot overwrite rows still consumed by the prior plan. Contiguous route tiles avoid padding top-k to a power of two. Remapping sorts each tile by (expert, original position), scans equal-expert runs to recover stable ordinals, and scatters runtime IDs back to their original positions. Alignment 1 retains the original quotas. Aligned placement balances rounded expert blocks, assigning remainder blocks to lower ranks. Tie breaks, router scores, and per-expert real-route order are unchanged.

core.transformer.moe.ops.virtual_expert_triton._scratch_layout(num_experts: int, ep_size: int) → tuple[dict, int]#: Fields of the planner’s int32 scratch arena as name -> (offset, shape) plus its total size; mirrors the field offsets in _plan_virtual_expert_routes_kernel.

class core.transformer.moe.ops.virtual_expert_triton.VirtualExpertPlannerWorkspace( *, num_experts: int, device: torch.device, group: torch.distributed.ProcessGroup, )#

Func:: _scratch_layout; :meth:field views one field for inspection.

Initialization

field(name: str) → torch.Tensor#: View one field of the scratch arena.

destroy() → None#: Drop the symmetric window while its process group is still alive.

core.transformer.moe.ops.virtual_expert_triton.launch_virtual_expert_planner( top_indices: torch.Tensor, workspace, alignment: int = 1, ) → tuple[torch.Tensor, torch.Tensor]#

Plan one layer’s routes in a fused 32-block cooperative launch.

top_indices are the router’s [num_tokens, topk] expert ids, workspace the planner scratch (VirtualExpertPlannerWorkspace) whose gathered_counts is this rank’s symmetric buffer of every rank’s histogram. Returns the int16 [num_tokens, topk] runtime ids and the int32 [ep_size, num_local_experts] slot table.

core.transformer.moe.ops.virtual_expert_triton._handshake( address, pending, dummy, COMPARE: triton.language.constexpr, VALUE: triton.language.constexpr, SEM: triton.language.constexpr, LABEL: triton.language.constexpr, )#

Flip one self-resetting signal per rank pair until every peer has flipped.

Every pair owns an independent word, so the system-scope atomics issue in parallel instead of serializing one NVLink round trip per peer. Lanes that have already flipped retarget a scratch word, because tl.atomic_cas takes no mask and re-flipping a live signal would forge a second arrival.

core.transformer.moe.ops.virtual_expert_triton._rendezvous( signal_bases, rank, peers, valid, dummy, LABEL: triton.language.constexpr, )#: Raise this rank’s word in every peer’s pad, then consume the word each peer raised in ours: a self-resetting release/acquire handshake over one int32 word per ordered rank pair. peers holds the peer ranks, valid the live lanes.

core.transformer.moe.ops.virtual_expert_triton._cross_rank_barrier( signal_bases, grid_barrier, dummy_signal, rank, WORLD: triton.language.constexpr, NUM_SMS: triton.language.constexpr, )#: Publish ordinary stores and acquire peer writes before subsequent TMA reads.

core.transformer.moe.ops.virtual_expert_triton._push_fc_layer( bases, plan, peer_bases, entry, mine, ordinal, active, block, rank, MEMBER_BYTES: triton.language.constexpr, ARENA_BYTES: triton.language.constexpr, TILE_BYTES: triton.language.constexpr, NUM_LOCAL_EXPERTS: triton.language.constexpr, NUM_SMS: triton.language.constexpr, )#

Push this block’s share of one component into every virtual-expert slot.

TMA reads feed coalesced uint64 stores to the peer without converting BF16/MXFP8 bytes. Build the source descriptor outside the pipeline: Triton cannot predicate descriptor construction. Four stages balance bandwidth and shared-memory use.

core.transformer.moe.ops.virtual_expert_triton._virtual_expert_weight_push_kernel( fc1_bases, fc2_bases, fc1_scale_bases, fc2_scale_bases, peer_bases, signal_bases, plan, grid_barrier, dummy_signal, rank, FC1_BYTES: triton.language.constexpr, FC2_BYTES: triton.language.constexpr, FC1_SCALE_BYTES: triton.language.constexpr, FC2_SCALE_BYTES: triton.language.constexpr, TILE_BYTES: triton.language.constexpr, SCALE_TILE_BYTES: triton.language.constexpr, NUM_LOCAL_EXPERTS: triton.language.constexpr, WORLD: triton.language.constexpr, NUM_SMS: triton.language.constexpr, )#

Push every owner-local expert into its virtual-expert slots and rendezvous.

plan holds the destination-major [world, num_local_experts] table of globally numbered experts the planner wants materialized, so the entries this rank owns are a sparse subset of it. Compacting them into a dense ordinal keeps the sweep free of idle iterations even when a rank owns 8 of 512 slots, and recovering each plan entry with one masked reduction avoids staging the compacted table through memory. The arena holds fc1 data, fc1 scales, fc2 data, fc2 scales, each section NUM_LOCAL_EXPERTS members long; the two scale sections are empty for BF16 weights.

core.transformer.moe.ops.virtual_expert_triton._staging_pointer( arena, address, ELEMENT_BYTES: triton.language.constexpr, )#

Return one native wgrad base as a pointer whose alignment Triton knows.

Reached as an offset from the arena pointer rather than by casting the address itself: a pointer cast from an integer carries no alignment, and Triton will not prefetch a load through one. The staging read then stops overlapping the peer transport and costs a third of the bandwidth. Staging is 16-byte aligned, which is what the transport has always assumed of it.

core.transformer.moe.ops.virtual_expert_triton._virtual_expert_grad_reduce_kernel( arena, fc1_bases, fc2_bases, peer_bases, signal_bases, plan, sources, grid_barrier, dummy_signal, rank, FC1_ROWS: triton.language.constexpr, FC2_ROWS: triton.language.constexpr, TILE_ROWS: triton.language.constexpr, ELEMENT_BYTES: triton.language.constexpr, TILE_BEGIN: triton.language.constexpr, TILE_END: triton.language.constexpr, NUM_LOCAL_EXPERTS: triton.language.constexpr, WORLD: triton.language.constexpr, NUM_SMS: triton.language.constexpr, )#

Reduce every peer’s virtual-expert gradients into native wgrad staging.

[TILE_BEGIN, TILE_END) selects the transport tiles of one launch: the FC1 members, the FC2 members, or both. The host launches the FC layers separately so each can start as soon as its wgrad GEMM has finished.

plan holds the destination-major [world, num_local_experts] table of globally numbered experts the planner materialized, so the sources of one owner-local expert are the entries naming it, one per peer that hosts it. Every block sweeps every materialized expert but only its own contiguous slice of the tiles, which splits the payload evenly however sparse the plan is, and each block starts at a different expert so the blocks running at any instant are spread over the peers instead of queued behind one of them.

The entry rendezvous proves every peer’s wgrad GEMM has finished writing; the exit rendezvous proves every owner has read, so a peer may rewrite its slots on the next backward.

core.transformer.moe.ops.virtual_expert_triton._transport_tile(limit: int, *components: int) → int#: Return the largest transport tile (in the components’ unit) dividing every component.

core.transformer.moe.ops.virtual_expert_triton._check_table( tensor: torch.Tensor, dtype: torch.dtype, shape: tuple[int, ...], what: str, ) → torch.Tensor#: Validate a kernel input table (pointer tables are int64 [L], plans int32 [W, L]).

core.transformer.moe.ops.virtual_expert_triton._barrier_scratch(device_index: int) → torch.Tensor#: The inert word that retargets peer lanes which already signalled.

core.transformer.moe.ops.virtual_expert_triton._source_scratch(device_index: int, entries: int) → torch.Tensor#: The table the reduction compacts each expert’s sources into; one per device and shape so its address is stable under CUDA-graph capture. Launches sharing this storage must be serialized, as the bridge does on its shared gradient stream.

core.transformer.moe.ops.virtual_expert_triton._launch_with_descriptor_allocator(kernel, *args, **kwargs)#: Scope Triton’s allocator override to this launch, including exception paths.

core.transformer.moe.ops.virtual_expert_triton.launch_virtual_expert_weight_prefetch( workspace, *, sources: tuple[torch.Tensor, torch.Tensor], experts_to_copy: torch.Tensor, scale_sources: tuple[torch.Tensor, torch.Tensor] | None = None, ) → None#

Push BF16 or native MXFP8 owner weights into destination virtual slots.

workspace is the host’s transport workspace: weight_arena with its symmetric-memory weight_handle and weight_grid_barrier, plus rank, world_size, num_local_experts, member_numels and num_sms. sources and scale_sources are int64 pointer tables with one FC1 or FC2 member base per local expert. A uint8 arena selects the MXFP8 layout and requires the matching orientation’s scale tables.

core.transformer.moe.ops.virtual_expert_triton.launch_virtual_expert_grad_reduce( workspace, *, native_grads: tuple[torch.Tensor, torch.Tensor], experts_to_copy: torch.Tensor, fc_layers: tuple[int, ...] = (0, 1), ) → None#

Accumulate every peer’s virtual-expert gradients into native wgrad staging.

workspace is the host’s transport workspace: grad_arena with its symmetric-memory grad_handle and grad_grid_barrier, plus rank, world_size, num_local_experts, member_numels and num_sms. native_grads are int64 pointer tables with one FC1 or FC2 staging base per local expert; fc_layers selects which of the two this launch reduces. Used virtual-expert slots are left holding their partials; the next wgrad GEMM overwrites them.

core.transformer.moe.ops.virtual_expert_triton#