flashinfer.fused_moe.trtllm_fused_moe_workspace_size

flashinfer.fused_moe.trtllm_fused_moe_workspace_size(max_num_tokens: int, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, *, quant_mode: str, local_num_experts: int | None = None, num_fused_shared_experts: int = 0, activation_type: int = 3, use_shuffled_weight: bool = True, weight_layout: int | None = None, use_routing_scales_on_input: bool = False, use_per_token_scale: bool = False, valid_hidden_size: int | None = None, valid_intermediate_size: int | None = None, gemm1_weights: Tensor | None = None, gemm1_weights_scale: Tensor | None = None, gemm2_weights: Tensor | None = None, gemm2_weights_scale: Tensor | None = None, device: device | None = None, tactic: Tuple[int, int] | None = None) → int

Return the workspace_buffer size in bytes for the trtllm-gen fused-MoE APIs.

The trtllm-gen counterpart of cutlass_fused_moe_workspace_size(). Allocate the returned number of bytes once (as a 1-D torch.uint8 CUDA tensor) and pass it as workspace_buffer= to trtllm_*_moe / trtllm_*_routed_moe. The launch then carves its routing maps, intermediate GEMM outputs and batched-GEMM workspaces from that buffer instead of allocating and freeing them on every call.

With tactic=None, the size is measured by the launcher’s own allocation code, over every tile size at max_num_tokens, so one buffer fits whichever tactic the autotuner or the default heuristic selects at that token count. The footprint grows with the token count, so a buffer sized for max_num_tokens also serves smaller batches.

With default sizing, autotuning profiles every tactic from the same buffer, so tuning needs no scratch beyond it. Size the buffer for max(max_num_tokens, tune_max_num_tokens): a launch whose scratch exceeds the buffer (a larger token count or tuning bucket) allocates the tensors that do not fit per call.

With an explicit tactic=(tile_tokens_dim, config_index), only that tactic’s scratch is measured. Use this after tuning, with the tactic and token count used by the launch. This size does not cover other tactics selected for other tuning buckets or during subsequent autotuning.

Parameters:

  • max_num_tokens (int) – Largest hidden_states.shape[0] the buffer must serve.

  • hidden_size (int) – Logical (unpacked) hidden dimension, padded as the weights are.

  • intermediate_size (int) – Logical intermediate dimension, as passed to the MoE API.

  • num_experts (int) – Global routed-expert count.

  • top_k (int) – Routed experts per token.

  • quant_mode (str) –

    The API and data types the buffer will serve:

  • local_num_experts (Optional[int]) – Experts on this rank. Defaults to num_experts.

  • num_fused_shared_experts (int) – Fused shared experts (FP4 and FP8 block-scale only). Default 0.

  • activation_type (int) – Activation, as passed to the MoE API. Default Swiglu.

  • use_shuffled_weight (bool, Optional[int]) – Weight preprocessing, as passed to the MoE API ("bf16" and "fp8_block_scale" only; the other modes fix them). weight_layout defaults to WeightLayout.BlockMajorK for "bf16" and WeightLayout.MajorK otherwise.

  • weight_layout (bool, Optional[int]) – Weight preprocessing, as passed to the MoE API ("bf16" and "fp8_block_scale" only; the other modes fix them). weight_layout defaults to WeightLayout.BlockMajorK for "bf16" and WeightLayout.MajorK otherwise.

  • use_routing_scales_on_input (bool) – As passed to the FP8 per-tensor / per-channel APIs. Default False.

  • use_per_token_scale (bool) – Whether per_token_scale will be passed ("nvfp4" only).

  • valid_hidden_size (Optional[int]) – As passed to the MoE API, when the weights are padded.

  • valid_intermediate_size (Optional[int]) – As passed to the MoE API, when the weights are padded.

  • gemm1_weights (Optional[torch.Tensor]) – The prepared weights (and block or per-channel weight scales) exactly as they will be passed to the MoE API. Optional: each one omitted is described to the native query by a one-element stand-in of the shape and dtype quant_mode implies, so only pass them when the real tensors are at hand (for example at layer load time). "bf16" and "fp8_per_tensor" have no weight scales.

  • gemm1_weights_scale (Optional[torch.Tensor]) – The prepared weights (and block or per-channel weight scales) exactly as they will be passed to the MoE API. Optional: each one omitted is described to the native query by a one-element stand-in of the shape and dtype quant_mode implies, so only pass them when the real tensors are at hand (for example at layer load time). "bf16" and "fp8_per_tensor" have no weight scales.

  • gemm2_weights (Optional[torch.Tensor]) – The prepared weights (and block or per-channel weight scales) exactly as they will be passed to the MoE API. Optional: each one omitted is described to the native query by a one-element stand-in of the shape and dtype quant_mode implies, so only pass them when the real tensors are at hand (for example at layer load time). "bf16" and "fp8_per_tensor" have no weight scales.

  • gemm2_weights_scale (Optional[torch.Tensor]) – The prepared weights (and block or per-channel weight scales) exactly as they will be passed to the MoE API. Optional: each one omitted is described to the native query by a one-element stand-in of the shape and dtype quant_mode implies, so only pass them when the real tensors are at hand (for example at layer load time). "bf16" and "fp8_per_tensor" have no weight scales.

  • device (Optional[torch.device]) – CUDA device the buffer will be used on. Defaults to the device of the given weights, else the current device.

  • tactic (Optional[Tuple[int, int]]) – Explicit (tile_tokens_dim, config_index) from the autotuner. A config index of -1 selects the default valid config for that tile. Defaults to None, which measures the maximum over supported tiles. Unsupported tiles raise an error. Configurations that cannot run for this problem return zero.

Returns:

Required workspace_buffer size in bytes.

Return type:

int

Note

Share one buffer only among launches that are stream-ordered with each other, such as the MoE layers of one forward pass. Overlapping micro-batches on separate streams each need their own buffer. workspace_buffer requires do_finalize=True.