Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all \u003cpre\u003e\u003ccode\u003e blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks"); } } catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); } })(); (function(){ try { var __m = "github.com"; var __re = new RegExp('^' + "github\\.com" + '
Skip to content

Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length \u003e 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); })();
Skip to content

Add large CUDA DFlash target prefill and small-query SDPA (#21628)#21628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meta-codesync merged 1 commit into from
Aug 7, 2026
Merged

Add large CUDA DFlash target prefill and small-query SDPA (#21628) #21628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions backends/cuda/runtime/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -35,16 +35,28 @@ def define_common_targets(is_fbcode = False):
name = "runtime_shims",
srcs = [
"shims/cuda_guard.cpp",
"shims/int4_plain_mm.cu",
"shims/int4mm.cu",
"shims/int5_plain_mm.cu",
"shims/int6_plain_mm.cu",
"shims/int8_plain_mm.cu",
"shims/memory.cpp",
"shims/rand.cu",
"shims/sort.cu",
"shims/tensor_attribute.cpp",
],
headers = [
"shims/cuda_guard.h",
"shims/int4_plain_mm.cuh",
"shims/int4_plain_mm.h",
"shims/int4mm.cuh",
"shims/int4mm.h",
"shims/int5_plain_mm.cuh",
"shims/int5_plain_mm.h",
"shims/int6_plain_mm.cuh",
"shims/int6_plain_mm.h",
"shims/int8_plain_mm.cuh",
"shims/int8_plain_mm.h",
"shims/memory.h",
"shims/rand.h",
"shims/sort.h",
Expand Down
19 changes: 19 additions & 0 deletions backends/cuda/tests/targets.bzl
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,25 @@ def define_common_targets(is_fbcode = False):
),
)

python_unittest_remote_gpu(
name = "test_triton_sdpa_splitk",
srcs = [
"test_triton_sdpa_splitk.py",
],
visibility = [
"//executorch/...",
],
deps = [
"//caffe2:torch",
"//executorch/backends/cuda:triton_kernels",
],
keep_gpu_sections = True,
remote_execution = re_test_utils.remote_execution(
platform = "gpu-remote-execution",
subplatform = "A100-exclusive",
),
)

python_unittest(
name = "test_cuda_partitioner",
srcs = [
Expand Down
192 changes: 170 additions & 22 deletions backends/cuda/tests/test_triton_sdpa_splitk.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,8 +10,10 @@
expanded KV heads in float32.
"""

import importlib
import itertools
import unittest
from unittest import mock

import torch
import torch.nn.functional as F
Expand All@@ -24,16 +26,20 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("BF16 not supported on this GPU")


def _import_splitk():
def _import_legacy_splitk():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa_decode_splitk

return sdpa_decode_splitk


def _import_sdpa():
from executorch.backends.cuda.triton.kernels.sdpa import sdpa
def _import_small_query_splitk():
from executorch.backends.cuda.triton.kernels import sdpa_small_query_splitk

return sdpa
return sdpa_small_query_splitk


def _import_sdpa_module():
return importlib.import_module("executorch.backends.cuda.triton.kernels.sdpa")


def _reference_sdpa(q, k, v, attn_mask=None, scale=None):
Expand DownExpand Up@@ -85,8 +91,10 @@ class TestTritonSdpaSplitK(unittest.TestCase):
@classmethod
def setUpClass(cls):
_skip_if_no_cuda()
cls.splitk = _import_splitk()
cls.sdpa = _import_sdpa()
cls.legacy_splitk = _import_legacy_splitk()
cls.small_query_splitk = _import_small_query_splitk()
cls.sdpa_module = _import_sdpa_module()
cls.sdpa = cls.sdpa_module.sdpa

# ------------------------------------------------------------------
# Correctness
Expand All@@ -106,7 +114,7 @@ def test_decode_basic(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -130,7 +138,7 @@ def test_decode_with_mask(self):
mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
mask[:, :, :, :200] = True

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -146,12 +154,132 @@ def test_decode_mha(self):
k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_small_query_blocks(self):
"""Split-K supports each small-query length through the verifier size."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 512, 128
for Lq in [2, 3, 4]:
with self.subTest(Lq=Lq):
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_non_power_of_two_gqa_ratio_with_three_queries(self):
"""Padded group tiles must retain every query row."""
B, H_q, H_kv, Lq, Lk, D = 1, 10, 2, 3, 512, 128
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.small_query_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_public_dispatch_isolates_kernel_families(self):
"""Lq1 and verifier blocks must use disjoint split-K launchers."""
B, H_q, H_kv, Lk, D = 1, 8, 2, 256, 64
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
kv_len = torch.tensor(Lk, dtype=torch.int64, device="cuda")

for Lq, legacy_calls, small_query_calls in [(1, 1, 0), (4, 0, 1)]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
with mock.patch.object(
self.sdpa_module, "_launch_decode_splitk"
) as legacy_launcher, mock.patch.object(
self.sdpa_module, "_launch_small_query_splitk"
) as small_query_launcher:
self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)

self.assertEqual(legacy_launcher.call_count, legacy_calls)
self.assertEqual(small_query_launcher.call_count, small_query_calls)

def test_high_gqa_small_query_with_runtime_kv_len(self):
"""Exercise Lq=4, 16:1 GQA, and a bottom-right causal mask."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 4096, 128
valid_len = 4089
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_runtime_kv_len_ignores_garbage_tail(self):
"""The public split-K dispatch must not read past the valid KV prefix."""
B, H_q, H_kv, Lq, Lk, D = 1, 32, 2, 4, 131072, 128
valid_len = 509
torch.manual_seed(42)
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
k[:, :, valid_len:] = 1000
v[:, :, valid_len:] = 1000
q_pos = torch.arange(valid_len - Lq, valid_len, device="cuda")
k_pos = torch.arange(Lk, device="cuda")
mask = (q_pos[:, None] >= k_pos[None, :])[None, None]
kv_len = torch.tensor(valid_len, dtype=torch.int64, device="cuda")

out = self.sdpa(
q,
k,
v,
attn_mask=mask,
enable_gqa=True,
kv_len=kv_len,
)
ref = _reference_sdpa(
q,
k[:, :, :valid_len],
v[:, :, :valid_len],
attn_mask=mask[:, :, :, :valid_len],
)

self.assertTrue(torch.isfinite(out).all())
self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL)

def test_qwen35_config(self):
"""Exact Qwen3.5 MoE config: H_q=16, H_kv=2, D=256."""
B, H_q, H_kv, D = 1, 16, 2, 256
Expand All@@ -165,7 +293,7 @@ def test_qwen35_config(self):

mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)
ref = _reference_sdpa(q, k, v, attn_mask=mask)

self.assertEqual(out.shape, (B, H_q, Lq, D))
Expand All@@ -181,7 +309,7 @@ def test_custom_scale(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

scale = 0.05
out = self.splitk(q, k, v, scale=scale)
out = self.legacy_splitk(q, k, v, scale=scale)
ref = _reference_sdpa(q, k, v, scale=scale)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -199,7 +327,7 @@ def test_cross_validate_with_sdpa(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
mask = torch.ones(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")

out_splitk = self.splitk(q, k, v, attn_mask=mask)
out_splitk = self.legacy_splitk(q, k, v, attn_mask=mask)
out_tiled = self.sdpa(q, k, v, attn_mask=mask, enable_gqa=True)

self.assertLess(
Expand All@@ -221,7 +349,7 @@ def test_all_masked(self):
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda")
out = self.splitk(q, k, v, attn_mask=mask)
out = self.legacy_splitk(q, k, v, attn_mask=mask)

self.assertFalse(torch.isnan(out).any(), "All-masked should not NaN")
self.assertFalse(torch.isinf(out).any(), "All-masked should not Inf")
Expand All@@ -234,7 +362,7 @@ def test_lk_1(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -250,7 +378,7 @@ def test_batch_size(self):
k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

out = self.splitk(q, k, v)
out = self.legacy_splitk(q, k, v)
ref = _reference_sdpa(q, k, v)

self.assertFalse(torch.isnan(out).any())
Expand All@@ -260,14 +388,34 @@ def test_batch_size(self):
# Validation errors
# ------------------------------------------------------------------

def test_lq_not_1_rejected(self):
"""L_q != 1 should raise RuntimeError."""
def test_legacy_lq_two_rejected(self):
"""The legacy decode op remains restricted to L_q == 1."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 2, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.legacy_splitk(q, k, v)

def test_small_query_lq_one_and_five_rejected(self):
"""The small-query op accepts only L_q values 2 through 4."""
B, H_q, H_kv, D = 1, 8, 2, 64
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
for Lq in [1, 5]:
with self.subTest(Lq=Lq):
q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.small_query_splitk(q, k, v)

def test_multi_query_implicit_causal_rejected(self):
"""Cached multi-query attention requires an explicit aligned mask."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 4, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.small_query_splitk(q, k, v, is_causal=True)

def test_dropout_rejected(self):
"""dropout_p != 0 should raise RuntimeError."""
Expand All@@ -276,15 +424,15 @@ def test_dropout_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v, dropout_p=0.1)
self.legacy_splitk(q, k, v, dropout_p=0.1)

def test_is_causal_accepted(self):
"""is_causal=True is a no-op at L_q=1, should not raise."""
B, H_q, H_kv, D = 1, 8, 2, 64
q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
out = self.splitk(q, k, v, is_causal=True)
out = self.legacy_splitk(q, k, v, is_causal=True)
self.assertEqual(out.shape, (B, H_q, 1, D))

def test_hq_not_divisible_rejected(self):
Expand All@@ -294,7 +442,7 @@ def test_hq_not_divisible_rejected(self):
k = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, 3, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)

def test_non_pow2_d_rejected(self):
"""Non-power-of-2 D should raise RuntimeError."""
Expand All@@ -303,7 +451,7 @@ def test_non_pow2_d_rejected(self):
k = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
v = torch.randn(B, H_kv, 64, D, dtype=torch.bfloat16, device="cuda")
with self.assertRaises(RuntimeError):
self.splitk(q, k, v)
self.legacy_splitk(q, k, v)


if __name__ == "__main__":
Expand Down
Loading
Loading