"""单元测试: 用理论文档中的典型边界 case 固化分支判定与仲裁逻辑. 覆盖: - v0.98 §十 典型边界 case (MergeBatch/IterBatch 分界) - v1.1 §4.5 统一分界条件 (K 截断 + b_core 阈值) - MergeBatch 五条进入条件逐条触发 - 时延模型自洽性 (MergeBatch 胜时确实更优) 运行: python -m unittest discover -s tests -v """ import unittest from bmm_theory.models import BmmCase from bmm_theory.router import BranchRouter from bmm_theory.branches.merge_batch import MergeBatchBranch from bmm_theory.branches.iter_batch import IterBatchBranch def mkcase(b, m, n, k, **kw): return BmmCase(case_id=f"B{b}_M{m}_N{n}_K{k}", batch_a=b, batch_b=b, m=m, n=n, k=k, **kw) class TestBranchEntry(unittest.TestCase): """v0.98 §十 典型边界 case 的分支归属.""" def setUp(self): self.router = BranchRouter() def test_mergebatch_boundary_case(self): # 文档: B=128 M=N=64 K=512 -> MergeBatch 五条全过 r = self.router.route(mkcase(128, 64, 64, 512)) self.assertEqual(r["candidates"].get("MergeBatch"), True) def test_iterbatch_when_datamount_insufficient(self): # 文档: 同上但 K=256 -> 条件 3 不满足 -> IterBatch r = self.router.route(mkcase(128, 64, 64, 256)) self.assertEqual(r["branch"], "IterBatch") mb = MergeBatchBranch().analyze(mkcase(128, 64, 64, 256)) self.assertFalse(mb.capable) def test_iterbatch_when_l0c_too_small(self): # 文档: B=512 M=N=128 K=128 -> 条件 2 不满足 (MN=16384>8192) -> IterBatch mb = MergeBatchBranch().analyze(mkcase(512, 128, 128, 128)) self.assertFalse(mb.capable) r = self.router.route(mkcase(512, 128, 128, 128)) self.assertEqual(r["branch"], "IterBatch") def test_iterbatch_form_d(self): # 文档: B=64 M=N=64 K=8192 -> IterBatch 形态 d r = self.router.route(mkcase(64, 64, 64, 8192)) self.assertEqual(r["branch"], "IterBatch") self.assertIn("d_", r["plan"].l1_form) def test_to_matmul(self): r = self.router.route(BmmCase(case_id="t", batch_a=1, batch_b=1, m=2048, n=2048, k=2048)) self.assertEqual(r["branch"], "转Matmul") def test_special_k1(self): r = self.router.route(mkcase(128, 256, 256, 1)) self.assertEqual(r["branch"], "特殊分支") def test_special_k0(self): r = self.router.route(mkcase(128, 256, 256, 0)) self.assertEqual(r["branch"], "特殊分支") def test_streamk_entry(self): # 文档: B=4 M=N=128 K=10240 -> StreamK (P=1 < C/2, K>=8192) r = self.router.route(mkcase(4, 128, 128, 10240)) self.assertEqual(r["branch"], "StreamK") self.assertGreaterEqual(r["plan"].grid_k, 2) def test_streamk_partial_sum_4byte(self): # StreamK 中间部分和按 4B (L0C dtype), 不随 C 的 fp16 转换 r = self.router.route(mkcase(4, 128, 128, 10240, dtype_c="fp16")) self.assertEqual(r["plan"].out_dtype_bytes, 4) def test_asw_basic_fallback(self): # 文档: B=32 M=N=4096 K=4096 -> ASW_Basic (L1 四形态不满足) r = self.router.route(mkcase(32, 4096, 4096, 4096)) self.assertEqual(r["branch"], "ASW_Basic") def test_asw_reduced_core(self): # 文档: B=16 M=N=256 K=128 -> 降核 ASW (P=4 < 32 且 K 不满足 StreamK) r = self.router.route(mkcase(16, 256, 256, 128)) self.assertEqual(r["branch"], "ASW_Basic_降核") self.assertLess(r["plan"].used_core_num, 32) def test_asw_tail_strategy_embedded(self): # ASW_Basic 的尾轮策略是必要组成, r>0 时不应为 A0 r = self.router.route(mkcase(8, 512, 512, 512)) # N_blk 不整除 if r["branch"] == "ASW_Basic" and r["plan"].tail_block_cnt > 0: self.assertIn(r["plan"].tail_strategy, ("A1b", "方案B")) class TestArbitration(unittest.TestCase): """v1.1 §4.5: MergeBatch 仅 K 截断且 b_core 足够大时胜.""" def setUp(self): self.mb = MergeBatchBranch() def test_l1_bound_mergebatch_loses(self): # L1 绑定 (k_L1 < K): MergeBatch 恒劣 win, detail = self.mb.beats_iterbatch(mkcase(256, 128, 128, 4096)) self.assertFalse(win) self.assertIn("L1绑定", detail) def test_large_batch_mergebatch_wins(self): # 大 B + 小 MN + K 截断: MergeBatch 应胜 (v1.1 §4.5) case = mkcase(2048, 32, 32, 256) self.assertTrue(MergeBatchBranch().analyze(case).capable) win, detail = self.mb.beats_iterbatch(case) self.assertTrue(win, detail) class TestTimingSanity(unittest.TestCase): """时延模型自洽性.""" def test_mergebatch_dma_cmd_saved(self): # K 截断时 MergeBatch 的 DMA 命令数 = IterBatch 的 1/b0 # B=2048 M=N=32 K=256 是 K 截断 case (b0=4, k_l1=256=K) case = mkcase(2048, 32, 32, 256) mb = MergeBatchBranch().analyze(case) ib = IterBatchBranch().analyze(case) self.assertTrue(mb.capable) self.assertTrue(ib.capable) self.assertGreaterEqual(mb.plan.k_l1, case.k) # 确认 K 截断前提 ratio = mb.timing.dma_cmd_count / ib.timing.dma_cmd_count self.assertAlmostEqual(ratio, 1.0 / mb.plan.merge_b0, places=1) def test_bottleneck_memory_bound_for_small_mn(self): # MergeBatch case 必为访存 Bound (进入条件 5) case = mkcase(128, 64, 64, 512) mb = MergeBatchBranch().analyze(case) self.assertIn(mb.timing.bottleneck, ("MTE2", "MMAD")) def test_fixpipe_dtype_conversion(self): # C 指定 fp16 输出时, 写出量按 2B 而非 L0C 的 4B; # 字节列为整芯片口径 (issue#29): fixpipe_bytes = B*MN*outB case = mkcase(128, 64, 64, 512, dtype_c="fp16") ib = IterBatchBranch().analyze(case) expect = case.batch_c * 64 * 64 * 2 self.assertAlmostEqual(ib.timing.fixpipe_bytes, expect) self.assertAlmostEqual(ib.timing.fixpipe_bytes, case.output_bytes) class TestIssueRegression(unittest.TestCase): """issue 复现 case 固化 (测评报告 §5).""" def setUp(self): self.router = BranchRouter() def test_issue4_k1_small_batch_real_plan(self): # issue#4/#12/#17: 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_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") class TestIssueRegression2(unittest.TestCase): """第三轮复评问题 (#17-#20) + 恢复 #11-#15 回归.""" def setUp(self): self.router = BranchRouter() def test_issue11_streamk_fixpipe_no_double_count(self): # issue#11/#17: 部分和写出只经 t_reduce 计账一次; 稳态 fixpipe 不得再计 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 self.assertAlmostEqual(t.t_fixpipe, 0.0) # 归约串行口径下无稳态 fixpipe 账 self.assertAlmostEqual(t.fixpipe_bytes, 0.0) # 端到端 = 稳态(MTE2搬移链) + 归约, fixpipe 无重复账 self.assertEqual(t.bottleneck, "MTE2") self.assertAlmostEqual(t.t_total, t.t_steady + t.t_drain) def test_issue12_k1_pingpong_still_ok(self): # issue#12/#17: 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) 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), []) def test_issue13_router_fallback_when_winner_infeasible(self): # issue#13: 仲裁胜出的 MergeBatch 自检违规时, 回退到可行候选 IterBatch # (issue#31 后该样例 MergeBatch 时延已不再胜出, 用 mock 压低其时延以固定 # "胜者违规 -> 回退" 机制路径; MergeBatch 方案 dValue=112B<128 仍违规) from unittest import mock from bmm_theory.branches.merge_batch import MergeBatchBranch from bmm_theory.constraints import check_plan_constraints case = mkcase(256, 1, 256, 4096, dtype_a="int8", dtype_b="int8") mb_plan = MergeBatchBranch().analyze(case).plan self.assertTrue(check_plan_constraints(case, mb_plan), "该样例 MergeBatch 方案应仍自检违规 (dValue)") real_eval = MergeBatchBranch.evaluate def fake_win(self, c, p): t = real_eval(self, c, p) t.t_total = 1e-9 # 强制 MergeBatch 时延胜出 -> 仲裁选它 -> 走自检/回退 return t with mock.patch.object(MergeBatchBranch, "evaluate", fake_win): 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") def test_issue18_placeholder_plan_infeasible_in_evaluate(self): # issue#18: 占位方案(used_core_num=0)在 evaluate 中必须不可行, # 不得被当作可行方案给出正常时延 from bmm_theory.evaluator import PlanEvaluator from bmm_theory.models import ImplPlan # K=1 且 B<128 已恢复真实单缓冲方案 (#17), 故直接构造占位 plan 验证约束层 case = mkcase(64, 8192, 32, 1, dtype_a="int8", dtype_b="int8") r = self.router.route(case) self.assertGreater(r["plan"].used_core_num, 0) # 真实方案 ph = ImplPlan(case_id="ph", branch="特殊分支", used_core_num=0, note="该区域暂无理论方案(进入条件不满足)") er = PlanEvaluator().evaluate(case, ph) self.assertFalse(er.feasible, "占位方案应判不可行") self.assertIn("used_core_num", er.violations) def test_issue19_transpose_dvalue_guard_effective(self): # issue#19: 转置感知 dValue 判据在生成守卫/条件4/约束三处同源后真正生效. # 判别形状: B=64 M=4096 N=64 K=4096 bf16 —— d 形态 k_l1=16, # 均不转置时 dv_a=k_l1*2=32B <128 挡下 (B 侧 dv_b=N*2=128B 恰好达标也不放行, # 因为两侧切 K 两侧都要高效); # A 转置后 dv_a=M*2=8192B, 应能走 IterBatch 形态 d. from bmm_theory.constraints import check_plan_constraints c_not = BmmCase(case_id="x", batch_a=64, batch_b=64, m=4096, n=64, k=4096, trans_a=False, trans_b=False) r_not = self.router.route(c_not) self.assertNotEqual(r_not["branch"], "IterBatch") # 非转置被 dValue 守卫挡下 c_tr = BmmCase(case_id="x", batch_a=64, batch_b=64, m=4096, n=64, k=4096, trans_a=True, trans_b=False) r_tr = self.router.route(c_tr) self.assertEqual(r_tr["branch"], "IterBatch") # A 转置 M 向连续, 守卫放行 self.assertIn("d_", r_tr["plan"].l1_form) self.assertEqual(check_plan_constraints(c_tr, r_tr["plan"]), []) class TestFp4Support(unittest.TestCase): """fp4 (0.5B) dtype 支持 (对齐 bmmv3, 2026-09-03).""" 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_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_fp4_dvalue_threshold(self): # fp4 (0.5B) 时 dValue 128B 需要 k_l1 >= 256 from bmm_theory.constraints import check_plan_constraints 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}") 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}") class TestZeroCmdHandling(unittest.TestCase): """T_cmd=0 (无命令时延/未标定) 时整链路不得除零/崩溃, 按策略优先 MergeBatch.""" def test_beats_iterbatch_policy(self): # T_cmd<=0: 阈值 +inf 不可除零; K截断按策略判 MergeBatch 胜 (指令级收益未量化), # L1 绑定仍恒劣 from bmm_theory.hardware import NpuSpec from bmm_theory.branches.merge_batch import MergeBatchBranch spec0 = NpuSpec(t_cmd_ns=0.0) mb = MergeBatchBranch(spec0) # (b, m, n, k, K截断与否) cases = [(2048, 32, 32, 256, True), (128, 64, 64, 512, True), (256, 128, 128, 4096, False)] for b, m, n, k, truncated in cases: win, detail = mb.beats_iterbatch(mkcase(b, m, n, k)) self.assertEqual(win, truncated, f"{b},{m},{n},{k}: {detail}") self.assertIn("T_cmd=0", detail) def test_route_with_zero_cmd_prefers_merge(self): # T_cmd=0 且 K截断时路由应优先 MergeBatch (命令/指令次数少 b0 倍, 结构性收益) from bmm_theory.hardware import NpuSpec router = BranchRouter(NpuSpec(t_cmd_ns=0.0)) r = router.route(mkcase(2048, 32, 32, 256)) # 大 B 小 MN 典型合并场景 self.assertEqual(r["branch"], "MergeBatch") self.assertIn("策略", r["arbitration"]) self.assertIsNotNone(r["timing"]) shapes = [(128, 64, 64, 512), (64, 64, 64, 8192), (512, 128, 128, 128), (128, 128, 128, 1024), (32, 4096, 4096, 4096)] for b, m, n, k in shapes: rr = router.route(mkcase(b, m, n, k)) t = rr["timing"] self.assertIsNotNone(t, f"{b},{m},{n},{k} 应有 timing") self.assertGreater(t.t_total, 0.0) self.assertEqual(t.t_total, t.t_total) # 非 NaN def test_zero_cmd_random_smoke(self): # 随机小样本冒烟: 无异常/无 NaN import random from bmm_theory.hardware import NpuSpec router = BranchRouter(NpuSpec(t_cmd_ns=0.0)) rng = random.Random(9) for _ in range(300): b = rng.choice([2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]) m = rng.choice([8, 16, 32, 64, 128, 256, 512, 1024, 4096]) n = rng.choice([8, 16, 32, 64, 128, 256, 512, 1024, 4096]) k = rng.choice([16, 32, 64, 128, 256, 512, 1024, 4096, 16384]) c = BmmCase(case_id="z", batch_a=b, batch_b=b, m=m, n=n, k=k) rr = router.route(c) self.assertIsNotNone(rr["plan"]) t = rr["timing"] self.assertIsNotNone(t) self.assertGreater(t.t_total, 0.0) self.assertEqual(t.t_total, t.t_total) class TestIssue27to30(unittest.TestCase): """第四轮评审 (issue#27-#30, 设计文档 docs/05). #27 MergeBatch Cube 公式复核 (每步 2*(b0M)(b0N)K x b_core/b0 步); #28 dtype 感知算力 (Cube/AIV 速率表); #29 GM 首读下限 GM>=V_in + 字节列整芯片口径; #30 Fixpipe 输出落点 (整case驻留 -> L2, 否则直写 GM). """ def setUp(self): from bmm_theory.hardware import ASCEND950PR self.s = ASCEND950PR self.router = BranchRouter() # ---------------- #27 ---------------- def test_issue27_merge_cube_flops_matches_formula(self): from bmm_theory.branches.merge_batch import MergeBatchBranch case = mkcase(2048, 32, 32, 256) mb = MergeBatchBranch().analyze(case) self.assertTrue(mb.capable) p, t = mb.plan, mb.timing b0, bc = p.merge_b0, p.b_core # 整芯片列 = B*b0*2MNK; 等价 (b_core/b0) 步 x 每步 2*(b0M)(b0N)K x C 核 step_flops = 2.0 * (b0 * 32) * (b0 * 32) * 256 self.assertAlmostEqual(t.cube_flops, step_flops * (bc / b0) * self.s.aic_num) self.assertAlmostEqual(t.cube_flops, case.batch_c * b0 * 2.0 * 32 * 32 * 256) # t_mmad = 每核 flops / 单核算力 self.assertAlmostEqual(t.t_mmad, (bc * b0 * 2.0 * 32 * 32 * 256) / self.s.q16) def test_issue27_merge_gm_chip_col_equals_input(self): # K 截断合并: GM 每字节一次 -> gm 整芯片列 = V_in (issue#29 口径) from bmm_theory.branches.merge_batch import MergeBatchBranch case = mkcase(2048, 32, 32, 256) t = MergeBatchBranch().analyze(case).timing self.assertAlmostEqual(t.gm_read_bytes, case.input_bytes) self.assertAlmostEqual(t.l2_read_bytes, 0.0) # ---------------- #28 ---------------- def test_issue28_cube_dtype_rate(self): from bmm_theory.branches.iter_batch import IterBatchBranch base = mkcase(128, 64, 64, 512) # IterBatch 形态 b, 各 dtype 均可行 t16 = IterBatchBranch().analyze(base).timing t32 = IterBatchBranch().analyze( mkcase(128, 64, 64, 512, dtype_a="fp32", dtype_b="fp32")).timing t8 = IterBatchBranch().analyze( mkcase(128, 64, 64, 512, dtype_a="fp8", dtype_b="fp8")).timing # 同计算量, 时延反比于 dtype 算力: fp32=1/2, fp8=2x (相对 bf16) self.assertAlmostEqual(t32.t_mmad / t16.t_mmad, 2.0, places=6) self.assertAlmostEqual(t8.t_mmad / t16.t_mmad, 0.5, places=6) def test_issue28_aiv_dtype_rate_k1(self): # K=1 AIV 逐元素: bf16 通量 2x fp32 -> fp32 时延为 bf16 的 2 倍 base = mkcase(64, 8192, 512, 1) t_bf16 = BranchRouter().route(base)["timing"] t_fp32 = BranchRouter().route( mkcase(64, 8192, 512, 1, dtype_a="fp32", dtype_b="fp32"))["timing"] self.assertAlmostEqual(t_fp32.t_mmad / t_bf16.t_mmad, 2.0, places=6) # ---------------- #29 ---------------- def test_issue29_gm_floor_and_chip_columns(self): # 各分支代表 case: GM 整芯片列 >= V_in (R1), 且迭代/合并/转matmul/special # 等于 V_in (无重复读结构) from bmm_theory.branches.iter_batch import IterBatchBranch from bmm_theory.branches.merge_batch import MergeBatchBranch from bmm_theory.branches.stream_k import StreamKBranch from bmm_theory.branches.asw_basic import AswBasicBranch checks = [ (IterBatchBranch().analyze(mkcase(128, 64, 64, 512)).timing, mkcase(128, 64, 64, 512)), (MergeBatchBranch().analyze(mkcase(2048, 32, 32, 256)).timing, mkcase(2048, 32, 32, 256)), (StreamKBranch().analyze(mkcase(4, 128, 128, 10240)).timing, mkcase(4, 128, 128, 10240)), (AswBasicBranch().analyze(mkcase(8, 4096, 4096, 1024)).timing, mkcase(8, 4096, 4096, 1024)), ] for t, case in checks: self.assertGreaterEqual(t.gm_read_bytes + 1e-6, case.input_bytes, f"{case.case_id} GM < V_in") # 转Matmul / 特殊分支 (K=1) 亦等于 V_in r = self.router.route(mkcase(1, 2048, 2048, 2048)) t = r["timing"] self.assertGreaterEqual(t.gm_read_bytes + 1e-6, BmmCase(case_id="x", batch_a=1, batch_b=1, m=2048, n=2048, k=2048).input_bytes) r = self.router.route(mkcase(128, 256, 256, 1)) self.assertAlmostEqual(r["timing"].gm_read_bytes, 128 * (256 + 256) * 2) def test_issue29_gm_floor_random_smoke(self): # 随机多 dtype 冒烟: 所有可行方案 GM 整芯片列 >= V_in, 无 NaN import random rng = random.Random(20260609) dtypes = ["bf16", "fp16", "fp32", "int8", "fp8", "fp4"] for _ in range(400): dt = rng.choice(dtypes) kw = dict(dtype_a=dt, dtype_b=dt, dtype_c=rng.choice(["bf16", "fp16"])) b = rng.choice([1, 2, 4, 8, 16, 32, 64, 128, 256, 2048]) m = rng.choice([1, 8, 32, 64, 128, 256, 1024, 4096]) n = rng.choice([1, 8, 32, 64, 128, 256, 1024, 4096]) k = rng.choice([0, 1, 32, 64, 128, 256, 512, 1024, 4096]) c = BmmCase(case_id="gm", batch_a=b, batch_b=b, m=m, n=n, k=k, **kw) r = self.router.route(c) t = r["timing"] self.assertIsNotNone(t, c.case_id) self.assertEqual(t.t_total, t.t_total) # 非 NaN self.assertGreaterEqual(t.gm_read_bytes + 1e-6, c.input_bytes, f"{dt} b={b} m={m} n={n} k={k} gm={t.gm_read_bytes}") # ---------------- #30 ---------------- def test_issue30_iter_output_l2_when_whole_case_fits(self): from bmm_theory.branches.iter_batch import IterBatchBranch # 整 case 输入+输出 16.8MB+1.0MB <= 128MB -> 输出写 L2 写口, GM 写=0 case = mkcase(128, 64, 64, 512) self.assertLess(case.input_bytes + case.output_bytes, self.s.l2_bytes) t = IterBatchBranch().analyze(case).timing self.assertAlmostEqual(t.fixpipe_bytes, case.output_bytes) self.assertAlmostEqual(t.t_fixpipe, case.output_bytes / self.s.bw_l2) def test_issue30_iter_output_gm_when_case_overflows(self): from bmm_theory.branches.iter_batch import IterBatchBranch # 输出 268MB > L2 余量 -> 直写 GM (输入优先驻留) case = mkcase(128, 1024, 1024, 64) self.assertGreater(case.input_bytes + case.output_bytes, self.s.l2_bytes) t = IterBatchBranch().analyze(case).timing self.assertAlmostEqual(t.fixpipe_bytes, case.output_bytes) self.assertAlmostEqual(t.t_fixpipe, case.output_bytes / self.s.bw_gm) def test_issue30_asw_scene_whole_case_labels(self): # S_A (整case全驻留): l2_policy_out=resident, t_fixpipe 走 L2 写口 case_sa = mkcase(2, 4096, 4096, 512) self.assertLess(case_sa.input_bytes + case_sa.output_bytes, self.s.l2_bytes) r = self.router.route(case_sa) self.assertEqual(r["branch"], "ASW_Basic") self.assertTrue(r["plan"].l2_policy_out.startswith("resident"), r["plan"].l2_policy_out) self.assertIn("A_整case全驻留", r["plan"].note) self.assertAlmostEqual(r["timing"].t_fixpipe, case_sa.output_bytes / self.s.bw_l2) # S_B (整case超L2, 单batch输入可驻留): 输出直写 GM case_sb = mkcase(64, 4096, 4096, 1024) self.assertGreater(case_sb.input_bytes + case_sb.output_bytes, self.s.l2_bytes) self.assertLess(4096 * 1024 * 2 * 2, self.s.l2_bytes) # 单batch输入 16.8MB r = self.router.route(case_sb) self.assertIn("ASW", r["branch"]) self.assertTrue(r["plan"].l2_policy_out.startswith("direct_gm"), r["plan"].l2_policy_out) self.assertAlmostEqual(r["timing"].t_fixpipe, case_sb.output_bytes / self.s.bw_gm) def test_issue30_to_matmul_output_l2_toggle(self): # 转Matmul 粗估: 整case可驻留 -> L2; 大 case -> GM small = mkcase(1, 2048, 2048, 2048) self.assertLess(small.input_bytes + small.output_bytes, self.s.l2_bytes) r = self.router.route(small) self.assertEqual(r["branch"], "转Matmul") self.assertAlmostEqual(r["timing"].t_fixpipe, small.output_bytes / self.s.bw_l2) big = mkcase(1, 8192, 8192, 4096) self.assertGreater(big.input_bytes + big.output_bytes, self.s.l2_bytes) r = self.router.route(big) self.assertEqual(r["branch"], "转Matmul") self.assertAlmostEqual(r["timing"].t_fixpipe, big.output_bytes / self.s.bw_gm) class TestIssue31to32(unittest.TestCase): """第五轮评审: #31 切B类 GM 每字节恰读一次 = V_in; #32 ASW 场景升级 (单侧全驻留 -> S_B 使 GM=V_in; 双侧超 L2 -> S_C 最小替换 2D 分组 + 窗口 L2 计账).""" def setUp(self): from bmm_theory.hardware import ASCEND950PR self.s = ASCEND950PR self.router = BranchRouter() # ---------------- #31 ---------------- def test_issue31_iter_form_c_gm_equals_vin(self): # b64_m16_n256_k512: 形态 c, K=512 切 n_K=3 (k_L1=240) 非整除 -> GM 仍 = V_in from bmm_theory.branches.iter_batch import IterBatchBranch case = mkcase(64, 16, 256, 512) ib = IterBatchBranch().analyze(case) self.assertIn("c_", ib.plan.l1_form) self.assertLess(ib.plan.k_l1, case.k) t = ib.timing self.assertAlmostEqual(t.gm_read_bytes, case.input_bytes) self.assertAlmostEqual(t.l2_read_bytes, 0.0) self.assertAlmostEqual(t.t_mte2_gm, case.input_bytes / self.s.bw_gm) def test_issue31_iter_form_d_gm_equals_vin(self): # K=5000 非整除切段 (k_L1=1024, n_K=5, 末段剩 904): GM 仍恰一次 = V_in from bmm_theory.branches.iter_batch import IterBatchBranch case = mkcase(128, 64, 64, 5000) ib = IterBatchBranch().analyze(case) self.assertIn("d_", ib.plan.l1_form) self.assertLess(ib.plan.k_l1, case.k) self.assertAlmostEqual(ib.timing.gm_read_bytes, case.input_bytes) self.assertAlmostEqual(ib.timing.t_mte2_gm, case.input_bytes / self.s.bw_gm) def test_issue31_merge_l1_bound_gm_equals_vin(self): # MergeBatch L1 绑定 (k_L1 L2, 但 B=58.7MB 可全驻留 # + A 行块滑窗 (58.7+2x2.5 <= 128) -> S_B: GM = V_in, 重复读全命中 L2 case = mkcase(32, 8192, 4096, 7168) r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") p, t = r["plan"], r["timing"] a_b = case.m * case.k * case.dtype_in_bytes bb_b = case.k * case.n * case.dtype_in_bytes m_cnt, n_cnt = p.m_cnt, p.n_cnt self.assertGreater(a_b + bb_b, self.s.l2_bytes) # 单 batch 确实超 L2 self.assertLessEqual(bb_b + 2 * (a_b / m_cnt), self.s.l2_bytes) # 单侧全驻留成立 self.assertAlmostEqual(t.gm_read_bytes, case.input_bytes) exp_l2 = case.batch_c * ((n_cnt - 1) * a_b + (m_cnt - 1) * bb_b) self.assertAlmostEqual(t.l2_read_bytes, exp_l2) def test_issue32_asw_scene_c_min_gm(self): # b8_m131072_n8192_k8192: 双侧均不可全驻留 -> S_C 最小替换分组; # gm == 容量约束最小解 (测试内复算), 且 >= V_in case = mkcase(8, 131072, 8192, 8192) r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") p, t = r["plan"], r["timing"] self.assertIn("C_", p.note) m_cnt, n_cnt = p.m_cnt, p.n_cnt dt = case.dtype_in_bytes a_b = case.m * case.k * dt bb_b = case.k * case.n * dt blk_a = a_b / m_cnt blk_b = bb_b / n_cnt best = None for mg in range(1, m_cnt + 1): for ng in range(1, n_cnt + 1): if mg * blk_a + ng * blk_b > self.s.l2_bytes: continue gm = (n_cnt + ng - 1) // ng * a_b + (m_cnt + mg - 1) // mg * bb_b if best is None or gm < best: best = gm self.assertIsNotNone(best) self.assertAlmostEqual(t.gm_read_bytes, case.batch_c * best) self.assertGreaterEqual(t.gm_read_bytes, case.input_bytes) # 最小替换解应严格优于"每块独立落 GM"的保守上界 self.assertLessEqual(t.gm_read_bytes, case.batch_c * (n_cnt * a_b + m_cnt * bb_b)) class TestIssue33(unittest.TestCase): """ASW_Basic tile 选择 (issue#33, v1.91 §5.1/§5.2 + 尾轮 v1.5 §2.1 修正口径): BaseM/N = 256 方形 (UnitFlag 单缓冲 L0C 用满 65536 元素), SingleCoreM/N 有界枚举.""" def setUp(self): from bmm_theory.hardware import ASCEND950PR self.s = ASCEND950PR self.router = BranchRouter() def test_base_tile_square_single_buffer(self): # 默认 BaseM=BaseN=256 (非 176 双缓冲口径), L0C 恰用满 (256*256*4=256KB) from bmm_theory.branches.asw_basic import AswBasicBranch p = AswBasicBranch().make_plan(mkcase(32, 4096, 4096, 4096)) self.assertEqual((p.base_m, p.base_n), (256, 256)) self.assertEqual(p.base_m * p.base_n * 4, self.s.l0c_bytes) self.assertEqual(p.base_k, 64) # 64KB/(2*256*2) = 64 self.assertTrue(p.fixpipe_unitflag) # UnitFlag 单缓冲 def test_base_tile_exception_small_m(self): # M=128 < 256: BaseM=128 (被迫跟随), BaseN = min(65536/128=512, N)=512 from bmm_theory.branches.asw_basic import AswBasicBranch p = AswBasicBranch().make_plan(mkcase(64, 128, 4096, 1024)) self.assertEqual((p.base_m, p.base_n), (128, 512)) self.assertEqual(p.base_k, 32) # min(64K/(2*128*2)=128, 64K/(2*512*2)=32) def test_v191_example_singlecore(self): # v1.91 §5.2 完整实例: B=8 M=N=2048 K=1024 bf16 -> (4,4) 512x512, k_l1=128, r=0 case = mkcase(8, 2048, 2048, 1024) r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") p = r["plan"] self.assertEqual((p.single_core_m, p.single_core_n), (512, 512)) self.assertEqual((p.m_cnt, p.n_cnt), (4, 4)) self.assertEqual(p.k_l1, 128) self.assertEqual(p.tail_block_cnt, 0) # 8*4*4=128 % 32 = 0 完美整除 def test_singlecore_not_constant_176(self): # 回归原 bug: SingleCoreM/N 恒为 176 (双缓冲 Base 兜底); 现在自适应 from bmm_theory.branches.asw_basic import AswBasicBranch p = AswBasicBranch().make_plan(mkcase(8, 4096, 4096, 1024)) self.assertNotEqual((p.single_core_m, p.single_core_n), (176, 176)) self.assertEqual((p.single_core_m, p.single_core_n), (512, 512)) # 且为 Base 的整数倍 (约束 4) self.assertEqual(p.single_core_m % p.base_m, 0) self.assertEqual(p.single_core_n % p.base_n, 0) def test_l0c_single_buffer_constraint_pass(self): # ASW 单缓冲: base 256x256 通过约束 (无违规) from bmm_theory.constraints import check_plan_constraints from bmm_theory.branches.asw_basic import AswBasicBranch case = mkcase(8, 4096, 4096, 1024) p = AswBasicBranch().make_plan(case) self.assertEqual(check_plan_constraints(case, p), []) def test_p1_fallback_to_enum_when_no_split_infeasible(self): # B>=C 时不切分 k_l1 过 dValue 下限失败 -> 强制切分枚举 -> 512x512 (8,8) case = mkcase(128, 4096, 4096, 8192) r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") p = r["plan"] self.assertEqual((p.single_core_m, p.single_core_n), (512, 512)) self.assertEqual((p.m_cnt, p.n_cnt), (8, 8)) def test_issue34_k_l1_decomposition_derived(self): # b32_m16_n8192_k7168: 分解 = Base 16x1024 (例外: M=16<256; N 侧受 L0B 单边 # 上限收敛), SingleCore 16x1024 (mCnt=1,nCnt=8), k_l1 = L1 双缓冲反推 # ⌊L1/(2·(sM+sN)·dt)⌋16 = ⌊524288/(2·1040·2)⌋16 = 112 (v1.91 §5.2) from bmm_theory.branches.asw_basic import AswBasicBranch case = mkcase(32, 16, 8192, 7168) p = AswBasicBranch().make_plan(case) self.assertEqual((p.base_m, p.base_n), (16, 1024)) self.assertEqual((p.single_core_m, p.single_core_n), (16, 1024)) self.assertEqual(p.k_l1, 112) def test_pathological_shape_degraded_warning_always_plan(self): # issue#34: 兜底分支恒出方案 —— 极端形状 (N=8 int8 大K, B 侧 dValue=8B # 物理不可满足) 照常给方案 + 标注效率降级 (warning), 不判违规/不判不可行 from bmm_theory.constraints import check_plan_constraints from bmm_theory.evaluator import PlanEvaluator case = mkcase(4096, 2048, 8, 1024, dtype_a="int8", dtype_b="int8") r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") self.assertEqual(check_plan_constraints(case, r["plan"]), []) self.assertIn("效率降级", r["plan"].note) er = PlanEvaluator().evaluate(case, r["plan"]) self.assertTrue(er.feasible) self.assertIn("效率降级", er.advice) self.assertGreater(er.timing.t_total, 0.0) def test_pathological_n8_hard_floor_fallback(self): # 更极端: N=8 int8 连 128B 硬下限也不满足 -> 仍恒出方案, 效率降级标注 from bmm_theory.evaluator import PlanEvaluator case = mkcase(512, 4096, 1, 8192, dtype_a="int8", dtype_b="int8") r = self.router.route(case) self.assertEqual(r["branch"], "ASW_Basic") self.assertIsNotNone(r["plan"]) self.assertIn("效率降级", r["plan"].note) er = PlanEvaluator().evaluate(case, r["plan"]) self.assertTrue(er.feasible) if __name__ == "__main__": unittest.main()