Update BMM_Theory: tests/test_branches.py (fix review issues #4-#10)

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()