89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
|
"""Correctness test for all v2 operators.
|
||
|
|
|
||
|
|
Run with: python -m pytest tests/test_correctness.py -v
|
||
|
|
or: python tests/test_correctness.py
|
||
|
|
"""
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def test_gelu_and_mul_v2_correctness():
|
||
|
|
import sys
|
||
|
|
sys.path.insert(0, "src")
|
||
|
|
from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline
|
||
|
|
from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2
|
||
|
|
|
||
|
|
torch.manual_seed(42)
|
||
|
|
for shape in [(1024, 1024), (4096, 4096), (2048, 11008)]:
|
||
|
|
for dtype in [torch.float16, torch.bfloat16, torch.float32]:
|
||
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
y = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
out_b = baseline(x, y)
|
||
|
|
out_v = v2(x, y)
|
||
|
|
abs_diff = (out_b - out_v).abs().max().item()
|
||
|
|
assert abs_diff < 1e-2, (
|
||
|
|
f"gelu_and_mul v2 diverges: shape={shape} dtype={dtype} abs_diff={abs_diff}"
|
||
|
|
)
|
||
|
|
print(f"gelu_and_mul shape={shape} dtype={dtype} abs_diff={abs_diff:.2e} ✓")
|
||
|
|
|
||
|
|
|
||
|
|
def test_rotary_embedding_v2_correctness():
|
||
|
|
import sys
|
||
|
|
sys.path.insert(0, "src")
|
||
|
|
from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb as baseline
|
||
|
|
from flag_gems_local.fused.rotary_embedding_v2 import apply_rotary_pos_emb as v2
|
||
|
|
|
||
|
|
torch.manual_seed(42)
|
||
|
|
for B, S, H, D in [(1, 1024, 32, 128), (2, 2048, 16, 64), (4, 512, 8, 128)]:
|
||
|
|
for interleaved in [False, True]:
|
||
|
|
q = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda")
|
||
|
|
k = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda")
|
||
|
|
cos = torch.randn(S, D // 2, dtype=torch.float16, device="cuda")
|
||
|
|
sin = torch.randn(S, D // 2, dtype=torch.float16, device="cuda")
|
||
|
|
qe_b, ke_b = baseline(q, k, cos, sin, rotary_interleaved=interleaved)
|
||
|
|
qe_v, ke_v = v2(q, k, cos, sin, rotary_interleaved=interleaved)
|
||
|
|
abs_diff_q = (qe_b - qe_v).abs().max().item()
|
||
|
|
abs_diff_k = (ke_b - ke_v).abs().max().item()
|
||
|
|
assert abs_diff_q < 1e-2, f"q diverges: {abs_diff_q}"
|
||
|
|
assert abs_diff_k < 1e-2, f"k diverges: {abs_diff_k}"
|
||
|
|
print(
|
||
|
|
f"rotary_emb B={B} S={S} H={H} D={D} "
|
||
|
|
f"interleaved={interleaved} abs_diff=({abs_diff_q:.2e}, {abs_diff_k:.2e}) ✓"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fused_moe_router_v2_correctness():
|
||
|
|
import sys
|
||
|
|
sys.path.insert(0, "src")
|
||
|
|
from flag_gems_local.ops.fused_moe_router_v2 import (
|
||
|
|
bitonic_sort_topk,
|
||
|
|
prepare_moe_inputs_v2,
|
||
|
|
)
|
||
|
|
|
||
|
|
torch.manual_seed(42)
|
||
|
|
for num_tokens, top_k in [(64, 4), (256, 8), (1024, 4)]:
|
||
|
|
weights = torch.rand(num_tokens, top_k, device="cuda")
|
||
|
|
ids = torch.randint(0, 8, (num_tokens, top_k), device="cuda")
|
||
|
|
|
||
|
|
sorted_w, sorted_ids = bitonic_sort_topk(weights, ids)
|
||
|
|
|
||
|
|
# Verify: weights should be descending
|
||
|
|
for i in range(num_tokens):
|
||
|
|
for j in range(top_k - 1):
|
||
|
|
assert sorted_w[i, j] >= sorted_w[i, j + 1] - 1e-5, (
|
||
|
|
f"Not sorted: token {i} weight[{j}]={sorted_w[i,j]} > weight[{j+1}]={sorted_w[i,j+1]}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Verify: expert ids are a permutation of original
|
||
|
|
orig_ids = torch.sort(ids, dim=1).values
|
||
|
|
new_ids = torch.sort(sorted_ids, dim=1).values
|
||
|
|
assert torch.equal(orig_ids, new_ids), "ids not a permutation!"
|
||
|
|
|
||
|
|
print(f"fused_moe_router tokens={num_tokens} top_k={top_k} ✓")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
test_gelu_and_mul_v2_correctness()
|
||
|
|
test_rotary_embedding_v2_correctness()
|
||
|
|
test_fused_moe_router_v2_correctness()
|
||
|
|
print("\nAll tests passed ✓")
|