File size: 4,338 Bytes
fd6abd3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 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 | import triton
import triton.language as tl
import torch
@triton.jit
def _moe_kernel(
X, W_gate, W_up, W_down, W_gate_scale, W_up_scale, W_down_scale,
top_k_indices, top_k_weights,
O,
stride_xb, stride_xd,
stride_og, stride_od,
num_experts: tl.constexpr,
top_k: tl.constexpr,
dim: tl.constexpr,
intermediate_dim: tl.constexpr,
use_mxfp4: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
BLOCK_INTER: tl.constexpr,
):
pid = tl.program_id(0)
offs_d = tl.arange(0, BLOCK_SIZE)
offs_inter = tl.arange(0, BLOCK_INTER)
x = tl.load(X + pid * stride_xb + offs_d, mask=offs_d < dim, other=0.0).to(tl.float32)
acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for k in range(top_k):
expert_idx = tl.load(top_k_indices + pid * top_k + k)
expert_weight = tl.load(top_k_weights + pid * top_k + k)
w_gate_ptrs = W_gate + expert_idx * dim * dim + offs_inter[:, None] * dim + offs_d[None, :]
w_gate = tl.load(w_gate_ptrs, mask=(offs_inter[:, None] < intermediate_dim) & (offs_d[None, :] < dim), other=0.0)
if use_mxfp4:
scale = tl.load(W_gate_scale + expert_idx)
w_gate = w_gate.to(tl.float32) * scale
gate_out = tl.dot(x, tl.trans(w_gate))
gate_out = tl.sigmoid(gate_out) * gate_out
w_up_ptrs = W_up + expert_idx * dim * dim + offs_inter[:, None] * dim + offs_d[None, :]
w_up = tl.load(w_up_ptrs, mask=(offs_inter[:, None] < intermediate_dim) & (offs_d[None, :] < dim), other=0.0)
if use_mxfp4:
scale = tl.load(W_up_scale + expert_idx)
w_up = w_up.to(tl.float32) * scale
up_out = tl.dot(x, tl.trans(w_up))
hidden = gate_out * up_out
w_down_ptrs = W_down + expert_idx * dim * dim + offs_d[:, None] * offs_inter[None, :]
w_down = tl.load(w_down_ptrs, mask=(offs_d[:, None] < dim) & (offs_inter[None, :] < intermediate_dim), other=0.0)
if use_mxfp4:
scale = tl.load(W_down_scale + expert_idx)
w_down = w_down.to(tl.float32) * scale
expert_out = tl.dot(hidden, tl.trans(w_down))
acc += expert_out * expert_weight
tl.store(O + pid * stride_og + offs_d, acc, mask=offs_d < dim)
def fused_moe_triton(
X, W_gate, W_up, W_down,
top_k_indices, top_k_weights,
W_gate_scale=None, W_up_scale=None, W_down_scale=None,
use_mxfp4=False
):
batch_size, dim = X.shape
num_experts = W_gate.shape[0]
intermediate_dim = W_gate.shape[1]
top_k = top_k_indices.shape[1]
BLOCK_SIZE = triton.next_power_of_2(dim)
BLOCK_INTER = triton.next_power_of_2(intermediate_dim)
O = torch.zeros_like(X)
grid = (batch_size,)
_moe_kernel[grid](
X, W_gate, W_up, W_down,
W_gate_scale if W_gate_scale is not None else X,
W_up_scale if W_up_scale is not None else X,
W_down_scale if W_down_scale is not None else X,
top_k_indices, top_k_weights,
O,
X.stride(0), X.stride(1),
O.stride(0), O.stride(1),
num_experts=num_experts,
top_k=top_k,
dim=dim,
intermediate_dim=intermediate_dim,
use_mxfp4=use_mxfp4,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_INTER=BLOCK_INTER,
)
return O
def test_moe():
torch.manual_seed(42)
batch, dim, intermediate, num_experts, top_k = 32, 512, 1024, 8, 4
X = torch.randn(batch, dim, dtype=torch.float16, device='cuda')
W_gate = torch.randn(num_experts, dim, intermediate, dtype=torch.float16, device='cuda')
W_up = torch.randn(num_experts, dim, intermediate, dtype=torch.float16, device='cuda')
W_down = torch.randn(num_experts, intermediate, dim, dtype=torch.float16, device='cuda')
gate_logits = torch.randn(batch, num_experts, device='cuda')
top_k_weights, top_k_indices = torch.topk(F.softmax(gate_logits, dim=-1), top_k)
top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
O = fused_moe_triton(X, W_gate, W_up, W_down, top_k_indices, top_k_weights)
assert O.shape == X.shape
print(f"MoE output shape: {O.shape}")
print(f"MoE test passed!")
if __name__ == '__main__':
test_moe() |