Skip to content

`vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl` ¶

CuteDSL BF16x3 router GEMM.

Computes X @ W.T for BF16 X with shape [N, K] and FP32 router weights W with shape [M, K] by decomposing each FP32 weight value into three BF16 residual terms inside the kernel, then accumulating the three BF16 MMA results into FP32 TMEM output.

Functions:

bf16x3_router_gemm –

Return X @ W.T using the SM100 BF16x3 router GEMM kernel.

`bf16x3_router_gemm(X, W)` ¶

Return X @ W.T using the SM100 BF16x3 router GEMM kernel.

Source code in vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py

def bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
    """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel."""
    N, K = X.shape
    M, _ = W.shape
    num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count

    # next power of 2 within 8 and 128
    BN = triton.next_power_of_2(N)
    BN = min(max(BN, 8), 128)

    BM = 128
    BK = 64
    k_tiles = math_utils.cdiv(K, BK)
    grid_m = math_utils.cdiv(M, BM)
    grid_n = math_utils.cdiv(N, BN)

    base_ctas = grid_m * grid_n
    split_k = min(k_tiles, max(1, num_sms // base_ctas))

    partials = X.new_empty(split_k, N, M, dtype=torch.float32)
    Sm100BF16x3RouterGemm.compile(BN, K)(X, W, partials, split_k)

    if split_k == 1:
        return partials.squeeze(0)

    out = X.new_empty(N, M, dtype=torch.float32)
    splitk_reduce_triton(partials, out)
    return out