diff --git a/BMM/BMM_Theory/tests/test_branches.py b/BMM/BMM_Theory/tests/test_branches.py index c6156c5..195d25e 100644 --- a/BMM/BMM_Theory/tests/test_branches.py +++ b/BMM/BMM_Theory/tests/test_branches.py @@ -144,5 +144,65 @@ class TestTimingSanity(unittest.TestCase): self.assertAlmostEqual(ib.timing.fixpipe_bytes, expect) +class TestIssueRegression(unittest.TestCase): + """issue 复现 case 固化 (测评报告 §5).""" + + def setUp(self): + self.router = BranchRouter() + + 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 溢出 + from bmm_theory.constraints import check_plan_constraints + case = mkcase(51, 255, 42, 682, dtype_a="fp32", dtype_b="fp32", dtype_c="fp32") + r = self.router.route(case) + self.assertEqual(r["plan"].branch, "StreamK") # 该 case 满足 StreamK + v = check_plan_constraints(case, r["plan"]) + self.assertEqual(v, [], f"应无违规: {v}") + + def test_issue5_asw_reduced_core_fp16_l0c(self): + # issue#5: ASW 降核 fp32 场景 base tile 不越 L0C + from bmm_theory.constraints import check_plan_constraints + case = mkcase(2, 64, 16, 32, dtype_a="fp16", dtype_b="fp16", dtype_c="fp16") + r = self.router.route(case) + self.assertIn("降核", r["plan"].branch) + # base_k 应由 L0A/L0B 反推, 不再硬编码 64 + self.assertLessEqual( + r["plan"].base_m * r["plan"].base_k * 2 * 2, 64 * 1024) # L0A 上限 + + def test_issue6_iterbatch_small_k_whole_resident(self): + # issue#6: IterBatch a/b 形态 K 整驻留, dValue 下限对 K 向豁免 + from bmm_theory.constraints import check_plan_constraints + case = mkcase(256, 128, 1024, 8) # K=8 bf16, 整驻留 + r = self.router.route(case) + self.assertEqual(r["branch"], "IterBatch") + v = check_plan_constraints(case, r["plan"]) + self.assertEqual(v, [], f"整驻留形态不应报 dValue 违规: {v}") + + def test_issue9_streamk_reduce_not_double_counted(self): + # issue#9: StreamK 归约串行追加, 不进稳态 max (否则双倍计账) + from bmm_theory.branches.stream_k import StreamKBranch + case = mkcase(4, 128, 128, 10240) + sk = StreamKBranch().analyze(case) + self.assertTrue(sk.capable) + t = sk.timing + # t_total = max(稳态) + drain(=reduce), 不应是 max(...,reduce)+reduce + self.assertAlmostEqual(t.t_total, t.t_steady + t.t_drain) + self.assertAlmostEqual(t.t_drain, t.t_reduce) + + def test_out_nd_disables_streamk(self): + # issue#7: out_nd=False 禁用 StreamK (经 CLI 层解析) + case = mkcase(4, 128, 128, 10240) + case.out_nd = False + r = self.router.route(case) + self.assertNotEqual(r["branch"], "StreamK") + + if __name__ == "__main__": unittest.main()