Skip to content

Commit 3847f1e

Browse files
author
103yiran
committed
style(ascend): format code
1 parent 891a00d commit 3847f1e

5 files changed

Lines changed: 3816 additions & 1244 deletions

File tree

‎src/flag_gems/runtime/backend/_ascend/fused/mhc/mhc_post.py‎

Lines changed: 103 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757

5858
try:
5959
import triton.experimental.tle as tle
60+
6061
_HAS_DSA = hasattr(tle, "dsa") and hasattr(tle.dsa, "alloc")
6162
except (ImportError, AttributeError):
6263
tle = None
@@ -101,10 +102,10 @@ def _mhc_post_kernel_tle(
101102
hp = tle.dsa.to_tensor(hpost_ub) # [4] f32
102103

103104
hr_vec = tle.dsa.to_tensor(hres_ub) # [16] f32
104-
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0, ), (4, ), (1, )), [4])
105-
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4, ), (4, ), (1, )), [4])
106-
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8, ), (4, ), (1, )), [4])
107-
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12, ), (4, ), (1, )), [4])
105+
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0,), (4,), (1,)), [4])
106+
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4,), (4,), (1,)), [4])
107+
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8,), (4,), (1,)), [4])
108+
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12,), (4,), (1,)), [4])
108109

109110
# ---- Allocate double-buffered UB for x[4,BLOCK_D] and h_out[BLOCK_D] ------
110111
ho_ub = tle.dsa.alloc([BLOCK_D], dtype=x_dt, mem_addr_space=tle.dsa.ascend.UB)
@@ -125,10 +126,18 @@ def _mhc_post_kernel_tle(
125126
# -- Compute stage: parallel FMA across 4 heads -----------------------
126127
ho = tle.dsa.to_tensor(ho_ub).to(tl.float32) # [BLOCK_D]
127128
x2d = tle.dsa.to_tensor(x_ub).to(tl.float32) # [4, BLOCK_D]
128-
x0 = tl.reshape(tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
129-
x1 = tl.reshape(tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
130-
x2 = tl.reshape(tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
131-
x3 = tl.reshape(tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
129+
x0 = tl.reshape(
130+
tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
131+
)
132+
x1 = tl.reshape(
133+
tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
134+
)
135+
x2 = tl.reshape(
136+
tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
137+
)
138+
x3 = tl.reshape(
139+
tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
140+
)
132141

133142
with tle.dsa.hint(inter_no_alias=True):
134143
Y = tl.expand_dims(hp, 1) * tl.expand_dims(ho, 0) # [4, BLOCK_D]
@@ -138,10 +147,18 @@ def _mhc_post_kernel_tle(
138147
Y += tl.expand_dims(hr3, 1) * tl.expand_dims(x3, 0)
139148

140149
# -- Store: to_buffer + dsa.copy (matching AscendC CopyOutTile) -------
141-
y0 = tl.reshape(tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
142-
y1 = tl.reshape(tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
143-
y2 = tl.reshape(tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
144-
y3 = tl.reshape(tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
150+
y0 = tl.reshape(
151+
tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
152+
)
153+
y1 = tl.reshape(
154+
tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
155+
)
156+
y2 = tl.reshape(
157+
tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
158+
)
159+
y3 = tl.reshape(
160+
tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
161+
)
145162

146163
y0b = tle.dsa.to_buffer(y0.to(x_dt), tle.dsa.ascend.UB)
147164
y1b = tle.dsa.to_buffer(y1.to(x_dt), tle.dsa.ascend.UB)
@@ -200,10 +217,18 @@ def _mhc_post_kernel_tle_rows(
200217
# Row j of h_res[j,i] contains the coefficients of x[j] for all 4 output heads.
201218
# Extract contiguous [4] slices — one per input x row:
202219
hr_vec = tle.dsa.to_tensor(hres_ub) # [16] f32
203-
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0, ), (4, ), (1, )), [4]) # [r00,r01,r02,r03]
204-
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4, ), (4, ), (1, )), [4]) # [r10,r11,r12,r13]
205-
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8, ), (4, ), (1, )), [4]) # [r20,r21,r22,r23]
206-
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12, ), (4, ), (1, )), [4]) # [r30,r31,r32,r33]
220+
hr0 = tl.reshape(
221+
tle.dsa.extract_slice(hr_vec, (0,), (4,), (1,)), [4]
222+
) # [r00,r01,r02,r03]
223+
hr1 = tl.reshape(
224+
tle.dsa.extract_slice(hr_vec, (4,), (4,), (1,)), [4]
225+
) # [r10,r11,r12,r13]
226+
hr2 = tl.reshape(
227+
tle.dsa.extract_slice(hr_vec, (8,), (4,), (1,)), [4]
228+
) # [r20,r21,r22,r23]
229+
hr3 = tl.reshape(
230+
tle.dsa.extract_slice(hr_vec, (12,), (4,), (1,)), [4]
231+
) # [r30,r31,r32,r33]
207232

