Is this page helpful?

core.transformer.moe.inference_routing_mask_kernel#

Triton kernel for masking CUDA-graph padding rows of a local routing map.

Under CUDA-graph capture the local token count is padded up to a captured graph size; those padding rows have garbage routing indices and, if left alone, would dispatch padding tokens to real experts. This kernel zeroes that out by writing -1 into every topk slot of rows in [real_token_count, local_tokens).

The kernel reads real_token_count from a fixed-address int32[1] GPU tensor, so it is safe to call from inside a captured graph: only the value behind the pointer changes between replays.

Module Contents#

Functions#

_mask_routing_padding_kernel

Fill routing_map[real_token_count:, :] with -1, BLOCK_M rows per program.

mask_routing_padding

In-place fill -1 into routing_map[real_token_count:, :].

API#

core.transformer.moe.inference_routing_mask_kernel._mask_routing_padding_kernel(
routing_map_ptr,
real_token_count_ptr,
total_rows: triton.language.int32,
tp_rank: triton.language.int32,
TOPK: triton.language.constexpr,
BLOCK_M: triton.language.constexpr,
BLOCK_TOPK: triton.language.constexpr,
)#

Fill routing_map[real_token_count:, :] with -1, BLOCK_M rows per program.

core.transformer.moe.inference_routing_mask_kernel.mask_routing_padding(
routing_map: torch.Tensor,
real_token_count_tensor: torch.Tensor,
tp_rank: int = 0,
) → None#

In-place fill -1 into routing_map[real_token_count:, :].

Parameters:

  • routing_map – [N, topk] int64 local routing map. N is the (possibly CUDA-graph-padded) local token count.

  • real_token_count_tensor – [1] int32 GPU tensor holding the real (unpadded) token count for this step, in the global (pre-SP-shard) frame. Read inside the kernel so the mask boundary moves correctly across CUDA-graph replays.

  • tp_rank – This rank’s index in the SP/TP group. Local row r is row r + tp_rank * N in the global frame; the kernel uses this offset to compare against real_token_count_tensor.