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

flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py

@@ -119,7 +119,9 @@
   (1740 on sm_107a) the long forward (3 CTAs per SM, at most 6 blocks of 32
   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.  The backward launches
+  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
   ``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

@@ -359,6 +359,29 @@
         "csrc/cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258_binding.cu"
       ]
     },
+    {
+      "arch": "sm_103a",
+      "closure_sha256": "1806bf4a6b51f1fd913a007d6dd9b5e2e36522517fd30854adc6e47866d06ab1",
+      "entry": "run",
+      "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_7f014401d5cc855ab883_sm_103a.so",
+      "sha256": {
+        "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_7f014401d5cc855ab883_kernel.cu",
+        "csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_binding.cu"
+      ]
+    },
     {
       "arch": "sm_100a",
       "closure_sha256": "904302c855a0071256abffe8ffeb5b410e2caf4d55526a7d064c0df4e5716148",

flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py

@@ -39,40 +39,44 @@
 (``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): 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_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 record; tried before the long and
-mid roles): on every architecture the long forward's kernel on 256-thread
-CTAs with 16 threads to a row, one 16-row block per CTA, that computes the
-packed SiLU gate before the row's sum of squares (the scalar gate after rstd
-where the range check fails), ``y`` the same bits.
-An architecture may also route the backward's short roles
-(``bwd_bf16_short``, ``bwd_f32_short``), whose module runs in the role's
+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.
+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
+record; tried before the extra-long, long and mid roles): on every
+architecture the long forward's kernel on 256-thread CTAs with 16 threads to a
+row, one 16-row block per CTA, that computes the packed SiLU gate before the
+row's sum of squares (the scalar gate after rstd where the range check fails),
+``y`` the same bits.  An architecture may also route the backward's short
+roles (``bwd_bf16_short``, ``bwd_f32_short``), whose module runs in the role's
 place up to ``max_rows`` rows of its grid rule: on every architecture one
 thread-block cluster of 16 CTAs of 256 threads with the backward's row
 arithmetic, ``dx`` and ``dz`` the same bits, that sums the weight gradient on
 chip (CTA ``r`` of the cluster adds every CTA's FP32 partial of its eight
 columns in rank order in FP64 and rounds once), with no workspace, atomic or
 arrival ticket, so its weight gradient is the same bits on every call.  A
 short role's rule bounds the rows per SM (``max_rows_per_sm``: the short
-forward, whose grid of one CTA per row block grows with the rows) or the
-rows themselves (``max_rows``: the short backward, whose grid is one cluster
+forward, whose grid of one CTA per row block grows with the rows) or the rows
+themselves (``max_rows``: the short backward, whose grid is one cluster
 whatever the rows and the SMs).
 
 Grids (:func:`launch_grid`): the forward and the short forward launch one
 CTA per block of ``rows_per_cta`` rows; the short backward one cluster of
 ``cluster`` CTAs (``launch.cluster``) whatever the rows, CTA ``r`` taking
 the 16-row blocks ``r``, ``r + 16``, ... (at most one up to its ``max_rows``).
-The long and mid forwards launch the fewest whole waves of
+The extra-long, long and mid forwards launch the fewest whole waves of
 ``ctas_per_sm * SMs`` CTAs (the CTAs per SM its launch bounds are compiled
 for) whose CTAs walk at most ``max_blocks_per_cta`` blocks of
 ``rows_per_block`` rows each, and never more CTAs than blocks; CTA ``b``
@@ -211,8 +215,8 @@
     # computes the current one, the gate on packed FP32 pairs (build {"lanes":
     # 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 and from 1740 (four whole waves) on