208233
# ---- Load x, h_out with mask (handles non-power-of-2 D correctly) ---------
209234
ho = tl.load(h_out_ptr + hout_base + d_off, mask=d_mask, other=0.0).to(tl.float32)
@@ -276,10 +301,10 @@ def _mhc_post_kernel_tle_rows_pipeline(
276301
hp = tle.dsa.to_tensor(hpost_ub) # [4] f32
277302

278303
hr_vec = tle.dsa.to_tensor(hres_ub) # [16] f32
279-
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0, ), (4, ), (1, )), [4])
280-
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4, ), (4, ), (1, )), [4])
281-
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8, ), (4, ), (1, )), [4])
282-
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12, ), (4, ), (1, )), [4])
304+
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0,), (4,), (1,)), [4])
305+
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4,), (4,), (1,)), [4])
306+
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8,), (4,), (1,)), [4])
307+
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12,), (4,), (1,)), [4])
283308

284309
# ---- Allocate double-buffered UB for x[4,BLOCK_D] and h_out[BLOCK_D] ------
285310
# tle.dsa.pipeline with num_stages=2 enables MTE2/V overlap via double buffer.
@@ -304,10 +329,18 @@ def _mhc_post_kernel_tle_rows_pipeline(
304329
# -- Compute stage: parallel FMA across 4 heads -----------------------
305330
ho = tle.dsa.to_tensor(ho_ub).to(tl.float32) # [BLOCK_D]
306331
x2d = tle.dsa.to_tensor(x_ub).to(tl.float32) # [4, BLOCK_D]
307-
x0 = tl.reshape(tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
308-
x1 = tl.reshape(tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
309-
x2 = tl.reshape(tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
310-
x3 = tl.reshape(tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
332+
x0 = tl.reshape(
333+
tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
334+
)
335+
x1 = tl.reshape(
336+
tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
337+
)
338+
x2 = tl.reshape(
339+
tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
340+
)
341+
x3 = tl.reshape(
342+
tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
343+
)
311344

312345
# Compute all 5 terms separately, then sum once
313346
term0 = tl.expand_dims(hp, 1) * tl.expand_dims(ho, 0) # [4, BLOCK_D]
@@ -318,10 +351,18 @@ def _mhc_post_kernel_tle_rows_pipeline(
318351
Y = term0 + term1 + term2 + term3 + term4 # [4, BLOCK_D]
319352

320353
# -- Extract rows and store back to GM --------------------------------
321-
y0 = tl.reshape(tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
322-
y1 = tl.reshape(tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
323-
y2 = tl.reshape(tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
324-
y3 = tl.reshape(tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
354+
y0 = tl.reshape(
355+
tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
356+
)
357+
y1 = tl.reshape(
358+
tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
359+
)
360+
y2 = tl.reshape(
361+
tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
362+
)
363+
y3 = tl.reshape(
364+
tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
365+
)
325366

326367
tl.store(out_ptr + x_base + 0 * D + d_off, y0.to(x_dt), mask=d_mask)
327368
tl.store(out_ptr + x_base + 1 * D + d_off, y1.to(x_dt), mask=d_mask)
@@ -368,10 +409,10 @@ def _mhc_post_kernel_tle_concat_reduce(
368409
# ---- Extract coefficient vectors (keep in FP32) --------------------------------------
369410
hp = tle.dsa.to_tensor(hpost_ub) # [4] f32
370411
hr_vec = tle.dsa.to_tensor(hres_ub) # [16] f32
371-
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0, ), (4, ), (1, )), [4])
372-
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4, ), (4, ), (1, )), [4])
373-
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8, ), (4, ), (1, )), [4])
374-
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12, ), (4, ), (1, )), [4])
412+
hr0 = tl.reshape(tle.dsa.extract_slice(hr_vec, (0,), (4,), (1,)), [4])
413+
hr1 = tl.reshape(tle.dsa.extract_slice(hr_vec, (4,), (4,), (1,)), [4])
414+
hr2 = tl.reshape(tle.dsa.extract_slice(hr_vec, (8,), (4,), (1,)), [4])
415+
hr3 = tl.reshape(tle.dsa.extract_slice(hr_vec, (12,), (4,), (1,)), [4])
375416

