diff --git a/flashinfer/cake_rmsnorm_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_train_backend/cake_jit.py index 001b210cf92..ecc0cb52bc9 100644 --- a/flashinfer/cake_rmsnorm_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_train_backend/cake_jit.py @@ -4540,6 +4540,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 64, "ctas_per_sm": 2, "grid_factor": 2, "kind": "claimed_chunks", @@ -4573,6 +4574,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 64, "ctas_per_sm": 2, "grid_factor": 2, "kind": "claimed_chunks", @@ -4606,6 +4608,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 128, "ctas_per_sm": 2, "grid_factor": 2, "kind": "claimed_chunks", @@ -4639,6 +4642,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 128, "ctas_per_sm": 2, "grid_factor": 2, "kind": "claimed_chunks", @@ -4672,6 +4676,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 80, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -4705,6 +4710,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -4738,6 +4744,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -4771,6 +4778,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -4804,6 +4812,7 @@ "sm_100a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -5392,78 +5401,12 @@ "sum_bytes": 98304, }, }, - # sm_103a H5120 bf16 weight: (320, 896] backward time 1.033-1.087x the det pair's at T 512 and 640 in all four timing sessions (GB300) - "cake_rmsnorm_train_080e0a6c3b18e9bad790__320_896": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas_per_sm": 1, - "grid_factor": 1, - "kind": "claimed_chunks", - "launch_bound_ctas_per_sm": 1, - "max_halving": 5, - "rows_per_issue": 1, - "tail_div": 12, - }, - "hidden": 5120, - "kernel": "cake_rmsnorm_train_080e0a6c3b18e9bad790", - "launch": { - "cooperative": False, - "use_pdl": False, - }, - "library": "cake_rmsnorm_train_080e0a6c3b18e9bad790", - "role": "bwd_main", - "select_rule": { - "above": 320, - "at_most": 896, - "kind": "rows_between", - }, - "workspace_rule": { - "counters": 1, - "kind": "claim_sum_f64", - "sum_bytes": 0, - }, - }, - # sm_103a H6144 bf16 weight: (448, 1152] backward time 1.029-1.043x the det pair's at T 1024 in all four timing sessions (GB300) - "cake_rmsnorm_train_44508df7dab2ab688f85__448_1152": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas_per_sm": 1, - "grid_factor": 1, - "kind": "claimed_chunks", - "launch_bound_ctas_per_sm": 1, - "max_halving": 5, - "rows_per_issue": 1, - "tail_div": 12, - }, - "hidden": 6144, - "kernel": "cake_rmsnorm_train_44508df7dab2ab688f85", - "launch": { - "cooperative": False, - "use_pdl": False, - }, - "library": "cake_rmsnorm_train_44508df7dab2ab688f85", - "role": "bwd_main", - "select_rule": { - "above": 448, - "at_most": 1152, - "kind": "rows_between", - }, - "workspace_rule": { - "counters": 1, - "kind": "claim_sum_f64", - "sum_bytes": 0, - }, - }, - # sm_103a H6144 f32 weight: (320, 1152] backward time 0.997-1.003x the det pair's at T 1024 over four timing sessions (GB300) - "cake_rmsnorm_train_b3e6531e46667901ab91__320_1152": { + "cake_rmsnorm_train_b3e6531e46667901ab91__320_448": { "arches": [ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -5482,7 +5425,7 @@ "role": "bwd_main_w32", "select_rule": { "above": 320, - "at_most": 1152, + "at_most": 448, "kind": "rows_between", }, "workspace_rule": { @@ -5497,6 +5440,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -5530,6 +5474,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -5695,6 +5640,7 @@ "sm_107a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -5728,6 +5674,7 @@ "sm_107a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -6982,6 +6929,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 80, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -7180,6 +7128,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -7213,6 +7162,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -7543,6 +7493,7 @@ "sm_107a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -7609,6 +7560,7 @@ "sm_103a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -7675,6 +7627,7 @@ "sm_107a", ], "grid_rule": { + "col_tasks": 96, "ctas_per_sm": 1, "grid_factor": 1, "kind": "claimed_chunks", @@ -11387,7 +11340,7 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_train_2fc4932fbde0bc848b32__22_320": { + "cake_rmsnorm_train_2fc4932fbde0bc848b32__22_5120": { "arches": [ "sm_103a", ], @@ -11410,37 +11363,6 @@ "role": "bwd_main", "select_rule": { "above": 22, - "at_most": 320, - "kind": "rows_between", - }, - "workspace_rule": { - "counters": 128, - "kind": "chunk_partials", - }, - }, - "cake_rmsnorm_train_2fc4932fbde0bc848b32__896_5120": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas_per_sm": 1, - "grid_factor": 2, - "kind": "claimed_chunks", - "launch_bound_ctas_per_sm": 1, - "max_halving": 2, - "rows_per_issue": 1, - "tail_div": 4, - }, - "hidden": 5120, - "kernel": "cake_rmsnorm_train_2fc4932fbde0bc848b32", - "launch": { - "cooperative": False, - "use_pdl": False, - }, - "library": "cake_rmsnorm_train_2fc4932fbde0bc848b32", - "role": "bwd_main", - "select_rule": { - "above": 896, "at_most": 5120, "kind": "rows_between", }, @@ -11480,52 +11402,6 @@ "kind": "chunk_partials", }, }, - "cake_rmsnorm_train_c5831b29d17867842e5f__22_320": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas": 160, - "kind": "fixed", - }, - "hidden": 5120, - "kernel": "cake_rmsnorm_train_c5831b29d17867842e5f", - "launch": { - "cooperative": False, - "use_pdl": True, - }, - "library": "cake_rmsnorm_train_c5831b29d17867842e5f", - "role": "bwd_reduce", - "select_rule": { - "above": 22, - "at_most": 320, - "kind": "rows_between", - }, - "workspace_rule": None, - }, - "cake_rmsnorm_train_c5831b29d17867842e5f__896_5120": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas": 160, - "kind": "fixed", - }, - "hidden": 5120, - "kernel": "cake_rmsnorm_train_c5831b29d17867842e5f", - "launch": { - "cooperative": False, - "use_pdl": True, - }, - "library": "cake_rmsnorm_train_c5831b29d17867842e5f", - "role": "bwd_reduce", - "select_rule": { - "above": 896, - "at_most": 5120, - "kind": "rows_between", - }, - "workspace_rule": None, - }, "cake_rmsnorm_train_c5831b29d17867842e5f__5120_61952": { "arches": [ "sm_103a", @@ -11634,7 +11510,7 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_train_5c804aed7ad14b99e1d0__3_448": { + "cake_rmsnorm_train_5c804aed7ad14b99e1d0__3_4096": { "arches": [ "sm_103a", ], @@ -11657,37 +11533,6 @@ "role": "bwd_small", "select_rule": { "above": 3, - "at_most": 448, - "kind": "rows_between", - }, - "workspace_rule": { - "counters": 192, - "kind": "chunk_partials", - }, - }, - "cake_rmsnorm_train_5c804aed7ad14b99e1d0__1152_4096_103": { - "arches": [ - "sm_103a", - ], - "grid_rule": { - "ctas_per_sm": 1, - "grid_factor": 2, - "kind": "claimed_chunks", - "launch_bound_ctas_per_sm": 1, - "max_halving": 5, - "rows_per_issue": 1, - "tail_div": 12, - }, - "hidden": 6144, - "kernel": "cake_rmsnorm_train_5c804aed7ad14b99e1d0", - "launch": { - "cooperative": False, - "use_pdl": False, - }, - "library": "cake_rmsnorm_train_5c804aed7ad14b99e1d0", - "role": "bwd_small", - "select_rule": { - "above": 1152, "at_most": 4096, "kind": "rows_between", }, @@ -11727,7 +11572,7 @@ "kind": "chunk_partials", }, }, - "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__3_448": { + "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__3_4096": { "arches": [ "sm_103a", ], @@ -11745,12 +11590,12 @@ "role": "bwd_reduce", "select_rule": { "above": 3, - "at_most": 448, + "at_most": 4096, "kind": "rows_between", }, "workspace_rule": None, }, - "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__1152_4096_0_4096": { + "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__448_4096": { "arches": [ "sm_103a", ], @@ -11767,7 +11612,7 @@ "library": "cake_rmsnorm_train_5f5cdc5151e3c0620bcd", "role": "bwd_reduce", "select_rule": { - "above": 1152, + "above": 448, "at_most": 4096, "kind": "rows_between", }, @@ -11827,7 +11672,7 @@ "kind": "chunk_partials", }, }, - "cake_rmsnorm_train_41b80b6af97974feb181__1152_4096": { + "cake_rmsnorm_train_41b80b6af97974feb181__448_4096": { "arches": [ "sm_103a", ], @@ -11849,7 +11694,7 @@ "library": "cake_rmsnorm_train_41b80b6af97974feb181", "role": "bwd_small_w32", "select_rule": { - "above": 1152, + "above": 448, "at_most": 4096, "kind": "rows_between", }, @@ -14600,13 +14445,10 @@ ], "bwd_h5120": [ "cake_rmsnorm_train_080e0a6c3b18e9bad790__0_22", - "cake_rmsnorm_train_2fc4932fbde0bc848b32__22_320", - "cake_rmsnorm_train_080e0a6c3b18e9bad790__320_896", - "cake_rmsnorm_train_2fc4932fbde0bc848b32__896_5120", + "cake_rmsnorm_train_2fc4932fbde0bc848b32__22_5120", "cake_rmsnorm_train_7ffb443a63a92ab6b0e0__5120_61952", "cake_rmsnorm_train_9e4c0825df92af02074f", - "cake_rmsnorm_train_c5831b29d17867842e5f__22_320", - "cake_rmsnorm_train_c5831b29d17867842e5f__896_5120", + "cake_rmsnorm_train_c5831b29d17867842e5f__22_5120", "cake_rmsnorm_train_c5831b29d17867842e5f__5120_61952", ], "bwd_w32_h5120": [ @@ -14619,13 +14461,10 @@ ], "bwd_h6144": [ "cake_rmsnorm_train_44508df7dab2ab688f85__0_3", - "cake_rmsnorm_train_5c804aed7ad14b99e1d0__3_448", - "cake_rmsnorm_train_44508df7dab2ab688f85__448_1152", - "cake_rmsnorm_train_5c804aed7ad14b99e1d0__1152_4096_103", + "cake_rmsnorm_train_5c804aed7ad14b99e1d0__3_4096", "cake_rmsnorm_train_322574af704b51aef0bb__4096_49152", "cake_rmsnorm_train_582b2a20f9be648af8ec", - "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__3_448", - "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__1152_4096_0_4096", + "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__3_4096", "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__4096_49152", ], "bwd_residual_h6144": [ @@ -14641,12 +14480,12 @@ "bwd_w32_h6144": [ "cake_rmsnorm_train_b3e6531e46667901ab91__0_20", "cake_rmsnorm_train_41b80b6af97974feb181__20_320", - "cake_rmsnorm_train_b3e6531e46667901ab91__320_1152", - "cake_rmsnorm_train_41b80b6af97974feb181__1152_4096", + "cake_rmsnorm_train_b3e6531e46667901ab91__320_448", + "cake_rmsnorm_train_41b80b6af97974feb181__448_4096", "cake_rmsnorm_train_9a8c0396edd85d4d545e__4096_49152", "cake_rmsnorm_train_4f9a76e68671f213ab7c", "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__20_320", - "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__1152_4096_0_4096", + "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__448_4096", "cake_rmsnorm_train_5f5cdc5151e3c0620bcd__4096_49152", ], "bwd_w32_residual_h6144": [ @@ -15387,7 +15226,10 @@ def launch_grid( ctas_per_sm * SM count`` persistent CTAs claim the chunks of the table of :func:`chunk_rounds` (``rows_per_issue`` rows per pipeline step of the kernel, under the taper its record spells: :func:`chunk_taper`), so the - grid is ``min(slots, n_chunks)``; the launch is a plain + grid is ``min(slots, n_chunks)``; a column-task kernel's rule spells its + ``col_tasks`` (``H / col`` column tasks, tickets beside the chunks), so its + grid is ``min(slots, n_chunks + col_tasks)`` and a small ``T`` runs every + task at once. The launch is a plain one -- no co-residency is required -- and ``ctas_per_sm`` is the production grid factor capped at the launch bound the kernel declares. ``rows`` must be positive; the public API returns empty outputs for zero rows without @@ -15413,7 +15255,8 @@ def launch_grid( slots = int(rule["ctas_per_sm"]) * sm_count(device_index) grid_x = min( slots, - n_chunks(rows, slots, int(rule["rows_per_issue"]), *chunk_taper(rule)), + n_chunks(rows, slots, int(rule["rows_per_issue"]), *chunk_taper(rule)) + + int(rule.get("col_tasks", 0)), ) else: raise ValueError(f"unknown Cake RMSNorm training grid rule: {kind!r}") diff --git a/flashinfer/cake_rmsnorm_train_backend/cake_windows.py b/flashinfer/cake_rmsnorm_train_backend/cake_windows.py index 12dfec9a333..b8fdd2a4269 100644 --- a/flashinfer/cake_rmsnorm_train_backend/cake_windows.py +++ b/flashinfer/cake_rmsnorm_train_backend/cake_windows.py @@ -417,9 +417,7 @@ ), "bwd_h5120": ( (0, 22, "cake_rmsnorm_train_080e0a6c3b18e9bad790"), - (22, 320, "cake_rmsnorm_train_2fc4932fbde0bc848b32"), - (320, 896, "cake_rmsnorm_train_080e0a6c3b18e9bad790"), - (896, 5120, "cake_rmsnorm_train_2fc4932fbde0bc848b32"), + (22, 5120, "cake_rmsnorm_train_2fc4932fbde0bc848b32"), (5120, 61952, "cake_rmsnorm_train_7ffb443a63a92ab6b0e0"), (61952, None, "cake_rmsnorm_train_9e4c0825df92af02074f"), ), @@ -431,9 +429,7 @@ ), "bwd_h6144": ( (0, 3, "cake_rmsnorm_train_44508df7dab2ab688f85"), - (3, 448, "cake_rmsnorm_train_5c804aed7ad14b99e1d0"), - (448, 1152, "cake_rmsnorm_train_44508df7dab2ab688f85"), - (1152, 4096, "cake_rmsnorm_train_5c804aed7ad14b99e1d0"), + (3, 4096, "cake_rmsnorm_train_5c804aed7ad14b99e1d0"), (4096, 49152, "cake_rmsnorm_train_322574af704b51aef0bb"), (49152, None, "cake_rmsnorm_train_582b2a20f9be648af8ec"), ), @@ -448,8 +444,8 @@ "bwd_w32_h6144": ( (0, 20, "cake_rmsnorm_train_b3e6531e46667901ab91"), (20, 320, "cake_rmsnorm_train_41b80b6af97974feb181"), - (320, 1152, "cake_rmsnorm_train_b3e6531e46667901ab91"), - (1152, 4096, "cake_rmsnorm_train_41b80b6af97974feb181"), + (320, 448, "cake_rmsnorm_train_b3e6531e46667901ab91"), + (448, 4096, "cake_rmsnorm_train_41b80b6af97974feb181"), (4096, 49152, "cake_rmsnorm_train_9a8c0396edd85d4d545e"), (49152, None, "cake_rmsnorm_train_4f9a76e68671f213ab7c"), ), @@ -868,21 +864,9 @@ (4096, "same"), (9216, "same"), ), - "bwd_h5120": ( - (22, "same"), - (320, "same"), - (896, "same"), - (5120, "same"), - (61952, "same"), - ), + "bwd_h5120": ((22, "same"), (5120, "same"), (61952, "same")), "bwd_w32_h5120": ((22, "same"), (5120, "same"), (61952, "same")), - "bwd_h6144": ( - (3, "same"), - (448, "same"), - (1152, "same"), - (4096, "same"), - (49152, "same"), - ), + "bwd_h6144": ((3, "same"), (4096, "same"), (49152, "same")), "bwd_residual_h6144": ( (2304, "same"), (2816, "same"), @@ -893,7 +877,7 @@ "bwd_w32_h6144": ( (20, "same"), (320, "same"), - (1152, "same"), + (448, "same"), (4096, "same"), (49152, "same"), ), @@ -1168,17 +1152,13 @@ ), "bwd_h5120": ( (0, 22, None), - (22, 320, 132), - (320, 896, None), - (896, 5120, 132), + (22, 5120, 132), (5120, 61952, 132), (61952, None, 132), ), "bwd_h6144": ( (0, 3, 132), - (3, 448, 132), - (448, 1152, 132), - (1152, 4096, 132), + (3, 4096, 132), (4096, 49152, 132), (49152, None, 100), ), @@ -1242,8 +1222,8 @@ "bwd_w32_h6144": ( (0, 20, 132), (20, 320, 132), - (320, 1152, 132), - (1152, 4096, 132), + (320, 448, 132), + (448, 4096, 132), (4096, 49152, 132), (49152, None, 100), ), @@ -1654,8 +1634,8 @@ * ``sm_103a`` ``bwd_h256``: (0, 1080] -, (1080, 4608] -, (4608, 5120] -, (5120, 12288] -, (12288, 40960] -, (40960, 93184] -, (93184, 139264] -, (139264, 1572864] -, (1572864, inf) - * ``sm_103a`` ``bwd_h4096``: (0, 3] -, (3, 256] -, (256, 1024] -, (1024, 4096] >= 132, (4096, 50176] >= 132, (50176, inf) - * ``sm_103a`` ``bwd_h512``: (0, 26] -, (26, 1024] -, (1024, 4096] -, (4096, 32768] -, (32768, inf) - -* ``sm_103a`` ``bwd_h5120``: (0, 22] -, (22, 320] >= 132, (320, 896] -, (896, 5120] >= 132, (5120, 61952] >= 132, (61952, inf) >= 132 -* ``sm_103a`` ``bwd_h6144``: (0, 3] >= 132, (3, 448] >= 132, (448, 1152] >= 132, (1152, 4096] >= 132, (4096, 49152] >= 132, (49152, inf) >= 100 +* ``sm_103a`` ``bwd_h5120``: (0, 22] -, (22, 5120] >= 132, (5120, 61952] >= 132, (61952, inf) >= 132 +* ``sm_103a`` ``bwd_h6144``: (0, 3] >= 132, (3, 4096] >= 132, (4096, 49152] >= 132, (49152, inf) >= 100 * ``sm_103a`` ``bwd_residual_h6144``: (0, 2304] >= 100, (2304, 2816] >= 100, (2816, 4608] >= 100, (4608, 5632] >= 100, (5632, 32768] >= 100, (32768, inf) >= 100 * ``sm_103a`` ``bwd_w32_h1536``: (0, 320] -, (320, 1024] >= 132, (1024, 4096] >= 100, (4096, 14928] >= 132, (14928, 15566] >= 100, (15566, 37888] >= 100, (37888, inf) - * ``sm_103a`` ``bwd_w32_h2048``: (0, 42] >= 132, (42, 3072] >= 100, (3072, 4096] >= 100, (4096, 16384] >= 100, (16384, 32768] >= 100, (32768, 61952] >= 100, (61952, inf) >= 132 @@ -1663,7 +1643,7 @@ * ``sm_103a`` ``bwd_w32_h4096``: (0, 15] -, (15, 256] -, (256, 1024] -, (1024, 4096] >= 132, (4096, 9216] >= 132, (9216, inf) - * ``sm_103a`` ``bwd_w32_h512``: (0, 292] -, (292, 960] -, (960, 32768] -, (32768, inf) - * ``sm_103a`` ``bwd_w32_h5120``: (0, 22] >= 132, (22, 5120] >= 132, (5120, 61952] >= 132, (61952, inf) >= 132 -* ``sm_103a`` ``bwd_w32_h6144``: (0, 20] >= 132, (20, 320] >= 132, (320, 1152] >= 132, (1152, 4096] >= 132, (4096, 49152] >= 132, (49152, inf) >= 100 +* ``sm_103a`` ``bwd_w32_h6144``: (0, 20] >= 132, (20, 320] >= 132, (320, 448] >= 132, (448, 4096] >= 132, (4096, 49152] >= 132, (49152, inf) >= 100 * ``sm_103a`` ``bwd_w32_residual_h6144``: (0, 6] >= 100, (6, 320] >= 100, (320, 2304] >= 100, (2304, 3584] >= 100, (3584, 4608] >= 100, (4608, 5632] >= 100, (5632, 15872] >= 100, (15872, inf) >= 100 * ``sm_107a`` ``bwd_h1536``: (0, 14] -, (14, 1024] -, (1024, 16384] >= 100, (16384, 53248] >= 100, (53248, inf) - * ``sm_107a`` ``bwd_h2048``: (0, 6] -, (6, 256] -, (256, 1024] >= 100, (1024, 16384] >= 100, (16384, 50176] >= 100, (50176, inf) - @@ -1763,19 +1743,17 @@ * ``sm_103a`` ``bwd_w32_h4096``: (0, 15] 131a495c = (15, 256] 3fd77c8c (det pair) = (256, 1024] 16545056 (det pair) = (1024, 4096] 78c26a53 (det pair) = (4096, 9216] 626aed74 (det pair) = (9216, inf) 131a495c -* ``sm_103a`` ``bwd_h5120``: (0, 22] 080e0a6c = (22, 320] 2fc4932f (det pair) = - (320, 896] 080e0a6c = (896, 5120] 2fc4932f (det pair) = +* ``sm_103a`` ``bwd_h5120``: (0, 22] 080e0a6c = (22, 5120] 2fc4932f (det pair) = (5120, 61952] 7ffb443a (det pair) = (61952, inf) 9e4c0825 * ``sm_103a`` ``bwd_w32_h5120``: (0, 22] f165933c = (22, 5120] 79811a76 (det pair) = (5120, 61952] f9f7d4cd (det pair) = (61952, inf) f165933c -* ``sm_103a`` ``bwd_h6144``: (0, 3] 44508df7 = (3, 448] 5c804aed (det pair) = - (448, 1152] 44508df7 = (1152, 4096] 5c804aed (det pair) = +* ``sm_103a`` ``bwd_h6144``: (0, 3] 44508df7 = (3, 4096] 5c804aed (det pair) = (4096, 49152] 322574af (det pair) = (49152, inf) 582b2a20 * ``sm_103a`` ``bwd_residual_h6144``: (0, 2304] 864e84f3 = (2304, 2816] 0dcdc563 (det pair) = (2816, 4608] 115dcc71 = (4608, 5632] 0dcdc563 (det pair) = (5632, 32768] 115dcc71 = (32768, inf) 52831494 * ``sm_103a`` ``bwd_w32_h6144``: (0, 20] b3e6531e = (20, 320] 41b80b6a (det pair) = - (320, 1152] b3e6531e = (1152, 4096] 41b80b6a (det pair) = + (320, 448] b3e6531e = (448, 4096] 41b80b6a (det pair) = (4096, 49152] 9a8c0396 (det pair) = (49152, inf) 4f9a76e6 * ``sm_103a`` ``bwd_w32_residual_h6144``: (0, 6] 71d36617 = (6, 320] 56f988dd (det pair) = (320, 2304] 71d36617 = (2304, 3584] 56f988dd (det pair) = (3584, 4608] 3c390339 = diff --git a/tests/norm/test_cake_rmsnorm_train.py b/tests/norm/test_cake_rmsnorm_train.py index 70efe7b0438..71b6448c114 100644 --- a/tests/norm/test_cake_rmsnorm_train.py +++ b/tests/norm/test_cake_rmsnorm_train.py @@ -871,7 +871,7 @@ def test_repeated_launches_against_fp64_and_det( # (on sm_103a its twosum form) above it; a single kernel also runs from T 1 up # to T 22 with FP32 weight on sm_100a and with both weight dtypes on sm_103a, # up to T 18 (BF16 weight) and 13 (FP32 weight) on sm_107a, and with BF16 -# weight up to T 896 on sm_100a and from T 320 to 896 on sm_103a. The plain +# weight up to T 896 on sm_100a. The plain # edge is the deterministic package's own switch to its compensated pair. The # accuracy edge is T 65536, the top of the qualified range, on # sm_100a, T 61952 on sm_103a and T 49152 on sm_107a, for both weight dtypes: diff --git a/tests/norm/test_cake_rmsnorm_train_pairs.py b/tests/norm/test_cake_rmsnorm_train_pairs.py index c614cbf3487..447fa1d7947 100644 --- a/tests/norm/test_cake_rmsnorm_train_pairs.py +++ b/tests/norm/test_cake_rmsnorm_train_pairs.py @@ -279,8 +279,15 @@ def _check_family_layout(arch: str, hidden: int, residual: bool, sms: int) -> in layout["total_bytes"] >= layout["partial_offset"] + layout["partial_bytes"] ) + # a column-task single (no FP64 sums) runs its ``col_tasks`` tickets + # beside the chunks; no other record spells them + tasks = rule.get("col_tasks", 0) + column = single and record["workspace_rule"]["sum_bytes"] == 0 + assert column == ("col_tasks" in rule), (name, rule) + if column: + assert tasks > 0 and hidden % tasks == 0, (name, tasks) grid = loader.launch_grid(record, rows, device_index=0) - assert grid == (min(slots, len(table)), 1, 1) + assert grid == (min(slots, len(table) + tasks), 1, 1) return singles diff --git a/tests/norm/test_cake_rmsnorm_train_stall.py b/tests/norm/test_cake_rmsnorm_train_stall.py index 151942aeaa7..7e223eeab05 100644 --- a/tests/norm/test_cake_rmsnorm_train_stall.py +++ b/tests/norm/test_cake_rmsnorm_train_stall.py @@ -260,7 +260,7 @@ def call(): ) -@pytest.mark.parametrize("rows", (129, 4096, 16384)) +@pytest.mark.parametrize("rows", (1, 3, 129, 4096, 16384)) @pytest.mark.parametrize("hidden,residual", ROUTES, ids=lambda v: str(v)) def test_outputs_do_not_depend_on_the_launch_grid( hidden: int, residual: bool, rows: int, monkeypatch: pytest.MonkeyPatch