Skip to content

Commit 8106b17

Browse files
apstenku123apstenku123
committed
perf: optimize Mamba3 TileLang forward pass by hoisting D and using grid intrinsic lane IDs
1 parent 7ac9e6f commit 8106b17

7 files changed

Lines changed: 594 additions & 91 deletions

File tree

‎.DS_Store‎

0 Bytes
Binary file not shown.

‎bench/tilelang_ports/mamba3_path_c.json‎

Lines changed: 78 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,15 @@
88
"machine": "arm64",
99
"python_version": "3.13.12",
1010
"mlx_version": "0.32.0",
11-
"tilelang_version": "0.1.9+gitd669bcee"
11+
"tilelang_version": "0.1.9"
1212
},
1313
"shape": {
14-
"batch": 2,
15-
"seq": 512,
16-
"heads": 4,
17-
"headdim": 32,
14+
"batch": 1,
15+
"seq": 2048,
16+
"heads": 112,
17+
"headdim": 64,
1818
"state": 64,
19-
"dtype": "float32"
19+
"dtype": "bfloat16"
2020
},
2121
"path_b_status": {
2222
"available": true,
@@ -27,75 +27,84 @@
2727
"reason": "Path C TileLang DSL ready"
2828
},
2929
"parity": {
30-
"y_max_abs": 2.9802322387695312e-08,
31-
"h_max_abs": 2.384185791015625e-07
30+
"y_max_abs": 0.000244140625,
31+
"h_max_abs": 0.0009765625
3232
},
3333
"timings": {
3434
"fwd_path_b": {
3535
"label": "fwd_path_b",
36-
"mean_ms": 1.1066609038971364,
37-
"median_ms": 1.0936045437119901,
38-
"min_ms": 1.044750097207725,
39-
"max_ms": 1.2936670100316405,
40-
"iters": 50,
41-
"warmup": 5,
36+
"mean_ms": 5.522979200759437,
37+
"median_ms": 5.502729502040893,
38+
"min_ms": 5.456540995510295,
39+
"max_ms": 5.6501670042052865,
40+
"iters": 10,
41+
"warmup": 3,
4242
"measurement": "paired_alternating"
4343
},
4444
"fwd_path_c": {
4545
"label": "fwd_path_c",
46-
"mean_ms": 1.1934833158738911,
47-
"median_ms": 1.1894790222868323,
48-
"min_ms": 1.1205419432371855,
49-
"max_ms": 1.3904999941587448,
50-
"iters": 50,
51-
"warmup": 5,
46+
"mean_ms": 4.618183099955786,
47+
"median_ms": 4.619500003173016,
48+
"min_ms": 4.545582996797748,
49+
"max_ms": 4.678625002270564,
50+
"iters": 10,
51+
"warmup": 3,
5252
"measurement": "paired_alternating"
5353
},
5454
"fwd_bwd_path_b": {
5555
"label": "fwd_bwd_path_b",
56-
"mean_ms": 7.528950059786439,
57-
"median_ms": 7.537187484558672,
58-
"min_ms": 7.227167021483183,
59-
"max_ms": 7.930250023491681,
60-
"iters": 50,
61-
"warmup": 5,
56+
"mean_ms": 75.87897090124898,
57+
"median_ms": 75.31389550422318,
58+
"min_ms": 74.12695899256505,
59+
"max_ms": 79.35158400505316,
60+
"iters": 10,
61+
"warmup": 3,
6262
"measurement": "paired_alternating"
6363
},
6464
"fwd_bwd_path_c": {
6565
"label": "fwd_bwd_path_c",
66-
"mean_ms": 11.153560862876475,
67-
"median_ms": 11.149812547955662,
68-
"min_ms": 10.769542073830962,
69-
"max_ms": 11.691749910824,
70-
"iters": 50,
71-
"warmup": 5,
66+
"mean_ms": 81.04167500277981,
67+
"median_ms": 80.02666699758265,
68+
"min_ms": 76.515417007613,
69+
"max_ms": 84.50462500331923,
70+
"iters": 10,
71+
"warmup": 3,
7272
"measurement": "paired_alternating"
7373
},
74-
"bwd_path_b_median_ms": 6.443582940846682,
75-
"bwd_path_c_median_ms": 9.96033352566883
74+
"bwd_path_b_median_ms": 69.81116600218229,
75+
"bwd_path_c_median_ms": 75.40716699440964
7676
},
7777
"bwd_profile": {
78+
"generated_partial_kernel": {
79+
"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,
84+
"iters": 10,
85+
"warmup": 3
86+
},
7887
"simd_p_reduce_kernel": {
79-
"label": "bwd_path_c_simd_p_reduce",
80-
"mean_ms": 9.195864284411073,
81-
"median_ms": 9.17252094950527,
82-
"min_ms": 9.00504202581942,
83-
"max_ms": 9.53387503977865,
84-
"iters": 50,
85-
"warmup": 5
88+
"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,
93+
"iters": 10,
94+
"warmup": 3
8695
}
8796
},
8897
"gflops": {
89-
"fwd_path_b": 39.5515547632831,
90-
"fwd_path_c": 36.363617339667336,
91-
"fwd_bwd_path_b": 22.954854228378085,
92-
"fwd_bwd_path_c": 15.517304820673656
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
93102
},
94103
"memory_bytes_peak": {
95-
"fwd_path_b": 5013536,
96-
"fwd_path_c": 5013536,
97-
"fwd_bwd_path_b": 111543876,
98-
"fwd_bwd_path_c": 43450980
104+
"fwd_path_b": 509215408,
105+
"fwd_path_c": 210108656,
106+
"fwd_bwd_path_b": 12643240030,
107+
"fwd_bwd_path_c": 12049613984
99108
},
100109
"strict_policy": {
101110
"phase": "fwd",
@@ -106,51 +115,51 @@
106115
},
107116
"scheduler_decision": {
108117
"source": "rule_z3_test_receipt",
109-
"mode": "path_b",
110-
"selected_forward_kernel": "metal_kernel_fwd_v1",
118+
"mode": "path_c_fwd_path_b_bwd",
119+
"selected_forward_kernel": "path_c_tilelang_dsl",
111120
"selected_backward_kernel": "metal_kernel_bwd_v1",
112121
"rule_z3_plan": {
113-
"batch": 2,
114-
"seq": 512,
115-
"heads": 4,
116-
"headdim": 32,
122+
"batch": 1,
123+
"seq": 2048,
124+
"heads": 112,
125+
"headdim": 64,
117126
"state": 64,
118-
"dtype": "float32",
119-
"lanes": 256,
127+
"dtype": "bfloat16",
128+
"lanes": 7168,
120129
"threads": 256,
121-
"grid_blocks": 1,
130+
"grid_blocks": 28,
122131
"fwd_path_c_candidate": true,
123132
"bwd_path_c_candidate": true,
124133
"mode": "path_c_fwd_bwd",
125134
"z3_used": true,
126135
"z3_proved": true,
127-
"reason": "rule: fp32-accumulating float32 per-lane scan with 256 threads over 1 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 thread_allreduce_sum over P-axis grads"
136+
"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"
128137
},
129138
"ratios": {
130-
"fwd_path_c_over_path_b": 1.0876683250139199,
131-
"bwd_path_c_over_path_b": 1.5457756371116176,
132-
"fwd_bwd_path_c_over_path_b": 1.4793067799889712
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
133142
},
134143
"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.",
135144
"blocked_path_c_codegen_gaps": [
136145
"lowered Path C fwd still recomputes some lane-derived indices inside the t loop",
137146
"Path C bwd still performs the reverse recurrence serially over T and N per lane",
138-
"unsupported bwd reduction shapes now fail closed until semantic reduction lowering covers them"
147+
"final-gradient in-kernel P-axis allreduce remains diagnostic only for recurrent bwd hot loops"
139148
],
140149
"remembered_optimizations": [
141150
"TileLang local.var scalar y_acc instead of thread float[1]",
142151
"TileLang Metal local.var PrintExpr statement-order fix",
143152
"Bench harness uses paired alternating samples to avoid order/warmup drift",
144153
"Path C bwd consumes explicit state snapshot boundaries instead of unsafe inverse-state reconstruction",
145-
"Path C bwd uses TileLang thread_allreduce P-axis reductions instead of public dB/dC partial buffers",
146-
"Bench harness profiles the final-gradient SIMD/snapshot bwd route directly",
154+
"Path C bwd writes generated per-lane partial gradients and moves P-axis reductions outside the recurrent hot loop",
155+
"Bench harness profiles generated-partial production bwd and final-gradient SIMD diagnostic bwd separately",
147156
"AUTO selects full Path C only when fwd, bwd, and fwd+bwd receipts are no-worse",
148157
"AUTO can still select Path C forward with Path B backward when only fwd is no-worse"
149158
]
150159
},
151-
"ratio_path_c_over_path_b_fwd": 1.0876683250139199,
152-
"ratio_path_c_over_path_b_bwd": 1.5457756371116176,
153-
"ratio_path_c_over_path_b_fwd_bwd": 1.4793067799889712,
154-
"verdict": "Path C is >20% slower than Path B; recommend keeping Path B as primary, KEEPING Path C as a reproducibility artifact (proves the DSL path works), and flagging the perf gap as a future TileLang scheduler bug.",
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,
163+
"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.",
155164
"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."
156165
}