-    # sm_107a.
+    # 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.
     "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258": {
         "arches": [
             "sm_100a",
@@ -272,6 +276,67 @@
         },
         "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.
+    "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883": {
+        "arches": [
+            "sm_103a",
+        ],
+        "arg_plan": [
+            ["buffer", "x"],
+            ["buffer", "z"],
+            ["buffer", "w"],
+            ["buffer", "y"],
+            ["parameter", "rows"],
+            ["parameter", "heads"],
+            ["parameter", "x_tok"],
+            ["parameter", "x_head"],
+            ["parameter", "z_tok"],
+            ["parameter", "z_head"],
+            ["parameter", "eps"],
+            ["grid", "grid_x"],
+            ["grid", "grid_y"],
+            ["grid", "grid_z"],
+        ],
+        "ffi_entry": "run",
+        "grid_rule": {
+            "ctas_per_sm": 3,
+            "kind": "persistent_row_blocks",
+            "max_blocks_per_cta": 6,
+            "min_rows_per_sm": {
+                "sm_103a": 3000,
+            },
+            "rows_per_block": 32,
+        },
+        "kernel": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883",
+        "kernel_symbol": "kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883",
+        "launch": {
+            "block": (256, 1, 1),
+            "cluster": (1, 1, 1),
+            "cooperative": False,
+            "dynamic_smem_bytes": 0,
+            "use_pdl": False,
+        },
+        "library": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883",
+        "logical": {
+            "eps": "eps",
+            "heads": "heads",
+            "rows": "rows",
+            "w": "w",
+            "x": "x",
+            "x_head": "x_head",
+            "x_tok": "x_tok",
+            "y": "y",
+            "z": "z",
+            "z_head": "z_head",
+            "z_tok": "z_tok",
+        },
+        "role": "fwd_xlong",
+    },
     # The mid forward: the long forward's kernel on 128 threads, 8 lanes per
     # row (16-row blocks), at most 6 CTAs per SM, whose thread 0 also warms in
     # L2, with one bulk prefetch a tensor where x and z hold their rows packed,
@@ -1029,6 +1094,7 @@
     },
     "sm_103a": {
         "fwd": "cake_rmsnorm_silu_gate_train_05ffc12001cf81513436",
+        "fwd_xlong": "cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883",
         "fwd_long": "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258",
         "fwd_mid": "cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a",
         "fwd_short": "cake_rmsnorm_silu_gate_train_de5df3c9a4a5a804fcf2",
@@ -1059,16 +1125,18 @@
     (10, 3): "sm_103a",
     (10, 7): "sm_107a",
 }
-# The kernels of a call, routed on every architecture; ``<role>_long`` is a
-# role's long role, ``<role>_mid`` its mid role and ``<role>_short`` its short
-# role, routed on some.
+# The kernels of a call, routed on every architecture; ``<role>_xlong`` is a
+# role's extra-long role, ``<role>_long`` its long role, ``<role>_mid`` its mid
+# role and ``<role>_short`` its short role, routed on some.
 ROLES = ("fwd", "bwd_bf16", "bwd_f32")
+XLONG_ROLES = ("fwd_xlong",)
 LONG_ROLES = ("fwd_long", "bwd_bf16_long", "bwd_f32_long")
 MID_ROLES = ("fwd_mid",)
 SHORT_ROLES = ("fwd_short", "bwd_bf16_short", "bwd_f32_short")
-# The suffixes of a role's long and mid roles in the order
-# :func:`module_for_rows` tries them (a long role runs from more rows per SM).
-TIERS = ("_long", "_mid")
+# 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).
+TIERS = ("_xlong", "_long", "_mid")
 GRID_KINDS = ("row_blocks", "persistent_row_blocks", "row_claims", "cluster")
 # Elements of one normalized row.
 HEAD_DIM = 128
