Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

CUDA: XOR swizzle flash attn K,V smem fp16 tiles (#25635)

Browse Source 
* 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>
This commit is contained in:
ynankani
2026-08-31 22:18:01 +02:00
committed by GitHub
parent ab0b3bd3c8
commit e4b9af007b
3 changed files with 194 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test-backend-ops.cpp
+18 -4
View File
@@ -10183,6 +10183,16 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false));
// FLASH_ATTN_EXT MMA: non-pow2 head size and MLA K/V view.
test_cases.emplace_back(new test_flash_attn_ext(192, 128, 8, {8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
// FLASH_ATTN_EXT MMA, swizzled K/V tiles, power-of-two stride: nbatch_K2 = 32, 64, 128, 256.
test_cases.emplace_back(new test_flash_attn_ext( 64, 64, 8, {8, 1}, 4096, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 4096, 8, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3}));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1}));
test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3}));
@@ -10581,10 +10591,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
for (int kv : { 4096, 8192, 16384, }) {
for (int hs : { 64, 128, }) {
for (int nr : { 1, 4, }) {
test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
for (int kv : { 4096, 8192, 16384,32768, 65536, }) {
for (int hs : { 64, 128, 256, 576, }) {
const int hsv = hs == 576 ? 512 : hs;
const bool v_view = hs == 576;
for (int nr : { 1, 4, 8, }) {
for (int nb : { 1, 4096, }) {
test_cases.emplace_back(new test_flash_attn_ext(hs, hsv, 8, {nr, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, v_view));
}
}
}
}