| """Exact native post-TP combine into caller-owned BF16 output. |
| |
| Channel partials must already have been reduced across both ranks. No routing |
| weight is applied here: it was applied before the down-input QDQ. Arithmetic |
| matches native.ops.moe_combine_official including ascending-ID FP32 additions. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| from ...native.ops import _bf16_rne |
|
|
|
|
| @triton.jit |
| def _combine(P, IDS, O, H: tl.constexpr, PS: tl.constexpr, OS: tl.constexpr): |
| row = tl.program_id(0) |
| c = tl.program_id(1)*256 + tl.arange(0, 256) |
| a = tl.load(IDS + row*3) |
| b = tl.load(IDS + row*3 + 1) |
| d = tl.load(IDS + row*3 + 2) |
| i, j, k = tl.full((), 0, tl.int32), tl.full((), 1, tl.int32), tl.full((), 2, tl.int32) |
| swap = a > b |
| a, b = tl.minimum(a, b), tl.maximum(a, b) |
| i, j = tl.where(swap, j, i), tl.where(swap, i, j) |
| swap = b > d |
| b, d = tl.minimum(b, d), tl.maximum(b, d) |
| j, k = tl.where(swap, k, j), tl.where(swap, j, k) |
| swap = a > b |
| i, j = tl.where(swap, j, i), tl.where(swap, i, j) |
| x = _bf16_rne(tl.load(P + row*PS + i*H + c, c < H, other=0).to(tl.float32)) |
| y = _bf16_rne(tl.load(P + row*PS + j*H + c, c < H, other=0).to(tl.float32)) |
| z = _bf16_rne(tl.load(P + row*PS + k*H + c, c < H, other=0).to(tl.float32)) |
| shared = _bf16_rne(tl.load(P + row*PS + 3*H + c, c < H, other=0).to(tl.float32)) |
| total = _bf16_rne(((x + y) + z) + shared) |
| tl.store(O + row*OS + c, total.to(tl.bfloat16), c < H) |
|
|
|
|
| def combine_official_into(parts: torch.Tensor, ids: torch.Tensor, out: torch.Tensor) -> None: |
| if parts.ndim != 3 or parts.shape[1] != 4 or parts.dtype != torch.float32: |
| raise ValueError('TP-reduced FP32 parts [T,4,H] required') |
| rows, _, channels = parts.shape |
| if ids.shape != (rows, 3) or ids.dtype not in (torch.int32, torch.int64) or not ids.is_contiguous(): |
| raise ValueError('contiguous int32/int64 IDs [T,3] required') |
| if out.shape != (rows, channels) or out.dtype != torch.bfloat16 or out.stride(1) != 1: |
| raise ValueError('BF16 output [T,H] with contiguous columns required') |
| if parts.device != ids.device or parts.device != out.device: |
| raise ValueError('combine buffers must be on the same device') |
| if parts.stride(2) != 1 or parts.stride(1) != channels or parts.stride(0) < 4*channels or out.stride(0) < channels: |
| raise ValueError('nonoverlapping parts/output rows and contiguous slots required') |
| if out.untyped_storage().data_ptr() in (parts.untyped_storage().data_ptr(), ids.untyped_storage().data_ptr()): |
| raise ValueError('combine output must not alias inputs') |
| _combine[(rows, triton.cdiv(channels, 256))](parts, ids, out, channels, parts.stride(0), out.stride(0), |
| num_warps=4, enable_fp_fusion=False) |
|
|