376417
# ---- Allocate double-buffered UB ----------------------------------------
377418
ho_ub = tle.dsa.alloc([BLOCK_D], dtype=x_dt, mem_addr_space=tle.dsa.ascend.UB)
@@ -393,10 +434,18 @@ def _mhc_post_kernel_tle_concat_reduce(
393434
# -- Compute: compute 5 separate outer products [4, BLOCK_D] each in FP32 ---
394435
ho = tle.dsa.to_tensor(ho_ub).to(tl.float32) # [BLOCK_D]
395436
x2d = tle.dsa.to_tensor(x_ub).to(tl.float32) # [4, BLOCK_D]
396-
x0 = tl.reshape(tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
397-
x1 = tl.reshape(tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
398-
x2 = tl.reshape(tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
399-
x3 = tl.reshape(tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
437+
x0 = tl.reshape(
438+
tle.dsa.extract_slice(x2d, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
439+
)
440+
x1 = tl.reshape(
441+
tle.dsa.extract_slice(x2d, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
442+
)
443+
x2 = tl.reshape(
444+
tle.dsa.extract_slice(x2d, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
445+
)
446+
x3 = tl.reshape(
447+
tle.dsa.extract_slice(x2d, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
448+
)
400449

401450
# Compute 5 terms, each is [4, BLOCK_D] in FP32
402451
term0 = tl.expand_dims(hp, 1) * tl.expand_dims(ho, 0) # [4, BLOCK_D]
@@ -409,10 +458,18 @@ def _mhc_post_kernel_tle_concat_reduce(
409458
Y = term0 + term1 + term2 + term3 + term4 # [4, BLOCK_D]
410459

411460
# -- Extract rows, cast once, and store ---
412-
y0 = tl.reshape(tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
413-
y1 = tl.reshape(tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
414-
y2 = tl.reshape(tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
415-
y3 = tl.reshape(tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D])
461+
y0 = tl.reshape(
462+
tle.dsa.extract_slice(Y, (0, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
463+
)
464+
y1 = tl.reshape(
465+
tle.dsa.extract_slice(Y, (1, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
466+
)
467+
y2 = tl.reshape(
468+
tle.dsa.extract_slice(Y, (2, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
469+
)
470+
y3 = tl.reshape(
471+
tle.dsa.extract_slice(Y, (3, 0), (1, BLOCK_D), (1, 1)), [BLOCK_D]
472+
)
416473

417474
tl.store(out_ptr + x_base + 0 * D + d_off, y0.to(x_dt), mask=d_mask)
418475
tl.store(out_ptr + x_base + 1 * D + d_off, y1.to(x_dt), mask=d_mask)
@@ -465,8 +522,10 @@ def mhc_post(
465522
Takes precedence over use_pipeline if both are True.
466523
"""
467524
if not _HAS_DSA:
468-
raise RuntimeError("This mhc_post implementation requires the TLE DSA surface "
469-
"(triton.experimental.tle.dsa.*).")
525+
raise RuntimeError(
526+
"This mhc_post implementation requires the TLE DSA surface "
527+
"(triton.experimental.tle.dsa.*)."
528+
)
470529

471530
xf, hres, hout, hpost, shape = _flatten_bsn(x, h_res, h_out, h_post)
472531
T, N, D = xf.shape
@@ -483,7 +542,7 @@ def mhc_post(
483542

484543
if use_concat_reduce:
485544
# Concat+reduce variant: 1D grid over T
486-
grid = (T, )
545+
grid = (T,)
487546
_mhc_post_kernel_tle_concat_reduce[grid](
488547
xf,
489548
hres.to(torch.float32),
@@ -497,7 +556,7 @@ def mhc_post(
497556
elif use_pipeline:
498557
# Pipeline kernel: 1D grid over T, iterates D-chunks internally
499558
# with tle.dsa.pipeline(num_stages=2) for MTE2/Vector overlap.
500-
grid = (T, )
559+
grid = (T,)
501560
_mhc_post_kernel_tle_rows_pipeline[grid](
502561
xf,
503562
hres.to(torch.float32),

0 commit comments

Comments
 (0)