Skip to content

vllm.models.muse_spark.amd.ops.router_bf16x3

FP32-exact router GEMM on bf16 matrix cores.

Three bf16 planes sum exactly to the fp32 weight; split-K sums in fixed order.

Functions:

router_bf16x3_gemm(x, planes)

x @ (planes[0] + planes[1] + planes[2]).T in fp32.

Source code in vllm/models/muse_spark/amd/ops/router_bf16x3.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def router_bf16x3_gemm(x: torch.Tensor, planes: torch.Tensor) -> torch.Tensor:
    """``x @ (planes[0] + planes[1] + planes[2]).T`` in fp32."""
    assert x.dtype == torch.bfloat16 and x.is_contiguous()
    m, k = x.shape
    n = planes.shape[1]
    assert planes.shape == (3, n, k)
    out = torch.empty((m, n), dtype=torch.float32, device=x.device)
    if m == 0:
        return out
    block_m, block_n, block_k, splits, num_warps = _config(m)
    part = torch.empty((splits, m, n), dtype=torch.float32, device=x.device)
    _router_bf16x3_kernel[(n // block_n, splits, triton.cdiv(m, block_m))](
        x,
        planes,
        part,
        m,
        N=n,
        K=k,
        K_PER_SPLIT=k // splits,
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        BLOCK_K=block_k,
        num_warps=num_warps,
    )
    block = 1024
    _reduce_splits_kernel[(triton.cdiv(m * n, block),)](
        part, out, m * n, SPLITS=splits, BLOCK=block, num_warps=4
    )
    return out

split_fp32_to_bf16x3(weight)

[3, N, K] bf16 planes whose sum is weight (fp32) exactly.

Source code in vllm/models/muse_spark/amd/ops/router_bf16x3.py 
39
40
41
42
43
44
45
46
def split_fp32_to_bf16x3(weight: torch.Tensor) -> torch.Tensor:
    """``[3, N, K]`` bf16 planes whose sum is ``weight`` (fp32) exactly."""
    assert weight.dtype == torch.float32
    hi = weight.to(torch.bfloat16)
    rest = weight - hi.float()
    mid = rest.to(torch.bfloat16)
    lo = (rest - mid.float()).to(torch.bfloat16)
    return torch.stack([hi, mid, lo]).contiguous()