perf(cake_rmsnorm_silu_gate_train): L2-prefetching extra-long forward on SM107

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

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"

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.

tests/norm/test_cake_rmsnorm_silu_gate_train.py

@@ -1,4 +1,4 @@
 """
 Copyright (c) 2026 by FlashInfer team.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -1079,8 +1079,9 @@
         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 @@
         "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 @@
             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 @@
         (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,18 +1274,21 @@
 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]
     assert record["grid_rule"] == {
         "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_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)

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"),

tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py

@@ -1,4 +1,4 @@
 """
 Copyright (c) 2026 by FlashInfer team.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -705,10 +705,13 @@
     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":