Skip to content

Commit 3bcc95e

Browse files
apstenku123apstenku123
committed
fix: refactor backward reduction logic to use explicit thread binding for improved kernel performance
1 parent 238cbf4 commit 3bcc95e

5 files changed

Lines changed: 89 additions & 68 deletions

File tree

‎bench/tilelang_ports/mamba3_path_c.json‎

Lines changed: 41 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -33,72 +33,72 @@
3333
"timings": {
3434
"fwd_path_b": {
3535
"label": "fwd_path_b",
36-
"mean_ms": 5.522979200759437,
37-
"median_ms": 5.502729502040893,
38-
"min_ms": 5.456540995510295,
39-
"max_ms": 5.6501670042052865,
36+
"mean_ms": 5.45246660039993,
37+
"median_ms": 5.445874492579605,
38+
"min_ms": 5.397458007792011,
39+
"max_ms": 5.517917001270689,
4040
"iters": 10,
4141
"warmup": 3,
4242
"measurement": "paired_alternating"
4343
},
4444
"fwd_path_c": {
4545
"label": "fwd_path_c",
46-
"mean_ms": 4.618183099955786,
47-
"median_ms": 4.619500003173016,
48-
"min_ms": 4.545582996797748,
49-
"max_ms": 4.678625002270564,
46+
"mean_ms": 4.607933398801833,
47+
"median_ms": 4.617666992999148,
48+
"min_ms": 4.51483299548272,
49+
"max_ms": 4.65354100742843,
5050
"iters": 10,
5151
"warmup": 3,
5252
"measurement": "paired_alternating"
5353
},
5454
"fwd_bwd_path_b": {
5555
"label": "fwd_bwd_path_b",
56-
"mean_ms": 75.87897090124898,
57-
"median_ms": 75.31389550422318,
58-
"min_ms": 74.12695899256505,
59-
"max_ms": 79.35158400505316,
56+
"mean_ms": 74.2388666010811,
57+
"median_ms": 74.38910450582625,
58+
"min_ms": 73.41383300081361,
59+
"max_ms": 74.8265000001993,
6060
"iters": 10,
6161
"warmup": 3,
6262
"measurement": "paired_alternating"
6363
},
6464
"fwd_bwd_path_c": {
6565
"label": "fwd_bwd_path_c",
66-
"mean_ms": 81.04167500277981,
67-
"median_ms": 80.02666699758265,
68-
"min_ms": 76.515417007613,
69-
"max_ms": 84.50462500331923,
66+
"mean_ms": 78.66660819854587,
67+
"median_ms": 78.39095799863571,
68+
"min_ms": 74.74725000793114,
69+
"max_ms": 85.78729200235102,
7070
"iters": 10,
7171
"warmup": 3,
7272
"measurement": "paired_alternating"
7373
},
74-
"bwd_path_b_median_ms": 69.81116600218229,
75-
"bwd_path_c_median_ms": 75.40716699440964
74+
"bwd_path_b_median_ms": 68.94323001324665,
75+
"bwd_path_c_median_ms": 73.77329100563657
7676
},
7777
"bwd_profile": {
7878
"generated_partial_kernel": {
7979
"label": "bwd_path_c_generated_partials",
80-
"mean_ms": 71.74030839523766,
81-
"median_ms": 71.4143749937648,
82-
"min_ms": 70.90587499260437,
83-
"max_ms": 72.96049999422394,
80+
"mean_ms": 70.47153330058791,
81+
"median_ms": 70.15604100160999,
82+
"min_ms": 69.36979199235793,
83+
"max_ms": 74.5268339960603,
8484
"iters": 10,
8585
"warmup": 3
8686
},
8787
"simd_p_reduce_kernel": {
8888
"label": "bwd_path_c_simd_p_reduce_diagnostic",
89-
"mean_ms": 92.45353759761201,
90-
"median_ms": 92.33262499037664,
91-
"min_ms": 91.85583400540054,
92-
"max_ms": 93.4950830123853,
89+
"mean_ms": 88.6143419003929,
90+
"median_ms": 88.66360449610511,
91+
"min_ms": 88.15495899762027,
92+
"max_ms": 89.18954200635199,
9393
"iters": 10,
9494
"warmup": 3
9595
}
9696
},
9797
"gflops": {
98-
"fwd_path_b": 880.3669375722112,
99-
"fwd_path_c": 1048.6894938137225,
100-
"fwd_bwd_path_b": 257.2922878343666,
101-
"fwd_bwd_path_c": 242.14034155121487
98+
"fwd_path_b": 889.5579812940733,
99+
"fwd_path_c": 1049.1057772993665,
100+
"fwd_bwd_path_b": 260.4908959279422,
101+
"fwd_bwd_path_c": 247.19285201664766
102102
},
103103
"memory_bytes_peak": {
104104
"fwd_path_b": 509215408,
@@ -136,19 +136,21 @@
136136
"reason": "rule: fp32-accumulating bfloat16 per-lane scan with 256 threads over 28 blocks; z3 proved per-lane index decomposition and buffer bounds; bwd reverse pass consumes explicit state snapshot tensor boundaries instead of inverse h_prev reconstruction; bwd emits TileLang per-lane partial gradients and reduces P-axis outside the hot scan kernel"
137137
},
138138
"ratios": {
139-
"fwd_path_c_over_path_b": 0.8394924739549167,
140-
"bwd_path_c_over_path_b": 1.08015911082265,
141-
"fwd_bwd_path_c_over_path_b": 1.0625750595133563
139+
"fwd_path_c_over_path_b": 0.8479202007484841,
140+
"bwd_path_c_over_path_b": 1.070058524839377,
141+
"fwd_bwd_path_c_over_path_b": 1.0537962315771126
142142
},
143143
"optimization_policy": "AUTO only promotes a Path C phase when the rule/Z3 plan is safe and paired bench receipt median is no-worse than Path B.",
144144
"blocked_path_c_codegen_gaps": [
145-
"lowered Path C fwd still recomputes some lane-derived indices inside the t loop",
146-
"Path C bwd still performs the reverse recurrence serially over T and N per lane",
147-
"final-gradient in-kernel P-axis allreduce remains diagnostic only for recurrent bwd hot loops"
145+
"Path C bwd materializes explicit state snapshots as a separate safety kernel",
146+
"Path C bwd returns per-lane partial buffers and relies on external MLX reductions",
147+
"final-gradient in-kernel P-axis allreduce remains diagnostic for recurrent bwd hot loops"
148148
],
149149
"remembered_optimizations": [
150150
"TileLang local.var scalar y_acc instead of thread float[1]",
151151
"TileLang Metal local.var PrintExpr statement-order fix",
152+
"TileLang Mamba3 fwd/bwd uses Metal thread_position_in_grid for absolute lane ids",
153+
"TileLang Mamba3 fwd/bwd uses Path-B-equivalent single-expression SiLU",
152154
"Bench harness uses paired alternating samples to avoid order/warmup drift",
153155
"Path C bwd consumes explicit state snapshot boundaries instead of unsafe inverse-state reconstruction",
154156
"Path C bwd writes generated per-lane partial gradients and moves P-axis reductions outside the recurrent hot loop",
@@ -157,9 +159,9 @@
157159
"AUTO can still select Path C forward with Path B backward when only fwd is no-worse"
158160
]
159161
},
160-
"ratio_path_c_over_path_b_fwd": 0.8394924739549167,
161-
"ratio_path_c_over_path_b_bwd": 1.08015911082265,
162-
"ratio_path_c_over_path_b_fwd_bwd": 1.0625750595133563,
162+
"ratio_path_c_over_path_b_fwd": 0.8479202007484841,
163+
"ratio_path_c_over_path_b_bwd": 1.070058524839377,
164+
"ratio_path_c_over_path_b_fwd_bwd": 1.0537962315771126,
163165
"verdict": "Path C forward is no-worse than Path B and is eligible for AUTO promotion, but Path C backward is slower; scheduler selects Path C forward with Path B backward.",
164166
"matched_run_guard": "Compare Path B and Path C only when both rows were collected on the same hardware with identical kernel inputs. This is a local_only receipt."
165167
}

