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 (TrtllmFp4Config, CuteDslConfig, CutlassNvfp4Config, Cake): packed uint8 [M, H/2] values with linear float8_e4m3fn [M, H/16] block scales and unit global scale. b12x / cuTile NVFP4 take raw bfloat16 [M, H] and quantize in-kernel.

  • MXFP4×MXFP8 (W4A8; TrtllmFp4Config, CuteDslConfig, CutlassMxfp8Mxfp4Config): 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 x MXFP8 with SM12xMxfp8Mxfp4Config: raw bfloat16 [M, H] values without an activation scale; K128 MXFP8 quantization runs internally.

  • DeepSeek FP8 with SM12xFp8Config: raw bfloat16 [M, H] values without an activation scale; K128 FP8 quantization runs internally.

  • 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 (TrtllmFp8BlockConfig, CutlassMxfp8Config): float8_e4m3fn [M, H] values with token-major uint8 [M, H/32] UE8M0 scales. CUTLASS MXFP8 runners accept the flat swizzled 1-D input_sf instead when QuantConfig(swizzled_scale_factors=True).

  • FP8 per-tensor (TrtllmFp8PerTensorConfig, CutlassFp8PerTensorConfig): float8_e4m3fn [M, H] values with no activation scale; the static calibration multipliers live in each backend’s weight view.

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