Hansung Kim
d0ef06cec1
flash: Complete Q_IS_K_MAJOR code for GEMM II
2024-09-19 20:36:03 -07:00
Hansung Kim
be15cffbf3
flash: Revert to gemmini config, remove DEBUG and unnecessary checks
2024-09-12 14:25:33 -07:00
Hansung Kim
b5916f3f07
flash: Fix hardcoded barrier for tcore; move tcore-specific flags
2024-09-11 22:08:44 -07:00
Hansung Kim
d69707f686
flash: Enable GEMM II fence; Pull 1st KV move out of the loop
2024-09-11 19:24:06 -07:00
Hansung Kim
18cf0e73cd
flash: Add early return for warp-indivisible row iter
2024-09-11 00:56:09 -07:00
Hansung Kim
068d48534e
flash: Swap S1/S0 to avoid GEMM II - softmax bank conflict
...
+ remove spurrious fences to better overlap GEMM I and DMA
2024-09-11 00:55:36 -07:00
Hansung Kim
dc746272fb
flash: Conditionally enable GEMM II fence code, fix tile_k for DEBUG
2024-09-10 22:53:35 -07:00
Hansung Kim
90e03894fc
flash: Add flag in SMEM for dependency check on O
...
TODO: results unverified.
Stalls O rescale until GEMM II finishes.
2024-09-10 13:42:47 -07:00
Hansung Kim
88760596cb
flash: Remove bogus mvout to SMEM code
2024-09-09 17:18:59 -07:00
Hansung Kim
a17edac875
flash: Fix barrier stall with DEBUG
...
Verified for up to P_expected on 2nd iter; O_before_PV is partially
correct
2024-09-09 17:02:05 -07:00
Hansung Kim
b652e25945
flash: Warp-specialize between warp 0 and 1-7
...
Finishes without stalls; No dependency check between O rescale and
GEMM-II.
2024-09-09 16:42:30 -07:00
Hansung Kim
d31c8ffd7d
flash: Fix grid size to hw cluster size
...
Verified fast config, minus the barrier stall at the end.
2024-09-09 15:44:03 -07:00
Hansung Kim
829af5d429
flash: Comment out mvout to smem
...
Verified up to O_before_PV; still stalls without DEBUG
2024-09-09 15:21:49 -07:00
Hansung Kim
ecc800964a
flash: Change smem alloc for less bank conflicts; noskip stc
2024-09-09 13:47:18 -07:00
Hansung Kim
714b9f501e
flash: Restructure to do delayed fence for better concurrency
...
Verified up to O_before_PV of 2nd iteration; O_after_PV needs preload
fix.
FIXME: Stalls at barrier without DEBUG set.
2024-09-08 22:16:17 -07:00
Hansung Kim
6911843a82
flash: Remove unnecessary dmem preload, fix rowmax/rowsum dependency
2024-09-08 21:11:59 -07:00
Hansung Kim
a4dd45bc1b
flash: Replace CISC with RISC
...
spadQuartile in hw does not match spad addresses in kernel; match them
later for optimization.
2024-09-08 20:52:28 -07:00
Hansung Kim
6547e92757
flash: Load Q to both quartiles; preload O for acc
2024-09-08 19:47:55 -07:00
Hansung Kim
8efa6868ea
flash: Restructure for full software pipelining
...
Verified up to P and O before PV; need to fix iteration for V load.
2024-09-08 18:45:32 -07:00
Hansung Kim
cdb8377b62
flash: Do GEMM II in Gemmini; verify 1st iteration
2024-09-08 16:09:06 -07:00
Hansung Kim
3f50ac57ee
flash: use 12bit dma interface
2024-09-08 15:29:56 -07:00
Hansung Kim
c51dc4902d
flash: Fix online softmax for DMA layout
2024-09-07 23:21:28 -07:00
Hansung Kim
2e1485877d
flash: Add Gemmini-accelerated kernel
2024-09-07 22:40:58 -07:00
Hansung Kim
b3be271b88
flash: Split impl to header file
2024-09-07 21:16:35 -07:00
Hansung Kim
03308f8033
flash: Write fast config for DMA
...
MAC utilization is 20-25% for the loop.
2024-09-07 20:47:24 -07:00
Hansung Kim
8d32a03d09
flash: Write DMA code for warp-specialized
...
TODO: result unverified
2024-09-07 20:32:08 -07:00
Hansung Kim
33bc084c37
flash: Fix DMA layout for GEMM II
2024-09-07 19:51:05 -07:00
Hansung Kim
e02892ab7d
flash: Fix DMA for up to GEMM II
...
yeah
2024-09-07 19:31:41 -07:00
Hansung Kim
d2f086344d
flash: Fix DMA addr stride, stop at S=Q*K
2024-09-07 15:48:37 -07:00
Hansung Kim
137df9bee2
WIP: flash: Use Gemmini DMA
...
Paused by the barrier bug in warp-divergent branches (tmask not being
considered)
2024-09-05 16:23:32 -07:00
Hansung Kim
bfb414c4eb
flash: Add DMA config logic
2024-09-05 16:23:32 -07:00
Hansung Kim
f028a97f75
sgemm_tcore: Verify wo DMA; warn untested against K-major A + DMA
2024-09-03 14:42:19 -07:00
Hansung Kim
70273fd00d
flash: Cleanup debug code
2024-09-02 00:40:05 -07:00
Hansung Kim
8125192846
flash: Specify leading_dim for split QK GEMM; fix uninit'd RF before GEMM
2024-09-02 00:15:57 -07:00
Hansung Kim
602fe4a400
flash: Change timing for QKV move
...
Verified with warp_specialized false; true remains to be fixed.
2024-09-01 22:06:46 -07:00
Hansung Kim
aea257349a
flash: Correct schedule with inter-warpgroup barriers
2024-09-01 20:40:51 -07:00
Hansung Kim
e5e65312d2
flash: Restructure to inter-warpgroup parallelism
...
This is similar to
https://tridao.me/blog/2024/flash3/#inter-warpgroup-overlapping-with-pingpong-scheduling
2024-09-01 19:58:33 -07:00
Hansung Kim
6cc1b5ca37
flash: Reduce smem_scratchpad alloc size
2024-09-01 16:02:06 -07:00
Hansung Kim
817cc9a5a5
flash: Fix overlap in smem alloc for P tile
2024-08-31 16:24:28 -07:00
Hansung Kim
bdd6e6a9ce
flash: Double-buffer between online softmax and GEMM II
...
TODO: O_after_PV at the last stage is incorrect.
2024-08-30 22:47:55 -07:00
Hansung Kim
042b47ff19
flash: Restructure for warp-specialization
2024-08-30 22:14:45 -07:00
Hansung Kim
1cfab40711
flash: Do Oi rescale with PV
...
Since Oi rescale has data dependency to previous Oi which gets produced
at the PV GEMM, both rescale+GEMM needs to be in a single pipeline stage
or otherwise it requires a stall. So instead, compute only the
rescale factor in the online softmax stage and apply rescaling right
before PV.
2024-08-30 20:11:07 -07:00
Hansung Kim
986d507223
flash: Fix single-tile GEMM for warp-specialized
...
With 4 warps, we can only do 32x64 GEMM; serialize 64x64 into 2 32x64
GEMM calls by split by the row.
2024-08-30 17:12:46 -07:00
Hansung Kim
72b6004e24
flash: Fix online softmax for warp-specialized
...
Note: now that threads_per_threadblock is passed as compile-time
constant, the compiler likes to completely loop unroll which can cause a
lot of stack spills.
todo fix GEMM part.
2024-08-29 21:50:02 -07:00
Hansung Kim
fd1ab358fa
flash: Add DOUBLE_BUF compile-time param (wip)
2024-08-29 14:18:56 -07:00
Hansung Kim
5ba06dfd9d
flash: Incomplete parallel stage-2 rowmax
2024-08-29 13:29:00 -07:00
Hansung Kim
3f20dd59c0
flash: Supply correct tile dims to single_tile
2024-08-20 19:50:45 -07:00
Hansung Kim
dde0372769
flash: Enable skipping Q*K for larger dimensions
2024-08-20 19:15:16 -07:00
Hansung Kim
615d36a5c2
flash: Reduce smem use for rowmax; verify result
2024-08-20 14:45:34 -07:00
Hansung Kim
d8d5df64e6
flash: Fix load addr for V tile; test with seqlen=128
2024-08-20 14:34:09 -07:00
