Add ROCm HIP small-seq fused attention via crossattn_hip_kernel

README.rst

@@ -3,7 +3,7 @@
     Copyright (c) 2023-2026, Advanced Micro Devices, Inc. All rights reserved.
     Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-    See LICENSE for license information.
+    See LICENSE for license information.Please also add one subsection in readme to tell our customers who to use small_seq attn
 
 |License|
 
@@ -259,6 +259,38 @@ ROCm TE provides the compile-time env NVTE_CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAU
 * 3 - standard asm, default;
 * 4 - rta_asm.
 
+Small-Sequence Attention in CK Backend (gfx942/gfx950)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+For workloads with very short sequences (up to 17 tokens), ROCm TE provides dedicated CK small-sequence
+attention kernels that are more efficient than the general fused-attention path for these shapes. A common
+use case is cross-attention where a single query token attends over a short key/value sequence.
+
+This path is part of the CK backend and is opt-in at runtime:
+
+* NVTE_FUSED_ATTN_CK_SMALLSEQ - by default 0 (disabled); set to 1 to route eligible problems through the small-seq kernels.
+
+It requires the CK backend to be enabled (NVTE_FUSED_ATTN_CK=1, the default). When enabled, a problem is
+routed to the small-seq kernels only when all of the following hold; otherwise TE transparently falls back
+to the regular CK/AITER fused-attention path:
+
+* GPU architecture is gfx942 or gfx950;
+* data type is BF16 (FP16 is not supported on this path yet);
+* head dimension is 128 or 256, with matching Q/K and V head dimensions;
+* number of attention heads is 16 or 32, with no GQA/MQA (num_heads == num_gqa_groups);
+* no attention bias and no dropout;
+* mask type is padding mask or no mask;
+* the (runtime) maximum sequence length of both Q and KV is at most 17.
+
+Both ``THD`` (variable-length, e.g. cross-attention) and ``BSHD`` (dense self-attention with s_q == s_kv)
+layouts are supported.
+
+When using the JAX integration, the small-seq path requires XLA GPU graph capture (command buffers) to be
+disabled, and XLA_FLAGS must be set before the process starts, for example:
+
+.. code-block:: bash
+
+    XLA_FLAGS='--xla_gpu_enable_command_buffer=' NVTE_FUSED_ATTN_CK_SMALLSEQ=1 python your_script.py
+
 Experimental Triton Kernels on ROCm
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Most CUDA kernels in Transformer Engine are hipified to run on ROCm. While the hipifiled CUDA kernels are functional, they are not necessarily optimal on ROCm.

ci/jax.sh

