From 10c601753cf14e9c7e0bef03e579d88be6f1103f Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 16:55:31 +0000 Subject: [PATCH] perf(cake_rmsnorm_silu_gate_train): warm the next row block in L2 in the long forward on SM107 With deterministic=False, on sm_107a from 1740 rows per SM (the long forward's start there) the forward runs the long forward's program whose thread 0 also warms in L2, with one bulk prefetch per tensor where x and z hold their rows packed, the CTA's next row block: the extra-long tier (role fwd_xlong) that sm_103a runs from 3000 rows per SM, in place of the long forward on sm_107a. Every other route, grid, threshold and architecture is unchanged, and y stays bitwise equal to the deterministic forward (a prefetch writes no register and no memory). cake_jit: the extra-long forward record on sm_103a and sm_107a (from 3000 and 1740 rows per SM), its sm_107a route in place of the long forward's, and the long forward record without sm_107a. The tests cover the extra-long tier's start, routes, grid and stall behavior on both architectures; the docs name it. --- .../__init__.py | 5 +- .../build.json | 32 +++++------ .../cake_jit.py | 53 ++++++++++--------- .../norm/test_cake_rmsnorm_silu_gate_train.py | 53 ++++++++++++------- .../test_cake_rmsnorm_silu_gate_train_api.py | 5 +- ...test_cake_rmsnorm_silu_gate_train_stall.py | 11 ++-- 6 files changed, 93 insertions(+), 66 deletions(-) diff --git a/flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py b/flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py index 8d04aa5c058..260fad07ac8 100644 --- a/flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py +++ b/flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py @@ -120,8 +120,9 @@ rows each). The mid forward's thread 0 also warms in L2, with one bulk prefetch a tensor where x and z hold their rows packed, the CTA's next block and the first block of the CTA one wave later; from 3000 rows per SM on - sm_103a the long forward's thread 0 also warms the CTA's next block the same - way. The backward launches + sm_103a and from 1740 on sm_107a (its whole range there) the long + forward's thread 0 also warms the CTA's next block the same way. The + backward launches ``cake_jit.backward_ctas`` CTAs, ``G = min(row blocks, ctas_per_sm * SMs)``: 256-thread CTAs of 16-row blocks with ``ctas_per_sm = 4`` (one wave); on sm_100a and sm_103a from 1280 rows per diff --git a/flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json b/flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json index 5b8db197a6d..89c3488cde0 100644 --- a/flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json +++ b/flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json @@ -337,42 +337,42 @@ ] }, { - "arch": "sm_107a", - "closure_sha256": "ba3c9d3213c685563270e20f81d1ef0d5537b589166b005b79ab1d0d7dc81089", + "arch": "sm_103a", + "closure_sha256": "1806bf4a6b51f1fd913a007d6dd9b5e2e36522517fd30854adc6e47866d06ab1", "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_sm_107a", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_103a", "nvcc_flags": [ "--prec-div=true", "--prec-sqrt=true", "--ftz=false" ], "nvcc_flags_if": [], - "path": "aot/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_sm_107a.so", + "path": "aot/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_103a.so", "sha256": { - "csrc/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_binding.cu": "d42c4185d610b49935462de9065cffafd54e32cd024e2203e5562f0177fb7185", - "csrc/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_kernel.cu": "a2c7d04cb91bef61cfac35362daa358b15d385504cfa003f1ae471ed017af1a7" + "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_binding.cu": "322eef384dabf9e7a4c58b7a32e8dc3b113845458687313adde05a21027cd70a", + "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_kernel.cu": "35202de8c110a6f8570ef3522e676e96a1bcbb16868ec3b55dc6a9739fae3fdd" }, "sources": [ - "csrc/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_kernel.cu", - "csrc/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_binding.cu" + "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_kernel.cu", + "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_binding.cu" ] }, { - "arch": "sm_103a", - "closure_sha256": "1806bf4a6b51f1fd913a007d6dd9b5e2e36522517fd30854adc6e47866d06ab1", + "arch": "sm_107a", + "closure_sha256": "7f059892d2c1adec9d9a1ad24e1460dee34c859b6640b9383d1ad5066cbdb81a", "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_103a", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_107a", "nvcc_flags": [ "--prec-div=true", "--prec-sqrt=true", "--ftz=false" ], "nvcc_flags_if": [], - "path": "aot/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_103a.so", + "path": "aot/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_sm_107a.so", "sha256": { "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_binding.cu": "322eef384dabf9e7a4c58b7a32e8dc3b113845458687313adde05a21027cd70a", "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_kernel.cu": "35202de8c110a6f8570ef3522e676e96a1bcbb16868ec3b55dc6a9739fae3fdd" diff --git a/flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py index b621d6595b8..4e14fcd7b75 100644 --- a/flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py @@ -39,19 +39,21 @@ (``fwd_mid``), whose module runs in the role's place from ``min_rows_per_sm`` rows per SM of its grid rule (:func:`min_rows_per_sm`: one count, or one per architecture of the record; :func:`module_for_rows`; the long role first, from -more rows), and sm_103a the forward's extra-long role (``fwd_xlong``), tried -before its long role, from more rows still: on every architecture a forward -that walks several row blocks per CTA, loading each next block while it -computes the current one, ``y`` the same bits (the long forward 32-row blocks -on 256-thread CTAs, the mid forward 16-row blocks on 128-thread CTAs whose -thread 0 also warms later row blocks of packed rows in L2 with one bulk -prefetch a tensor; on sm_103a the extra-long forward is the long forward whose -thread 0 also warms the CTA's next row block of packed rows in L2 with one -bulk prefetch a tensor, from 3000 rows per SM, the long forward keeping the -rows from its own start up to there); on sm_100a and sm_103a from 1280 rows -per SM a backward whose thread 0 warms the next row block of packed rows in L2 -with one bulk prefetch a tensor, launched with 5 CTAs per SM, and on sm_107a -from 1024 the backward with 128-thread CTAs, ``dx`` and ``dz`` the same bits. +more rows), and sm_103a and sm_107a the forward's extra-long role +(``fwd_xlong``), tried before the long role (on sm_107a, which routes no long +forward, in its place): on every architecture a forward that walks several row +blocks per CTA, loading each next block while it computes the current one, +``y`` the same bits (the long forward 32-row blocks on 256-thread CTAs, the mid +forward 16-row blocks on 128-thread CTAs whose thread 0 also warms later row +blocks of packed rows in L2 with one bulk prefetch a tensor; the extra-long +forward is the long forward whose thread 0 also warms the CTA's next row block +of packed rows in L2 with one bulk prefetch a tensor, on sm_103a from 3000 rows +per SM, the long forward keeping the rows from its own start up to there, and +on sm_107a from the long forward's 1740 rows per SM); on sm_100a and sm_103a +from 1280 rows per SM a backward whose thread 0 warms the next row block of +packed rows in L2 with one bulk prefetch a tensor, launched with 5 CTAs per SM, +and on sm_107a from 1024 the backward with 128-thread CTAs, ``dx`` and ``dz`` +the same bits. An architecture may also route the forward's short role (``fwd_short``), whose module runs in the role's place up to ``max_rows_per_sm`` rows per SM of its grid rule (:func:`max_rows_per_sm`: one count, or one per architecture of the @@ -216,12 +218,11 @@ # 8, "math": "x2", "min_blocks": 3, "prefetch_ahead": 0, "prefetch_wave": # 0, "threads": 256}); whole waves of at most 6 blocks per CTA, from 2400 # rows per SM on sm_100a and sm_103a (on sm_103a up to the extra-long - # forward's 3000) and from 1740 (four whole waves) on sm_107a. + # forward's 3000; sm_107a runs the extra-long forward instead). "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258": { "arches": [ "sm_100a", "sm_103a", - "sm_107a", ], "arg_plan": [ ["buffer", "x"], @@ -247,7 +248,6 @@ "min_rows_per_sm": { "sm_100a": 2400, "sm_103a": 2400, - "sm_107a": 1740, }, "rows_per_block": 32, }, @@ -276,15 +276,17 @@ }, "role": "fwd_long", }, - # The extra-long forward (sm_103a): the long forward's program whose thread - # 0 also warms in L2, with one bulk prefetch a tensor where x and z hold - # their rows packed, the CTA's next row block (build {"lanes": 8, "math": - # "x2", "min_blocks": 3, "prefetch_ahead": 1, "prefetch_wave": 0, - # "threads": 256}); whole waves of at most 6 blocks per CTA, from 3000 rows - # per SM. + # The extra-long forward (sm_103a and sm_107a): the long forward's program + # whose thread 0 also warms in L2, with one bulk prefetch a tensor where x + # and z hold their rows packed, the CTA's next row block (build {"lanes": + # 8, "math": "x2", "min_blocks": 3, "prefetch_ahead": 1, "prefetch_wave": + # 0, "threads": 256}); whole waves of at most 6 blocks per CTA, from 3000 + # rows per SM on sm_103a and from 1740 (four whole waves, the long + # forward's start there) on sm_107a. "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883": { "arches": [ "sm_103a", + "sm_107a", ], "arg_plan": [ ["buffer", "x"], @@ -309,6 +311,7 @@ "max_blocks_per_cta": 6, "min_rows_per_sm": { "sm_103a": 3000, + "sm_107a": 1740, }, "rows_per_block": 32, }, @@ -1107,7 +1110,7 @@ }, "sm_107a": { "fwd": "cake_rmsnorm_silu_gate_train_05ffc12001cf81513436", - "fwd_long": "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258", + "fwd_xlong": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883", "fwd_mid": "cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a", "fwd_short": "cake_rmsnorm_silu_gate_train_de5df3c9a4a5a804fcf2", "bwd_bf16": "cake_rmsnorm_silu_gate_train_58f1ac96c476c73ec5c8", @@ -1134,8 +1137,8 @@ MID_ROLES = ("fwd_mid",) SHORT_ROLES = ("fwd_short", "bwd_bf16_short", "bwd_f32_short") # The suffixes of a role's extra-long, long and mid roles in the order -# :func:`module_for_rows` tries them (an extra-long role runs from more rows -# per SM than a long one, a long role from more than a mid one). +# :func:`module_for_rows` tries them (an extra-long role runs from at least as +# many rows per SM as a long one, a long role from more than a mid one). TIERS = ("_xlong", "_long", "_mid") GRID_KINDS = ("row_blocks", "persistent_row_blocks", "row_claims", "cluster") # Elements of one normalized row. diff --git a/tests/norm/test_cake_rmsnorm_silu_gate_train.py b/tests/norm/test_cake_rmsnorm_silu_gate_train.py index 7e4c6dcfcc6..04fb9e55226 100644 --- a/tests/norm/test_cake_rmsnorm_silu_gate_train.py +++ b/tests/norm/test_cake_rmsnorm_silu_gate_train.py @@ -1079,8 +1079,9 @@ def test_backward_cta_rule(rows, sm_count): assert grid == (ctas, 1, 1) -# The long and mid roles each architecture routes, with the rows per SM they -# run from. +# The extra-long, long and mid roles each architecture routes, with the rows +# per SM they run from (sm_107a: the extra-long forward from the long forward's +# 1740, no long forward). TIER_ROUTES = { "sm_100a": { "fwd_long": 2400, @@ -1096,7 +1097,7 @@ def test_backward_cta_rule(rows, sm_count): "bwd_f32_long": 1280, }, "sm_107a": { - "fwd_long": 1740, + "fwd_xlong": 1740, "fwd_mid": 800, "bwd_bf16_long": 1024, "bwd_f32_long": 1024, @@ -1234,8 +1235,8 @@ def test_short_long_and_mid_modules_run_from_their_rows_per_sm(sm_count): got = loader.module_for_rows("sm_107a", role, rows, sm_count) assert got == loader.ROUTES["sm_107a"][role + tier], (role, rows) # the forward on sm_107a: the short forward up to 59 rows per SM, the - # forward above, then the mid forward from 800 rows per SM, then the long - # forward from 1740 + # forward above, then the mid forward from 800 rows per SM, then the + # extra-long forward from 1740 (sm_107a routes no long forward) routes = loader.ROUTES["sm_107a"] for rows, role in ( (1, "fwd_short"), @@ -1245,9 +1246,11 @@ def test_short_long_and_mid_modules_run_from_their_rows_per_sm(sm_count): (800 * sm_count - 1, "fwd"), (800 * sm_count, "fwd_mid"), (1740 * sm_count - 1, "fwd_mid"), - (1740 * sm_count, "fwd_long"), + (1740 * sm_count, "fwd_xlong"), + (2**31 - 1, "fwd_xlong"), ): assert loader.module_for_rows("sm_107a", "fwd", rows, sm_count) == routes[role] + assert "fwd_long" not in routes # (rows, SMs, CTAs) of the long forward: 32-row blocks, the fewest whole waves @@ -1271,7 +1274,7 @@ def test_short_long_and_mid_modules_run_from_their_rows_per_sm(sm_count): def test_long_forward_grid_is_whole_waves_of_at_most_six_blocks_a_cta( rows, sm_count, ctas ): - for arch in ("sm_100a", "sm_103a", "sm_107a"): + for arch in ("sm_100a", "sm_103a"): name = loader.ROUTES[arch]["fwd_long"] assert name == loader.ROUTES["sm_100a"]["fwd_long"] record = loader.MODULES[name] @@ -1279,10 +1282,13 @@ def test_long_forward_grid_is_whole_waves_of_at_most_six_blocks_a_cta( "ctas_per_sm": 3, "kind": "persistent_row_blocks", "max_blocks_per_cta": 6, - "min_rows_per_sm": {"sm_100a": 2400, "sm_103a": 2400, "sm_107a": 1740}, + "min_rows_per_sm": {"sm_100a": 2400, "sm_103a": 2400}, "rows_per_block": 32, } - assert record["arches"] == ["sm_100a", "sm_103a", "sm_107a"] + assert record["arches"] == ["sm_100a", "sm_103a"] + # sm_107a's extra-long forward launches the same grid + xlong = loader.MODULES[loader.ROUTES["sm_107a"]["fwd_xlong"]] + assert loader.launch_grid(xlong, rows=rows, sm_count=sm_count) == (ctas, 1, 1) grid = loader.launch_grid(record, rows=rows, sm_count=sm_count) assert grid == (ctas, 1, 1) blocks, wave = -(-rows // 32), 3 * sm_count @@ -1299,21 +1305,32 @@ def test_long_forward_grid_is_whole_waves_of_at_most_six_blocks_a_cta( ) -def test_extra_long_forward_is_the_long_forwards_grid_from_3000_rows_per_sm(): - # sm_103a's extra-long forward: the long forward's grid rule (32-row blocks, +def test_extra_long_forward_is_the_long_forwards_grid_from_its_rows_per_sm(): + # the extra-long forward: the long forward's grid rule (32-row blocks, # whole waves of 3 CTAs per SM walking at most 6 blocks) from 3000 rows per - # SM, 256 threads, no PDL + # SM on sm_103a and from 1740 on sm_107a (the long forward's start there), + # 256 threads, no PDL name = loader.ROUTES["sm_103a"]["fwd_xlong"] + assert loader.ROUTES["sm_107a"]["fwd_xlong"] == name assert "fwd_xlong" not in loader.ROUTES["sm_100a"] - assert "fwd_xlong" not in loader.ROUTES["sm_107a"] record = loader.MODULES[name] - assert record["arches"] == ["sm_103a"] and record["role"] == "fwd_xlong" - long_rule = loader.MODULES[loader.ROUTES["sm_103a"]["fwd_long"]]["grid_rule"] - assert record["grid_rule"] == {**long_rule, "min_rows_per_sm": {"sm_103a": 3000}} + assert record["arches"] == ["sm_103a", "sm_107a"] + assert record["role"] == "fwd_xlong" + long_record = loader.MODULES[loader.ROUTES["sm_103a"]["fwd_long"]] + assert record["grid_rule"] == { + **long_record["grid_rule"], + "min_rows_per_sm": {"sm_103a": 3000, "sm_107a": 1740}, + } assert record["launch"]["block"] == (256, 1, 1) assert record["launch"]["use_pdl"] is False - for rows, sm_count in ((3000 * 152, 152), (32737 * 48, 152), (2**31 - 1, 152)): - long_record = loader.MODULES[loader.ROUTES["sm_103a"]["fwd_long"]] + for rows, sm_count in ( + (3000 * 152, 152), + (32737 * 48, 152), + (2**31 - 1, 152), + (1740 * 212, 212), + (32737 * 48, 212), + (2**31 - 1, 216), + ): assert loader.launch_grid( record, rows=rows, sm_count=sm_count ) == loader.launch_grid(long_record, rows=rows, sm_count=sm_count) diff --git a/tests/norm/test_cake_rmsnorm_silu_gate_train_api.py b/tests/norm/test_cake_rmsnorm_silu_gate_train_api.py index 2e449724868..c5d2a4fe5a3 100644 --- a/tests/norm/test_cake_rmsnorm_silu_gate_train_api.py +++ b/tests/norm/test_cake_rmsnorm_silu_gate_train_api.py @@ -1177,7 +1177,7 @@ def _sm_count() -> int: MID_FORWARD_ROWS_PER_SM = {"sm_100a": 976, "sm_103a": 976, "sm_107a": 800} LONG_FORWARD_ROWS_PER_SM = {"sm_100a": 2400, "sm_103a": 2400, "sm_107a": 1740} # The rows per SM from which an architecture runs the extra-long forward. -XLONG_FORWARD_ROWS_PER_SM = {"sm_103a": 3000} +XLONG_FORWARD_ROWS_PER_SM = {"sm_103a": 3000, "sm_107a": 1740} # Forward inputs around and past the short, the mid and the long forward's # thresholds: (tokens, heads, layout); ``heads=None`` takes ``threshold * SMs # + tokens`` rows (the short threshold for ``flat-short``, the mid one for @@ -1201,6 +1201,9 @@ def _sm_count() -> int: pytest.param(-1, None, "flat-xlong", id="xlong-threshold-1"), pytest.param(0, None, "flat-xlong", id="xlong-threshold"), pytest.param(1, None, "flat-xlong", id="xlong-threshold+1"), + pytest.param(7685, 48, "contiguous", id="7685x48"), + pytest.param(7685, 48, "projection", id="7685x48-projection"), + pytest.param(7685, 48, "head_major", id="7685x48-head_major"), pytest.param(9500, 48, "contiguous", id="9500x48"), pytest.param(9500, 48, "projection", id="9500x48-projection"), pytest.param(9500, 48, "head_major", id="9500x48-head_major"), diff --git a/tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py b/tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py index d6a90d0d5d5..97c278c3f70 100644 --- a/tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py +++ b/tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py @@ -705,10 +705,13 @@ def _persistent_forward_tokens(role: str, edge: str) -> tuple[int, int, float]: step = int(rule["ctas_per_sm"]) * rows_per_block blocks_per_slot = -(-loader.min_rows_per_sm(name, arch) // step) # the "last" edge: the most below the next tier's start where the - # architecture routes one (the mid forward below the long one, on sm_103a - # the long forward below the extra-long one), else four times "first" - upper = {"fwd_mid": "fwd_long", "fwd_long": "fwd_xlong"}.get(role) - if edge == "last" and upper in routes: + # architecture routes one (the mid forward below the long one, or on + # sm_107a, which routes no long forward, below the extra-long one; on + # sm_103a the long forward below the extra-long one), else four times + # "first" + uppers = {"fwd_mid": ("fwd_long", "fwd_xlong"), "fwd_long": ("fwd_xlong",)} + upper = next((u for u in uppers.get(role, ()) if u in routes), None) + if edge == "last" and upper is not None: upper_start = loader.min_rows_per_sm(routes[upper], arch) blocks_per_slot = (upper_start - 1) // step elif edge == "last":