CUDA: XOR swizzle flash attn K,V smem fp16 tiles (#25635)
* CUDA: XOR swizzle flash attn K,V smem fp16 tiles
Signed-off-by: ynankani <ynankani@nvidia.com>
* Fix use 64bit generic pointer instead of 32bit shared pointer
Signed-off-by: ynankani <ynankani@nvidia.com>
* fix shared memory race in FA on DGX Spark
* Handle corener case
Signed-off-by: ynankani <ynankani@nvidia.com>
* Add swizzle test cases and gate sync for swizzled path only
Signed-off-by: ynankani <ynankani@nvidia.com>
* gate CUDA PTX
Signed-off-by: ynankani <ynankani@nvidia.com>
* offset calculation specific for swizzle branch
Signed-off-by: ynankani <ynankani@nvidia.com>
* Reafctor code
Signed-off-by: ynankani <ynankani@nvidia.com>
* Refactor FA swizzle ldmatrix if/else into helpers (K row/col, V offset)
Signed-off-by: ynankani <ynankani@nvidia.com>
* rebase and update test case args
Signed-off-by: ynankani <ynankani@nvidia.com>
* Allow swizzle for non-pow2 shapes, for which nbatch_2%32==0
Signed-off-by: ynankani <ynankani@nvidia.com>
---------
Signed-off-by: ynankani <ynankani@nvidia.com>
