`core.inference.moe.flashinfer_mxfp8`#

FlashInfer MoE support for MCore MXFP8 expert weights.

Module Contents#

Classes#

`FlashInferRoutedMXFP8Weight`	An expert-weight stack in TRT-LLM Major-K MXFP8 layout.
`FlashInferCutlassMXFP8Weight`	An expert-weight stack in CUTLASS row-major MXFP8 layout.

Functions#

`require_flashinfer_cutlass_mxfp8`	Raise an error when the CUTLASS MXFP8 APIs are unavailable.
`require_flashinfer_routed_mxfp8`	Raise an error when routed MXFP8 APIs are unavailable.
`_round_up`
`_pack_cutlass_mxfp8_scale`	Repackage existing MXFP8 scales for FlashInfer’s CUTLASS interface.
`prepare_cutlass_mxfp8_weights`	View a canonical expert stack as CUTLASS weights, sharing its storage.
`_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.
`attach_routed_mxfp8_writeback`	Make an expert’s canonical MXFP8 weight write through to its routed slot.
`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_cutlass_mxfp8_moe`	Run full-row MXFP8 MoE through FlashInfer’s CUTLASS backend.
`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_cutlass_mxfp8() → None#: Raise an error when the CUTLASS MXFP8 APIs are unavailable.

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

class core.inference.moe.flashinfer_mxfp8.FlashInferCutlassMXFP8Weight#

An expert-weight stack in CUTLASS row-major MXFP8 layout.

data: torch.Tensor#: None

scale: torch.Tensor#: None

input_scale: torch.Tensor#: None

logical_rows: int#: None

logical_cols: int#: None

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

core.inference.moe.flashinfer_mxfp8._pack_cutlass_mxfp8_scale( scale: torch.Tensor, experts: int, rows: int, cols: int, ) → torch.Tensor#

Repackage existing MXFP8 scales for FlashInfer’s CUTLASS interface.

Each scale occupies one byte and applies to 32 weights. CUTLASS expects four adjacent scale bytes per int32 entry, with shape [experts, rows, cols // 128]. Only the tensor view changes: scale values and byte order are preserved.

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

View a canonical expert stack as CUTLASS weights, sharing its storage.

Callers must retain canonical per-expert views for in-place refits. Neither the data nor the scale bytes need a layout conversion for this backend.

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.attach_routed_mxfp8_writeback( weight: megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor, routed: core.inference.moe.flashinfer_mxfp8.FlashInferRoutedMXFP8Weight, expert: int, ) → None#

Make an expert’s canonical MXFP8 weight write through to its routed slot.

Refit’s quantize_/copy_ then requantizes with the canonical kernel and repacks into the routed stack in place (same bytes, same CUDA-graph addresses), so the canonical storage is released and replaced by broadcast placeholders: only the routed copy stays resident.

After this call weight.data and weight.scale no longer hold the weights: reading them (dequantize, torch.stack, .contiguous(), …) silently returns one uninitialized element broadcast to the full shape. Read the routed stack instead.

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_cutlass_mxfp8_moe( hidden_states: torch.Tensor, routing_map: torch.Tensor, probabilities: torch.Tensor, fc1_weight: core.inference.moe.flashinfer_mxfp8.FlashInferCutlassMXFP8Weight, fc2_weight: core.inference.moe.flashinfer_mxfp8.FlashInferCutlassMXFP8Weight, *, ep_size: int, ep_rank: int, activation_type, out: torch.Tensor | None = None, activation_clamp_limit: torch.Tensor | None = None, ) → torch.Tensor#: Run full-row MXFP8 MoE through FlashInfer’s CUTLASS backend.

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#