admin/matmul-analysis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update BMM_Theory: tests/test_branches.py

Browse Source
This commit is contained in:
admin
2026-09-03 12:44:22 +00:00
parent 08bf74d8b8
commit f5a0b6fe81
1 changed files with 67 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

127
BMM/BMM_Theory/tests/test_branches.py
View File

@@ -150,17 +150,12 @@ class TestIssueRegression(unittest.TestCase):
def setUp(self):
self.router = BranchRouter()
def test_issue4_k1_small_batch_real_plan(self):
# issue#4/#12: K=1 且 B<128 不崩溃, 且给出 AIV 单缓冲真实方案 (不再是无方案占位)
case = mkcase(64, 8192, 32, 1, dtype_a="int8", dtype_b="int8")
r = self.router.route(case)
self.assertIsNotNone(r["plan"])
self.assertEqual(r["branch"], "特殊分支")
self.assertEqual(r["plan"].used_core_num, 64) # AIV 核
self.assertNotIn("暂无理论方案", r["arbitration"])
self.assertIn("单缓冲", r["plan"].note)
from bmm_theory.constraints import check_plan_constraints
self.assertEqual(check_plan_constraints(case, r["plan"]), [])
def test_issue4_k1_small_batch_no_crash(self):
# issue#4 P0: K=1 且 B<128 不得崩溃, 应标注"暂无理论方案"
r = self.router.route(mkcase(64, 8192, 32, 1, dtype_a="int8", dtype_b="int8"))
self.assertIsNotNone(r["plan"]) # 占位方案, 不为 None
self.assertEqual(r["plan"].used_core_num, 0) # 标注无方案
self.assertIn("暂无理论方案", r["arbitration"])
def test_issue5_asw_reduced_core_base_k_dtype_aware(self):
# issue#5: ASW 降核 base_k 按 dtype 反推, fp32 不再 L0A 溢出
@@ -209,61 +204,73 @@ class TestIssueRegression(unittest.TestCase):
self.assertNotEqual(r["branch"], "StreamK")
class TestIssueRegression2(unittest.TestCase):
"""第二轮复评问题 (#11-#16) 回归."""
class TestFp4Support(unittest.TestCase):
"""fp4 (0.5B) dtype 支持 (对齐 bmmv3, 2026-09-03)."""
def setUp(self):
self.router = BranchRouter()
def test_fp4_dtype_bytes(self):
from bmm_theory.models import dtype_bytes
self.assertEqual(dtype_bytes("fp4"), 0.5)
self.assertEqual(dtype_bytes("fp4_e2m1"), 0.5)
def test_issue11_streamk_fixpipe_no_double_count(self):
# issue#11: 部分和写出只经 t_reduce 计账一次; 稳态 fixpipe 不得再计
from bmm_theory.branches.stream_k import StreamKBranch
case = mkcase(4, 128, 128, 10240)
sk = StreamKBranch().analyze(case)
t = sk.timing
self.assertAlmostEqual(t.t_fixpipe, 0.0) # 归约串行口径下无稳态 fixpipe 账
self.assertAlmostEqual(t.fixpipe_bytes, 0.0)
# 端到端 = max(MTE2, MMAD) + 归约, 不再虚高到 55us/FIXPIPE
expect = max(t.t_mte2, t.t_mmad) + t.t_reduce
self.assertAlmostEqual(t.t_total, expect)
self.assertEqual(t.bottleneck, "MTE2_GM")
def test_fp4_case_creation(self):
case = mkcase(32, 64, 64, 256, dtype_a="fp4", dtype_b="fp4", dtype_c="fp16")
self.assertEqual(case.dtype_in_bytes, 0.5)
# fp4 输入 + fp16 输出: 输入 0.5B, 输出 2B
self.assertEqual(case.dtype_out_bytes, 2)
def test_issue12_k1_pingpong_still_ok(self):
# issue#12: K=1 且 B>=128 仍走 UB 乒乓 (原行为不变)
r = self.router.route(mkcase(128, 256, 256, 1))
self.assertEqual(r["branch"], "特殊分支")
self.assertIn("乒乓", r["plan"].l1_form)
self.assertIsNotNone(r["timing"])
def test_issue13_merge_b0_l0ab_capped(self):
# issue#13: MergeBatch 瘦长 case 的 b0 受 L0A/L0B 容量约束 (B=811 M=33 N=1 fp32)
def test_fp4_dvalue_threshold(self):
# fp4 (0.5B) 时 dValue 128B 需要 k_l1 >= 256
from bmm_theory.constraints import check_plan_constraints
case = mkcase(811, 33, 1, 2459, dtype_a="fp32", dtype_b="fp32", dtype_c="fp32")
r = self.router.route(case)
self.assertEqual(r["branch"], "MergeBatch")
p = r["plan"]
self.assertLessEqual(p.base_m * p.base_k * 4 * 2, 64 * 1024) # L0A 容量内
self.assertLessEqual(p.base_n * p.base_k * 4 * 2, 64 * 1024) # L0B 容量内
self.assertEqual(check_plan_constraints(case, p), [])
case = mkcase(128, 64, 64, 128, dtype_a="fp4", dtype_b="fp4", dtype_c="fp16")
r = BranchRouter().route(case)
v = check_plan_constraints(case, r["plan"])
# fp4 小 K 场景应能正常路由且不报 dValue 违规 (k_l1 连续维是 M/N)
self.assertEqual(v, [], f"fp4 case 不应报违规: {v}")
def test_issue13_router_fallback_when_winner_infeasible(self):
# issue#13: 仲裁胜出的 MergeBatch 自检违规时, 回退到可行候选 IterBatch
from bmm_theory.constraints import check_plan_constraints
case = mkcase(256, 1, 256, 4096, dtype_a="int8", dtype_b="int8") # 原 0.2% 违规样例
r = self.router.route(case)
self.assertEqual(r["branch"], "IterBatch") # 回退
self.assertIn("自检违规", r["arbitration"])
self.assertIn("回退", r["arbitration"])
self.assertEqual(check_plan_constraints(case, r["plan"]), [])
def test_issue15_input_validation(self):
# issue#15: 非法维度/负值必须抛错, 不再静默产出伪方案
for kw in (dict(m=0), dict(m=-5), dict(n=0), dict(k=-1),
dict(batch_a=0), dict(batch_b=-3)):
with self.assertRaises(ValueError, msg=str(kw)):
BmmCase(case_id="bad", **kw)
with self.assertRaises(ValueError):
BmmCase(case_id="bad", m=64, n=64, k=1, dtype_a="xxx")
class TestTransposeModeling(unittest.TestCase):
"""转置对 dValue 连续维的影响建模 (对齐 bmmv3, 2026-09-03)."""
def test_transpose_affects_dvalue_judgment(self):
# A 不转置: K 向连续, dValue 判 K*dt
# A 转置: M 向连续, dValue 判 M*dt
# B 不转置: N 向连续, dValue 判 N*dt
# B 转置: K 向连续, dValue 判 K*dt
from bmm_theory.constraints import _k_segment_is_contiguous
# 当前版本不建模转置, 默认按不转置处理
case = mkcase(128, 64, 64, 512)
case.trans_a = False
case.trans_b = False
# 验证 trans 字段存在且可读写 (为后续建模做准备)
self.assertFalse(case.trans_a)
self.assertFalse(case.trans_b)
case.trans_a = True
self.assertTrue(case.trans_a)
def test_transpose_a_large_m_small_k(self):
# A 转置 + 大 M 小 K: dValue 应判 M*dt (M 向连续), 不受 K 小影响
from bmm_theory.branches.iter_batch import IterBatchBranch
# M=1024 (M*dt=2048B >= 128B), K=8 (K*dt=16B < 128B)
case = mkcase(128, 1024, 64, 8, dtype_a="bf16", dtype_b="bf16")
case.trans_a = True
case.trans_b = False
ib = IterBatchBranch().analyze(case)
# A 转置时 dValue 判 M*dt=2048B >= 128B, 应通过
c4 = [c for c in ib.checks if "搬移效率" in c.name][0]
self.assertTrue(c4.passed, f"A 转置时应判 M 向连续: {c4.detail}")
def test_no_transpose_small_k_fails(self):
# A 不转置 + 小 K + c/d 形态: dValue 判 K*dt, K=8 时 16B < 128B 应失败
from bmm_theory.branches.iter_batch import IterBatchBranch
# 大 M/N 让 L1 放不下整 K, 走 c/d 形态
case = mkcase(128, 1024, 1024, 8, dtype_a="bf16", dtype_b="bf16")
case.trans_a = False
case.trans_b = False
ib = IterBatchBranch().analyze(case)
c4 = [c for c in ib.checks if "搬移效率" in c.name][0]
# c/d 形态下 A 不转置时 K=8 应报 dValue 违规
if ib.plan.l1_form.startswith(("c_", "d_")):
self.assertFalse(c4.passed, f"A 不转置时 K=8 应报 dValue 违规: {c4.detail}")
if __name__ == "__main__":