How can I turn on the Warp specialization in the grouped_gemm op?

[linlang1837]

The below are my codes.

from fbgemm_gpu.experimental.gemm.triton_gemm.grouped_gemm import (
    grouped_gemm,
)    
a = torch.randn(M * G, K, dtype=torch.bfloat16, device=device)
b = torch.randn(N * G, K, dtype=torch.bfloat16, device=device)
m_sizes = torch.full((G, ), M, device=device, dtype=torch.int32)
for _ in range(num_warmup):
    result = grouped_gemm(
        a,
        b,
        m_sizes,
    )

When I run the codes, the following warning is displayed:
fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py:975: UserWarning: Warp specialization is disabled as the Triton build in current environment doesn't have such support. Please build from https://github.com/facebookexperimental/triton/tree/ws-3.2.x to enable it for best performance on Nvidia's SM90 GPUs.

I want to know if the only way to turn on warp specialization is to build Triton from source in the specific branch?
And the versions of the relevant packages are below:
fbgemm_gpu_genai 2025.6.10+cu126
pytorch-triton 3.3.1+gitc8757738
torch 2.8.0.dev20250611+cu126
triton 3.3.1

[linlang1837]

@q10 @jianyuh

[moderato]

Got a very similar question: #4192. Would appreciate any guidance on the best practice of setting up triton for FBGEMM.

@q10 @jianyuh

[levendlee]

Using https://github.com/facebookexperimental/triton/tree/ws-3.2.x is needed to use WS on Hopper.
Alternatively, one can use CUTLASS/ComposableKernel based GroupedGEMM, APIs are defined in https://github.com/pytorch/FBGEMM/blob/main/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize_defs.cpp, as ".*_grouped_cat".

[moderato]

Using https://github.com/facebookexperimental/triton/tree/ws-3.2.x is needed to use WS on Hopper. Alternatively, one can use CUTLASS/ComposableKernel based GroupedGEMM, APIs are defined in https://github.com/pytorch/FBGEMM/blob/main/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize_defs.cpp, as ".*_grouped_cat".

Thank you for the answer. I have tried ws-3.2.x and ws-3.3.x and got errors for grouped fp8 GEMM. Would you take a look at this issue? #4192