Skip to content

Commit b6a5ee8

Browse files
VeeraRajasekharVeeraRajasekhar
committed
Guard CK small-seq behind NVTE_FUSED_ATTN_CK_SMALLSEQ=1; add FP16 support to small-seq kernels
1 parent db685c4 commit b6a5ee8

8 files changed

Lines changed: 196 additions & 107 deletions

File tree

‎tests/jax/test_fused_attn.py‎

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from functools import partial
1010
from math import sqrt
1111
from typing import Tuple, Optional, Dict
12+
import os
1213
import random
1314

1415
import jax
@@ -329,7 +330,12 @@ class FusedAttnRunner:
329330
# generating zero-length ragged tensors. This setting adjusts the test to avoid the zero-length cases.
330331
def _get_max_segments_per_sequence(self):
331332
if self.qkv_layout.is_thd():
332-
if 90400 <= get_cudnn_version() < 90500 or self.max_seqlen_q == 1:
333+
if (
334+
90400 <= get_cudnn_version() < 90500
335+
or ( self.max_seqlen_q == 1 and
336+
is_hip_extension() and
337+
os.environ.get("NVTE_FUSED_ATTN_CK_SMALLSEQ", "0") == "1")
338+
):
333339
return self.num_segments_per_seq
334340
else:
335341
# +1 for testing runtime_segments < max_segments
@@ -539,7 +545,7 @@ def generate_random_segment_ids(
539545
return segment_ids, segment_pos, segment_pad
540546

541547
if self.qkv_layout.is_thd():
542-
if self.max_seqlen_q == 1:
548+
if self.max_seqlen_q == 1 and is_hip_extension() and os.environ.get("NVTE_FUSED_ATTN_CK_SMALLSEQ", "0") == "1":
543549
self.num_segments_per_seq = 1
544550
# Q: deterministic — one segment of length 1 per batch -> cu_seqlen [0,1,2,...,batch_size]
545551
self.segment_ids_q = jnp.ones((self.batch_size, self.max_seqlen_q), dtype=jnp.int32)
@@ -555,7 +561,6 @@ def generate_random_segment_ids(
555561
)
556562

557563
# KV: one segment per batch (num_segments_per_seq=1) to match smallseq kernel
558-
# expectations (batch_size == max_tokens_q, cu_seqlens of size batch_size+1).
559564
min_segment_len = None if self.window_size is None else self.seqlens_q
560565
self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = (
561566
generate_random_segment_ids(
@@ -1247,26 +1252,43 @@ def test_jax_new_rng():
12471252
runner.test_forward()
12481253

12491254

1250-
# ROCm CK internal small-seq (varlen unfused) branch tests.
1255+
# ROCm CK small-seq varlen tests.
12511256
# Uses THD_THD_THD with s_q=1, s_kv<=16 so the small-seq path is taken.
1257+
# Run only when NVTE_FUSED_ATTN_CK_SMALLSEQ=1.
1258+
@pytest.mark.skipif(
1259+
os.environ.get("NVTE_FUSED_ATTN_CK_SMALLSEQ", "0") != "1",
1260+
reason="CK unfused smallseq tests require NVTE_FUSED_ATTN_CK_SMALLSEQ=1",
1261+
)
12521262
@pytest.mark.skipif(
12531263
not is_hip_extension(), reason="CK unfused smallseq backend only available on AMD hardware"
12541264
)
12551265
@pytest.mark.parametrize(
12561266
"b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype",
12571267
[
1258-
pytest.param(30720, 1, 2, 16, 16, 128, 128, jnp.bfloat16,
1259-
id="30720-1-2-16-16-128-128-BF16"),
1260-
pytest.param(30720, 1, 4, 16, 16, 128, 128, jnp.bfloat16,
1261-
id="30720-1-4-16-16-128-128-BF16"),
1262-
pytest.param(30720, 1, 6, 16, 16, 128, 128, jnp.bfloat16,
1263-
id="30720-1-6-16-16-128-128-BF16"),
1264-
pytest.param(30720, 1, 8, 16, 16, 128, 128, jnp.bfloat16,
1265-
id="30720-1-8-16-16-128-128-BF16"),
1266-
pytest.param(30720, 1, 12, 16, 16, 128, 128, jnp.bfloat16,
1267-
id="30720-1-12-16-16-128-128-BF16"),
1268-
pytest.param(30720, 1, 16, 16, 16, 128, 128, jnp.bfloat16,
1269-
id="30720-1-16-16-16-128-128-BF16"),
1268+
pytest.param(4000, 1, 2, 16, 16, 128, 128, jnp.bfloat16,
1269+
id="4000-1-2-16-16-128-128-BF16"),
1270+
pytest.param(4000, 1, 4, 16, 16, 128, 128, jnp.bfloat16,
1271+
id="4000-1-4-16-16-128-128-BF16"),
1272+
pytest.param(4000, 1, 6, 16, 16, 128, 128, jnp.bfloat16,
1273+
id="4000-1-6-16-16-128-128-BF16"),
1274+
pytest.param(4000, 1, 8, 16, 16, 128, 128, jnp.bfloat16,
1275+
id="4000-1-8-16-16-128-128-BF16"),
1276+
pytest.param(4000, 1, 12, 16, 16, 128, 128, jnp.bfloat16,
1277+
id="4000-1-12-16-16-128-128-BF16"),
1278+
pytest.param(4000, 1, 16, 16, 16, 128, 128, jnp.bfloat16,
1279+
id="4000-1-16-16-16-128-128-BF16"),
1280+
pytest.param(4000, 1, 2, 16, 16, 128, 128, jnp.float16,
1281+
id="4000-1-2-16-16-128-128-FP16"),
1282+
pytest.param(4000, 1, 4, 16, 16, 128, 128, jnp.float16,
1283+
id="4000-1-4-16-16-128-128-FP16"),
1284+
pytest.param(4000, 1, 6, 16, 16, 128, 128, jnp.float16,
1285+
id="4000-1-6-16-16-128-128-FP16"),
1286+
pytest.param(4000, 1, 8, 16, 16, 128, 128, jnp.float16,
1287+
id="4000-1-8-16-16-128-128-FP16"),
1288+
pytest.param(4000, 1, 12, 16, 16, 128, 128, jnp.float16,
1289+
id="4000-1-12-16-16-128-128-FP16"),
1290+
pytest.param(4000, 1, 16, 16, 16, 128, 128, jnp.float16,
1291+
id="4000-1-16-16-16-128-128-FP16"),
12701292
],
12711293
)
12721294
def test_ck_unfused_smallseq_backend(b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype):

‎transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*************************************************************************
2-
* Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
2+
* Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
33
*
44
* License for AMD contributions = MIT. See LICENSE for more information
55
************************************************************************/

‎transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
************************************************************************/
66

77
#include <utility>
8-
98
#include "ck_fused_attn_utils.hpp"
109
#include "ck_fused_attn/ck_fused_attn.hpp"
1110
#include "mask.hpp"

‎transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*************************************************************************
2-
* Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
2+
* Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
33
*
44
* License for AMD contributions = MIT. See LICENSE for more information
55
************************************************************************/
@@ -615,9 +615,10 @@ void fused_attn_ck_fwd_impl(
615615
// denote the next available section of workspace from upstream
616616
void* workspace_next = workspace;
617617

618-
if (is_ragged) {
618+
const char* nvte_smallseq = std::getenv("NVTE_FUSED_ATTN_CK_SMALLSEQ");
619+
if (is_ragged && nvte_smallseq && std::string(nvte_smallseq) == "1") {
619620
void* max_seqlen_workspace = workspace;
620-
621+
621622
size_t runtime_max_seqlen_q = static_cast<size_t>(ck_fused_attn::get_runtime_max_seqlen(
622623
static_cast<uint64_t>(b), devPtrCuSeqlensQ, nullptr, max_seqlen_workspace, stream));
623624
size_t runtime_max_seqlen_kv = static_cast<size_t>(ck_fused_attn::get_runtime_max_seqlen(
@@ -630,7 +631,7 @@ void fused_attn_ck_fwd_impl(
630631
std::cout << "runtime_max_seqlen_kv: " << runtime_max_seqlen_kv << std::endl;
631632
}
632633

633-
if (runtime_max_seqlen_q==1 && runtime_max_seqlen_kv >= 2 && runtime_max_seqlen_kv <= 16) {
634+
if (runtime_max_seqlen_q == 1 && runtime_max_seqlen_kv >= 2 && runtime_max_seqlen_kv <= 16) {
634635
fused_attn_rocm::fused_attn_smallseq_fwd(
635636
b, h, hg, runtime_max_seqlen_kv, d_qk, d_v,
636637
is_training, scaling_factor, dropout_probability,
@@ -944,7 +945,8 @@ void fused_attn_ck_bwd_impl(
944945
// denote the next available section of workspace from upstream
945946
void* workspace_next = workspace;
946947

947-
if (is_ragged) {
948+
const char* nvte_smallseq = std::getenv("NVTE_FUSED_ATTN_CK_SMALLSEQ");
949+
if (is_ragged && nvte_smallseq && std::string(nvte_smallseq) == "1") {
948950
void* max_seqlen_workspace = workspace;
949951

950952
size_t runtime_max_seqlen_q = static_cast<size_t>(ck_fused_attn::get_runtime_max_seqlen(

‎transformer_engine/common/fused_attn_rocm/fused_attn_smallseq.cpp‎

Lines changed: 111 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
11
/*************************************************************************
2-
* Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
2+
* Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
33
*
44
* License for AMD contributions = MIT. See LICENSE for more information
55
************************************************************************/
66

77
/*! \file fused_attn_smallseq.cpp
88
* \brief Unfused small-seq (varlen) attention: seq_q=1, max_seqlen_kv<=16, THD only.
9-
* Ported from varlen_attn/attn_fwd.cpp and attn_bwd.cpp with runtime b, head_num.
109
*/
1110

1211
#include <hip/hip_runtime.h>
1312
#include <hip/hip_bfloat16.h>
13+
#include <hip/hip_fp16.h>
1414

1515
#include <cstdint>
1616
#include <cstdlib>
@@ -106,30 +106,30 @@ __global__ void compute_scores_kernel(const T* Q,
106106
for (int i = 0; i < seq_kv; i++)
107107
results[i] = 0.0f;
108108
for (int dim_offset = 0; dim_offset < head_dim; dim_offset += block_k) {
109-
if constexpr (std::is_same<T, hip_bfloat16>::value) {
109+
if constexpr (std::is_same<T, hip_bfloat16>::value || std::is_same<T, __half>::value) {
110110
for (int k = 0; k < block_k / 8; k++) {
111111
ls_dwordx4_tmp_var = *((uint4*)&Q_ptr[dim_offset + k * 8]);
112-
fetch_Q[k * 8 + 0] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[0];
113-
fetch_Q[k * 8 + 1] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[1];
114-
fetch_Q[k * 8 + 2] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[0];
115-
fetch_Q[k * 8 + 3] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[1];
116-
fetch_Q[k * 8 + 4] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[0];
117-
fetch_Q[k * 8 + 5] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[1];
118-
fetch_Q[k * 8 + 6] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[0];
119-
fetch_Q[k * 8 + 7] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[1];
112+
fetch_Q[k * 8 + 0] = ((T*)&ls_dwordx4_tmp_var.x)[0];
113+
fetch_Q[k * 8 + 1] = ((T*)&ls_dwordx4_tmp_var.x)[1];
114+
fetch_Q[k * 8 + 2] = ((T*)&ls_dwordx4_tmp_var.y)[0];
115+
fetch_Q[k * 8 + 3] = ((T*)&ls_dwordx4_tmp_var.y)[1];
116+
fetch_Q[k * 8 + 4] = ((T*)&ls_dwordx4_tmp_var.z)[0];
117+
fetch_Q[k * 8 + 5] = ((T*)&ls_dwordx4_tmp_var.z)[1];
118+
fetch_Q[k * 8 + 6] = ((T*)&ls_dwordx4_tmp_var.w)[0];
119+
fetch_Q[k * 8 + 7] = ((T*)&ls_dwordx4_tmp_var.w)[1];
120120
}
121121
for (int kv_idx = 0; kv_idx < seq_kv; kv_idx++) {
122122
for (int k = 0; k < block_k / 8; k++) {
123123
ls_dwordx4_tmp_var =
124124
*((uint4*)&K_ptr[kv_idx * head_num * head_dim + dim_offset + k * 8]);
125-
fetch_K[k * 8 + 0] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[0];
126-
fetch_K[k * 8 + 1] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[1];
127-
fetch_K[k * 8 + 2] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[0];
128-
fetch_K[k * 8 + 3] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[1];
129-
fetch_K[k * 8 + 4] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[0];
130-
fetch_K[k * 8 + 5] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[1];
131-
fetch_K[k * 8 + 6] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[0];
132-
fetch_K[k * 8 + 7] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[1];
125+
fetch_K[k * 8 + 0] = ((T*)&ls_dwordx4_tmp_var.x)[0];
126+
fetch_K[k * 8 + 1] = ((T*)&ls_dwordx4_tmp_var.x)[1];
127+
fetch_K[k * 8 + 2] = ((T*)&ls_dwordx4_tmp_var.y)[0];
128+
fetch_K[k * 8 + 3] = ((T*)&ls_dwordx4_tmp_var.y)[1];
129+
fetch_K[k * 8 + 4] = ((T*)&ls_dwordx4_tmp_var.z)[0];
130+
fetch_K[k * 8 + 5] = ((T*)&ls_dwordx4_tmp_var.z)[1];
131+
fetch_K[k * 8 + 6] = ((T*)&ls_dwordx4_tmp_var.w)[0];
132+
fetch_K[k * 8 + 7] = ((T*)&ls_dwordx4_tmp_var.w)[1];
133133
}
134134
#pragma unroll
135135
for (int k = 0; k < block_k; k++)
@@ -502,30 +502,30 @@ __global__ void compute_grad_attn_kernel(const T* grad_O,
502502
results[i] = 0.0f;
503503

504504
for (int dim_offset = 0; dim_offset < head_dim; dim_offset += block_k) {
505-
if constexpr (std::is_same<T, hip_bfloat16>::value) {
505+
if constexpr (std::is_same<T, hip_bfloat16>::value || std::is_same<T, __half>::value) {
506506
for (int k = 0; k < block_k / 8; k++) {
507507
ls_dwordx4_tmp_var = *((uint4*)&grad_O_ptr[dim_offset + k * 8]);
508-
fetch_grad_O[k * 8 + 0] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[0];
509-
fetch_grad_O[k * 8 + 1] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[1];
510-
fetch_grad_O[k * 8 + 2] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[0];
511-
fetch_grad_O[k * 8 + 3] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[1];
512-
fetch_grad_O[k * 8 + 4] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[0];
513-
fetch_grad_O[k * 8 + 5] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[1];
514-
fetch_grad_O[k * 8 + 6] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[0];
515-
fetch_grad_O[k * 8 + 7] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[1];
508+
fetch_grad_O[k * 8 + 0] = ((T*)&ls_dwordx4_tmp_var.x)[0];
509+
fetch_grad_O[k * 8 + 1] = ((T*)&ls_dwordx4_tmp_var.x)[1];
510+
fetch_grad_O[k * 8 + 2] = ((T*)&ls_dwordx4_tmp_var.y)[0];
511+
fetch_grad_O[k * 8 + 3] = ((T*)&ls_dwordx4_tmp_var.y)[1];
512+
fetch_grad_O[k * 8 + 4] = ((T*)&ls_dwordx4_tmp_var.z)[0];
513+
fetch_grad_O[k * 8 + 5] = ((T*)&ls_dwordx4_tmp_var.z)[1];
514+
fetch_grad_O[k * 8 + 6] = ((T*)&ls_dwordx4_tmp_var.w)[0];
515+
fetch_grad_O[k * 8 + 7] = ((T*)&ls_dwordx4_tmp_var.w)[1];
516516
}
517517
for (int kv_idx = 0; kv_idx < seq_kv; kv_idx++) {
518518
for (int k = 0; k < block_k / 8; k++) {
519519
ls_dwordx4_tmp_var =
520520
*((uint4*)&V_base[kv_idx * V_stride + dim_offset + k * 8]);
521-
fetch_V[k * 8 + 0] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[0];
522-
fetch_V[k * 8 + 1] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.x)[1];
523-
fetch_V[k * 8 + 2] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[0];
524-
fetch_V[k * 8 + 3] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.y)[1];
525-
fetch_V[k * 8 + 4] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[0];
526-
fetch_V[k * 8 + 5] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.z)[1];
527-
fetch_V[k * 8 + 6] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[0];
528-
fetch_V[k * 8 + 7] = ((hip_bfloat16*)&ls_dwordx4_tmp_var.w)[1];
521+
fetch_V[k * 8 + 0] = ((T*)&ls_dwordx4_tmp_var.x)[0];
522+
fetch_V[k * 8 + 1] = ((T*)&ls_dwordx4_tmp_var.x)[1];
523+
fetch_V[k * 8 + 2] = ((T*)&ls_dwordx4_tmp_var.y)[0];
524+
fetch_V[k * 8 + 3] = ((T*)&ls_dwordx4_tmp_var.y)[1];
525+
fetch_V[k * 8 + 4] = ((T*)&ls_dwordx4_tmp_var.z)[0];
526+
fetch_V[k * 8 + 5] = ((T*)&ls_dwordx4_tmp_var.z)[1];
527+
fetch_V[k * 8 + 6] = ((T*)&ls_dwordx4_tmp_var.w)[0];
528+
fetch_V[k * 8 + 7] = ((T*)&ls_dwordx4_tmp_var.w)[1];
529529
}
530530
#pragma unroll
531531
for (int k = 0; k < block_k; k++)
@@ -708,7 +708,7 @@ __global__ void compute_grad_qk_kernel(const T* grad_scores,
708708
head_dim +
709709
thread_head_offset + i * dwordx4_load_elt];
710710
for (int b = 0; b < dwordx4_load_elt; b++)
711-
grad_Q_ptr[b] = ((T*)&store_dwordx4_tmp_var[i])[b] * scale;
711+
grad_Q_ptr[b] = ((T*)&store_dwordx4_tmp_var[i])[b] * T(scale);
712712
}
713713
#pragma unroll
714714
for (int i = 0; i < block_k / dwordx4_load_elt; i++) {
@@ -833,8 +833,9 @@ void fused_attn_smallseq_fwd(size_t b,
833833
size_t* workspace_size,
834834
cudaStream_t stream)
835835
{
836-
if (std::getenv("NVTE_FUSED_ATTN_CK_SMALLSEQ")) {
837-
std::cout << std::endl << "attn_fwd(ck small-seq kernel): ";
836+
const char* nvte_smallseq = std::getenv("NVTE_LOG_CK_CONFIG");
837+
if (nvte_smallseq && std::string(nvte_smallseq) == "1") {
838+
std::cout << std::endl << "attn_fwd(small-seq kernel): ";
838839
std::cout << "b: " << b << ", ";
839840
std::cout << "h_q: " << h_q << ", ";
840841
std::cout << "h_kv: " << h_kv << ", ";
@@ -871,6 +872,38 @@ void fused_attn_smallseq_fwd(size_t b,
871872
int bi = static_cast<int>(b);
872873
int hi = static_cast<int>(h_q);
873874

875+
switch (max_seqlen_kv) {
876+
SMALLSEQ_DISPATCH_FWD_CASE(2)
877+
SMALLSEQ_DISPATCH_FWD_CASE(3)
878+
SMALLSEQ_DISPATCH_FWD_CASE(4)
879+
SMALLSEQ_DISPATCH_FWD_CASE(5)
880+
SMALLSEQ_DISPATCH_FWD_CASE(6)
881+
SMALLSEQ_DISPATCH_FWD_CASE(7)
882+
SMALLSEQ_DISPATCH_FWD_CASE(8)
883+
SMALLSEQ_DISPATCH_FWD_CASE(9)
884+
SMALLSEQ_DISPATCH_FWD_CASE(10)
885+
SMALLSEQ_DISPATCH_FWD_CASE(11)
886+
SMALLSEQ_DISPATCH_FWD_CASE(12)
887+
SMALLSEQ_DISPATCH_FWD_CASE(13)
888+
SMALLSEQ_DISPATCH_FWD_CASE(14)
889+
SMALLSEQ_DISPATCH_FWD_CASE(15)
890+
SMALLSEQ_DISPATCH_FWD_CASE(16)
891+
default:
892+
NVTE_ERROR("Unsupported max_seqlen_kv for small-seq: max_seqlen_kv <= 16.");
893+
}
894+
} else if (qkv_dtype == DType::kFloat16) {
895+
using T = __half;
896+
const T* Q_ptr = static_cast<const T*>(devPtrQ);
897+
const T* K_ptr = static_cast<const T*>(devPtrK);
898+
const T* V_ptr = static_cast<const T*>(devPtrV);
899+
T* O_ptr = static_cast<T*>(devPtrO);
900+
T* attn_workspace = static_cast<T*>(attn_weights_buffer);
901+
const int* cu_kv = static_cast<const int*>(devPtrCuSeqlensKV);
902+
const int* cu_kv_p = static_cast<const int*>(devPtrSeqOffsetsKV);
903+
const T* dropout_mask = nullptr;
904+
int bi = static_cast<int>(b);
905+
int hi = static_cast<int>(h_q);
906+
874907
switch (max_seqlen_kv) {
875908
SMALLSEQ_DISPATCH_FWD_CASE(2)
876909
SMALLSEQ_DISPATCH_FWD_CASE(3)
@@ -891,7 +924,7 @@ void fused_attn_smallseq_fwd(size_t b,
891924
NVTE_ERROR("Unsupported max_seqlen_kv for small-seq: max_seqlen_kv <= 16.");
892925
}
893926
} else {
894-
NVTE_ERROR("small-seq path supports only BF16 (and optionally FP16).");
927+
NVTE_ERROR("small-seq path supports only BF16 and FP16.");
895928
}
896929

897930
if (workspace_size) {
@@ -941,10 +974,6 @@ void fused_attn_smallseq_bwd(size_t b,
941974
(void)h_kv;
942975
(void)d_qk;
943976
(void)d_v;
944-
NVTE_CHECK(max_seqlen_kv >= 2 && max_seqlen_kv <= 16,
945-
"small-seq path requires 2 <= max_seqlen_kv <= 16.");
946-
NVTE_CHECK(d_qk == 128 && d_v == 128, "small-seq path currently supports head_dim 128 only.");
947-
NVTE_CHECK(workspace != nullptr, "small-seq bwd requires workspace.");
948977

949978
float sqr_dk_scale = attn_scale;
950979
hipStream_t hip_stream = reinterpret_cast<hipStream_t>(stream);
@@ -967,6 +996,43 @@ void fused_attn_smallseq_bwd(size_t b,
967996
int bi = static_cast<int>(b);
968997
int hi = static_cast<int>(h_q);
969998

999+
switch (max_seqlen_kv) {
1000+
SMALLSEQ_DISPATCH_BWD_CASE(2)
1001+
SMALLSEQ_DISPATCH_BWD_CASE(3)
1002+
SMALLSEQ_DISPATCH_BWD_CASE(4)
1003+
SMALLSEQ_DISPATCH_BWD_CASE(5)
1004+
SMALLSEQ_DISPATCH_BWD_CASE(6)
1005+
SMALLSEQ_DISPATCH_BWD_CASE(7)
1006+
SMALLSEQ_DISPATCH_BWD_CASE(8)
1007+
SMALLSEQ_DISPATCH_BWD_CASE(9)
1008+
SMALLSEQ_DISPATCH_BWD_CASE(10)
1009+
SMALLSEQ_DISPATCH_BWD_CASE(11)
1010+
SMALLSEQ_DISPATCH_BWD_CASE(12)
1011+
SMALLSEQ_DISPATCH_BWD_CASE(13)
1012+
SMALLSEQ_DISPATCH_BWD_CASE(14)
1013+
SMALLSEQ_DISPATCH_BWD_CASE(15)
1014+
SMALLSEQ_DISPATCH_BWD_CASE(16)
1015+
default:
1016+
NVTE_ERROR("Unsupported max_seqlen_kv for small-seq: max_seqlen_kv <= 16.");
1017+
}
1018+
} else if (qkv_dtype == DType::kFloat16) {
1019+
using T = __half;
1020+
const T* Q_ptr = static_cast<const T*>(devPtrQ);
1021+
const T* K_ptr = static_cast<const T*>(devPtrK);
1022+
const T* V_ptr = static_cast<const T*>(devPtrV);
1023+
const T* O_ptr = static_cast<const T*>(devPtrO);
1024+
const T* dO_ptr = static_cast<const T*>(devPtrdO);
1025+
const T* attn_ptr = static_cast<const T*>(attn_weights);
1026+
T* dQ_ptr = static_cast<T*>(devPtrdQ);
1027+
T* dK_ptr = static_cast<T*>(devPtrdK);
1028+
T* dV_ptr = static_cast<T*>(devPtrdV);
1029+
T* workspace_ptr = static_cast<T*>(workspace);
1030+
const int* cu_kv = static_cast<const int*>(devPtrCuSeqlensKV);
1031+
const int* cu_kv_p = static_cast<const int*>(devPtrSeqOffsetsKV);
1032+
const T* dropout_mask = nullptr;
1033+
int bi = static_cast<int>(b);
1034+
int hi = static_cast<int>(h_q);
1035+
9701036
switch (max_seqlen_kv) {
9711037
SMALLSEQ_DISPATCH_BWD_CASE(2)
9721038
SMALLSEQ_DISPATCH_BWD_CASE(3)
@@ -987,7 +1053,7 @@ void fused_attn_smallseq_bwd(size_t b,
9871053
NVTE_ERROR("Unsupported max_seqlen_kv for small-seq: max_seqlen_kv <= 16.");
9881054
}
9891055
} else {
990-
NVTE_ERROR("small-seq path supports only BF16 (and optionally FP16).");
1056+
NVTE_ERROR("small-seq path supports only BF16 and FP16.");
9911057
}
9921058

9931059
if (workspace_size)

0 commit comments

Comments
 (0)