flashinfer.comm.is_mnnvl_cutedsl_supported

flashinfer.comm.is_mnnvl_cutedsl_supported(tp_size: int, hidden_dim: int, dtype: dtype, *, max_token_num: int, top_k: int = 10, apply_rms_norm: bool = True, write_residual_output: bool = True, output_dtype: dtype = torch.bfloat16, dynamic_fp8: bool = False, write_norm_output: bool = False, config: MNNVLCuteDSLConfig | None = None, device: device | str | int | None = None) → bool

Whether this rank passes the backend="mnnvl-cutedsl" workspace checks.

Runs the checks of the MNNVL CuTe DSL workspace that neither allocate nor communicate, and returns False instead of raising: the CuTe DSL is importable, tp_size is 2, 4, 8 or 16, dtype is torch.bfloat16, max_token_num is positive, the output settings are a valid combination, config has a static profile for (tp_size, hidden_dim, top_k, dtype) on the device’s compute capability whose routes cover max_token_num and suit apply_rms_norm and dynamic_fp8, and the device is a Blackwell-class GPU with PyTorch symmetric memory and NVLink multicast. The workspace constructor runs the same checks with the same arguments.

True does not guarantee that construction succeeds: building the workspace also allocates symmetric and multicast memory, which can fail for lack of device memory or multicast resources, and compiles the kernels.

The answer describes this rank only. Workspace construction is collective, so combine the answers of all ranks (for example a MIN all-reduce of the flag) and build the workspace only when every rank supports it.

Parameters:

  • tp_size (int) – Number of ranks in the tensor-parallel group.

  • hidden_dim (int) – Hidden dimension of the communicated activations.

  • dtype (torch.dtype) – Activation dtype.

  • max_token_num (int) – Workspace capacity in tokens.

  • top_k (int) – Routed experts per token; part of the profile key. Default 10.

  • apply_rms_norm (bool) – Whether the workspace would compile the RMSNorm in. Default True.

  • write_residual_output (bool) – As the workspace argument. Default True.

  • output_dtype (torch.dtype) – As the workspace argument. Default torch.bfloat16.

  • dynamic_fp8 (bool) – As the workspace argument. Default False.

  • write_norm_output (bool) – As the workspace argument. Default False.

  • config (Optional[MNNVLCuteDSLConfig]) – Routing configuration. Defaults to DEFAULT_CONFIG with the norm and NO_NORM_CONFIG without it, as the workspace does.

  • device (Optional[Union[torch.device, int, str]]) – CUDA device to check. Defaults to the current device.

Returns:

True if this rank passes every check.

Return type:

bool