flashinfer.fused_moe.MoEActivationPack

class flashinfer.fused_moe.MoEActivationPack(hidden_states_q: Tensor, hidden_states_scale: Tensor | None, topk_ids: Tensor | None = None, topk_weights: Tensor | None = None, *, per_token_scale: Tensor | None = None, routing_input_mode: RoutingInputMode = RoutingInputMode.PackedPrecomputed, routing_logits: Tensor | None = None, routing_bias: Tensor | None = None)

Per-call backend-native activations plus routing inputs.

Activation encoding depends on the MMA pair on QuantConfig:

  • NVFP4×NVFP4 with TrtllmFp4Config or CuteDslConfig: packed uint8 [M, H/2] values with float8_e4m3fn [M, H/16] block scales.

  • NVFP4×NVFP4 with CutlassNvfp4Config: raw bfloat16 [M, H] values without an activation scale.

  • MXFP4×MXFP8 (W4A8): float8_e4m3fn [M, H] MXFP8 values with token-major float8_e4m3fn [M, H/32] tensors carrying UE8M0 scale bytes, matching the TRTLLM FP4 launcher ABI.

  • MXFP4×BF16 (TRTLLM W4A16): raw bfloat16 [M, H] values with no activation scale; weights use the MXFP4 preparation contract.

  • NVFP4×BF16 (CuTe-DSL / b12x W4A16): raw bfloat16 [M, H] values with no activation scale; weights use the NVFP4 preparation contract.

  • BF16: raw bfloat16 [M, H] values with no scale tensor.

  • MxInt4: raw bfloat16 [M, H] values with no scale tensor; weights are packed signed INT4 with BF16 block scales.

  • DeepSeek FP8: float8_e4m3fn [M, H] values with transposed float32 [H/128, M] block scales.

  • MXFP8: float8_e4m3fn [M, H] values with token-major uint8 [M, H/32] UE8M0 scales.

  • FP8 per-tensor with TrtllmFp8PerTensorConfig: float8_e4m3fn [M, H] values with no activation scale; the calibrated scalar is folded into the TRT-LLM weight view.

  • FP8 per-tensor with CutlassFp8PerTensorConfig: float8_e4m3fn [M, H] values with a scalar float32 dequantization scale. Obtain both tensors from CutlassFp8PerTensorConfig.prepare_activations. This pack cannot be shared with the TRT-LLM per-tensor backend.

routing_input_mode selects how routing reaches the kernel (the runner reads it directly):

  • PackedPrecomputed (default) — pre-routed: the caller computes expert selection on the host and passes topk_ids + topk_weights. The TRTLLM runners normally combine both fields into one packed int32 tensor before launch.

  • UnpackedPrecomputed — pre-routed, separate kernel inputs: supported by the TRTLLM runners. The caller supplies int32 ids and BF16 or FP32 weights directly, avoiding packed-id construction. The launcher consumes the weights in their native dtype.

  • FromLogits — in-kernel: the caller passes raw routing_logits (and, for bias-aware methods like DeepSeekV3/MiniMax2, routing_bias); the kernel computes the top-k selection itself per RoutingConfig.method. topk_ids / topk_weights stay None — the runner allocates internal kernel-filled buffers, and the routing result is not surfaced back through the pack (routing replay is a separate, future capability). TRTLLM FP4, BF16, block-FP8, per-tensor-FP8, and MxInt4 runners support this mode; MoELayer dispatches a logits pack only to capable backends (see each runner’s supported_routing_modes).

topk_ids / topk_weights follow the routed-MoE naming convention (gh #2425); they keep the field positions of the former selected_experts / final_scales, so positional construction of pre-routed packs is unchanged. Additional activation metadata and the in-kernel routing fields are keyword-only.

__init__(hidden_states_q: Tensor, hidden_states_scale: Tensor | None, topk_ids: Tensor | None = None, topk_weights: Tensor | None = None, *, per_token_scale: Tensor | None = None, routing_input_mode: RoutingInputMode = RoutingInputMode.PackedPrecomputed, routing_logits: Tensor | None = None, routing_bias: Tensor | None = None) → None

Methods

__init__(hidden_states_q, hidden_states_scale)

Attributes

num_tokens

per_token_scale

routing_bias

routing_input_mode

routing_logits

topk_ids

topk_weights

hidden_states_q

hidden_states_scale