flashinfer.comm.create_allreduce_fusion_workspace

flashinfer.comm.create_allreduce_fusion_workspace(backend: Literal['trtllm', 'mnnvl', 'mnnvl-cutedsl', 'auto'] = 'auto', world_size: int = None, rank: int = None, max_token_num: int = None, hidden_dim: int = None, dtype: dtype = None, gpus_per_node: int = None, comm_backend: CommBackend | None = None, force_oneshot_support: bool = False, group: ProcessGroup | None = None, *, top_k: int | None = None, rms_eps: float | None = None, weight_bias: float | None = None, routed_scaling_factor: float | None = None, include_shared_expert: bool | None = None, add_residual: bool | None = None, write_residual_output: bool | None = None, apply_rms_norm: bool | None = None, output_dtype: dtype | None = None, dynamic_fp8: bool | None = None, write_norm_output: bool | None = None, config: MNNVLCuteDSLConfig | None = None) → AllReduceFusionWorkspace

Create workspace for AllReduce fusion operations.

Backend selection uses topology-based checks and heuristics.

Important: Workspace Reusability The workspace is allocated based on the total size (max_token_num * hidden_dim * dtype_size). You can reuse the same workspace with different shapes as long as the total size fits.

Use workspace.is_buffer_size_sufficient(tp_size, num_tokens, hidden_dim, dtype) to check before reusing.

Parameters:

  • backend (Literal["trtllm", "mnnvl", "mnnvl-cutedsl", "auto"]) – Backend to use. "auto" uses a topology-based heuristic to pick between "trtllm" and "mnnvl"; it never selects "mnnvl-cutedsl", which must be requested explicitly and builds an MNNVLCuteDSLAllReduceFusionWorkspace.

  • world_size (int) – Number of ranks in the process group.

  • rank (int) – Current rank id.

  • max_token_num (int) – Maximum number of tokens the workspace must support.

  • hidden_dim (int) – Hidden dimension size.

  • dtype (torch.dtype) – Element dtype of the communication tensors.

  • gpus_per_node (int, optional) – Number of GPUs per node (used for multi-node topology decisions). Defaults to min(torch.cuda.device_count(), world_size).

  • comm_backend (Optional[CommBackend]) – Communication backend to use for rendezvous. Defaults to the process-group’s default.

  • force_oneshot_support (bool) – If True, allocate workspace for the oneshot strategy up to the largest problem size requested. If False (default), allocate workspace for the twoshot strategy across all problem sizes and for the oneshot strategy up to the heuristic threshold. Only the MNNVL backend needs to be initialized with the correct strategy; the TRT-LLM backend works for both.

  • group (Optional[ProcessGroup]) – Process group used for symmetric-memory rendezvous (TRT-LLM and MNNVL CuTe DSL backends). Defaults to torch.distributed.group.WORLD.

  • backend="mnnvl-cutedsl" (The keyword-only arguments below configure)

  • are (and are rejected by every other backend. The CuTe DSL kernels)

  • built (compiled for them when the workspace is)

  • use (so every call must)

  • shown. (the same values; None keeps the workspace default)

  • top_k (Optional[int]) – Routed experts per token for the MoE finalize patterns; part of the static profile key. Default 10.

  • rms_eps (Optional[float]) – RMSNorm epsilon compiled into the kernels. Default 1e-6.

  • weight_bias (Optional[float]) – Offset added to rms_gamma (1.0 for Gemma-style norms). Default 0.0.

  • routed_scaling_factor (Optional[float]) – Scale applied to the routed-expert sum of the finalize patterns. Default 1.0.

  • include_shared_expert (Optional[bool]) – Whether the finalize patterns add shared_expert_output. Default True.

  • add_residual (Optional[bool]) – Whether residual_in is added to the reduced sum. residual_in must then be identical on every rank. Default True.

  • write_residual_output (Optional[bool]) – Whether the pre-norm sum is written to residual_out. Default True.

  • apply_rms_norm (Optional[bool]) – Whether the RMSNorm is compiled in. With apply_rms_norm=False and add_residual=False the workspace serves kAllReduce. Default True.

  • output_dtype (Optional[torch.dtype]) – torch.bfloat16 (default) or torch.float8_e4m3fn for the FP8 quantized patterns.

  • dynamic_fp8 (Optional[bool]) – Per-token dynamic FP8 scales (patterns 10 and 11) instead of a static scale; requires output_dtype=torch.float8_e4m3fn. Default False.

  • write_norm_output (Optional[bool]) – Also write the BF16 norm for the *Out* FP8 patterns (4 and 11). Default False.

  • config (Optional[MNNVLCuteDSLConfig]) – Static profiles and token-count routes. Defaults to DEFAULT_CONFIG with the norm and NO_NORM_CONFIG without it. Construction fails if no profile matches the shape and the device’s compute capability; is_mnnvl_cutedsl_supported() checks this without raising.

Returns:

A TRTLLMAllReduceFusionWorkspace, MNNVLAllReduceFusionWorkspace or MNNVLCuteDSLAllReduceFusionWorkspace. The workspace type determines which backend allreduce_fusion() will dispatch to.

Return type:

AllReduceFusionWorkspace

Raises:

  • ValueError – If no suitable backend is available for the requested configuration, or if the problem size is not supported by the chosen backend.

  • RuntimeError – If an explicit backend argument is passed that does not match any known backend implementation.

Examples

>>> # Auto-select best backend
>>> workspace = create_allreduce_fusion_workspace(
...     backend="auto",
...     world_size=8,
...     rank=0,
...     max_token_num=2048,
...     hidden_dim=4096,
...     dtype=torch.bfloat16,
... )
>>> print(workspace.backend)  # "trtllm"

>>> # Explicit backend selection
>>> workspace = create_allreduce_fusion_workspace(
...     backend="mnnvl",
...     world_size=16,
...     rank=0,
...     max_token_num=2048,
...     hidden_dim=4096,
...     dtype=torch.bfloat16,
... )

>>> # MNNVL CuTe DSL plain all-reduce (collective; every rank calls it)
>>> workspace = create_allreduce_fusion_workspace(
...     backend="mnnvl-cutedsl",
...     world_size=4,
...     rank=rank,
...     max_token_num=2048,
...     hidden_dim=8192,
...     dtype=torch.bfloat16,
...     apply_rms_norm=False,
...     add_residual=False,
... )
>>> out = allreduce_fusion(
...     input=x, workspace=workspace, pattern=AllReduceFusionPattern.kAllReduce
... )