Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

SDNQ add causal atten support

Browse Source
This commit is contained in:
Disty0
2026-06-28 20:08:40 +03:00
parent 0eb7d7032b
commit e89f1a1fe8
2 changed files with 20 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/attention.py
+1 -1
View File
@@ -25,7 +25,7 @@ def set_sdnq_attention():
@wraps(sdpa_pre_sdnq_atten)
def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
if (
not is_causal and query.device.type != "cpu"
query.device.type != "cpu"
and (query.shape[-1] >= 32 and key.shape[-1] >= 32 and value.shape[-1] >= 32) # Dim < 32 is unsupported by Matrix Cores
and (query.shape[-2] >= 512 or key.shape[-2] >= 512) # Skip TE
and query.shape[-3] > 1 # Skip VAE
modules/sdnq/kernels/triton_atten.py
+19 -5
View File
@@ -62,10 +62,12 @@ def quantize_attn(
return q_q, q_scale, k_q, k_scale, v_q, v_scale
@triton.autotune(configs=matmul_configs, key=["QN_AT", "QHD", "VN_AT", "VHD", "qk_is_quantized", "pv_is_quantized", "q_dtype", "v_dtype", "out_dtype", "mask_dtype"], cache_results=True)
@triton.autotune(configs=matmul_configs, key=["QN_AT", "QHD", "VN_AT", "VHD", "qk_is_quantized", "pv_is_quantized", "q_dtype", "v_dtype", "out_dtype", "mask_dtype", "is_causal"], cache_results=True)
@triton.jit
def sdnq_attn_kernel(
q_ptr, k_ptr, v_ptr, q_scale_ptr, k_scale_ptr, v_scale_ptr, out_ptr, mask_ptr,
q_ptr, k_ptr, v_ptr,
q_scale_ptr, k_scale_ptr, v_scale_ptr,
out_ptr, mask_ptr, is_causal: tl.constexpr,
QZ: tl.constexpr, QH: tl.constexpr, QN: tl.constexpr, QHD: tl.constexpr,
KZ: tl.constexpr, KH: tl.constexpr, KN: tl.constexpr, KHD: tl.constexpr,
VZ: tl.constexpr, VH: tl.constexpr, VN: tl.constexpr, VHD: tl.constexpr,
@@ -88,6 +90,8 @@ def sdnq_attn_kernel(
start_m = tl.program_id(0)
off_z = tl.program_id(2)
off_h = tl.program_id(1)
offs_m = start_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_n = tl.arange(0, BLOCK_SIZE_N)
num_kv_groups = QH // VH
off_h_kv = off_h // num_kv_groups
@@ -110,7 +114,12 @@ def sdnq_attn_kernel(
q = Q_desc.load([start_m * BLOCK_SIZE_M, 0])
lo, hi = 0, KN
lo = 0
if is_causal:
hi = tl.minimum(KN, (start_m + 1) * BLOCK_SIZE_M + KN - QN)
else:
hi = KN
for start_n in range(lo, hi, BLOCK_SIZE_N):
start_n = tl.multiple_of(start_n, BLOCK_SIZE_N)
skip = False
@@ -131,6 +140,10 @@ def sdnq_attn_kernel(
else:
qk = tl.dot(q, k, out_dtype=tl.float32)
if is_causal:
causal_mask = (offs_m[:, None] + (KN - QN)) >= (start_n + offs_n[None, :])
qk = tl.where(causal_mask, qk, -float('inf'))
if mask_ptr is not None:
if mask.dtype == tl.int1:
qk = tl.where(mask, qk, -float('inf'))
@@ -190,7 +203,6 @@ def sdnq_triton_atten_forward(
do_quantize: bool = True,
out_dtype: torch.dtype | None = None,
) -> torch.FloatTensor:
assert not is_causal
if out_dtype is None:
out_dtype = query.dtype
QZ, QH, QN, QHD = query.shape
@@ -221,7 +233,9 @@ def sdnq_triton_atten_forward(
pv_matmul_dtype=pv_matmul_dtype if do_quantize else "no",
)
sdnq_attn_kernel[grid](
query, key, value, query_scale, key_scale, value_scale, out, attn_mask,
query, key, value,
query_scale, key_scale, value_scale,
out, attn_mask, is_causal,
*query.shape, *key.shape, *value.shape, *out.shape,
*(attn_mask.shape if attn_mask is not None else (0, 0, 0, 0)),
*query.stride(), *key.stride(), *value.stride(), *out.stride(),