`core.inference.moe.flashinfer_mxfp8`#

FlashInfer routed-MoE support for MCore MXFP8 expert weights.

Module Contents#

Classes#

FlashInferRoutedMXFP8Weight

An expert-weight stack in TRT-LLM Major-K MXFP8 layout.

Functions#

`require_flashinfer_routed_mxfp8`	Raise an error when routed MXFP8 APIs are unavailable.
`_round_up`
`_unshuffle_cublas_scale`	Undo MCore’s 128x4 cuBLAS scale swizzle into a logical scale matrix.
`_block_scale_interleave`	Apply TRT-LLM’s architecture-independent 128x4 scale interleave.
`_shuffle_routed_scale`	Convert a logical UE8M0 scale matrix to TRT-LLM Major-K byte layout.
`prepare_routed_mxfp8_weights`	Repack a stacked MCore MXFP8 expert weight without requantizing it.
`pack_routed_mxfp8_routing`	Pack global expert IDs and BF16 routing weights for TRT-LLM MoE.
`_unwrap_output`	Normalize FlashInfer’s 0.6.x tensor/list return variants.
`quantize_routed_mxfp8_input`	Pad and quantize routed-MoE input, returning a 2D uint8 scale matrix.
`select_routed_mxfp8_active_rows`	Select the graph-stable row count for one routed-MoE invocation.
`flashinfer_routed_mxfp8_moe_prequantized`	Launch routed MXFP8 MoE with prequantized input and packed routing.
`flashinfer_routed_mxfp8_moe`	Run the FlashInfer TRT-LLM routed MXFP8 MoE kernel.

Data#

`logger`
`_LOGGED_TOKEN_POLICIES`

API#

core.inference.moe.flashinfer_mxfp8.logger#: ‘getLogger(…)’

core.inference.moe.flashinfer_mxfp8._LOGGED_TOKEN_POLICIES: set[tuple[str, int, int]]#: ‘set(…)’

core.inference.moe.flashinfer_mxfp8.require_flashinfer_routed_mxfp8() → None#: Raise an error when routed MXFP8 APIs are unavailable.

class core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight#

An expert-weight stack in TRT-LLM Major-K MXFP8 layout.

data: torch.Tensor#: None

scale: torch.Tensor#: None

logical_rows: int#: None

logical_cols: int#: None

padded_rows: int#: None

padded_cols: int#: None

core.inference.moe.flashinfer_mxfp8._round_up(value: int, alignment: int) → int#

core.inference.moe.flashinfer_mxfp8._unshuffle_cublas_scale( scale: torch.Tensor, logical_rows: int, logical_cols: int, ) → torch.Tensor#: Undo MCore’s 128x4 cuBLAS scale swizzle into a logical scale matrix.

core.inference.moe.flashinfer_mxfp8._block_scale_interleave(scale: torch.Tensor) → torch.Tensor#: Apply TRT-LLM’s architecture-independent 128x4 scale interleave.

core.inference.moe.flashinfer_mxfp8._shuffle_routed_scale(scale: torch.Tensor) → torch.Tensor#: Convert a logical UE8M0 scale matrix to TRT-LLM Major-K byte layout.

core.inference.moe.flashinfer_mxfp8.prepare_routed_mxfp8_weights( weight: megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor, out: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight | None = None, ) → core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight#

Repack a stacked MCore MXFP8 expert weight without requantizing it.

When out is provided, its data and scale tensors are refreshed in place. This preserves addresses captured by CUDA graphs across model refits.

core.inference.moe.flashinfer_mxfp8.pack_routed_mxfp8_routing( expert_ids: torch.Tensor, probabilities: torch.Tensor, ) → torch.Tensor#: Pack global expert IDs and BF16 routing weights for TRT-LLM MoE.

core.inference.moe.flashinfer_mxfp8._unwrap_output(output) → torch.Tensor#: Normalize FlashInfer’s 0.6.x tensor/list return variants.

core.inference.moe.flashinfer_mxfp8.quantize_routed_mxfp8_input( hidden_states: torch.Tensor, padded_hidden_size: int, ) → tuple[torch.Tensor, torch.Tensor]#: Pad and quantize routed-MoE input, returning a 2D uint8 scale matrix.

core.inference.moe.flashinfer_mxfp8.select_routed_mxfp8_active_rows( full_rows: int, *, token_capacity: int | None, use_bounded_rows: bool, ) → tuple[int, str]#

Select the graph-stable row count for one routed-MoE invocation.

The controller enables a bounded prefix only for a decode-only graph whose static EP-wide token bound fits the configured capacity. All other graphs retain the full dispatcher buffer so prompt tokens are never truncated.

core.inference.moe.flashinfer_mxfp8.flashinfer_routed_mxfp8_moe_prequantized( quantized_hidden: torch.Tensor, hidden_scale: torch.Tensor, packed_routing: torch.Tensor, fc1_weight: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight, fc2_weight: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight, *, num_experts: int, local_expert_offset: int, top_k: int, activation_type: int, ) → torch.Tensor#: Launch routed MXFP8 MoE with prequantized input and packed routing.

core.inference.moe.flashinfer_mxfp8.flashinfer_routed_mxfp8_moe( hidden_states: torch.Tensor, routing_map: torch.Tensor, probabilities: torch.Tensor, fc1_weight: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight, fc2_weight: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight, *, num_experts: int, local_expert_offset: int, activation_type: int, out: torch.Tensor | None = None, token_capacity: int | None = None, use_bounded_rows: bool = False, ) → torch.Tensor#

Run the FlashInfer TRT-LLM routed MXFP8 MoE kernel.

When token_capacity is set and the controller marks the current graph safe, only that fixed prefix is processed. Prefill, mixed steps, and too-large decode configurations process the full input. Invalid rows in the bounded prefix must already have expert ID -1.

The row choice is made when each CUDA graph is built, so graph replay sees fixed shapes/addresses.

core.inference.moe.flashinfer_mxfp8#