Skip to content

Commit 15fdc5d

Browse files
apstenku123apstenku123
committed
refactor: decouple TileLang kernel definition from lowering logic and extend path C fusion runtime with FP8 training infrastructure.
1 parent 8875b99 commit 15fdc5d

7 files changed

Lines changed: 1059 additions & 57 deletions

File tree

‎cppmega_mlx/nn/_tilelang/m2rnn_path_c.py‎

Lines changed: 99 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1571,7 +1571,7 @@ def fwd(
15711571

15721572

15731573
@lru_cache(maxsize=128)
1574-
def _mapped_packed_post_fwd_k_parallel_kernel_for(
1574+
def _make_mapped_packed_post_fwd_k_parallel_prim_func(
15751575
batch: int,
15761576
seq: int,
15771577
total_heads: int,
@@ -1585,12 +1585,11 @@ def _mapped_packed_post_fwd_k_parallel_kernel_for(
15851585
v_dim: int,
15861586
projected_dim: int,
15871587
carrier_dtype: str,
1588-
*,
1589-
return_msl: bool = False,
1590-
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1588+
) -> Any:
1589+
"""Build the raw mapped packed-post k-parallel PrimFunc before lowering."""
1590+
15911591
import tilelang.language as T
15921592

1593-
del return_msl
15941593
conv_dim = _mapped_conv_dim(q_heads, k_heads, v_heads, k_dim, v_dim)
15951594
k_offset = q_heads * k_dim
15961595
v_offset = k_offset + k_heads * k_dim
@@ -1703,7 +1702,7 @@ def fwd(
17031702
for vv in T.serial(v_dim):
17041703
h_last[b, h, tid, vv] = T.cast(h_row[vv], carrier_dtype)
17051704

1706-
fwd = _with_global_symbol(
1705+
return _with_global_symbol(
17071706
fwd,
17081707
"m2rnn_mapped_packed_post_fwd_kp",
17091708
batch,
@@ -1720,6 +1719,42 @@ def fwd(
17201719
projected_dim,
17211720
carrier_dtype,
17221721
)
1722+
1723+
1724+
@lru_cache(maxsize=128)
1725+
def _mapped_packed_post_fwd_k_parallel_kernel_for(
1726+
batch: int,
1727+
seq: int,
1728+
total_heads: int,
1729+
q_heads: int,
1730+
k_heads: int,
1731+
v_heads: int,
1732+
g_heads: int,
1733+
w_heads: int,
1734+
f_heads: int,
1735+
k_dim: int,
1736+
v_dim: int,
1737+
projected_dim: int,
1738+
carrier_dtype: str,
1739+
*,
1740+
return_msl: bool = False,
1741+
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1742+
del return_msl
1743+
fwd = _make_mapped_packed_post_fwd_k_parallel_prim_func(
1744+
batch,
1745+
seq,
1746+
total_heads,
1747+
q_heads,
1748+
k_heads,
1749+
v_heads,
1750+
g_heads,
1751+
w_heads,
1752+
f_heads,
1753+
k_dim,
1754+
v_dim,
1755+
projected_dim,
1756+
carrier_dtype,
1757+
)
17231758
kernel, lowering = _compile_native_metal_kernel_with_lowering(
17241759
fwd,
17251760
out_idx=list(_PACKED_POST_FWD_OUTPUT_IDX),
@@ -1738,7 +1773,7 @@ def fwd(
17381773

17391774

17401775
@lru_cache(maxsize=128)
1741-
def _mapped_packed_post_fwd_kernel_for(
1776+
def _make_mapped_packed_post_fwd_prim_func(
17421777
batch: int,
17431778
seq: int,
17441779
total_heads: int,
@@ -1752,30 +1787,11 @@ def _mapped_packed_post_fwd_kernel_for(
17521787
v_dim: int,
17531788
projected_dim: int,
17541789
carrier_dtype: str,
1755-
*,
1756-
return_msl: bool = False,
1757-
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1758-
if k_dim >= _MAPPED_PACKED_FWD_K_PARALLEL_MIN_K and max(k_dim, v_dim) <= 256:
1759-
return _mapped_packed_post_fwd_k_parallel_kernel_for(
1760-
batch,
1761-
seq,
1762-
total_heads,
1763-
q_heads,
1764-
k_heads,
1765-
v_heads,
1766-
g_heads,
1767-
w_heads,
1768-
f_heads,
1769-
k_dim,
1770-
v_dim,
1771-
projected_dim,
1772-
carrier_dtype,
1773-
return_msl=return_msl,
1774-
)
1790+
) -> Any:
1791+
"""Build the raw mapped packed-post PrimFunc before lowering."""
17751792

17761793
import tilelang.language as T
17771794

1778-
del return_msl
17791795
conv_dim = _mapped_conv_dim(q_heads, k_heads, v_heads, k_dim, v_dim)
17801796
k_offset = q_heads * k_dim
17811797
v_offset = k_offset + k_heads * k_dim
@@ -1891,7 +1907,7 @@ def fwd(
18911907
carrier_dtype,
18921908
)
18931909

1894-
fwd = _with_global_symbol(
1910+
return _with_global_symbol(
18951911
fwd,
18961912
"m2rnn_mapped_packed_post_fwd",
18971913
batch,
@@ -1908,6 +1924,60 @@ def fwd(
19081924
projected_dim,
19091925
carrier_dtype,
19101926
)
1927+
1928+
1929+
@lru_cache(maxsize=128)
1930+
def _mapped_packed_post_fwd_kernel_for(
1931+
batch: int,
1932+
seq: int,
1933+
total_heads: int,
1934+
q_heads: int,
1935+
k_heads: int,
1936+
v_heads: int,
1937+
g_heads: int,
1938+
w_heads: int,
1939+
f_heads: int,
1940+
k_dim: int,
1941+
v_dim: int,
1942+
projected_dim: int,
1943+
carrier_dtype: str,
1944+
*,
1945+
return_msl: bool = False,
1946+
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1947+
if k_dim >= _MAPPED_PACKED_FWD_K_PARALLEL_MIN_K and max(k_dim, v_dim) <= 256:
1948+
return _mapped_packed_post_fwd_k_parallel_kernel_for(
1949+
batch,
1950+
seq,
1951+
total_heads,
1952+
q_heads,
1953+
k_heads,
1954+
v_heads,
1955+
g_heads,
1956+
w_heads,
1957+
f_heads,
1958+
k_dim,
1959+
v_dim,
1960+
projected_dim,
1961+
carrier_dtype,
1962+
return_msl=return_msl,
1963+
)
1964+
1965+
del return_msl
1966+
fwd = _make_mapped_packed_post_fwd_prim_func(
1967+
batch,
1968+
seq,
1969+
total_heads,
1970+
q_heads,
1971+
k_heads,
1972+
v_heads,
1973+
g_heads,
1974+
w_heads,
1975+
f_heads,
1976+
k_dim,
1977+
v_dim,
1978+
projected_dim,
1979+
carrier_dtype,
1980+
)
19111981
kernel, lowering = _compile_native_metal_kernel_with_lowering(
19121982
fwd,
19131983
out_idx=list(_PACKED_POST_FWD_OUTPUT_IDX),

‎cppmega_mlx/nn/_tilelang/mamba3_path_c.py‎

Lines changed: 96 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -887,7 +887,7 @@ def _bwd_scan_plan_for(
887887

888888

889889
@lru_cache(maxsize=128)
890-
def _fwd_kernel_for(
890+
def _make_fwd_prim_func(
891891
BATCH: int,
892892
SEQ: int,
893893
HEADS: int,
@@ -903,10 +903,8 @@ def _fwd_kernel_for(
903903
h0_dtype: str = "float32",
904904
y_dtype: str = "float32",
905905
h_last_dtype: str = "float32",
906-
*,
907-
return_msl: bool = False,
908-
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
909-
"""Build & cache the Path C TileLang fwd kernel for a given (B, T, H, P, N)."""
906+
) -> Any:
907+
"""Build the raw Path C Mamba3 forward PrimFunc before lowering."""
910908

911909
import tilelang.language as T
912910

@@ -942,7 +940,7 @@ def fwd(
942940
y: T.Tensor((BATCH, SEQ, HEADS, HEADDIM), y_dtype),
943941
h_last: T.Tensor((BATCH, HEADS, HEADDIM, STATE), h_last_dtype),
944942
):
945-
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
943+
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
946944
# Metal already exposes the absolute 1-D lane id. Use it directly
947945
# so the lowered hot loop matches Path B and does not repeatedly
948946
# carry blockIdx * THREADS + threadIdx address arithmetic.
@@ -981,6 +979,50 @@ def fwd(
981979
for n in T.serial(STATE):
982980
h_last[b, h, p, n] = T.cast(h_state[n], h_last_dtype)
983981

982+
return fwd
983+
984+
985+
@lru_cache(maxsize=128)
986+
def _fwd_kernel_for(
987+
BATCH: int,
988+
SEQ: int,
989+
HEADS: int,
990+
HEADDIM: int,
991+
STATE: int,
992+
x_dtype: str = "float32",
993+
B_dtype: str = "float32",
994+
C_dtype: str = "float32",
995+
z_dtype: str = "float32",
996+
A_dtype: str = "float32",
997+
dt_dtype: str = "float32",
998+
D_dtype: str = "float32",
999+
h0_dtype: str = "float32",
1000+
y_dtype: str = "float32",
1001+
h_last_dtype: str = "float32",
1002+
*,
1003+
return_msl: bool = False,
1004+
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1005+
"""Build & cache the Path C TileLang fwd kernel for a given (B, T, H, P, N)."""
1006+
1007+
del return_msl
1008+
fwd = _make_fwd_prim_func(
1009+
BATCH,
1010+
SEQ,
1011+
HEADS,
1012+
HEADDIM,
1013+
STATE,
1014+
x_dtype,
1015+
B_dtype,
1016+
C_dtype,
1017+
z_dtype,
1018+
A_dtype,
1019+
dt_dtype,
1020+
D_dtype,
1021+
h0_dtype,
1022+
y_dtype,
1023+
h_last_dtype,
1024+
)
1025+
9841026
artifact = dispatch_lower(fwd, target="metal", return_msl=True)
9851027
if hasattr(artifact, "_tilelang_engine_target"):
9861028
raise MSLDispatchUnsupported("Mamba3 Path C requires TileLang MSL extraction metadata")
@@ -1005,7 +1047,7 @@ def fwd(
10051047

10061048

10071049
@lru_cache(maxsize=128)
1008-
def _fwd_with_snapshots_kernel_for(
1050+
def _make_fwd_with_snapshots_prim_func(
10091051
BATCH: int,
10101052
SEQ: int,
10111053
HEADS: int,
@@ -1022,8 +1064,8 @@ def _fwd_with_snapshots_kernel_for(
10221064
y_dtype: str = "float32",
10231065
h_last_dtype: str = "float32",
10241066
h_snap_dtype: str = "float32",
1025-
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1026-
"""Build Path C fwd that also materializes per-step states for training."""
1067+
) -> Any:
1068+
"""Build the raw Path C Mamba3 training PrimFunc before lowering."""
10271069

10281070
import tilelang.language as T
10291071

@@ -1101,6 +1143,49 @@ def fwd_with_snapshots(
11011143
for n in T.serial(STATE):
11021144
h_last[b, h, p, n] = T.cast(h_state[n], h_last_dtype)
11031145

1146+
return fwd_with_snapshots
1147+
1148+
1149+
@lru_cache(maxsize=128)
1150+
def _fwd_with_snapshots_kernel_for(
1151+
BATCH: int,
1152+
SEQ: int,
1153+
HEADS: int,
1154+
HEADDIM: int,
1155+
STATE: int,
1156+
x_dtype: str = "float32",
1157+
B_dtype: str = "float32",
1158+
C_dtype: str = "float32",
1159+
z_dtype: str = "float32",
1160+
A_dtype: str = "float32",
1161+
dt_dtype: str = "float32",
1162+
D_dtype: str = "float32",
1163+
h0_dtype: str = "float32",
1164+
y_dtype: str = "float32",
1165+
h_last_dtype: str = "float32",
1166+
h_snap_dtype: str = "float32",
1167+
) -> tuple[Any, _msl_transform.TileLangMSLLowering]:
1168+
"""Build Path C fwd that also materializes per-step states for training."""
1169+
1170+
fwd_with_snapshots = _make_fwd_with_snapshots_prim_func(
1171+
BATCH,
1172+
SEQ,
1173+
HEADS,
1174+
HEADDIM,
1175+
STATE,
1176+
x_dtype,
1177+
B_dtype,
1178+
C_dtype,
1179+
z_dtype,
1180+
A_dtype,
1181+
dt_dtype,
1182+
D_dtype,
1183+
h0_dtype,
1184+
y_dtype,
1185+
h_last_dtype,
1186+
h_snap_dtype,
1187+
)
1188+
11041189
artifact = dispatch_lower(fwd_with_snapshots, target="metal", return_msl=True)
11051190
if hasattr(artifact, "_tilelang_engine_target"):
11061191
raise MSLDispatchUnsupported("Mamba3 Path C requires TileLang MSL extraction metadata")
@@ -1176,7 +1261,6 @@ def snapshots(
11761261
with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
11771262
global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
11781263
h_state = T.alloc_local((STATE,), accum_dtype)
1179-
dh = T.alloc_local((STATE,), accum_dtype)
11801264
if LANES % THREADS == 0 or global_lane < LANES:
11811265
p = global_lane % HEADDIM
11821266
if BATCH == 1:
@@ -1258,6 +1342,7 @@ def _bwd_simd_reduce_kernel_for(
12581342
A_dtype: str = "float32",
12591343
dt_dtype: str = "float32",
12601344
D_dtype: str = "float32",
1345+
h0_dtype: str = "float32",
12611346
dx_dtype: str = "float32",
12621347
dz_dtype: str = "float32",
12631348
dB_dtype: str = "float32",
@@ -1322,6 +1407,7 @@ def bwd_simd(
13221407
tid = T.get_thread_binding(0)
13231408
global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
13241409
h_state = T.alloc_local((STATE,), accum_dtype)
1410+
dh = T.alloc_local((STATE,), accum_dtype)
13251411
if LANES % THREADS == 0 or global_lane < LANES:
13261412
p = global_lane % HEADDIM
13271413
reduce_lane = tid % HEADDIM

0 commit comments

Comments
 (0)