@@ -57,9 +57,10 @@ run_test_config() {
     export NVTE_JAX_UNITTEST_LEVEL=L0 # this env variable controls parameters set for some tests
     run_default_fa 1 test_custom_call_compute.py
     run_default_fa 1 test_functions.py
-    run 1 test_fused_attn.py
+    run 1 test_fused_attn.py -k 'not TestFusedAttnCkSmallseq' # skip smallseq in normal flow
+    XLA_FLAGS='--xla_gpu_enable_command_buffer=' run 1 test_fused_attn.py -k 'TestFusedAttnCkSmallseq' # CK small-seq path; requires GPU graph capture disabled
     NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 run_default_fa_lbl "deterministic" 3 test_fused_attn.py -k "TestFusedAttnWithDeterminism"
-    NVTE_CK_USES_FWD_V3=0 NVTE_CK_USES_BWD_V3=0 run_default_fa_lbl "v2" 3 test_fused_attn.py # Using FAv2 for forward and backward pass
+    NVTE_CK_USES_FWD_V3=0 NVTE_CK_USES_BWD_V3=0 run_default_fa_lbl "v2" 3 test_fused_attn.py -k 'not TestFusedAttnCkSmallseq' # Using FAv2 for forward and backward pass
     run_default_fa 1 test_layer.py # it effectively always uses unfused attention
     run_default_fa 1 test_sanity_import.py
     run_default_fa 1 test_softmax.py

tests/cpp/small_seq_kernels/CMakeLists.txt

@@ -0,0 +1,224 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+cmake_minimum_required(VERSION 3.21)
+
+# Declare project with both CXX and HIP languages.
+# Requires hip-lang CMake package (available under ${ROCM_PATH}/lib/cmake).
+project(crossattn_hip_kernel LANGUAGES CXX HIP)
+
+# ---------------------------------------------------------------------------
+# ROCm / HIP setup
+# ---------------------------------------------------------------------------
+
+if(NOT DEFINED ROCM_PATH)
+    set(ROCM_PATH "/opt/rocm" CACHE PATH "Path to ROCm installation")
+endif()
+
+list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}/lib/cmake")
+find_package(hip REQUIRED CONFIG)
+
+# GPU architecture — override with -DGPU_TARGETS=gfx906 etc.
+# set(GPU_TARGETS "gfx950" CACHE STRING "GPU architecture targets")
+set(GPU_TARGETS "gfx942" CACHE STRING "GPU architecture targets")
+set(CMAKE_HIP_ARCHITECTURES "${GPU_TARGETS}")
+
+# ---------------------------------------------------------------------------
+# Language standards
+# ---------------------------------------------------------------------------
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_HIP_STANDARD 17)
+set(CMAKE_HIP_STANDARD_REQUIRED ON)
+
+# ---------------------------------------------------------------------------
+# Kernel headers
+#
+# The MFMA kernel headers (attn_*.h) are vendored once for the Transformer
+# Engine build under transformer_engine/common/ck_fused_attn/small_seq_kernels.
+# These reference tests include them directly from that canonical location
+# instead of keeping a duplicate copy here.
+# ---------------------------------------------------------------------------
+
+set(KERNEL_INCLUDE_DIR
+    "${CMAKE_SOURCE_DIR}/../../../transformer_engine/common/ck_fused_attn/small_seq_kernels"
+    CACHE PATH "Path to the vendored small-seq MFMA kernel headers")
+
+# ---------------------------------------------------------------------------
+# CPU reference library
+#
+# Compiled as plain C++ (no HIP device code) and linked by both test
+# executables. Avoids duplicating CPU reference compilation.
+# ---------------------------------------------------------------------------
+
+add_library(attn_ref STATIC
+    ref/attn_fwd_ref.cpp
+    ref/attn_bwd_ref.cpp
+)
+# Mark as HIP so clang++ handles hip_bfloat16 conversions (float<->bfloat16)
+# correctly. No device kernels are compiled; clang just enables the host-side
+# HIP type support.
+set_source_files_properties(
+    ref/attn_fwd_ref.cpp
+    ref/attn_bwd_ref.cpp
+    PROPERTIES LANGUAGE HIP
+)
+target_include_directories(attn_ref PUBLIC
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+)
+target_compile_options(attn_ref PRIVATE --offload-arch=${GPU_TARGETS})
+target_link_libraries(attn_ref PUBLIC hip::host)
+
+# ---------------------------------------------------------------------------
+# test_fwd executable
+#
+# test_fwd.cpp includes attn_fwd.h which contains __global__ kernels.
+# Mark it as LANGUAGE HIP so clang++ uses -x hip and sees <<< >>> syntax.
+# ---------------------------------------------------------------------------
+
+add_executable(test_fwd tests/test_fwd.cpp)
+set_source_files_properties(tests/test_fwd.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_fwd PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_fwd PRIVATE attn_ref hip::host)
+target_compile_options(test_fwd PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_bwd executable
+# ---------------------------------------------------------------------------
+
+add_executable(test_bwd tests/test_bwd.cpp)
+set_source_files_properties(tests/test_bwd.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_bwd PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_bwd PRIVATE attn_ref hip::host)
+target_compile_options(test_bwd PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_fwd_mfma executable
+#
+# Tests the fused MFMA forward kernel (attn_fwd_mfma.h).
+# ---------------------------------------------------------------------------
+
+add_executable(test_fwd_mfma tests/test_fwd_mfma.cpp)
+set_source_files_properties(tests/test_fwd_mfma.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_fwd_mfma PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_fwd_mfma PRIVATE attn_ref hip::host)
+target_compile_options(test_fwd_mfma PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_fwd_mfma_16x16 executable
+#
+# Tests the fused MFMA 16x16x16 forward kernel (attn_fwd_mfma_16x16.h).
+# ---------------------------------------------------------------------------
+
+add_executable(test_fwd_mfma_16x16 tests/test_fwd_mfma_16x16.cpp)
+set_source_files_properties(tests/test_fwd_mfma_16x16.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_fwd_mfma_16x16 PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_fwd_mfma_16x16 PRIVATE attn_ref hip::host)
+target_compile_options(test_fwd_mfma_16x16 PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_mfma_head_dims executable
+#
+# Forward MFMA 16x16 correctness for head_dim 128, 256, 512 (small config).
+# ---------------------------------------------------------------------------
+
+add_executable(test_mfma_head_dims tests/test_mfma_head_dims.cpp)
+set_source_files_properties(tests/test_mfma_head_dims.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_mfma_head_dims PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_mfma_head_dims PRIVATE attn_ref hip::host)
+target_compile_options(test_mfma_head_dims PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_fwd_mfma_multiq executable
+#
+# Tests multi-Q dispatch across 4x4x4 and 16x16x16 MFMA kernels.
+# ---------------------------------------------------------------------------
+
+add_executable(test_fwd_mfma_multiq tests/test_fwd_mfma_multiq.cpp)
+set_source_files_properties(tests/test_fwd_mfma_multiq.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_fwd_mfma_multiq PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_fwd_mfma_multiq PRIVATE attn_ref hip::host)
+target_compile_options(test_fwd_mfma_multiq PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_bwd_mfma_16x16 executable
+#
+# Tests the MFMA 16x16x16 backward kernels (attn_bwd_mfma_16x16.h).
+# ---------------------------------------------------------------------------
+
+add_executable(test_bwd_mfma_16x16 tests/test_bwd_mfma_16x16.cpp)
+set_source_files_properties(tests/test_bwd_mfma_16x16.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_bwd_mfma_16x16 PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_bwd_mfma_16x16 PRIVATE attn_ref hip::host)
+target_compile_options(test_bwd_mfma_16x16 PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_varlen_mfma_16x16 executable
+#
+# Unified varlen test for MFMA 16x16x16 forward + backward kernels.
+# ---------------------------------------------------------------------------
+
+add_executable(test_varlen_mfma_16x16 tests/test_varlen_mfma_16x16.cpp)
+set_source_files_properties(tests/test_varlen_mfma_16x16.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_varlen_mfma_16x16 PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_varlen_mfma_16x16 PRIVATE attn_ref hip::host)
+target_compile_options(test_varlen_mfma_16x16 PRIVATE -O3 --offload-arch=${GPU_TARGETS})
+
+# ---------------------------------------------------------------------------
+# test_small_seq_sweep executable
+#
+# Small-sequence sweep benchmark (seqlen 1..17, bs=2048, fwd+bwd, TE format).
+# ---------------------------------------------------------------------------
+
+add_executable(test_small_seq_sweep tests/test_small_seq_sweep.cpp)
+set_source_files_properties(tests/test_small_seq_sweep.cpp PROPERTIES LANGUAGE HIP)
+
+target_include_directories(test_small_seq_sweep PRIVATE
+    ${KERNEL_INCLUDE_DIR}
+    ${CMAKE_SOURCE_DIR}/ref
+    ${CMAKE_SOURCE_DIR}/tests
+)
+target_link_libraries(test_small_seq_sweep PRIVATE attn_ref hip::host)
+target_compile_options(test_small_seq_sweep PRIVATE -O3 --offload-arch=${GPU_TARGETS})