61855506ed
3 道题参赛(全部基于队友参考实现 + 国产 NPU 套路): - Task 30 interleaved_rope (M-RoPE) 跨芯片通用版: 35.79× 平均 (5 款跑通) · 天数 86.81× · 海光 42.86× · 国通A 26.91× · 沐曦 15.54× · 华为 6.84× (燧原/昆仑芯 Failed) - Task 29 gelu_and_mul 跨芯片通用版: 3.02× 平均 (7 款跑通) · 燧原 0.96× + 华为 1.09× 拖后腿 - Task 35 rotary_embedding 跨芯片通用版: 待传(等 0 点提交次数重置) 3 个 zip (results/) + 1 个 README (D245 总览) + 1 个 LESSONS_LEARNED.md (5 作品问题 + 3 过程问题 + 协作模式 + 3 条硬规则) + BATTLECARDS + PR 模板 含 D243 失败版本 silu_and_mul_masked.py 留作复盘 作者: 阿念(Mavis) · ICE-GL-AN-001 · Code · 为 甄静(8592_apivqhj)· 之之的家 · 2026-09-02 D245 23:35 CST
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 ✓")
|