Star20
31
代码
代码
Issues5
Pull Requests2
项目讨论
Wiki
分析
Star20
31
已合并
perf: conv1d_and_silu consume/produce [B,S,C] directly, drop transposes #476
Igor Prusov创建于 6 天前
perf: conv1d_and_silu consume/produce [B,S,C] directly, drop transposes #476
已合并
Igor Prusov创建于 6 天前
从igor-prusov/GVirteliminate_transpose合入到openeuler/GVirtmaster
共 6 个文件变更+377-352
Mxlite/csrc/kernels/conv1d_and_silu.h+340-298
@@ -8,22 +8,20 @@
8 * if updateState: state = concat[..., -K:]8 * if updateState: state = concat[..., -K:]
9 *9 *
10 * No GM workspace concat. State/input are loaded into aligned UB float10 * No GM workspace concat. State/input are loaded into aligned UB float
11- * buffers. Outputs whose window lies entirely inside input (i >= K-1) are11+ * buffers.
12- * produced in full-width 64-lane blocks: each of the K taps is fetched with a
13- * vgather (which accepts arbitrary byte-offset bases, unlike plain vector ops
14- * that need 32B-aligned sources) and accumulated with vmuls+vadd, so 64
15- * outputs cost ~2K+ vector ops with fully-utilized lanes. Only the first K-1
16- * outputs (window straddling state and input) are computed via scalar moves.
17- * GM copies use CopyGmToUbufAligned / CopyUbufToGmAligned so lengths need not
18- * be 32B-aligned.
19- *
20- * Limits: kernelDim <= 16, seqLen (or per-request lens) <= 4096.
21 *12 *
22 * Layouts:13 * Layouts:
23- * Uniform (seqLen != 0): input/output [B, C, S], state [B, C, K]14+ * Uniform (seqLen != 0): input/output [B, S, C] (channel is the inner,
24- * Packed (seqLen == 0): input/output [T, C] token-major, state still [B, C, K];15+ * contiguous dim), state [B, C, K], weight [C, 1, K].
25- * per-request tokens are input[queryStartLoc[b] : queryStartLoc[b]+queryLens[b]].16+ * Packed (seqLen == 0): input/output [T, C] token-major, state still [B, C, K].
26- * Host sets seqLen=0 for packed; do not test GM pointers against nullptr.17+ *
18+ * The 64-lane (vectorized) dimension is the CHANNEL, not the sequence. Each tap
19+ * of the K-tap convolution is a contiguous, 32B-aligned 64-channel vector load
20+ * (vmul + vadd), instead of the per-tap vgather the [B,C,S] layout required.
21+ * Consuming/producing [B,S,C] directly also removes the two [B,S,C]<->[B,C,S]
22+ * Transpose passes around the kernel.
23+ *
24+ * Limits: kernelDim <= 16, seqLen (or per-request lens) <= 4096.
27 */25 */
28#pragma once26#pragma once
29#include "kernel_operator.h"27#include "kernel_operator.h"
@@ -41,9 +39,9 @@ public:
41 static constexpr int kMaxKernel = 16;39 static constexpr int kMaxKernel = 16;
42 static constexpr int kMaxInputF = 4096;40 static constexpr int kMaxInputF = 4096;
43 static constexpr int kMaxBatchMeta = 256;41 static constexpr int kMaxBatchMeta = 256;
44- // Extra pad so full-width (64-lane) tail reads stay inside input_f.42+ static constexpr int kBlock = VECTOR_MAX_NUM_OF_FP32; // 64 lanes
45- static constexpr int kGatherPad = 128;43+ static constexpr int kTile = 128;
46- static constexpr int kBlock = VECTOR_MAX_NUM_OF_FP32;44+ static constexpr int kWin = kMaxKernel - 1 + kTile;
47 45 
48 __aicore__ inline XliteCausalConv1dSiLU()46 __aicore__ inline XliteCausalConv1dSiLU()
49 {47 {
@@ -70,44 +68,63 @@ public:
70 this->updateState = updateState;68 this->updateState = updateState;
71 69 
72 uint64_t off = 0;70 uint64_t off = 0;
73- // Aligned staging for GM<->UB dtype copies (must stay at offset 0 of a 32B region).71+ // Aligned staging for GM<->UB dtype conversion round trips.
74- // Sized for a single fp16/bf16 conversion round trip of the whole input.
75 stage_buf = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);72 stage_buf = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);
76- off += kMaxInputF * sizeof(Dtype); // 8KB (fp16/bf16)73+ off += kMaxInputF * sizeof(Dtype);
77- tmp_kernel_buf = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);
78- off += 8 * 32;
79 74 
80- // fp32 weights (K floats) read into scalar registers once per channel.75+ w_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
81- kernel_buf = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);76+ off += kBlock * kMaxKernel * sizeof(float);
82- off += 8 * 32;
83- 
84- state_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
85- off += kMaxKernel * sizeof(float); // 64B, keep 32B-aligned
86 if (off % 32 != 0) {77 if (off % 32 != 0) {
87 off = (off + 31) / 32 * 32;78 off = (off + 31) / 32 * 32;
88 }79 }
89- input_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);80+ w_reorg = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
90- // Extra pad so full-width tap reads stay in bounds at the sequence tail.81+ off += kBlock * kMaxKernel * sizeof(float);
91- off += (kMaxInputF + kGatherPad) * sizeof(float);82+ if (off % 32 != 0) {
83+ off = (off + 31) / 32 * 32;
84+ }
85+ state_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
86+ off += kBlock * kMaxKernel * sizeof(float);
87+ if (off % 32 != 0) {
88+ off = (off + 31) / 32 * 32;
89+ }
90+ state_reorg = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
91+ off += kBlock * kMaxKernel * sizeof(float);
92+ if (off % 32 != 0) {
93+ off = (off + 31) / 32 * 32;
94+ }
95+ 
96+ if constexpr (std::is_same<Dtype, float>::value) {
97+ win_raw = nullptr;
98+ } else {
99+ win_raw = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);
100+ off += kWin * kBlock * sizeof(Dtype);
101+ if (off % 32 != 0) {
102+ off = (off + 31) / 32 * 32;
103+ }
104+ }
105+ window_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
106+ off += kWin * kBlock * sizeof(float);
107+ if (off % 32 != 0) {
108+ off = (off + 31) / 32 * 32;
109+ }
110+ 
111+ state_win = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
112+ off += kMaxKernel * kBlock * sizeof(float);
92 if (off % 32 != 0) {113 if (off % 32 != 0) {
93 off = (off + 31) / 32 * 32;114 off = (off + 31) / 32 * 32;
94 }115 }
95 new_state_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);116 new_state_f = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
96- off += kMaxKernel * sizeof(float);117+ off += kMaxKernel * kBlock * sizeof(float);
97 if (off % 32 != 0) {118 if (off % 32 != 0) {
98 off = (off + 31) / 32 * 32;119 off = (off + 31) / 32 * 32;
99 }120 }
100 121 
101- off_ramp = reinterpret_cast<__ubuf__ uint32_t *>((uintptr_t)off);
102- off += 8 * 32;
103- qkv_tmp = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
104- off += 8 * 32;
105 acc_buf = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);122 acc_buf = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
106 off += 8 * 32;123 off += 8 * 32;
107 calc_buf = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);124 calc_buf = reinterpret_cast<__ubuf__ float *>((uintptr_t)off);
108 off += 8 * 32;125 off += 8 * 32;
109- out_buf = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);126+ out_tile = reinterpret_cast<__ubuf__ Dtype *>((uintptr_t)off);
110- off += 64 * sizeof(Dtype);127+ off += kTile * kBlock * sizeof(Dtype);
111 if (off % 32 != 0) {128 if (off % 32 != 0) {
112 off = (off + 31) / 32 * 32;129 off = (off + 31) / 32 * 32;
113 }130 }
@@ -116,26 +133,62 @@ public:
116 meta_lens = reinterpret_cast<__ubuf__ int32_t *>((uintptr_t)off);133 meta_lens = reinterpret_cast<__ubuf__ int32_t *>((uintptr_t)off);
117 }134 }
118 135 
119- __aicore__ inline float ReadFloat(__ubuf__ float *buf, int idx)136+ __aicore__ inline void SetLaneMask(int w)
137+ {
138+ if (w >= kBlock) {
139+ set_vector_mask((uint64_t)-1, (uint64_t)-1);
140+ } else {
141+ SetMask(w);
142+ }
143+ }
144+ 
145+ // pipe_barrier(PIPE_X) alone does NOT order UB data produced/consumed by
146+ // *different* pipes (MTE2/V/MTE3/S); explicit event flags are required.
147+ // EVENT_ID0: MTE2 <-> V and V <-> MTE3 for the window/out_tile pipeline.
148+ // EVENT_ID1: scalar (S) pipe fences.
149+ // EVENT_ID2: stage_buf MTE2 <-> V inside LoadGmToFloat/StoreFloatToGm.
150+ __aicore__ inline void FlagMTE2V()
151+ {
152+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
153+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
154+ }
155+ __aicore__ inline void FlagVMTE2()
156+ {
157+ set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
158+ wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
159+ }
160+ __aicore__ inline void FlagVMTE3()
161+ {
162+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
163+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
164+ }
165+ __aicore__ inline void FlagMTE3V()
166+ {
167+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
168+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
169+ }
170+ __aicore__ inline void FlagMTE2S()
171+ {
172+ set_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);
173+ wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);
174+ }
175+ __aicore__ inline void FlagVS()
120 {176 {
121 set_flag(PIPE_V, PIPE_S, EVENT_ID1);177 set_flag(PIPE_V, PIPE_S, EVENT_ID1);
122 wait_flag(PIPE_V, PIPE_S, EVENT_ID1);178 wait_flag(PIPE_V, PIPE_S, EVENT_ID1);
123- float val = buf[idx];
124- set_flag(PIPE_S, PIPE_V, EVENT_ID1);
125- wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
126- return val;
127 }179 }
128- 180+ __aicore__ inline void FlagSV()
129- __aicore__ inline void WriteFloat(__ubuf__ float *buf, int idx, float val)
130 {181 {
131- set_flag(PIPE_V, PIPE_S, EVENT_ID1);
132- wait_flag(PIPE_V, PIPE_S, EVENT_ID1);
133- buf[idx] = val;
134 set_flag(PIPE_S, PIPE_V, EVENT_ID1);182 set_flag(PIPE_S, PIPE_V, EVENT_ID1);
135 wait_flag(PIPE_S, PIPE_V, EVENT_ID1);183 wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
136 }184 }
185+ __aicore__ inline void FlagSMTE3()
186+ {
187+ set_flag(PIPE_S, PIPE_MTE3, EVENT_ID1);
188+ wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID1);
189+ }
137 190 
138- __aicore__ inline void SiLU()191+ __aicore__ inline void SiLU(__ubuf__ Dtype *dst)
139 {192 {
140 vmuls(calc_buf, acc_buf, (float)-1.0, 1, 1, 1, 8, 8);193 vmuls(calc_buf, acc_buf, (float)-1.0, 1, 1, 1, 8, 8);
141 pipe_barrier(PIPE_V);194 pipe_barrier(PIPE_V);
@@ -144,21 +197,18 @@ public:
144 vadds(calc_buf, calc_buf, (float)1.0, 1, 1, 1, 8, 8);197 vadds(calc_buf, calc_buf, (float)1.0, 1, 1, 1, 8, 8);
145 pipe_barrier(PIPE_V);198 pipe_barrier(PIPE_V);
146 199 
147- wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
148 if constexpr (std::is_same<Dtype, float>::value) {200 if constexpr (std::is_same<Dtype, float>::value) {
149- vdiv(out_buf, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);201+ vdiv(dst, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);
150 } else {202 } else {
151 vdiv(calc_buf, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);203 vdiv(calc_buf, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);
152 pipe_barrier(PIPE_V);204 pipe_barrier(PIPE_V);
153 if constexpr (std::is_same<Dtype, float16_t>::value) {205 if constexpr (std::is_same<Dtype, float16_t>::value) {
154- vconv_f322f16(out_buf, calc_buf, 1, 1, 1, 4, 8);206+ vconv_f322f16(dst, calc_buf, 1, 1, 1, 4, 8);
155 } else {207 } else {
156- vconv_f322bf16r(out_buf, calc_buf, 1, 1, 1, 4, 8);208+ vconv_f322bf16r(dst, calc_buf, 1, 1, 1, 4, 8);
157 }209 }
158 }210 }
159 pipe_barrier(PIPE_V);211 pipe_barrier(PIPE_V);
160- set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
161- wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
162 }212 }
163 213 
164 // Load nElem dtype values from GM into dstF (float). One round trip per214 // Load nElem dtype values from GM into dstF (float). One round trip per
@@ -182,9 +232,6 @@ public:
182 pipe_barrier(PIPE_MTE2);232 pipe_barrier(PIPE_MTE2);
183 set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);233 set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
184 wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);234 wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
185- 
186- // The 128-bit vector mask cycles across repeats (64 lanes each);
187- // SetMask only covers up to two repeats, so use a full mask above.
188 if (take >= 2 * VECTOR_MAX_NUM_OF_FP32) {235 if (take >= 2 * VECTOR_MAX_NUM_OF_FP32) {
189 set_vector_mask((uint64_t)-1, (uint64_t)-1);236 set_vector_mask((uint64_t)-1, (uint64_t)-1);
190 } else {237 } else {
@@ -197,6 +244,8 @@ public:
197 vconv_bf162f32(dstF + done, stage_buf, repeat, 1, 1, 8, 4);244 vconv_bf162f32(dstF + done, stage_buf, repeat, 1, 1, 8, 4);
198 }245 }
199 pipe_barrier(PIPE_V);246 pipe_barrier(PIPE_V);
247+ set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2);
248+ wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2);
200 done += take;249 done += take;
201 }250 }
202 set_vector_mask((uint64_t)-1, (uint64_t)-1);251 set_vector_mask((uint64_t)-1, (uint64_t)-1);
@@ -242,6 +291,216 @@ public:
242 set_vector_mask((uint64_t)-1, (uint64_t)-1);291 set_vector_mask((uint64_t)-1, (uint64_t)-1);
243 }292 }
244 293 
294+ // Load a [nPos, w] tile of channel-block c0 (w channels starting at c0) from
295+ // GM positions [tokenBase+loPos, tokenBase+loPos+nPos) into UB at row dstIdx.
296+ // Each UB row is kBlock elements wide. When both the row and the channel
297+ // width are 32B-aligned, use one multi-burst DMA; otherwise fall back to a
298+ // per-position copy.
299+ __aicore__ inline void LoadTileInto(__ubuf__ Dtype *dst, int tokenBase, int c0, int loPos,
300+ int nPos, int dstIdx)
301+ {
302+ int C = (int)channels;
303+ int w = MIN(kBlock, C - c0);
304+ __gm__ Dtype *src = input + ((tokenBase + loPos) * C + c0);
305+ int burstBytes = w * (int)sizeof(Dtype);
306+ int burstBlocks = burstBytes / BLOCK_SIZE;
307+ int rowBytes = C * (int)sizeof(Dtype);
308+ if (w == kBlock && burstBlocks * BLOCK_SIZE == burstBytes && rowBytes % BLOCK_SIZE == 0) {
309+ uint64_t cfg =
310+ __set_dmi_config(0, nPos, burstBlocks, rowBytes / BLOCK_SIZE - burstBlocks, 0);
311+ copy_gm_to_ubuf(dst + dstIdx * kBlock, src, cfg);
312+ } else {
313+ for (int i = 0; i < nPos; i++) {
314+ CopyGmToUbufAligned(dst + (dstIdx + i) * kBlock, src + i * C,
315+ static_cast<uint32_t>(burstBytes));
316+ }
317+ }
318+ pipe_barrier(PIPE_MTE2);
319+ }
320+ 
321+ __aicore__ inline void ConvertTile(int nPos, int dstIdx)
322+ {
323+ if constexpr (std::is_same<Dtype, float>::value) {
324+ return;
325+ }
326+ int off = dstIdx * kBlock;
327+ set_vector_mask((uint64_t)-1, (uint64_t)-1);
328+ if constexpr (std::is_same<Dtype, float16_t>::value) {
329+ vconv_f162f32(window_f + off, win_raw + off, nPos, 1, 1, 8, 4);
330+ } else {
331+ vconv_bf162f32(window_f + off, win_raw + off, nPos, 1, 1, 8, 4);
332+ }
333+ pipe_barrier(PIPE_V);
334+ }
335+ 
336+ // Load [w, K] weights and re-layout them to [K, kBlock] so tap j lives at a
337+ // contiguous 64-lane vector (w_reorg + j*kBlock).
338+ __aicore__ inline void LoadWeights(int c0, int w)
339+ {
340+ int K = (int)kernelDim;
341+ LoadGmToFloat(weight + c0 * K, w * K, w_f);
342+ if constexpr (std::is_same<Dtype, float>::value) {
343+ FlagMTE2S();
344+ } else {
345+ FlagVS();
346+ }
347+ for (int l = 0; l < w; l++) {
348+ for (int j = 0; j < K; j++) {
349+ w_reorg[j * kBlock + l] = w_f[l * K + j];
350+ }
351+ }
352+ FlagSV();
353+ }
354+ 
355+ // Load [w, K] state and re-layout it to [K, kBlock].
356+ __aicore__ inline void LoadState(int b, int c0, int w)
357+ {
358+ int K = (int)kernelDim;
359+ int stateBase = (b * (int)channels + c0) * K;
360+ LoadGmToFloat(state + stateBase, w * K, state_f);
361+ if constexpr (std::is_same<Dtype, float>::value) {
362+ FlagMTE2S();
363+ } else {
364+ FlagVS();
365+ }
366+ for (int l = 0; l < w; l++) {
367+ for (int p = 0; p < K; p++) {
368+ state_reorg[p * kBlock + l] = state_f[l * K + p];
369+ }
370+ }
371+ FlagSV();
372+ }
373+ 
374+ // Compute one output position (channel block c0, w lanes) and store the
375+ // SiLU result into out_tile row (s - s0).
376+ __aicore__ inline void ComputePosition(int s, int s0, int w)
377+ {
378+ int K = (int)kernelDim;
379+ SetLaneMask(w);
380+ vector_dup(acc_buf, 0.0f, 1, 1, 1, 8, 8);
381+ pipe_barrier(PIPE_V);
382+ for (int j = 0; j < K; j++) {
383+ int p = s + 1 + j;
384+ __ubuf__ float *tap =
385+ (p < K) ? (state_reorg + p * kBlock) : (window_f + (p - s0 - 1) * kBlock);
386+ vmul(calc_buf, tap, w_reorg + j * kBlock, 1, 1, 1, 1, 8, 8, 8);
387+ pipe_barrier(PIPE_V);
388+ vadd(acc_buf, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);
389+ pipe_barrier(PIPE_V);
390+ }
391+ SiLU(out_tile + (s - s0) * kBlock);
392+ }
393+ 
394+ // Store out_tile rows [0, nPos) to GM positions [tokenBase+s0, tokenBase+s0+nPos)
395+ // with a channel-stride C. Bulk multi-burst DMA when the full block is used.
396+ __aicore__ inline void StoreTile(int tokenBase, int c0, int s0, int nPos, int w)
397+ {
398+ int C = (int)channels;
399+ __gm__ Dtype *dst = output + (tokenBase + s0) * C + c0;
400+ int burstBytes = w * (int)sizeof(Dtype);
401+ int burstBlocks = burstBytes / BLOCK_SIZE;
402+ int rowBlocks = C * (int)sizeof(Dtype) / BLOCK_SIZE;
403+ if (w == kBlock && burstBlocks * BLOCK_SIZE == burstBytes &&
404+ rowBlocks * BLOCK_SIZE == C * (int)sizeof(Dtype)) {
405+ uint64_t cfg = __set_dmi_config(0, nPos, burstBlocks, 0, rowBlocks - burstBlocks);
406+ copy_ubuf_to_gm(dst, out_tile, cfg);
407+ } else {
408+ for (int s = 0; s < nPos; s++) {
409+ CopyUbufToGmAligned(dst + s * C, out_tile + s * kBlock,
410+ static_cast<uint32_t>(burstBytes));
411+ }
412+ }
413+ pipe_barrier(PIPE_MTE3);
414+ }
415+ 
416+ __aicore__ inline void ProcessSequence(int tokenBase, int c0, int w, int S)
417+ {
418+ int K = (int)kernelDim;
419+ for (int s0 = 0; s0 < S; s0 += kTile) {
420+ int loPos = (s0 == 0) ? 0 : (s0 - K + 1);
421+ int hiPos = MIN(s0 + kTile - 1, S - 1);
422+ if (hiPos < loPos) {
423+ break;
424+ }
425+ int nPos = hiPos - loPos + 1;
426+ int dstIdx = (s0 == 0) ? (K - 1) : 0;
427+ 
428+ if constexpr (std::is_same<Dtype, float>::value) {
429+ LoadTileInto(window_f, tokenBase, c0, loPos, nPos, dstIdx);
430+ } else {
431+ LoadTileInto(win_raw, tokenBase, c0, loPos, nPos, dstIdx);
432+ }
433+ FlagMTE2V();
434+ ConvertTile(nPos, dstIdx);
435+ FlagVMTE2();
436+ 
437+ int sEnd = MIN(s0 + kTile, S);
438+ for (int s = s0; s < sEnd; s++) {
439+ ComputePosition(s, s0, w);
440+ }
441+ FlagVMTE3();
442+ StoreTile(tokenBase, c0, s0, sEnd - s0, w);
443+ FlagMTE3V();
444+ }
445+ }
446+ 
447+ // Load the trailing input window used to derive the new state.
448+ __aicore__ inline void LoadStateWin(int tokenBase, int c0, int S)
449+ {
450+ int K = (int)kernelDim;
451+ int firstPos = MAX(0, S - K);
452+ int nPos = S - firstPos;
453+ if (nPos <= 0) {
454+ return;
455+ }
456+ if constexpr (std::is_same<Dtype, float>::value) {
457+ LoadTileInto(state_win, tokenBase, c0, firstPos, nPos, 0);
458+ } else {
459+ LoadTileInto(win_raw, tokenBase, c0, firstPos, nPos, 0);
460+ FlagMTE2V();
461+ set_vector_mask((uint64_t)-1, (uint64_t)-1);
462+ if constexpr (std::is_same<Dtype, float16_t>::value) {
463+ vconv_f162f32(state_win, win_raw, nPos, 1, 1, 8, 4);
464+ } else {
465+ vconv_bf162f32(state_win, win_raw, nPos, 1, 1, 8, 4);
466+ }
467+ pipe_barrier(PIPE_V);
468+ FlagVMTE2();
469+ }
470+ }
471+ 
472+ __aicore__ inline void WriteBackState(int b, int tokenBase, int c0, int w, int S)
473+ {
474+ int K = (int)kernelDim;
475+ int C = (int)channels;
476+ int firstPos = MAX(0, S - K);
477+ 
478+ LoadStateWin(tokenBase, c0, S);
479+ if constexpr (std::is_same<Dtype, float>::value) {
480+ FlagMTE2S();
481+ } else {
482+ FlagVS();
483+ }
484+ for (int t = 0; t < K; t++) {
485+ int concatIdx = S + t;
486+ for (int l = 0; l < w; l++) {
487+ float v;
488+ if (concatIdx < K) {
489+ v = state_reorg[concatIdx * kBlock + l];
490+ } else {
491+ v = state_win[(concatIdx - K - firstPos) * kBlock + l];
492+ }
493+ new_state_f[l * K + t] = v;
494+ }
495+ }
496+ if constexpr (std::is_same<Dtype, float>::value) {
497+ FlagSMTE3();
498+ } else {
499+ FlagSV();
500+ }
501+ StoreFloatToGm(new_state_f, state + (b * C + c0) * K, w * K);
502+ }
503+ 
245 __aicore__ inline bool Packed()504 __aicore__ inline bool Packed()
246 {505 {
247 // Host sets seqLen=0 for packed mixed-length. GM pointer != nullptr is506 // Host sets seqLen=0 for packed mixed-length. GM pointer != nullptr is
@@ -249,98 +508,6 @@ public:
249 return seqLen == 0;508 return seqLen == 0;
250 }509 }
251 510 
252- // Packed [T,C]: tokens of one channel are stride-C apart.
253- // align_b16 bursts are placed 32B apart in UB; compact the first Dtype of
254- // each slot, then convert. Never DMA from a non-32B UB pointer (ADDR_MISALIGN).
255- __aicore__ inline void LoadPackedChannel(__gm__ Dtype *base, int S, int stride,
256- __ubuf__ float *dstF)
257- {
258- constexpr int kSlot = BLOCK_SIZE / sizeof(Dtype);
259- constexpr int kMaxTake = kMaxInputF / kSlot;
260- uint32_t elemBytes = static_cast<uint32_t>(sizeof(Dtype));
261- uint32_t srcGap = static_cast<uint32_t>(stride - 1) * elemBytes;
262- int done = 0;
263- while (done < S) {
264- int take = S - done;
265- if (take > kMaxTake) {
266- take = kMaxTake;
267- }
268- if constexpr (std::is_same<Dtype, float>::value) {
269- copy_gm_to_ubuf_align_b32(stage_buf, base + done * stride, 0, (uint16_t)take,
270- elemBytes, 0, 0, srcGap, 0);
271- } else {
272- copy_gm_to_ubuf_align_b16(stage_buf, base + done * stride, 0, (uint16_t)take,
273- elemBytes, 0, 0, srcGap, 0);
274- }
275- pipe_barrier(PIPE_MTE2);
276- set_flag(PIPE_MTE2, PIPE_S, EVENT_ID2);
277- wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID2);
278- for (int i = 0; i < take; ++i) {
279- stage_buf[i] = stage_buf[i * kSlot];
280- }
281- set_flag(PIPE_S, PIPE_V, EVENT_ID2);
282- wait_flag(PIPE_S, PIPE_V, EVENT_ID2);
283- if constexpr (std::is_same<Dtype, float>::value) {
284- for (int i = 0; i < take; ++i) {
285- WriteFloat(dstF, done + i, ReadFloat(stage_buf, i));
286- }
287- } else {
288- if (take >= 2 * VECTOR_MAX_NUM_OF_FP32) {
289- set_vector_mask((uint64_t)-1, (uint64_t)-1);
290- } else {
291- SetMask(take);
292- }
293- int repeat = DIV_ROUND_UP(take, VECTOR_MAX_NUM_OF_FP32);
294- if constexpr (std::is_same<Dtype, float16_t>::value) {
295- vconv_f162f32(dstF + done, stage_buf, repeat, 1, 1, 8, 4);
296- } else {
297- vconv_bf162f32(dstF + done, stage_buf, repeat, 1, 1, 8, 4);
298- }
299- pipe_barrier(PIPE_V);
300- set_vector_mask((uint64_t)-1, (uint64_t)-1);
301- }
302- done += take;
303- }
304- }
305- 
306- __aicore__ inline void StorePackedChannel(__gm__ Dtype *base, int S, int stride,
307- __ubuf__ Dtype *src, int nElem)
308- {
309- (void)S;
310- uint32_t elemBytes = static_cast<uint32_t>(sizeof(Dtype));
311- // Bulk scatter: expand contiguous src[0..nElem) into the 32B-slot layout
312- // that align_b16 expects on the UB side (each burst lives in its own
313- // 32B slot), then one copy_ubuf_to_gm_align_b16 writes every burst to a
314- // GM position stride elements apart (dstGap). nElem <= kBlock=64, so
315- // slot staging fits stage_buf (8KB for fp16/bf16).
316- constexpr int kSlot = BLOCK_SIZE / sizeof(Dtype);
317- // SiLU (vdiv) writes out_buf on V pipe
318- set_flag(PIPE_V, PIPE_S, EVENT_ID0);
319- wait_flag(PIPE_V, PIPE_S, EVENT_ID0);
320- for (int s = 0; s < nElem; ++s) {
321- stage_buf[s * kSlot] = src[s];
322- }
323- // S→MTE3: slot scatter must be visible before align_b16 DMA reads stage_buf.
324- set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
325- wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
326- uint32_t dstGap = static_cast<uint32_t>(stride - 1) * elemBytes;
327- copy_ubuf_to_gm_align_b16(base, stage_buf, 0, static_cast<uint16_t>(nElem), elemBytes, 0, 0,
328- 0, dstGap);
329- pipe_barrier(PIPE_MTE3);
330- }
331- 
332- __aicore__ inline void WriteBackState(int batchIdx, int channel, int S)
333- {
334- int K = (int)kernelDim;
335- for (int j = 0; j < K; ++j) {
336- int absIdx = S + j;
337- float v = (absIdx < K) ? ReadFloat(state_f, absIdx) : ReadFloat(input_f, absIdx - K);
338- WriteFloat(new_state_f, j, v);
339- }
340- int stateBase = (batchIdx * (int)channels + channel) * K;
341- StoreFloatToGm(new_state_f, state + stateBase, K);
342- }
343- 
344 __aicore__ inline void Process()511 __aicore__ inline void Process()
345 {512 {
346 int K = (int)kernelDim;513 int K = (int)kernelDim;
@@ -357,165 +524,39 @@ public:
357 pipe_barrier(PIPE_MTE2);524 pipe_barrier(PIPE_MTE2);
358 CopyGmToUbufAligned(meta_lens, queryLens, n * sizeof(int32_t));525 CopyGmToUbufAligned(meta_lens, queryLens, n * sizeof(int32_t));
359 pipe_barrier(PIPE_MTE2);526 pipe_barrier(PIPE_MTE2);
360- set_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);527+ FlagMTE2S();
361- wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);
362 }528 }
363 529 
364- // Byte offsets for vgather: lane p reads base + off_ramp[p].530+ int C = (int)channels;
365- for (int k = 0; k < kBlock; k++) {531+ int nBlocks = DIV_ROUND_UP(C, kBlock);
366- off_ramp[k] = static_cast<uint32_t>(k * (int)sizeof(float));532+ for (int cb = 0; cb < nBlocks; cb++) {
367- }533+ if (cb % GetBlockNum() != GetBlockIdx()) {
368- set_flag(PIPE_S, PIPE_V, EVENT_ID1);
369- wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
370- 
371- for (int channel = 0; channel < (int)channels; channel++) {
372- if (channel % GetBlockNum() != GetBlockIdx())
373 continue;534 continue;
374- 
375- int wOffset = channel * K;
376- if constexpr (std::is_same<Dtype, float>::value) {
377- CopyGmToUbufAligned(kernel_buf, weight + wOffset,
378- static_cast<uint32_t>(K * sizeof(Dtype)));
379- pipe_barrier(PIPE_MTE2);
380- } else {
381- CopyGmToUbufAligned(tmp_kernel_buf, weight + wOffset,
382- static_cast<uint32_t>(K * sizeof(Dtype)));
383- pipe_barrier(PIPE_MTE2);
384- set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
385- wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
386- uint64_t wmask = (K >= 64) ? (uint64_t)-1 : ((1ull << K) - 1ull);
387- set_vector_mask(0, wmask);
388- if constexpr (std::is_same<Dtype, float16_t>::value) {
389- vconv_f162f32(kernel_buf, tmp_kernel_buf, 1, 1, 1, 8, 4);
390- } else {
391- vconv_bf162f32(kernel_buf, tmp_kernel_buf, 1, 1, 1, 8, 4);
392- }
393- pipe_barrier(PIPE_V);
394- set_vector_mask((uint64_t)-1, (uint64_t)-1);
395 }535 }
536+ int c0 = cb * kBlock;
537+ int w = MIN(kBlock, C - c0);
396 538 
397- // Read the K fp32 weights into scalar registers (once per channel).539+ LoadWeights(c0, w);
398- set_flag(PIPE_V, PIPE_S, EVENT_ID1);540+ for (int b = 0; b < (int)batch; b++) {
399- wait_flag(PIPE_V, PIPE_S, EVENT_ID1);
400- float w[kMaxKernel];
401- for (int j = 0; j < K; ++j) {
402- w[j] = kernel_buf[j];
403- }
404- set_flag(PIPE_S, PIPE_V, EVENT_ID1);
405- wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
406- 
407- set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
408- for (int batchIdx = 0; batchIdx < (int)batch; ++batchIdx) {
409 int S;541 int S;
410- int start;542+ int tokenBase;
411 if (packed) {543 if (packed) {
412- // meta_* live in UB; must not read them from V without a drain.544+ S = (int)meta_lens[b];
413- pipe_barrier(PIPE_ALL);545+ tokenBase = (int)meta_start[b];
414- S = (int)meta_lens[batchIdx];
415- start = (int)meta_start[batchIdx];
416 } else {546 } else {
417 S = (int)seqLen;547 S = (int)seqLen;
418- start = 0;548+ tokenBase = b * S;
419 }549 }
420 if (S <= 0 || S > kMaxInputF) {550 if (S <= 0 || S > kMaxInputF) {
421 continue;551 continue;
422 }552 }
423- int stateBase = (batchIdx * (int)channels + channel) * K;553+ LoadState(b, c0, w);
424- LoadGmToFloat(state + stateBase, K, state_f);554+ ProcessSequence(tokenBase, c0, w, S);
425- if (packed) {
426- LoadPackedChannel(input + start * (int)channels + channel, S, (int)channels,
427- input_f);
428- } else {
429- int inputBase = (batchIdx * (int)channels + channel) * S;
430- LoadGmToFloat(input + inputBase, S, input_f);
431- }
432- 
433- // ---- Straddle outputs [0, K-1): window spans state and input ----
434- int scalarEnd = MIN(K - 1, S);
435- if (scalarEnd > 0) {
436- float st[kMaxKernel];
437- float in[kMaxKernel];
438- set_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);
439- wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID1);
440- set_flag(PIPE_V, PIPE_S, EVENT_ID1);
441- wait_flag(PIPE_V, PIPE_S, EVENT_ID1);
442- for (int j = 0; j < K; ++j) {
443- st[j] = state_f[j];
444- }
445- for (int j = 0; j < MIN(K, S); ++j) {
446- in[j] = input_f[j];
447- }
448- set_flag(PIPE_S, PIPE_V, EVENT_ID1);
449- wait_flag(PIPE_S, PIPE_V, EVENT_ID1);
450- for (int pos = 0; pos < scalarEnd; ++pos) {
451- float dot = 0.0f;
452- for (int j = 0; j < K; ++j) {
453- int t = pos + 1 + j;
454- float v = (t < K) ? st[t] : in[t - K];
455- dot += w[j] * v;
456- }
457- WriteFloat(acc_buf, pos, dot);
458- }
459- SiLU();
460- set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
461- wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
462- if (packed) {
463- StorePackedChannel(output + start * (int)channels + channel, S,
464- (int)channels, out_buf, scalarEnd);
465- } else {
466- int outOffset = (batchIdx * (int)channels + channel) * S;
467- CopyUbufToGmAligned(output + outOffset, out_buf, scalarEnd * sizeof(Dtype));
468- pipe_barrier(PIPE_MTE3);
469- }
470- set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
471- }
472- 
473- // ---- Windows fully inside input: full-width 64-lane blocks ----
474- for (int i = K - 1; i < S; i += kBlock) {
475- int realLen = MIN(kBlock, S - i);
476- 
477- vector_dup(acc_buf, 0.0f, 1, 1, 1, 8, 8);
478- pipe_barrier(PIPE_V);
479- for (int j = 0; j < K; ++j) {
480- uint32_t baseAddr = static_cast<uint32_t>(
481- (uint64_t)input_f + (i + 1 - K + j) * (int)sizeof(float));
482- vgather((__ubuf__ uint32_t *)qkv_tmp, off_ramp, baseAddr, 8, 1);
483- // vgather UB writeback is NOT covered by pipe_barrier(PIPE_V)
484- // (proven on this NPU). Use an S-roundtrip event fence before
485- // vmuls consumes qkv_tmp, otherwise the dependent vector op can
486- // be issued while the gather is still writing -> AIV hazard.
487- set_flag(PIPE_V, PIPE_S, EVENT_ID3);
488- wait_flag(PIPE_V, PIPE_S, EVENT_ID3);
489- set_flag(PIPE_S, PIPE_V, EVENT_ID3);
490- wait_flag(PIPE_S, PIPE_V, EVENT_ID3);
491- vmuls(calc_buf, qkv_tmp, w[j], 1, 1, 1, 8, 8);
492- pipe_barrier(PIPE_V);
493- vadd(acc_buf, acc_buf, calc_buf, 1, 1, 1, 1, 8, 8, 8);
494- pipe_barrier(PIPE_V);
495- }
496- SiLU();
497- 
498- set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
499- wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
500- if (packed) {
501- StorePackedChannel(output + (start + i) * (int)channels + channel, S,
502- (int)channels, out_buf, realLen);
503- } else {
504- int outOffset = (batchIdx * (int)channels + channel) * S + i;
505- CopyUbufToGmAligned(output + outOffset, out_buf, realLen * sizeof(Dtype));
506- pipe_barrier(PIPE_MTE3);
507- }
508- set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
509- }
510- 
511 if (updateState) {555 if (updateState) {
512- wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);556+ WriteBackState(b, tokenBase, c0, w, S);
513- WriteBackState(batchIdx, channel, S);
514- set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
515 }557 }
516 pipe_barrier(PIPE_ALL);558 pipe_barrier(PIPE_ALL);
517 }559 }
518- wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
519 }560 }
520 }561 }
521 562 
@@ -528,16 +569,17 @@ private:
528 __gm__ int32_t *queryLens;569 __gm__ int32_t *queryLens;
529 570 
530 __ubuf__ Dtype *stage_buf;571 __ubuf__ Dtype *stage_buf;
531- __ubuf__ Dtype *tmp_kernel_buf;572+ __ubuf__ float *w_f;
532- __ubuf__ float *kernel_buf;573+ __ubuf__ float *w_reorg;
533 __ubuf__ float *state_f;574 __ubuf__ float *state_f;
534- __ubuf__ float *input_f;575+ __ubuf__ float *state_reorg;
576+ __ubuf__ Dtype *win_raw;
577+ __ubuf__ float *window_f;
578+ __ubuf__ float *state_win;
535 __ubuf__ float *new_state_f;579 __ubuf__ float *new_state_f;
536- __ubuf__ uint32_t *off_ramp;
537- __ubuf__ float *qkv_tmp;
538 __ubuf__ float *acc_buf;580 __ubuf__ float *acc_buf;
539 __ubuf__ float *calc_buf;581 __ubuf__ float *calc_buf;
540- __ubuf__ Dtype *out_buf;582+ __ubuf__ Dtype *out_tile;
541 __ubuf__ int32_t *meta_start;583 __ubuf__ int32_t *meta_start;
542 __ubuf__ int32_t *meta_lens;584 __ubuf__ int32_t *meta_lens;
543 585 
Mxlite/csrc/model.cpp+9-20
@@ -801,20 +801,15 @@ void XModel::ForwardAttnLinear(XRuntime &rt, uint32_t layer,
801 }801 }
802 XTensor convStateBatch;802 XTensor convStateBatch;
803 convStateBatch.Init({batch, convDim, _c.linearConvKernelDim}, convState.dtype, convState.ptr);803 convStateBatch.Init({batch, convDim, _c.linearConvKernelDim}, convState.dtype, convState.ptr);
804- // Decode (seqlen==1) and small tasks use the packed 2D path:804+ // Decode (seqlen==1) uses the packed 2D path: [B,1,C] and [B,C,1] are
805- // mixQkv [m, qkvDim] token-major (matmul output)805+ // memory-identical for contiguous tensors, so the two Transpose_1_2 are
806- // convPacked [m, convDim] token-major (Step4 SplitCol consumes directly)806+ // identity ops. Skip them and use the packed 2D path.
807- // queryStartLoc/lens are populated every step by PrepareAttn.
808- //
809- // decode: seqlen==1 (any batch). [B,1,C] and [B,C,1] are memory-identical
810- // for contiguous tensors, so the two Transpose_1_2 are identity ops even
811- // with multiple requests in flight. Skip them and use the packed 2D path.
812 //807 //
813 // Prefill (seqlen>1, uniform) uses the token-row-parallel kernel when its808 // Prefill (seqlen>1, uniform) uses the token-row-parallel kernel when its
814- // constraints hold (K in {1,2,4}, seqlen >= K, convDim % 1024 == 0): it809+ // constraints hold (K in {1,2,4}, seqlen >= K, convDim % 1024 == 0).
815- // consumes token-major [m,C] mixQkv directly, removing the two810+ // Otherwise the fused conv kernel consumes/produces the token-major
816- // XliteOpTranspose_1_2 of the old 3D path.811+ // [B,S,C] layout directly, so no Transpose passes are needed either.
817- // Constraints unmet -> old 3D transpose path. Decode/mixed -> packed.812+ // Decode/mixed -> packed.
818 bool decodeStep = (seqlen == 1);813 bool decodeStep = (seqlen == 1);
819 uint32_t convK = _c.linearConvKernelDim;814 uint32_t convK = _c.linearConvKernelDim;
820 bool useToken = uniform && !decodeStep && seqlen >= convK &&815 bool useToken = uniform && !decodeStep && seqlen >= convK &&
@@ -824,18 +819,12 @@ void XModel::ForwardAttnLinear(XRuntime &rt, uint32_t layer,
824 seqlen, /*updateState=*/true);819 seqlen, /*updateState=*/true);
825 rt.PutTensor(mixQkv);820 rt.PutTensor(mixQkv);
826 } else if (uniform && !decodeStep) {821 } else if (uniform && !decodeStep) {
827- XTensor &mixTrans = rt.GetTensor({batch, qkvDim, seqlen}, hiddenState.dtype, DBG_LOC);
828- XTensor &convOut = rt.GetTensor({batch, convDim, seqlen}, hiddenState.dtype, DBG_LOC);
829 XTensor mix3d;822 XTensor mix3d;
830 mix3d.Init({batch, seqlen, qkvDim}, mixQkv.dtype, mixQkv.ptr);823 mix3d.Init({batch, seqlen, qkvDim}, mixQkv.dtype, mixQkv.ptr);
831- XliteOpTranspose_1_2(rt, mix3d, mixTrans);
832- XliteOpConv1dAndSiLU(rt, convStateBatch, mixTrans, linearConv1d[layer], convOut,
833- /*updateState=*/true);
834- rt.PutTensor(mixTrans);
835 XTensor convSeq3d;824 XTensor convSeq3d;
836 convSeq3d.Init({batch, seqlen, convDim}, convPacked.dtype, convPacked.ptr);825 convSeq3d.Init({batch, seqlen, convDim}, convPacked.dtype, convPacked.ptr);
837- XliteOpTranspose_1_2(rt, convOut, convSeq3d);826+ XliteOpConv1dAndSiLU(rt, convStateBatch, mix3d, linearConv1d[layer], convSeq3d,
838- rt.PutTensor(convOut);827+ /*updateState=*/true);
839 rt.PutTensor(mixQkv);828 rt.PutTensor(mixQkv);
840 } else {829 } else {
841 XliteOpConv1dAndSiLU(rt, convStateBatch, mixQkv, linearConv1d[layer], convPacked,830 XliteOpConv1dAndSiLU(rt, convStateBatch, mixQkv, linearConv1d[layer], convPacked,
Mxlite/csrc/op.cpp+12-8
@@ -1896,22 +1896,26 @@ void XliteOpConv1dAndSiLU(XRuntime &rt, XTensor &state, XTensor &input, XTensor
1896 }1896 }
1897 } else {1897 } else {
1898 if (state.shape.size() != 3 || input.shape.size() != 3 || output.shape.size() != 3) {1898 if (state.shape.size() != 3 || input.shape.size() != 3 || output.shape.size() != 3) {
1899- throw std::runtime_error("XliteOpConv1dAndSiLU: state/input/output must be 3D [B,C,*]");1899+ throw std::runtime_error("XliteOpConv1dAndSiLU: state/input/output must be 3D [B,*,C]");
1900 }1900 }
1901- if (state.shape[0] != input.shape[0] || state.shape[1] != input.shape[1] ||1901+ if (state.shape[0] != input.shape[0] || output.shape[0] != input.shape[0]) {
1902- output.shape[0] != input.shape[0] || output.shape[1] != input.shape[1] ||1902+ throw std::runtime_error("XliteOpConv1dAndSiLU: batch shape mismatch");
1903- output.shape[2] != input.shape[2]) {1903+ }
1904- throw std::runtime_error("XliteOpConv1dAndSiLU: batch/channel/seq shape mismatch");1904+ if (state.shape[1] != input.shape[2] || output.shape[2] != input.shape[2]) {
1905+ throw std::runtime_error("XliteOpConv1dAndSiLU: channel shape mismatch");
1906+ }
1907+ if (output.shape[1] != input.shape[1]) {
1908+ throw std::runtime_error("XliteOpConv1dAndSiLU: seq shape mismatch");
1905 }1909 }
1906 if (state.shape[2] != kernelDim) {1910 if (state.shape[2] != kernelDim) {
1907 throw std::runtime_error("XliteOpConv1dAndSiLU: state last dim != kernelDim");1911 throw std::runtime_error("XliteOpConv1dAndSiLU: state last dim != kernelDim");
1908 }1912 }
1909- if (weight.shape[0] != input.shape[1]) {1913+ if (weight.shape[0] != input.shape[2]) {
1910 throw std::runtime_error("XliteOpConv1dAndSiLU: weight channels mismatch");1914 throw std::runtime_error("XliteOpConv1dAndSiLU: weight channels mismatch");
1911 }1915 }
1912 batch = input.shape[0];1916 batch = input.shape[0];
1913- channels = input.shape[1];1917+ channels = input.shape[2];
1914- seqLen = input.shape[2];1918+ seqLen = input.shape[1];
1915 if (kernelDim > 16 || seqLen > 4096) {1919 if (kernelDim > 16 || seqLen > 4096) {
1916 throw std::runtime_error(1920 throw std::runtime_error(
1917 "XliteOpConv1dAndSiLU: require kernelDim<=16 and seqLen<=4096 for fused kernel");1921 "XliteOpConv1dAndSiLU: require kernelDim<=16 and seqLen<=4096 for fused kernel");
Mxlite/tests/kernels/conv1d_and_silu.py+2-9
@@ -10,21 +10,14 @@
10import torch10import torch
11import time11import time
12import torch.nn.functional as F12import torch.nn.functional as F
13-from xlite._C import Runtime, transpose_1_2, linear_att_conv_and_silu13+from xlite._C import Runtime, linear_att_conv_and_silu
14 14 
15channels = 614415channels = 6144
16kernel_dim = 416kernel_dim = 4
17 17 
18 18 
19def my_impl(rt, input, weight, conv_state, output, batch, seq_len):19def my_impl(rt, input, weight, conv_state, output, batch, seq_len):
20- mix_qkv = torch.empty(batch, channels, seq_len)20+ linear_att_conv_and_silu(rt, input, conv_state, weight, output)
21- torch.npu.synchronize()
22- transpose_1_2(rt, input, mix_qkv)
23- torch.npu.synchronize()
24- out = torch.empty(batch, channels, seq_len)
25- linear_att_conv_and_silu(rt, mix_qkv, conv_state, weight, out)
26- torch.npu.synchronize()
27- transpose_1_2(rt, out, output)
28 torch.npu.synchronize()21 torch.npu.synchronize()
29 22 
30 23 
Mxlite/tests/kernels/conv1d_and_silu_token.py+12-15
@@ -12,7 +12,7 @@ import time
12import torch12import torch
13import torch.nn.functional as F13import torch.nn.functional as F
14from xlite._C import (Runtime, linear_att_conv_and_silu,14from xlite._C import (Runtime, linear_att_conv_and_silu,
15- linear_att_conv_and_silu_token, transpose_1_2)15+ linear_att_conv_and_silu_token)
16 16 
17kernel_dim = 417kernel_dim = 4
18 18 
@@ -69,34 +69,31 @@ def run_perf(rt, batch=1, seq_len=512, channels=10240, dtype=torch.bfloat16):
69 input = torch.randn(batch, seq_len, channels, dtype=dtype, device="npu:0")69 input = torch.randn(batch, seq_len, channels, dtype=dtype, device="npu:0")
70 weight = torch.randn(channels, 1, kernel_dim, dtype=dtype, device="npu:0")70 weight = torch.randn(channels, 1, kernel_dim, dtype=dtype, device="npu:0")
71 state_t = torch.randn(batch, channels, kernel_dim, dtype=dtype, device="npu:0")71 state_t = torch.randn(batch, channels, kernel_dim, dtype=dtype, device="npu:0")
72- state_3d = state_t.clone()72+ state_chan = state_t.clone()
73 73 
74 mix_qkv = input.reshape(batch * seq_len, channels).contiguous()74 mix_qkv = input.reshape(batch * seq_len, channels).contiguous()
75 out_tok = torch.zeros(batch * seq_len, channels, dtype=dtype, device="npu:0")75 out_tok = torch.zeros(batch * seq_len, channels, dtype=dtype, device="npu:0")
76 76 
77- # old 3D path: transpose [B,S,C]->[B,C,S], channel-parallel conv, transpose back77+ # channel-parallel path: the fused conv consumes/produces [B,S,C] directly
78- mix_trans = torch.empty(batch, channels, seq_len, dtype=dtype, device="npu:0")78+ # (no Transpose passes) and updates the state in-kernel.
79- conv_out = torch.empty(batch, channels, seq_len, dtype=dtype, device="npu:0")79+ out_chan = torch.empty(batch, seq_len, channels, dtype=dtype, device="npu:0")
80- out_3d = torch.empty(batch, seq_len, channels, dtype=dtype, device="npu:0")
81 torch.npu.synchronize() # torch stream -> rt stream barrier80 torch.npu.synchronize() # torch stream -> rt stream barrier
82 81 
83 def token_path():82 def token_path():
84 linear_att_conv_and_silu_token(rt, mix_qkv, state_t, weight, out_tok, seq_len)83 linear_att_conv_and_silu_token(rt, mix_qkv, state_t, weight, out_tok, seq_len)
85 84 
86- def old_3d_path():85+ def chan_path():
87- transpose_1_2(rt, input, mix_trans)86+ linear_att_conv_and_silu(rt, input, state_chan, weight, out_chan)
88- linear_att_conv_and_silu(rt, mix_trans, state_3d, weight, conv_out)
89- transpose_1_2(rt, conv_out, out_3d)
90 87 
91 t_tok = bench(token_path)88 t_tok = bench(token_path)
92- t_old = bench(old_3d_path)89+ t_chan = bench(chan_path)
93 # numerical cross-check while we are here90 # numerical cross-check while we are here
94- torch.testing.assert_close(out_tok.reshape(batch, seq_len, channels), out_3d,91+ torch.testing.assert_close(out_tok.reshape(batch, seq_len, channels), out_chan,
95 rtol=1e-2, atol=1e-3)92 rtol=1e-2, atol=1e-3)
96- torch.testing.assert_close(state_t, state_3d, rtol=1e-2, atol=1e-3)93+ torch.testing.assert_close(state_t, state_chan, rtol=1e-2, atol=1e-3)
97 print(f"[perf B={batch} S={seq_len} C={channels} {dtype}] "94 print(f"[perf B={batch} S={seq_len} C={channels} {dtype}] "
98- f"token={t_tok:.3f} ms/call, old3d={t_old:.3f} ms/call, "95+ f"token={t_tok:.3f} ms/call, chan={t_chan:.3f} ms/call, "
99- f"speedup={t_old / t_tok:.2f}x")96+ f"speedup={t_chan / t_tok:.2f}x")
100 97 
101 98 
102if __name__ == "__main__":99if __name__ == "__main__":
Mxlite/xlite/_C.pyi+2-2
@@ -2514,10 +2514,10 @@ def linear_att_conv_and_silu(
2514 2514 
2515 Args:2515 Args:
2516 rt (Runtime): Native runtime handle.2516 rt (Runtime): Native runtime handle.
2517- mix_qkv (torch.Tensor): Input mixed QKV tensor, shape [B, C, S].2517+ mix_qkv (torch.Tensor): Input mixed QKV tensor, shape [B, S, C].
2518 conv_state (torch.Tensor): Convolution state tensor, shape [B, C, K].2518 conv_state (torch.Tensor): Convolution state tensor, shape [B, C, K].
2519 weight (torch.Tensor): Kernel weight tensor, shape [C, 1, K] or [C, K].2519 weight (torch.Tensor): Kernel weight tensor, shape [C, 1, K] or [C, K].
2520- output (torch.Tensor): Output tensor, shape [B, C, S].2520+ output (torch.Tensor): Output tensor, shape [B, S, C].
2521 2521 
2522 Returns:2522 Returns:
2523 None: `output` is written in place. State is always updated.2523 None: `output` is written in place. State is always updated.