‎cppmega_mlx/nn/_tilelang/mamba3_path_c.py‎

Lines changed: 57 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -879,14 +879,17 @@ def fwd(
879879
h_last: T.Tensor((BATCH, HEADS, HEADDIM, STATE), h_last_dtype),
880880
):
881881
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
882-
tid_in_block = T.get_thread_binding(0)
883-
global_lane = bx * THREADS + tid_in_block
882+
# Metal already exposes the absolute 1-D lane id. Use it directly
883+
# so the lowered hot loop matches Path B and does not repeatedly
884+
# carry blockIdx * THREADS + threadIdx address arithmetic.
885+
global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
884886
# Per-lane state lives in registers (size N).
885887
h_state = T.alloc_local((STATE,), accum_dtype)
886888
if global_lane < LANES:
887889
p = global_lane % HEADDIM
888890
h = (global_lane // HEADDIM) % HEADS
889891
b = global_lane // (HEADDIM * HEADS)
892+
D_h = T.cast(D[h], accum_dtype)
890893
for n in T.serial(STATE):
891894
h_state[n] = T.cast(h0[b, h, p, n], accum_dtype)
892895
for t in T.serial(SEQ):
@@ -903,14 +906,8 @@ def fwd(
903906
)
904907
h_state[n] = new_h
905908
y_acc += new_h * T.cast(C[b, t, h, n], accum_dtype)
906-
D_h = T.cast(D[h], accum_dtype)
907909
y_skipped = y_acc + D_h * x_val
908-
sig_z = T.alloc_var(T.float32, init=0.0)
909-
if z_val >= 0.0:
910-
sig_z = 1.0 / (1.0 + T.exp(-z_val))
911-
else:
912-
sig_z = T.exp(z_val)
913-
sig_z = sig_z / (1.0 + sig_z)
910+
sig_z = 1.0 / (1.0 + T.exp(-z_val))
914911
y[b, t, h, p] = T.cast(z_val * sig_z * y_skipped, y_dtype)
915912
for n in T.serial(STATE):
916913
h_last[b, h, p, n] = T.cast(h_state[n], h_last_dtype)
@@ -2271,7 +2268,13 @@ def _mamba3_mimo_apply_with_state_path_c_fwd_path_b_bwd_vjp(
22712268
# Path B's hand-written MSL against Path C's machine-emitted MSL without
22722269
# having to re-run the lowering pipeline.
22732270
def dump_lowered_fwd_msl(
2274-
*, batch: int, seq: int, heads: int, headdim: int, state: int
2271+
*,
2272+
batch: int,
2273+
seq: int,
2274+
heads: int,
2275+
headdim: int,
2276+
state: int,
2277+
dtype: str = "float32",
22752278
) -> str:
22762279
"""Return the raw lowered MSL for the Path C forward kernel.
22772280
@@ -2280,19 +2283,60 @@ def dump_lowered_fwd_msl(
22802283
"""
22812284

22822285
kernel, lowering = _fwd_kernel_for(
2283-
batch, seq, heads, headdim, state, return_msl=True
2286+
batch,
2287+
seq,
2288+
heads,
2289+
headdim,
2290+
state,
2291+
dtype,
2292+
dtype,
2293+
dtype,
2294+
dtype,
2295+
dtype,
2296+
dtype,
2297+
dtype,
2298+
dtype,
2299+
dtype,
2300+
dtype,
2301+
return_msl=True,
22842302
)
22852303
del kernel
22862304
return _source_with_reused_scalar_bindings(lowering)
22872305

22882306

22892307
def dump_lowered_bwd_msl(
2290-
*, batch: int, seq: int, heads: int, headdim: int, state: int
2308+
*,
2309+
batch: int,
2310+
seq: int,
2311+
heads: int,
2312+
headdim: int,
2313+
state: int,
2314+
dtype: str = "float32",
22912315
) -> str:
22922316
"""Return the raw lowered MSL for the production Path C backward kernel."""
22932317

22942318
kernel, lowering = _bwd_partial_kernel_for_state_snapshots(
2295-
batch, seq, heads, headdim, state
2319+
batch,
2320+
seq,
2321+
heads,
2322+
headdim,
2323+
state,
2324+
dtype,
2325+
dtype,
2326+
dtype,
2327+
dtype,
2328+
dtype,
2329+
dtype,
2330+
dtype,
2331+
dtype,
2332+
dtype,
2333+
dtype,
2334+
dtype,
2335+
dtype,
2336+
dtype,
2337+
dtype,
2338+
dtype,
2339+
dtype,
22962340
)
22972341
del kernel
22982342
return _source_with_reused_scalar_bindings(lowering)

0 commit comments

Comments
 (0)