‎cppmega_mlx/nn/_tilelang/mamba3_path_c.py‎

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1120,11 +1120,13 @@ def bwd_simd(
11201120
dh0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), dh0_dtype),
11211121
):
11221122
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
1123+
tid = T.get_thread_binding(0)
11231124
global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
11241125
h_state = T.alloc_local((STATE,), accum_dtype)
11251126
dh = T.alloc_local((STATE,), accum_dtype)
11261127
if global_lane < LANES:
11271128
p = global_lane % HEADDIM
1129+
reduce_lane = tid % HEADDIM
11281130
h = (global_lane // HEADDIM) % HEADS
11291131
b = global_lane // (HEADDIM * HEADS)
11301132

@@ -1179,10 +1181,10 @@ def bwd_simd(
11791181
dC_sum = T.alloc_local((1,), accum_dtype)
11801182
dB_sum = T.alloc_local((1,), accum_dtype)
11811183
T.thread_allreduce_sum(
1182-
d_y_skipped * h_state[n], dC_sum[0], p
1184+
d_y_skipped * h_state[n], dC_sum[0], reduce_lane
11831185
)
1184-
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], p)
1185-
if p == 0:
1186+
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], reduce_lane)
1187+
if reduce_lane == 0:
11861188
dC[b, t, h, n] = T.cast(dC_sum[0], dC_dtype)
11871189
dB[b, t, h, n] = T.cast(dB_sum[0], dB_dtype)
11881190
dx_inp += dh_n * B_val
@@ -1196,10 +1198,10 @@ def bwd_simd(
11961198
dC_sum = T.alloc_local((1,), accum_dtype)
11971199
dB_sum = T.alloc_local((1,), accum_dtype)
11981200
T.thread_allreduce_sum(
1199-
d_y_skipped * h_state[n], dC_sum[0], p
1201+
d_y_skipped * h_state[n], dC_sum[0], reduce_lane
12001202
)
1201-
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], p)
1202-
if p == 0:
1203+
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], reduce_lane)
1204+
if reduce_lane == 0:
12031205
dC[b, t, h, n] = T.cast(dC_sum[0], dC_dtype)
12041206
dB[b, t, h, n] = T.cast(dB_sum[0], dB_dtype)
12051207
dx_inp += dh_n * B_val
@@ -1215,17 +1217,17 @@ def bwd_simd(
12151217
ddt_lane = d_logdecay * A_val
12161218
dA_sum = T.alloc_local((1,), accum_dtype)
12171219
ddt_sum = T.alloc_local((1,), accum_dtype)
1218-
T.thread_allreduce_sum(dA_lane, dA_sum[0], p)
1219-
T.thread_allreduce_sum(ddt_lane, ddt_sum[0], p)
1220-
if p == 0:
1220+
T.thread_allreduce_sum(dA_lane, dA_sum[0], reduce_lane)
1221+
T.thread_allreduce_sum(ddt_lane, ddt_sum[0], reduce_lane)
1222+
if reduce_lane == 0:
12211223
dA[b, t, h] = T.cast(dA_sum[0], dA_dtype)
12221224
ddt[b, t, h] = T.cast(ddt_sum[0], ddt_dtype)
12231225

12241226
for n in T.serial(STATE):
12251227
dh0[b, h, p, n] = T.cast(dh[n], dh0_dtype)
12261228
dD_sum = T.alloc_local((1,), accum_dtype)
1227-
T.thread_allreduce_sum(dD_acc, dD_sum[0], p)
1228-
if p == 0:
1229+
T.thread_allreduce_sum(dD_acc, dD_sum[0], reduce_lane)
1230+
if reduce_lane == 0:
12291231
dD_batch[b, h] = T.cast(dD_sum[0], dD_dtype)
12301232

12311233
artifact = dispatch_lower(bwd_simd, target="metal", return_msl=True)
@@ -1323,11 +1325,13 @@ def bwd_snap_simd(
13231325
dh0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), dh0_dtype),
13241326
):
13251327
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
1328+
tid = T.get_thread_binding(0)
13261329
global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
13271330
h_state = T.alloc_local((STATE,), accum_dtype)
13281331
dh = T.alloc_local((STATE,), accum_dtype)
13291332
if global_lane < LANES:
13301333
p = global_lane % HEADDIM
1334+
reduce_lane = tid % HEADDIM
13311335
h = (global_lane // HEADDIM) % HEADS
13321336
b = global_lane // (HEADDIM * HEADS)
13331337

@@ -1382,10 +1386,10 @@ def bwd_snap_simd(
13821386
T.thread_allreduce_sum(
13831387
d_y_skipped * h_state[n],
13841388
dC_sum[0],
1385-
p,
1389+
reduce_lane,
13861390
)
1387-
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], p)
1388-
if p == 0:
1391+
T.thread_allreduce_sum(dh_n * x_val, dB_sum[0], reduce_lane)
1392+
if reduce_lane == 0:
13891393
dC[b, t, h, n] = T.cast(dC_sum[0], dC_dtype)
13901394
dB[b, t, h, n] = T.cast(dB_sum[0], dB_dtype)
13911395
dx_inp += dh_n * B_val
@@ -1401,17 +1405,17 @@ def bwd_snap_simd(
14011405
ddt_lane = d_logdecay * A_val
14021406
dA_sum = T.alloc_local((1,), accum_dtype)
14031407
ddt_sum = T.alloc_local((1,), accum_dtype)
1404-
T.thread_allreduce_sum(dA_lane, dA_sum[0], p)
1405-
T.thread_allreduce_sum(ddt_lane, ddt_sum[0], p)
1406-
if p == 0:
1408+
T.thread_allreduce_sum(dA_lane, dA_sum[0], reduce_lane)
1409+
T.thread_allreduce_sum(ddt_lane, ddt_sum[0], reduce_lane)
1410+
if reduce_lane == 0:
14071411
dA[b, t, h] = T.cast(dA_sum[0], dA_dtype)
14081412
ddt[b, t, h] = T.cast(ddt_sum[0], ddt_dtype)
14091413

14101414
for n in T.serial(STATE):
14111415
dh0[b, h, p, n] = T.cast(dh[n], dh0_dtype)
14121416
dD_sum = T.alloc_local((1,), accum_dtype)
1413-
T.thread_allreduce_sum(dD_acc, dD_sum[0], p)
1414-
if p == 0:
1417+
T.thread_allreduce_sum(dD_acc, dD_sum[0], reduce_lane)
1418+
if reduce_lane == 0:
14151419
dD_batch[b, h] = T.cast(dD_sum[0], dD_dtype)
14161420

14171421
artifact = dispatch_lower(bwd_snap_simd, target="metal", return_msl=True)

‎reports/profiling/mamba3_path_c_forward_gap_2026-05-18.html‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ <h1>Mamba3 Path C Forward Gap</h1>
8080
<div class="head"><h2>Conclusion</h2><p>2026-05-18</p></div>
8181
<div class="body">
8282
<div class="callout">
83-
The forward regression is fixed for the BF16 1B-model shape: Path C forward moved from <span class="bad">1.29x slower</span> than Path B to <span class="good">0.84x of Path B</span>. The remaining default-training blocker is backward, not forward.
83+
The forward regression is fixed for the BF16 1B-model shape: Path C forward moved from <span class="bad">1.29x slower</span> than Path B to <span class="good">0.85x of Path B</span>. Carrying the same lane-id and SiLU fixes into backward reduced BF16 bwd C/B from <span class="bad">1.080</span> to <span class="warn">1.070</span>; the remaining blocker is the safer snapshot/reduction backward route, not forward.
8484
</div>
8585
</div>
8686
</section>
@@ -128,6 +128,15 @@ <h1>Mamba3 Path C Forward Gap</h1>
128128
<td><code>path_c_fwd_path_b_bwd</code></td>
129129
<td><a href="mamba3_model_shape_path_b_vs_c_after_grid_tid_bf16_2026-05-18.json">json</a> · <a href="mamba3_model_shape_path_b_vs_c_after_grid_tid_bf16_2026-05-18.metal">msl</a> · <a href="mamba3_model_shape_path_b_vs_c_after_grid_tid_bf16_2026-05-18.diff">diff</a></td>
130130
</tr>
131+
<tr>
132+
<td>Backward lane + SiLU carry-over</td>
133+
<td>bf16</td>
134+
<td class="num">5.446 ms</td>
135+
<td class="num">4.618 ms</td>
136+
<td class="num good">0.848</td>
137+
<td><code>path_c_fwd_path_b_bwd</code></td>
138+
<td><a href="mamba3_model_shape_path_b_vs_c_after_bwd_grid_tid_bf16_2026-05-18.json">json</a> · <a href="mamba3_model_shape_path_b_vs_c_after_bwd_grid_tid_bf16_2026-05-18.metal">msl</a> · <a href="mamba3_model_shape_path_b_vs_c_after_bwd_grid_tid_bf16_2026-05-18.diff">diff</a></td>
139+
</tr>
131140
<tr>
132141
<td>Grid lane intrinsic</td>
133142
<td>fp32</td>
@@ -143,18 +152,21 @@ <h1>Mamba3 Path C Forward Gap</h1>
143152
</section>
144153

145154
<section>
146-
<div class="head"><h2>What Changed</h2><p>forward codegen/schedule</p></div>
155+
<div class="head"><h2>What Changed</h2><p>codegen/schedule</p></div>
147156
<div class="body grid">
148-
<div class="metric"><span>BF16 forward C/B</span><strong class="good">0.839</strong></div>
157+
<div class="metric"><span>BF16 forward C/B</span><strong class="good">0.848</strong></div>
149158
<div class="metric"><span>FP32 forward C/B</span><strong>1.050</strong></div>
150-
<div class="metric"><span>BF16 fwd+bwd C/B</span><strong>1.063</strong></div>
159+
<div class="metric"><span>BF16 bwd C/B</span><strong>1.070</strong></div>
160+
<div class="metric"><span>BF16 fwd+bwd C/B</span><strong>1.054</strong></div>
161+
<div class="metric"><span>BF16 generated bwd</span><strong>70.156 ms</strong></div>
151162
<div class="metric"><span>BF16 Path C fwd+bwd peak</span><strong>11.49 GB</strong></div>
152163
</div>
153164
<div class="body">
154165
<ul>
155166
<li>Replaced the branchy generated forward sigmoid with the Path-B-equivalent single-expression SiLU.</li>
156167
<li>Hoisted <code>D[h]</code> out of the time loop.</li>
157-
<li>Changed forward lane id lowering to use <code>thread_position_in_grid</code>, removing the hot <code>blockIdx * THREADS + threadIdx</code> address arithmetic from Path C.</li>
168+
<li>Changed forward and backward lane id lowering to use <code>thread_position_in_grid</code>, removing the hot <code>blockIdx * THREADS + threadIdx</code> address arithmetic from Path C.</li>
169+
<li>Carried the Path-B-equivalent SiLU expression into generated backward; the old branchy backward sigmoid is gone from the lowered MSL.</li>
158170
<li>Reinstalled TileLang into <code>nanochat/.venv</code> with <code>--no-build-isolation</code> so the native MLX/TVM-FFI bridge is present in the benchmark environment.</li>
159171
</ul>
160172
</div>

‎scripts/bench_tilelang_mamba3_path_c.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -489,13 +489,15 @@ def _row(label: str, b_val: float, c_val: float, fmt: str = "{:.3f}", suffix: st
489489
"and paired bench receipt median is no-worse than Path B."
490490
),
491491
"blocked_path_c_codegen_gaps": [
492-
"lowered Path C fwd still recomputes some lane-derived indices inside the t loop",
493-
"Path C bwd still performs the reverse recurrence serially over T and N per lane",
494-
"final-gradient in-kernel P-axis allreduce remains diagnostic only for recurrent bwd hot loops",
492+
"Path C bwd materializes explicit state snapshots as a separate safety kernel",
493+
"Path C bwd returns per-lane partial buffers and relies on external MLX reductions",
494+
"final-gradient in-kernel P-axis allreduce remains diagnostic for recurrent bwd hot loops",
495495
],
496496
"remembered_optimizations": [
497497
"TileLang local.var scalar y_acc instead of thread float[1]",
498498
"TileLang Metal local.var PrintExpr statement-order fix",
499+
"TileLang Mamba3 fwd/bwd uses Metal thread_position_in_grid for absolute lane ids",
500+
"TileLang Mamba3 fwd/bwd uses Path-B-equivalent single-expression SiLU",
499501
"Bench harness uses paired alternating samples to avoid order/warmup drift",
500502
"Path C bwd consumes explicit state snapshot boundaries instead of unsafe inverse-state reconstruction",
501503
"Path C bwd writes generated per-lane partial gradients and moves P-axis reductions outside the recurrent hot loop",

‎tests/test_tilelang_mamba3_path_c.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -382,14 +382,15 @@ def test_raw_lowering_uses_tilelang_metal_scalar_pipeline() -> None:
382382
_require_mamba3_path_c()
383383

384384
_kernel, lowering = mamba3_path_c._bwd_simd_reduce_kernel_for_state_snapshots(
385-
1, 4, 1, 2, 4
385+
1, 4, 1, 64, 4
386386
)
387387
assert lowering.grid == (1, 1, 1)
388-
assert lowering.threadgroup == (2, 1, 1)
388+
assert lowering.threadgroup == (64, 1, 1)
389389
assert "h_snap" in lowering.body
390390
assert "float decay = exp(" in lowering.body
391391
assert "1.000000e+00 / decay" not in lowering.body
392392
assert "exp((A_val * dt_val))" not in lowering.body
393+
assert "T.metal_thread_position_in_grid_x() % 64" not in lowering.body
393394

394395

395396
def test_path_c_launch_geometry_comes_from_tilelang_lowering() -> None:

0 commit comments

Comments
 (0)