[FA2] fix tiling-qk misaligned address✔️ (#174)

* Update README.md

* fix misaligned address

* fix misaligned address

Author: DefTruth

kernels/flash-attn/README.md

@@ -49,9 +49,7 @@ python3 flash_attn_mma.py --B 1 --H 48 --D 64 --N 8192 --iters 10 --torch  # NVI
 ```
 - Example: B=1, H=48, N=8192, D=512 (NVIDIA RTX 3080 Laptop), FA2 not supported.
 ```bash
- python3 flash_attn_mma.py --B 1 --H 8 --N 8192 --iters 10 --show-all --sdpa --D 512 # NVIDIA RTX 3080 Laptop, Faster than SDPA
------------------------B: batch_size, H: n_head, N: seq_len, D: head_dim, seed: 1041, Warmup: 1, Iters: 10------------------------
-----------------------------------------------------------------------------------------------------------------------------------
+python3 flash_attn_mma.py --B 1 --H 8 --N 8192 --iters 10 --show-all --sdpa --D 512 # NVIDIA RTX 3080 Laptop, Faster than SDPA
 ------------------------------------------B=1, H=8, N=8192, D=512, Warmup: 1, Iters: 10-------------------------------------------
    mma(split-q+tiling-qk+stage1): ['-0.00433731 ', '0.02165222  ', '-0.01544189 '], time:48.775554ms, TFLOPS:22.60 (+0.00%)
    mma(split-q+tiling-qk+stage2): ['-0.00433731 ', '0.02165222  ', '-0.01544189 '], time:47.503424ms, TFLOPS:23.20 (+2.68%)

kernels/flash-attn/flash_attn_mma.py

@@ -66,6 +66,8 @@ def get_args():
 
 
 args = get_args()
+if args.D and args.D >= 256:
+    args.run_torch_sdpa = True
 pretty_print_line()
 print(args)
 pretty_print_line()

kernels/flash-attn/mma/flash_attn_mma_tiling_qk.cu

@@ -138,6 +138,8 @@ flash_attn_mma_stages_split_q_tiling_qk_kernel(half* Q,
   int load_gmem_Q_Br = Q_tile_id * Br + load_smem_Q_Br; 
   if (load_gmem_Q_Br >= QKV_seqlen) return;
   constexpr bool kIsVCanLoadIn128b = (kHeadDim / (kNumThreads / kMmaAtomK)) % 8 == 0;
+  constexpr bool kIsVCanLoadIn64b  = (kHeadDim / (kNumThreads / kMmaAtomK)) % 4 == 0;
+  static_assert(kIsVCanLoadIn128b || kIsVCanLoadIn64b, "V can't load in 128b or 64b."); // 32,64,128,192,256,...
 
   // Shared memory for Q,K,V, we don not need additional smem for O 
   // collective store which perform via registers reuse and warp shuffle.
@@ -763,17 +765,17 @@ flash_attn_mma_stages_split_q_tiling_qk_kernel(half* Q,
 template<const int kHeadDim, const int kStage>
 void launch_flash_attn_mma_stages_split_q_tiling_qk(
   torch::Tensor Q, torch::Tensor K, torch::Tensor V, torch::Tensor O) {
-  // Now: fixed tile BrxBc=128x128
+  // Now: fixed tile BrxBc=128x128 for d>= 128, 64x64 for d<128.
   // TODO: dynamic tile size for Br, Bc according to kHeadDim and shared memory size.
   constexpr int kMmaAtomM = 16;
   constexpr int kMmaAtomN = 8;
   constexpr int kMmaAtomK = 16;
-  constexpr int kMmaTileSeqLenQ  = 8;
+  constexpr int kMmaTileSeqLenQ  = (kHeadDim < 128) ? 4 : 8;
   constexpr int kMmaTileSeqLenK  = 1;
-  constexpr int kMmaTileSeqLenP  = 8;
+  constexpr int kMmaTileSeqLenP  = (kHeadDim < 128) ? 4 : 8;
   constexpr int kMmaTileHeadDimV = 1;
   constexpr int kWarpTileSeqLenQ = 1;
-  constexpr int kWarpTileSeqLenK = 16;
+  constexpr int kWarpTileSeqLenK = (kHeadDim < 128) ? 8 : 16;
   constexpr int kWarpTileSeqLenP = 1;
   constexpr int kWarpTileHeadDimV = (kHeadDim / (kMmaAtomN * kMmaTileHeadDimV)); // (d=64)8,(d=128)16,32,....
   constexpr int Br = kMmaAtomM * kMmaTileSeqLenQ * kWarpTileSeqLenQ; // 16*4*1=64