Skip to content

perf(cake_rmsnorm_silu_gate): L2-prefetching mid forward on SM107 #8656

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yyihuang merged 1 commit into from
Nov 4, 2026
Merged

perf(cake_rmsnorm_silu_gate): L2-prefetching mid forward on SM107 #8656

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,10 +117,10 @@
blocks: from 976 rows per SM on sm_100a and sm_103a (800 on sm_107a) the
mid forward (6 CTAs per SM, at most 2 blocks of 16 rows each), from 2400
(1740 on sm_107a) the long forward (3 CTAs per SM, at most 6 blocks of 32
rows each). On sm_100a and sm_103a 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 ``cake_jit.backward_ctas`` CTAs, ``G = min(row blocks,
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
``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
SM the same CTAs, whose thread 0 warms the next block of packed rows in L2
Expand Down
46 changes: 23 additions & 23 deletions flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,29 @@
"csrc/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_binding.cu"
]
},
{
"arch": "sm_107a",
"closure_sha256": "d258873b275d5546ed0e1c14b18e01237ea3c0e9d4611ad0f253d0665cf519a9",
"entry": "run",
"gencode": "-gencode=arch=compute_107a,code=sm_107a",
"min_cuda": "13.4",
"name": "cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_sm_107a",
"nvcc_flags": [
"--prec-div=true",
"--prec-sqrt=true",
"--ftz=false"
],
"nvcc_flags_if": [],
"path": "aot/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_sm_107a.so",
"sha256": {
"csrc/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_binding.cu": "0a446913ad2d58f2f88d18d689e3a3797499ae94f4af342d9b42bd8781ec6391",
"csrc/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_kernel.cu": "5bbba7afa0443ae2fb45bc01bb44f909310339e638d25dd36c2d3023004f6b46"
},
"sources": [
"csrc/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_kernel.cu",
"csrc/cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a_binding.cu"
]
},
{
"arch": "sm_107a",
"closure_sha256": "ceeafdf81e3a0e00d352f26a5418dea918bf79e36379003423ad58667ee1e6db",
Expand Down Expand Up @@ -474,29 +497,6 @@
"csrc/cake_rmsnorm_silu_gate_train_9f2f4d1f2b85a7469a7d_binding.cu"
]
},
{
"arch": "sm_107a",
"closure_sha256": "384ae3b705ef936cf80e84a9a50c59d7574df93913ee5f235b9e7bad4396f48a",
"entry": "run",
"gencode": "-gencode=arch=compute_107a,code=sm_107a",
"min_cuda": "13.4",
"name": "cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_sm_107a",
"nvcc_flags": [
"--prec-div=true",
"--prec-sqrt=true",
"--ftz=false"
],
"nvcc_flags_if": [],
"path": "aot/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_sm_107a.so",
"sha256": {
"csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_binding.cu": "1ac3efcc471d61950cfb1504d484e2b6724798d7cb545f4ca5d2b46fd884a249",
"csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_kernel.cu": "17af21a4aa2e5ea72db2c8d764c9a2745e10f7247f493fade3ee63f65cb49df5"
},
"sources": [
"csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_kernel.cu",
"csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_binding.cu"
]
},
{
"arch": "sm_100a",
"closure_sha256": "f74b9b06ec82436bc5940503192953fb7013e9f1e4dee522d72106b138e18865",
Expand Down
84 changes: 13 additions & 71 deletions flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,9 @@
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; on sm_100a and sm_103a the mid forward's
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
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
Expand Down Expand Up @@ -274,76 +273,18 @@
"role": "fwd_long",
},
# The mid forward: the long forward's kernel on 128 threads, 8 lanes per
# row (16-row blocks), at most 6 CTAs per SM (build {"lanes": 8, "math":
# "x2", "min_blocks": 6, "prefetch_ahead": 0, "prefetch_wave": 0,
# "threads": 128}); whole waves of at most 2 blocks per CTA, from 800 rows
# per SM on sm_107a, up to the long forward's.
"cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850": {
"arches": [
"sm_107a",
],
"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": 6,
"kind": "persistent_row_blocks",
"max_blocks_per_cta": 2,
"min_rows_per_sm": {
"sm_107a": 800,
},
"rows_per_block": 16,
},
"kernel": "cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850",
"kernel_symbol": "kernel_cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850",
"launch": {
"block": (128, 1, 1),
"cluster": (1, 1, 1),
"cooperative": False,
"dynamic_smem_bytes": 0,
"use_pdl": False,
},
"library": "cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850",
"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_mid",
},
# The mid forward with the L2 prefetch (sm_100a, sm_103a): the mid
# 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
# and the first row block of the CTA one wave of 6 CTAs per SM later
# (build {"lanes": 8, "math": "x2", "min_blocks": 6, "prefetch_ahead": 1,
# "prefetch_wave": 1, "threads": 128}); whole waves of at most 2 blocks per
# CTA, from 976 rows per SM up to the long forward's.
# 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,
# the CTA's next row block and the first row block of the CTA one wave of 6
# CTAs per SM later (build {"lanes": 8, "math": "x2", "min_blocks": 6,
# "prefetch_ahead": 1, "prefetch_wave": 1, "threads": 128}); whole waves of
# at most 2 blocks per CTA, from 976 rows per SM on sm_100a and sm_103a and
# from 800 on sm_107a, up to the long forward's.
"cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a": {
"arches": [
"sm_100a",
"sm_103a",
"sm_107a",
],
"arg_plan": [
["buffer", "x"],
Expand All @@ -369,6 +310,7 @@
"min_rows_per_sm": {
"sm_100a": 976,
"sm_103a": 976,
"sm_107a": 800,
},
"rows_per_block": 16,
},
Expand Down Expand Up @@ -1100,7 +1042,7 @@
"sm_107a": {
"fwd": "cake_rmsnorm_silu_gate_train_05ffc12001cf81513436",
"fwd_long": "cake_rmsnorm_silu_gate_train_7b7fac905aff1ddaa258",
"fwd_mid": "cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850",
"fwd_mid": "cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a",
"fwd_short": "cake_rmsnorm_silu_gate_train_de5df3c9a4a5a804fcf2",
"bwd_bf16": "cake_rmsnorm_silu_gate_train_58f1ac96c476c73ec5c8",
"bwd_bf16_long": "cake_rmsnorm_silu_gate_train_8b61372f6a32eeca4d10",
Expand Down
127 changes: 0 additions & 127 deletions ...silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_binding.cu
View file
Open in desktop

This file was deleted.

Loading
Loading