`core.transformer.moe.router_replay`#

Module Contents#

Classes#

`RouterReplayAction`	A Enum to define the actions for router replay.
`RouterReplay`	A class to manage the recording and replaying of MoE routing decisions. It holds all router instances and provides static methods to globally control recording and replaying.

API#

class core.transformer.moe.router_replay.RouterReplayAction(*args, **kwds)#

Bases: enum.Enum

A Enum to define the actions for router replay.

Initialization

RECORD#: ‘record’

REPLAY_FORWARD#: ‘replay_forward’

REPLAY_BACKWARD#: ‘replay_backward’

class core.transformer.moe.router_replay.RouterReplay(is_mtp_layer: bool = False)#

A class to manage the recording and replaying of MoE routing decisions. It holds all router instances and provides static methods to globally control recording and replaying.

Initialization

Initialize a RouterReplay instance for a specific layer.

Parameters:: is_mtp_layer (bool) – Whether this recorder belongs to an MTP layer.

global_router_replay_instances: List[core.transformer.moe.router_replay.RouterReplay]#: []

static get_instances( is_mtp_layer: bool | None = None, ) → List[core.transformer.moe.router_replay.RouterReplay]#

Get registered router replay instances, optionally filtered by layer type.

Parameters:: is_mtp_layer (bool | None) – Select MTP routers when true, base-model routers when false, or all routers when unset.
Returns:: Matching router replay instances in registration order.
Return type:: List[RouterReplay]

static set_replay_data(all_layers_topk_indices: List[torch.Tensor])#

Distributes the topk indices for all layers to their respective RouterReplay instances.

Parameters:: all_layers_topk_indices – A list of tensors, where each tensor contains the topk indices for a specific layer. The order must match the instantiation order of the routers.

static get_recorded_data( is_mtp_layer: bool | None = None, ) → List[torch.Tensor | None]#

Collect recorded top-k indices, optionally filtered by layer type.

Parameters:: is_mtp_layer (bool | None) – Select MTP routers when true, base-model routers when false, or all routers when unset.
Returns:: Recorded indices in router registration order.
Return type:: List[torch.Tensor | None]

static clear_global_indices()#: Clears the recorded and target topk indices in all instances.

static set_global_router_replay_action( router_replay_action: core.transformer.moe.router_replay.RouterReplayAction, *, is_mtp_layer: bool | None = None, ) → None#

Set the replay action, optionally filtered by layer type.

Parameters:

router_replay_action (RouterReplayAction) – Replay action to set.
is_mtp_layer (bool | None) – Select MTP routers when true, base-model routers when false, or all routers when unset.

static clear_global_router_replay_action()#: Clears the router replay action for all router instances.

static clear_global_router_replay_instances()#: Clear the global list of router replay instances to prevent memory leaks.

static set_global_static_buffers( static_buffer: torch.Tensor, *, is_mtp_layer: bool | None = None, buffer_index: torch.Tensor | None = None, layer_indices: list[torch.Tensor] | None = None, ) → None#

Set static buffers, optionally filtered by layer type.

Parameters:

static_buffer (torch.Tensor) – Tensor shaped [max_tokens, num_layers, topk], with a leading buffer dimension when buffer_index is provided.
is_mtp_layer (bool | None) – Select MTP routers when true, base-model routers when false, or all routers when unset.
buffer_index (torch.Tensor | None) – GPU selector for double-buffered recording.
layer_indices (list[torch.Tensor] | None) – Stable GPU layer indices for double-buffered recording, one scalar per router.

static clear_global_static_buffers( *, is_mtp_layer: bool | None = None, ) → None#

Clear static buffers, optionally filtered by layer type.

Parameters:: is_mtp_layer (bool | None) – Select MTP routers when true, base-model routers when false, or all routers when unset.

set_target_indices(topk_indices: torch.Tensor)#: Sets the target topk indices for replay.

get_recorded_indices() → Optional[torch.Tensor]#: Returns the recorded topk indices.

clear_indices()#: Clears the recorded and target topk indices.

set_router_replay_action( router_replay_action: core.transformer.moe.router_replay.RouterReplayAction, )#: Sets the router replay action for this layer.

clear_router_replay_action()#: Clears the router replay action for this layer.

get_replay_topk( scores: torch.Tensor, topk: int, num_groups: Optional[int] = None, group_topk: Optional[int] = None, default_compute_topk: Callable[[torch.Tensor, int, Optional[int], Optional[int]], Tuple[torch.Tensor, torch.Tensor]] = None, ) → Tuple[torch.Tensor, torch.Tensor]#

A wrapper for top-k computation that handles different replay actions.

Parameters:

scores (torch.Tensor) – The scores to compute top-k on.
topk (int) – The number of top elements to select.
num_groups (Optional[int]) – Number of expert groups for group-limited routing.
group_topk (Optional[int]) – Number of groups to select for each token.
default_compute_topk (Callable) – The default top-k computation function, which should return a tuple of (values, indices).

Returns:

A tuple containing the top-k values and indices.

Return type:

Tuple[torch.Tensor, torch.Tensor]

set_static_buffer( buffer: torch.Tensor, *, buffer_index: torch.Tensor | None = None, layer_index: torch.Tensor | None = None, ) → None#

Sets a static buffer for CUDA graph compatible recording.

Parameters:

buffer (torch.Tensor) – Per-layer destination shaped [max_tokens, topk] for legacy recording, or the full contiguous [2, max_tokens, num_layers, topk] buffer for async recording.
buffer_index (torch.Tensor | None) – GPU selector for double-buffered recording.
layer_index (torch.Tensor | None) – GPU layer index within the shared buffer.

clear_static_buffer()#: Clears the static buffer.

record_indices(topk_indices: torch.Tensor) → None#

Records the topk indices.

If a static buffer is set (for CUDA graph compatibility), copies into it. Otherwise, just stores the tensor reference.

Parameters:: topk_indices (torch.Tensor) – Selected expert indices for this layer’s tokens.

core.transformer.moe.router_replay#

Module Contents#

Classes#

API#

`core.transformer.moe.router_replay`#