@@ -1126,8 +1194,8 @@ def module_for_rows(arch: str, role: str, rows: int, sm_count: int) -> str:
     routes one and ``rows`` is at most its grid rule's ``max_rows`` (a grid
     that does not grow with the rows: the short backward's one cluster) or
     else ``max_rows_per_sm(module, arch) * sm_count`` (one CTA per row block:
-    the short forward), else the first of the role's long and mid modules
-    (``TIERS`` order) that ``arch`` routes with
+    the short forward), else the first of the role's extra-long, long and
+    mid modules (``TIERS`` order) that ``arch`` routes with
     ``rows >= min_rows_per_sm(module, arch) * sm_count``, else
     :func:`route_module`."""
 
@@ -1316,6 +1384,7 @@ def run_module(
     "ROUTES",
     "SHORT_ROLES",
     "TIERS",
+    "XLONG_ROLES",
     "arch_for_capability",
     "backward_ctas",
     "device_arch",

flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883_binding.cu

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883, for the device source that defines it. Do not edit.
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+
+#include "cake_ffi_checks.h"
+
+#include <cstdint>
+
+struct CakeFastDivmod { int32_t divisor; uint32_t multiplier; uint32_t shift_right; };
+static_assert(sizeof(CakeFastDivmod) == 12, "CakeFastDivmod CUDA ABI must be 12 bytes");
+static_assert(alignof(CakeFastDivmod) == 4, "CakeFastDivmod CUDA ABI must be 4-byte aligned");
+
+extern "C" __global__ void kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ z, __nv_bfloat16* __restrict__ w, __nv_bfloat16* __restrict__ y, int rows, CakeFastDivmod heads, long long x_tok, long long x_head, long long z_tok, long long z_head, float eps);
+
+namespace cake_binding_kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883 {
+
+using tvm::ffi::TensorView;
+
+constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883";
+
+void CheckCuda(cudaError_t status, const char* what) {
+  TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
+      << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status);
+}
+
+void* DataPtr(const TensorView& t) { return static_cast<char*>(t.data_ptr()) + t.byte_offset(); }
+
+[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) {
+  TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b
+                            << " has no int64 value in a host check";
+}
+
+int64_t CMod(int64_t a, int64_t b) {
+  if (b == 0) NoInt64(a, "%", b);
+  return b == -1 ? 0 : a % b;
+}
+
+CakeFastDivmod MakeFastDivmod(int32_t divisor) {
+  if (divisor == 1) return CakeFastDivmod{1, 0u, 0u};
+  uint32_t log2_ceil = 0;
+  for (uint32_t value = static_cast<uint32_t>(divisor) - 1u; value != 0u; value >>= 1) ++log2_ceil;
+  const uint64_t power = uint64_t{1} << (31u + log2_ceil);
+  const uint32_t multiplier =
+      static_cast<uint32_t>((power + static_cast<uint64_t>(divisor) - 1u) / static_cast<uint64_t>(divisor));
+  return CakeFastDivmod{divisor, multiplier, log2_ceil - 1u};
+}
+
+void Run(TensorView a_x, TensorView a_z, TensorView a_w, TensorView a_y, int64_t a_rows, int64_t a_heads, int64_t a_x_tok, int64_t a_x_head, int64_t a_z_tok, int64_t a_z_head, double a_eps, int64_t grid_x, int64_t grid_y, int64_t grid_z) {
+  DLDevice dev = a_x.device();
+  tvm::ffi::CUDADeviceGuard device_guard(dev.device_id);
+  check_cuda_tensor(a_x, "x");
+  check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x");
+  check_cuda_tensor(a_z, "z");
+  check_dtype(a_z, DLDataType{kDLBfloat, 16, 1}, "z");
+  check_cuda_tensor(a_w, "w");
+  check_dtype(a_w, DLDataType{kDLBfloat, 16, 1}, "w");
+  check_cuda_tensor(a_y, "y");
+  check_dtype(a_y, DLDataType{kDLBfloat, 16, 1}, "y");
+  TVM_FFI_CHECK(a_rows >= -2147483648LL && a_rows <= 2147483647LL, ValueError)
+      << KERNEL_NAME << ": 'rows' = " << a_rows
+      << " is outside the i32 range [-2147483648, 2147483647]";
+  TVM_FFI_CHECK(a_heads >= 1LL && a_heads <= 2147483647LL, ValueError)
+      << KERNEL_NAME << ": fast divisor 'heads' = " << a_heads
+      << " is outside [1, 2147483647]";
+  check_same_device(a_z, a_x, "z", "x");
+  check_contiguous(a_w, "w");
+  check_same_device(a_w, a_x, "w", "x");
+  check_contiguous(a_y, "y");
+  check_same_device(a_y, a_x, "y", "x");
+  TVM_FFI_CHECK((a_rows >= 1LL), ValueError)
+      << KERNEL_NAME << ": requires (rows >= 1)";
+  TVM_FFI_CHECK((a_heads >= 1LL), ValueError)
+      << KERNEL_NAME << ": requires (heads >= 1)";
+  TVM_FFI_CHECK((CMod(a_x_tok, 16LL) == 0LL), ValueError)
+      << KERNEL_NAME << ": requires (CMod(x_tok, 16) == 0)";
+  TVM_FFI_CHECK((CMod(a_x_head, 16LL) == 0LL), ValueError)
+      << KERNEL_NAME << ": requires (CMod(x_head, 16) == 0)";
+  TVM_FFI_CHECK((CMod(a_z_tok, 16LL) == 0LL), ValueError)
+      << KERNEL_NAME << ": requires (CMod(z_tok, 16) == 0)";
+  TVM_FFI_CHECK((CMod(a_z_head, 16LL) == 0LL), ValueError)
+      << KERNEL_NAME << ": requires (CMod(z_head, 16) == 0)";
+  TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError)
+      << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive";
+  TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError)
+      << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits";
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(dev.device_type, dev.device_id));
+  void* v_x = DataPtr(a_x);
+  void* v_z = DataPtr(a_z);
+  void* v_w = DataPtr(a_w);
+  void* v_y = DataPtr(a_y);
+  int32_t v_rows = static_cast<int32_t>(a_rows);
+  CakeFastDivmod v_heads = MakeFastDivmod(static_cast<int32_t>(a_heads));
+  int64_t v_x_tok = static_cast<int64_t>(a_x_tok);
+  int64_t v_x_head = static_cast<int64_t>(a_x_head);
+  int64_t v_z_tok = static_cast<int64_t>(a_z_tok);
+  int64_t v_z_head = static_cast<int64_t>(a_z_head);
+  float v_eps = static_cast<float>(a_eps);
+  void* kernel_args[] = {&v_x, &v_z, &v_w, &v_y, &v_rows, &v_heads, &v_x_tok, &v_x_head, &v_z_tok, &v_z_head, &v_eps};
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3(static_cast<unsigned>(grid_x), static_cast<unsigned>(grid_y), static_cast<unsigned>(grid_z));
+  config.blockDim = dim3(256, 1, 1);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883), kernel_args),
+            "cudaLaunchKernelExC");
+}
+
+}  // namespace cake_binding_kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_silu_gate_train_7f014401d5cc855ab883::Run);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);