From a2ad7a7ec77540eb861281276d5d31b72ce1c604 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Sun, 1 Nov 2026 22:35:11 -0800 Subject: [PATCH] cake_rmsnorm_det_train: tighten two large-row weight-gradient accuracy edges The compensated backward programs now start at lower row counts on two routes: the H = 4096 in-chunk compensated form above 44032 rows on SM100 (was 52736) and the compensated-dot H = 256 ring above 99840 rows on SM103 (was 114688). Below the old edges the plain programs' weight gradient came closer than the 0.8 routing margin to the per-element bound. cake_jit.py moves the window records of the two edges, cake_windows.py follows and records the measured class of the accuracy edges that stay, and the pairs test covers both sides of the new edges. No program changes. --- .../cake_jit.py | 68 ++++----- .../cake_windows.py | 135 +++++++----------- .../norm/test_cake_rmsnorm_det_train_pairs.py | 14 +- 3 files changed, 96 insertions(+), 121 deletions(-) diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py index a98b4b65778..1eb6e689d01 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py @@ -1821,7 +1821,7 @@ "role": "bwd_main", "select_rule": { "above": 4096, - "at_most": 52736, + "at_most": 44032, "kind": "rows_between", }, "workspace_rule": { @@ -1853,7 +1853,7 @@ "role": "bwd_main_w32", "select_rule": { "above": 4096, - "at_most": 52736, + "at_most": 44032, "kind": "rows_between", }, "workspace_rule": { @@ -1883,7 +1883,7 @@ "library": "cake_rmsnorm_det_train_7122dbcd697b3fbe1ee0", "role": "bwd_main", "select_rule": { - "above": 52736, + "above": 44032, "at_most": None, "kind": "rows_between", }, @@ -1915,7 +1915,7 @@ "library": "cake_rmsnorm_det_train_5ef5f1f13ed8f2b994c4", "role": "bwd_main_w32", "select_rule": { - "above": 52736, + "above": 44032, "at_most": None, "kind": "rows_between", }, @@ -2096,7 +2096,7 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_52736": { + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_44032": { "arches": [ "sm_100a", ], @@ -2114,12 +2114,12 @@ "role": "bwd_reduce", "select_rule": { "above": 4096, - "at_most": 52736, + "at_most": 44032, "kind": "rows_between", }, "workspace_rule": None, }, - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__52736_inf": { + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__44032_inf": { "arches": [ "sm_100a", ], @@ -2136,7 +2136,7 @@ "library": "cake_rmsnorm_det_train_fcbfa14d591a9babfd26", "role": "bwd_reduce", "select_rule": { - "above": 52736, + "above": 44032, "at_most": None, "kind": "rows_between", }, @@ -9490,7 +9490,7 @@ "kind": "chunk_partials", }, }, - "cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_114688": { + "cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_99840": { "arches": [ "sm_103a", ], @@ -9513,7 +9513,7 @@ "role": "bwd_main", "select_rule": { "above": 40960, - "at_most": 114688, + "at_most": 99840, "kind": "rows_between", }, "workspace_rule": { @@ -9521,8 +9521,8 @@ "kind": "chunk_partials", }, }, - # The FP32-weight twin of cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_114688 (the effective weight read at FP32; every other operand and the host plan the owner's). - "cake_rmsnorm_det_train_420788e663a2c3a810f6__40960_114688": { + # The FP32-weight twin of cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_99840 (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_420788e663a2c3a810f6__40960_99840": { "arches": [ "sm_103a", ], @@ -9545,7 +9545,7 @@ "role": "bwd_main_w32", "select_rule": { "above": 40960, - "at_most": 114688, + "at_most": 99840, "kind": "rows_between", }, "workspace_rule": { @@ -9679,7 +9679,7 @@ "kind": "chunk_partials", }, }, - "cake_rmsnorm_det_train_0291949ba535c646c3a8__114688_1572864": { + "cake_rmsnorm_det_train_0291949ba535c646c3a8__99840_1572864": { "arches": [ "sm_103a", ], @@ -9701,7 +9701,7 @@ "library": "cake_rmsnorm_det_train_0291949ba535c646c3a8", "role": "bwd_main", "select_rule": { - "above": 114688, + "above": 99840, "at_most": 1572864, "kind": "rows_between", }, @@ -9710,8 +9710,8 @@ "kind": "chunk_partials", }, }, - # The FP32-weight twin of cake_rmsnorm_det_train_0291949ba535c646c3a8__114688_1572864 (the effective weight read at FP32; every other operand and the host plan the owner's). - "cake_rmsnorm_det_train_89505762fa1cbf33cb8f__114688_1572864": { + # The FP32-weight twin of cake_rmsnorm_det_train_0291949ba535c646c3a8__99840_1572864 (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_89505762fa1cbf33cb8f__99840_1572864": { "arches": [ "sm_103a", ], @@ -9733,7 +9733,7 @@ "library": "cake_rmsnorm_det_train_89505762fa1cbf33cb8f", "role": "bwd_main_w32", "select_rule": { - "above": 114688, + "above": 99840, "at_most": 1572864, "kind": "rows_between", }, @@ -9993,7 +9993,7 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_114688": { + "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840": { "arches": [ "sm_103a", ], @@ -10011,7 +10011,7 @@ "role": "bwd_reduce", "select_rule": { "above": 40960, - "at_most": 114688, + "at_most": 99840, "kind": "rows_between", }, "workspace_rule": None, @@ -10085,7 +10085,7 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_det_train_0ac243f8385c32960e14__114688_1572864": { + "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864": { "arches": [ "sm_103a", ], @@ -10102,7 +10102,7 @@ "library": "cake_rmsnorm_det_train_0ac243f8385c32960e14", "role": "bwd_reduce", "select_rule": { - "above": 114688, + "above": 99840, "at_most": 1572864, "kind": "rows_between", }, @@ -11610,8 +11610,8 @@ "cake_rmsnorm_det_train_5d75ea5f3c25579af496", "cake_rmsnorm_det_train_5d75ea5f3c25579af496__256_1024", "cake_rmsnorm_det_train_5d75ea5f3c25579af496__1024_4096", - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_52736", - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__52736_inf", + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_44032", + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__44032_inf", ], "bwd_w32_h4096": [ "cake_rmsnorm_det_train_20c64c610d22c6db4c05", @@ -11622,8 +11622,8 @@ "cake_rmsnorm_det_train_5d75ea5f3c25579af496", "cake_rmsnorm_det_train_5d75ea5f3c25579af496__256_1024", "cake_rmsnorm_det_train_5d75ea5f3c25579af496__1024_4096", - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_52736", - "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__52736_inf", + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__4096_44032", + "cake_rmsnorm_det_train_fcbfa14d591a9babfd26__44032_inf", ], "bwd_h512": [ "cake_rmsnorm_det_train_09601d42628de24f6f87", @@ -11850,15 +11850,15 @@ "cake_rmsnorm_det_train_af005462cd6a301f31c4", "cake_rmsnorm_det_train_a385ae1681a2c078344e__5120_12288", "cake_rmsnorm_det_train_af005462cd6a301f31c4__12288_40960", - "cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_114688", - "cake_rmsnorm_det_train_0291949ba535c646c3a8__114688_1572864", + "cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_99840", + "cake_rmsnorm_det_train_0291949ba535c646c3a8__99840_1572864", "cake_rmsnorm_det_train_c18297d75c0d424978fd", "cake_rmsnorm_det_train_72138dce40d43da44aaf", "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_114688", - "cake_rmsnorm_det_train_0ac243f8385c32960e14__114688_1572864", + "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840", + "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_w32_h256": [ @@ -11866,15 +11866,15 @@ "cake_rmsnorm_det_train_b7d9577ff884de98ea06", "cake_rmsnorm_det_train_420788e663a2c3a810f6__5120_12288", "cake_rmsnorm_det_train_b7d9577ff884de98ea06__12288_40960", - "cake_rmsnorm_det_train_420788e663a2c3a810f6__40960_114688", - "cake_rmsnorm_det_train_89505762fa1cbf33cb8f__114688_1572864", + "cake_rmsnorm_det_train_420788e663a2c3a810f6__40960_99840", + "cake_rmsnorm_det_train_89505762fa1cbf33cb8f__99840_1572864", "cake_rmsnorm_det_train_f788ef57163552825d27", "cake_rmsnorm_det_train_72138dce40d43da44aaf", "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_114688", - "cake_rmsnorm_det_train_0ac243f8385c32960e14__114688_1572864", + "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840", + "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_h4096": [ diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py index 5df80702ae5..baf4af89d0a 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py @@ -81,15 +81,15 @@ (0, 256, "cake_rmsnorm_det_train_3e644ea19b7d14e61ab9"), (256, 1024, "cake_rmsnorm_det_train_59270fb249dcba1ff679"), (1024, 4096, "cake_rmsnorm_det_train_9a830956447d4e304914"), - (4096, 52736, "cake_rmsnorm_det_train_35a0b0ca6b2060313255"), - (52736, None, "cake_rmsnorm_det_train_7122dbcd697b3fbe1ee0"), + (4096, 44032, "cake_rmsnorm_det_train_35a0b0ca6b2060313255"), + (44032, None, "cake_rmsnorm_det_train_7122dbcd697b3fbe1ee0"), ), "bwd_w32_h4096": ( (0, 256, "cake_rmsnorm_det_train_20c64c610d22c6db4c05"), (256, 1024, "cake_rmsnorm_det_train_c9f55c0ae4b198c859d1"), (1024, 4096, "cake_rmsnorm_det_train_0c84999caf2b86c2e09e"), - (4096, 52736, "cake_rmsnorm_det_train_626aed743358c27513f1"), - (52736, None, "cake_rmsnorm_det_train_5ef5f1f13ed8f2b994c4"), + (4096, 44032, "cake_rmsnorm_det_train_626aed743358c27513f1"), + (44032, None, "cake_rmsnorm_det_train_5ef5f1f13ed8f2b994c4"), ), "bwd_h512": ( (0, 1024, "cake_rmsnorm_det_train_09601d42628de24f6f87"), @@ -268,8 +268,8 @@ (2048, 5120, "cake_rmsnorm_det_train_af005462cd6a301f31c4"), (5120, 12288, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), (12288, 40960, "cake_rmsnorm_det_train_af005462cd6a301f31c4"), - (40960, 114688, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), - (114688, 1572864, "cake_rmsnorm_det_train_0291949ba535c646c3a8"), + (40960, 99840, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), + (99840, 1572864, "cake_rmsnorm_det_train_0291949ba535c646c3a8"), (1572864, None, "cake_rmsnorm_det_train_c18297d75c0d424978fd"), ), "bwd_w32_h256": ( @@ -277,8 +277,8 @@ (2048, 5120, "cake_rmsnorm_det_train_b7d9577ff884de98ea06"), (5120, 12288, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), (12288, 40960, "cake_rmsnorm_det_train_b7d9577ff884de98ea06"), - (40960, 114688, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), - (114688, 1572864, "cake_rmsnorm_det_train_89505762fa1cbf33cb8f"), + (40960, 99840, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), + (99840, 1572864, "cake_rmsnorm_det_train_89505762fa1cbf33cb8f"), (1572864, None, "cake_rmsnorm_det_train_f788ef57163552825d27"), ), "bwd_h4096": ( @@ -627,29 +627,19 @@ # the residual route's edge 1024), "dx" (the backward kernel switch), "unmeasured". EDGES: dict[str, dict[str, tuple[tuple[int, str], ...]]] = { "sm_100a": { - "bwd_h1536": ( - (1024, "same"), - (4096, "same"), - (32768, "same"), - (36352, "unmeasured"), - ), + "bwd_h1536": ((1024, "same"), (4096, "same"), (32768, "same"), (36352, "same")), "bwd_w32_h1536": ( (1024, "same"), (4096, "same"), (32768, "same"), - (36352, "unmeasured"), - ), - "bwd_h2048": ( - (3072, "same"), - (4096, "ulp"), - (16384, "same"), - (58880, "unmeasured"), + (36352, "same"), ), + "bwd_h2048": ((3072, "same"), (4096, "ulp"), (16384, "same"), (58880, "same")), "bwd_w32_h2048": ( (3072, "same"), (4096, "ulp"), (16384, "same"), - (58880, "unmeasured"), + (58880, "same"), ), "bwd_h256": ( (2048, "same"), @@ -671,20 +661,20 @@ (256, "same"), (1024, "same"), (4096, "same"), - (52736, "unmeasured"), + (44032, "unmeasured"), ), "bwd_w32_h4096": ( (256, "same"), (1024, "same"), (4096, "same"), - (52736, "unmeasured"), + (44032, "unmeasured"), ), "bwd_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_w32_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_h5120": ((256, "unmeasured"), (6144, "unmeasured")), "bwd_w32_h5120": ((256, "unmeasured"), (6144, "unmeasured")), - "bwd_h6144": ((256, "same"), (4096, "same"), (42496, "unmeasured")), - "bwd_w32_h6144": ((256, "same"), (4096, "same"), (42496, "unmeasured")), + "bwd_h6144": ((256, "same"), (4096, "same"), (42496, "same")), + "bwd_w32_h6144": ((256, "same"), (4096, "same"), (42496, "same")), "bwd_residual_h6144": (), "bwd_w32_residual_h6144": (), "fwd_h1536": ((1536, "ulp"), (3584, "same"), (6144, "same")), @@ -724,35 +714,35 @@ (4096, "same"), (14928, "unmeasured"), (15566, "unmeasured"), - (40448, "unmeasured"), + (40448, "same"), ), "bwd_w32_h1536": ( (1024, "same"), (4096, "same"), (14928, "unmeasured"), (15566, "unmeasured"), - (40448, "unmeasured"), + (40448, "same"), ), "bwd_h2048": ( (3072, "same"), (4096, "ulp"), (16384, "same"), (32768, "same"), - (55296, "unmeasured"), + (55296, "same"), ), "bwd_w32_h2048": ( (3072, "same"), (4096, "ulp"), (16384, "same"), (32768, "same"), - (55296, "unmeasured"), + (55296, "same"), ), "bwd_h256": ( (2048, "same"), (5120, "same"), (12288, "same"), (40960, "same"), - (114688, "unmeasured"), + (99840, "unmeasured"), (1572864, "unmeasured"), ), "bwd_w32_h256": ( @@ -760,27 +750,22 @@ (5120, "same"), (12288, "same"), (40960, "same"), - (114688, "unmeasured"), + (99840, "unmeasured"), (1572864, "unmeasured"), ), - "bwd_h4096": ( - (256, "same"), - (1024, "same"), - (4096, "same"), - (48640, "unmeasured"), - ), + "bwd_h4096": ((256, "same"), (1024, "same"), (4096, "same"), (48640, "same")), "bwd_w32_h4096": ( (256, "same"), (1024, "same"), (4096, "same"), - (48640, "unmeasured"), + (48640, "same"), ), "bwd_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_w32_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_h5120": ((5120, "unmeasured"),), "bwd_w32_h5120": ((5120, "unmeasured"),), - "bwd_h6144": ((4096, "same"), (48640, "unmeasured")), - "bwd_w32_h6144": ((4096, "same"), (48640, "unmeasured")), + "bwd_h6144": ((4096, "same"), (48640, "same")), + "bwd_w32_h6144": ((4096, "same"), (48640, "same")), "bwd_residual_h6144": (), "bwd_w32_residual_h6144": (), "fwd_h1536": ((1536, "ulp"), (3584, "same"), (6144, "same")), @@ -819,19 +804,14 @@ "fwd_w32_residual_h6144": ((256, "same"), (1024, "ulp")), }, "sm_107a": { - "bwd_h1536": ((1024, "same"), (16384, "same"), (45056, "unmeasured")), - "bwd_w32_h1536": ((1024, "same"), (16384, "same"), (45056, "unmeasured")), - "bwd_h2048": ( - (256, "same"), - (1024, "ulp"), - (16384, "same"), - (50688, "unmeasured"), - ), + "bwd_h1536": ((1024, "same"), (16384, "same"), (45056, "same")), + "bwd_w32_h1536": ((1024, "same"), (16384, "same"), (45056, "same")), + "bwd_h2048": ((256, "same"), (1024, "ulp"), (16384, "same"), (50688, "same")), "bwd_w32_h2048": ( (256, "same"), (1024, "ulp"), (16384, "same"), - (50688, "unmeasured"), + (50688, "same"), ), "bwd_h256": ( (65536, "unmeasured"), @@ -843,24 +823,19 @@ (70656, "unmeasured"), (1572864, "unmeasured"), ), - "bwd_h4096": ( - (256, "same"), - (1024, "same"), - (4096, "same"), - (24064, "unmeasured"), - ), + "bwd_h4096": ((256, "same"), (1024, "same"), (4096, "same"), (24064, "same")), "bwd_w32_h4096": ( (256, "same"), (1024, "same"), (4096, "same"), - (24064, "unmeasured"), + (24064, "same"), ), "bwd_h512": ((1024, "same"), (4096, "same"), (16384, "unmeasured")), "bwd_w32_h512": ((1024, "same"), (4096, "same"), (16384, "unmeasured")), - "bwd_h5120": ((4096, "unmeasured"), (19968, "unmeasured")), - "bwd_w32_h5120": ((4096, "unmeasured"), (19968, "unmeasured")), - "bwd_h6144": ((19456, "unmeasured"),), - "bwd_w32_h6144": ((19456, "unmeasured"),), + "bwd_h5120": ((4096, "unmeasured"), (19968, "same")), + "bwd_w32_h5120": ((4096, "unmeasured"), (19968, "same")), + "bwd_h6144": ((19456, "same"),), + "bwd_w32_h6144": ((19456, "same"),), "bwd_residual_h6144": (), "bwd_w32_residual_h6144": (), "fwd_h1536": ((2048, "ulp"), (12288, "same")), @@ -926,8 +901,8 @@ "bwd_w32_h2048": 58880, "bwd_h256": 80896, "bwd_w32_h256": 80896, - "bwd_h4096": 52736, - "bwd_w32_h4096": 52736, + "bwd_h4096": 44032, + "bwd_w32_h4096": 44032, "bwd_h5120": 6144, "bwd_w32_h5120": 6144, "bwd_h6144": 42496, @@ -938,8 +913,8 @@ "bwd_w32_h1536": 40448, "bwd_h2048": 55296, "bwd_w32_h2048": 55296, - "bwd_h256": 114688, - "bwd_w32_h256": 114688, + "bwd_h256": 99840, + "bwd_w32_h256": 99840, "bwd_h4096": 48640, "bwd_w32_h4096": 48640, "bwd_h5120": 5120, @@ -1170,51 +1145,51 @@ its own windows. * ``sm_100a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 4096] c1a9efdf = - (4096, 32768] 21c8bb55 = (32768, 36352] b457c4d1 ? (36352, inf) b5bdabba + (4096, 32768] 21c8bb55 = (32768, 36352] b457c4d1 = (36352, inf) b5bdabba * ``sm_100a`` ``bwd_h2048``: (0, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ - (4096, 16384] c73074cf = (16384, 58880] ceeef043 ? (58880, inf) 0cb817d1 + (4096, 16384] c73074cf = (16384, 58880] ceeef043 = (58880, inf) 0cb817d1 * ``sm_100a`` ``bwd_h256``: (0, 2048] a385ae16 = (2048, 5120] af005462 = (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 80896] a385ae16 ? (80896, 1572864] 0291949b ? (1572864, inf) c18297d7 * ``sm_100a`` ``bwd_h4096``: (0, 256] 3e644ea1 = (256, 1024] 59270fb2 = - (1024, 4096] 9a830956 = (4096, 52736] 35a0b0ca ? (52736, inf) 7122dbcd + (1024, 4096] 9a830956 = (4096, 44032] 35a0b0ca ? (44032, inf) 7122dbcd * ``sm_100a`` ``bwd_h512``: (0, 1024] 09601d42 ~ (1024, 4096] 123db0b2 = (4096, 32768] ad1a5a3c = (32768, inf) 8877322a * ``sm_100a`` ``bwd_h5120``: (0, 256] db01cec1 ? (256, 6144] b482c84d ? (6144, inf) c145ef19 * ``sm_100a`` ``bwd_h6144``: (0, 256] b93bf060 = (256, 4096] 5c804aed = - (4096, 42496] b93bf060 ? (42496, inf) c440e57b + (4096, 42496] b93bf060 = (42496, inf) c440e57b * ``sm_100a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 * ``sm_103a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 4096] c1a9efdf = - (4096, 14928] 1df7db17 ? (14928, 15566] 4df08ce5 ? (15566, 40448] f522f5b8 ? + (4096, 14928] 1df7db17 ? (14928, 15566] 4df08ce5 ? (15566, 40448] f522f5b8 = (40448, inf) cbeae168 * ``sm_103a`` ``bwd_h2048``: (0, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ - (4096, 16384] caf6978e = (16384, 32768] caf6978e = (32768, 55296] 302f64db ? + (4096, 16384] caf6978e = (16384, 32768] caf6978e = (32768, 55296] 302f64db = (55296, inf) 018f0cd7 * ``sm_103a`` ``bwd_h256``: (0, 2048] a385ae16 = (2048, 5120] af005462 = - (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 114688] a385ae16 ? - (114688, 1572864] 0291949b ? (1572864, inf) c18297d7 + (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 99840] a385ae16 ? + (99840, 1572864] 0291949b ? (1572864, inf) c18297d7 * ``sm_103a`` ``bwd_h4096``: (0, 256] 50044ee3 = (256, 1024] ad2bb26e = - (1024, 4096] 24643edc = (4096, 48640] 35a0b0ca ? (48640, inf) 850eec3a + (1024, 4096] 24643edc = (4096, 48640] 35a0b0ca = (48640, inf) 850eec3a * ``sm_103a`` ``bwd_h512``: (0, 1024] 09601d42 ~ (1024, 4096] 8e7f4d81 = (4096, 32768] 274b0055 = (32768, inf) cfaae3ce * ``sm_103a`` ``bwd_h5120``: (0, 5120] 2fc4932f ? (5120, inf) 7ffb443a -* ``sm_103a`` ``bwd_h6144``: (0, 4096] 5c804aed = (4096, 48640] 322574af ? +* ``sm_103a`` ``bwd_h6144``: (0, 4096] 5c804aed = (4096, 48640] 322574af = (48640, inf) 289dc9ab * ``sm_103a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 * ``sm_107a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 16384] 6a428da2 = - (16384, 45056] 56d26996 ? (45056, inf) bb711cc7 + (16384, 45056] 56d26996 = (45056, inf) bb711cc7 * ``sm_107a`` ``bwd_h2048``: (0, 256] 84efa2af = (256, 1024] 845cdcd8 ~ - (1024, 16384] 81c3d8ab = (16384, 50688] 845cdcd8 ? (50688, inf) 694decfb + (1024, 16384] 81c3d8ab = (16384, 50688] 845cdcd8 = (50688, inf) 694decfb * ``sm_107a`` ``bwd_h256``: (0, 65536] af005462 ? (65536, 70656] a385ae16 ? (70656, 1572864] 0291949b ? (1572864, inf) c18297d7 * ``sm_107a`` ``bwd_h4096``: (0, 256] 0da5dd4b = (256, 1024] a647fccc = - (1024, 4096] d41a514f = (4096, 24064] ee4a9053 ? (24064, inf) 99a01f2a + (1024, 4096] d41a514f = (4096, 24064] ee4a9053 = (24064, inf) 99a01f2a * ``sm_107a`` ``bwd_h512``: (0, 1024] 3feed502 = (1024, 4096] 3feed502 = (4096, 16384] dfaf378a ? (16384, inf) eb141cc9 -* ``sm_107a`` ``bwd_h5120``: (0, 4096] 838c077d ? (4096, 19968] 6816fb80 ? +* ``sm_107a`` ``bwd_h5120``: (0, 4096] 838c077d ? (4096, 19968] 6816fb80 = (19968, inf) fbe20dbb -* ``sm_107a`` ``bwd_h6144``: (0, 19456] 5d292172 ? (19456, inf) 19e18371 +* ``sm_107a`` ``bwd_h6144``: (0, 19456] 5d292172 = (19456, inf) 19e18371 * ``sm_107a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 """ diff --git a/tests/norm/test_cake_rmsnorm_det_train_pairs.py b/tests/norm/test_cake_rmsnorm_det_train_pairs.py index 4e58ca6f621..056e63d184c 100644 --- a/tests/norm/test_cake_rmsnorm_det_train_pairs.py +++ b/tests/norm/test_cake_rmsnorm_det_train_pairs.py @@ -46,9 +46,9 @@ ROWS = ( 1, 2, 7, 63, 64, 65, 127, 128, 129, 255, 256, 257, 1023, 1024, 1025, 2048, 2049, 3072, 3073, 4096, 4097, 5120, 5121, 6144, 6145, 12288, 12289, 14928, 14929, 15566, 15567, 16231, 16263, 16384, 16385, 19456, 19457, 19968, - 19969, 20480, 20481, 24064, 24065, 32768, 32769, 36352, 36353, 40448, 40449, 40960, 40961, 42496, 42497, 45056, - 45057, 48640, 48641, 50688, 50689, 52736, 52737, 55296, 55297, 58880, 58881, 65536, 65537, 70656, 70657, 80896, - 80897, 81920, 81921, 114688, 114689, 139264, 139265, 1179648, 1179649, 1572864, 1572865, + 19969, 20480, 20481, 24064, 24065, 32768, 32769, 36352, 36353, 40448, 40449, 40960, 40961, 42496, 42497, 44032, + 44033, 45056, 45057, 48640, 48641, 50688, 50689, 52736, 52737, 55296, 55297, 58880, 58881, 65536, 65537, 70656, + 70657, 80896, 80897, 81920, 81921, 99840, 99841, 114688, 114689, 139264, 139265, 1179648, 1179649, 1572864, 1572865, ) # fmt: skip # The row count standing in for an unbounded window's upper edge. OPEN_WINDOW_ROWS = 65536 @@ -316,7 +316,7 @@ def test_rows_cover_every_backward_window_edge() -> None: # The H = 256 backward route: per sm_100a / sm_103a window the faster of the four-CTA pipeline (``pd4``) and the # five-CTA ring at issue distance three without the L2 hint (``a3``), both at ``tail_div=4``, in two sessions per # architecture (BF16 and FP32 weight alike); sm_107a runs the pipeline to 65536 rows and the ring above. Above the -# ring's dw accuracy edge of each architecture (80896 rows on sm_100a, 114688 on sm_103a, 70656 on sm_107a) the ring +# ring's dw accuracy edge of each architecture (80896 rows on sm_100a, 99840 on sm_103a, 70656 on sm_107a) the ring # with compensated dw sums (``dc``) runs instead, and above 1572864 rows its form with hi / lo partial rows # (``dchl``). H256_BWD_PROGRAMS = { @@ -338,7 +338,7 @@ def test_rows_cover_every_backward_window_edge() -> None: ) for arch, edge, hl in ( ("sm_100a", 80896, 1572864), - ("sm_103a", 114688, 1572864), + ("sm_103a", 99840, 1572864), ) }, "sm_107a": ( @@ -562,14 +562,14 @@ def test_fp32_weight_routes_mirror_their_bf16_routes(arch: str) -> None: "h1536": 36352, "h2048": 58880, "h256": 80896, - "h4096": 52736, + "h4096": 44032, "h5120": 6144, "h6144": 42496, }, "sm_103a": { "h1536": 40448, "h2048": 55296, - "h256": 114688, + "h256": 99840, "h4096": 48640, "h5120": 5120, "h6144": 48640,