diff --git a/README.rst b/README.rst index 98ba14f55..084f170ec 100644 --- a/README.rst +++ b/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| @@ -275,6 +275,53 @@ 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 (gfx950 and gfx942) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +For workloads with very short sequences, ROCm TE supports these shapes on both gfx950 +(MI350) and gfx942 (MI300/MI325). Both ``THD`` (variable-length / ragged, e.g. +cross-attention) and ``BSHD`` (dense self-attention with ``s_q == s_kv``) layouts are +supported. The CK backend must be enabled (``NVTE_FUSED_ATTN_CK=1``, the default). + +Sequence length limits depend on layout: + +* **``THD``** — eligibility is based on the **runtime** maximum sequence length per batch + (from ``cu_seqlens`` on device). On **gfx950**, both Q and KV runtime max seqlen must be + **at most 16**. + +* **``BSHD``** — eligibility uses the **static** sequence length ``s_q == s_kv`` (no + ``cu_seqlens`` probe). Self-attention with ``2 <= s_q == s_kv <= 17`` is supported, on both gfx950 and gfx942. + +**gfx950 — traditional CK/AITER path (default).** + On gfx950, short-sequence problems use the regular CK fused-attention backend. No extra environment variables + are required; ``NVTE_FUSED_ATTN_CK_SMALLSEQ`` is ignored on this architecture. Apply the + sequence length limits above: ``THD`` up to **16** per side, ``BSHD`` up to **17**. + +**gfx942 — dedicated MFMA small-seq path (opt-in).** + On gfx942 only, set ``NVTE_FUSED_ATTN_CK_SMALLSEQ=1`` to route eligible problems through + dedicated CK small-sequence MFMA kernels that are more efficient than the general + fused-attention path for these shapes. When enabled, a problem is routed to the MFMA + 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; + * 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 (``THD``) or no mask (``BSHD``); + * sequence length within layout limits above. + + When using the JAX integration with ``THD`` layouts on this path, both of the following must be set before the process starts: + + * ``NVTE_FUSED_ATTN_CK_SMALLSEQ=1`` + * ``XLA_FLAGS='--xla_gpu_enable_command_buffer='`` — disables XLA GPU graph capture (command buffers / cudagraphs), which is incompatible with the runtime ``cu_seqlens`` segment check performed by the MFMA small-seq ``THD`` path. + + 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. diff --git a/ci/jax.sh b/ci/jax.sh index 95bb2d71f..9a1808a10 100755 --- a/ci/jax.sh +++ b/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 diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b8c4fb9e2..4a6478f8f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -432,6 +432,8 @@ class FusedAttnRunner: stripe_size: int | None = None num_segments_per_seq: int | None = None use_old_rng: bool = True #ROCm may use new-style RNG + # THD ragged segment setup for small-seqlen CK tests + use_thd_smallseq_segments: bool = False # Specifies sharding resources for distributed tests number_of_devices: int = 1 @@ -482,6 +484,8 @@ def _get_max_segments_per_sequence(self): if self.qkv_layout.is_thd(): if 90400 <= get_cudnn_version() < 90500: return self.num_segments_per_seq + elif is_hip_extension() and self.use_thd_smallseq_segments: + return self.num_segments_per_seq else: # +1 for testing runtime_segments < max_segments return self.num_segments_per_seq + 1 @@ -602,6 +606,81 @@ def _check_configs(self): "the F16_arbitrary_seqlen backend." ) + def _setup_segments_ck_smallseq(self, generate_random_segment_ids): + """ + Segment ids / seqlens for NVTE_FUSED_ATTN_CK_SMALLSEQ + padded ragged layouts. + + num_segments_per_seq follows max_seqlen_q; max_seqlen_q==1 uses a fixed Q row and + corrected seqlens_q. KV always uses generate_random_segment_ids. + """ + num_segments_per_seq = self.max_seqlen_q + if self.max_seqlen_q == 1: + # Q: one length-1 segment per batch. Take offsets_q from get_seqlens_and_offsets: this + # integration's SequenceDescriptor expects the intra-sequence offset form the helper + # produces (all-zero), NOT PR #542's cumulative arange form -- the arange form makes + # get_runtime_max_seqlen underflow and faults the GPU. Only override seqlens_q to ones, + # since the helper's seqlens are wrong here (bincount(length=1) drops the id==1 segment). + segment_ids_q = jnp.ones((self.batch_size, self.max_seqlen_q), dtype=jnp.int32) + segment_pos_q = jnp.zeros((self.batch_size, self.max_seqlen_q), dtype=jnp.int32) + pad_q = jnp.zeros((self.batch_size, self.max_seqlen_q), dtype=jnp.int32) + seqlens_q, offsets_q = get_seqlens_and_offsets(segment_ids_q) + seqlens_q = jnp.ones((self.batch_size, 1), dtype=jnp.int32) + else: + segment_ids_q, segment_pos_q, pad_q = generate_random_segment_ids( + self.batch_size, self.max_seqlen_q, num_segments_per_seq, seed=42 + ) + # Compute seqlens/offsets directly instead of using get_seqlens_and_offsets. + # get_seqlens_and_offsets uses bincount(length=max_seqlen) which cannot capture + # segment IDs equal to max_seqlen (when num_segments == max_seqlen_q, segment + # IDs range from 1 to max_seqlen_q). The missing segment plus the appended + # sentinel causes _fix_len_take in impl() to leak entries across batches. + # Since each Q segment has exactly 1 token (max_segment_size = max_seqlen_q // + # num_segments_per_seq = 1), we build seqlens as all-ones with no sentinels. + seqlens_q = jnp.ones((self.batch_size, num_segments_per_seq), dtype=jnp.int32) + offsets_q = jnp.concatenate( + [ + jnp.tile( + jnp.arange(num_segments_per_seq, dtype=jnp.int32)[None, :], + (self.batch_size, 1), + ), + jnp.full((self.batch_size, 1), -1, dtype=jnp.int32), + ], + axis=1, + ) + + min_segment_len = None if self.window_size is None else seqlens_q + segment_ids_kv, segment_pos_kv, pad_kv = generate_random_segment_ids( + self.batch_size, + self.max_seqlen_kv, + num_segments_per_seq, + seed=2024, + min_segment_len=min_segment_len, + ) + seqlens_kv, offsets_kv = get_seqlens_and_offsets(segment_ids_kv) + # get_seqlens_and_offsets derives seqlens via bincount(length=max_seqlen_kv), which drops + # the segment whose ID == max_seqlen_kv. That happens whenever num_segments_per_seq == + # max_seqlen_kv (i.e. self-attention), corrupting KV seqlens and yielding empty-KV + # softmax NaNs. Recompute seqlens with a wide-enough bincount so the last segment is + # retained; the offsets from _find_offsets are already correct. + kv_counts = jax.vmap(partial(jnp.bincount, length=self.max_seqlen_kv + 1))( + segment_ids_kv.astype(jnp.int32) + ) + seqlens_kv = kv_counts[..., 1:] + seqlens_kv = jnp.where(seqlens_kv, seqlens_kv, -1) + return ( + num_segments_per_seq, + segment_ids_q, + segment_pos_q, + pad_q, + seqlens_q, + offsets_q, + segment_ids_kv, + segment_pos_kv, + pad_kv, + seqlens_kv, + offsets_kv, + ) + def _setup_inputs(self): self._check_configs() @@ -742,34 +821,51 @@ def generate_random_segment_ids_and_pos( return segment_ids, segment_pos, segment_pad if self.qkv_layout.is_thd(): - self.segment_ids_q, self.segment_pos_q, self.pad_q = ( + if is_hip_extension() and self.use_thd_smallseq_segments: + ( + self.num_segments_per_seq, + self.segment_ids_q, + self.segment_pos_q, + self.pad_q, + self.seqlens_q, + self.offsets_q, + self.segment_ids_kv, + self.segment_pos_kv, + self.pad_kv, + self.seqlens_kv, + self.offsets_kv, + ) = self._setup_segments_ck_smallseq(generate_random_segment_ids_and_pos) + else: + self.segment_ids_q, self.segment_pos_q, self.pad_q = ( generate_random_segment_ids_and_pos( - self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42 - ) + self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42 + ) ) - self.seqlens_q, self.offsets_q = get_seqlens_and_offsets(self.segment_ids_q) - # TODO(rewang): record only self attention and find the reason of cross attention - if self.qkv_layout == QKVLayout.T3HD or self.max_seqlen_q == self.max_seqlen_kv: - self.segment_ids_kv = self.segment_ids_q - self.segment_pos_kv = self.segment_pos_q - self.pad_kv = self.pad_q - else: - # Force kv_len >= q_len for swa, otherwise, cuDNN kernels don't support - min_segment_len = None - if ( - self.window_size is not None or self.attn_mask_type.is_bottom_right() - ): # SWA or BRCM requires kv_len >= q_len - min_segment_len = self.seqlens_q - self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = ( + self.seqlens_q, self.offsets_q = get_seqlens_and_offsets(self.segment_ids_q) + # TODO(rewang): record only self attention and find the reason of cross attention + if self.qkv_layout == QKVLayout.T3HD or self.max_seqlen_q == self.max_seqlen_kv: + self.segment_ids_kv = self.segment_ids_q + self.segment_pos_kv = self.segment_pos_q + self.pad_kv = self.pad_q + else: + # Force kv_len >= q_len for swa, otherwise, cuDNN kernels don't support + min_segment_len = None + if ( + self.window_size is not None or self.attn_mask_type.is_bottom_right() + ): # SWA or BRCM requires kv_len >= q_len + min_segment_len = self.seqlens_q + self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = ( + ( generate_random_segment_ids_and_pos( - self.batch_size, - self.max_seqlen_kv, - self.num_segments_per_seq, - seed=2024, - min_segment_len=min_segment_len, - ) + self.batch_size, + self.max_seqlen_kv, + self.num_segments_per_seq, + seed=2024, + min_segment_len=min_segment_len, + ) + ) ) - self.seqlens_kv, self.offsets_kv = get_seqlens_and_offsets(self.segment_ids_kv) + self.seqlens_kv, self.offsets_kv = get_seqlens_and_offsets(self.segment_ids_kv) else: self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_valid_segment_ids_and_pos( self.batch_size, self.max_seqlen_q, pad_ratio @@ -2089,3 +2185,115 @@ def fused_fn(q, k, v): for name, x, y in zip(("dQ", "dK", "dV"), grads1, grads2): # Bitwise reproducibility across consecutive runs assert_allclose(x, y, atol=0, rtol=0, err_msg=f"{name} not bitwise reproducible") + + +# ROCm CK small-seq tests. +@pytest.fixture +def ck_smallseq_env(monkeypatch): + """ROCm test env for small-sequence CK attention tests. + + On gfx942 only: enable NVTE_FUSED_ATTN_CK_SMALLSEQ and require + XLA command buffers disabled. + + On gfx950, normal CK/aiter kernels are used. + """ + if not is_hip_extension(): + pytest.skip("CK unfused small-seq tests only on ROCm") + # This test uses the dedicated small-seq CK path (NVTE_FUSED_ATTN_CK_SMALLSEQ), + # which requires XLA GPU graph capture (command buffers) disabled via an empty + # --xla_gpu_enable_command_buffer= + if get_device_compute_capability(0) == 94: + # gfx942-only MFMA small-seq path; requires command buffers disabled. + if "xla_gpu_enable_command_buffer=" not in os.environ.get("XLA_FLAGS", ""): + pytest.skip("Test must be run with XLA_FLAGS='--xla_gpu_enable_command_buffer='") + monkeypatch.setenv("NVTE_FUSED_ATTN_CK_SMALLSEQ", "1") + yield + + +@pytest.mark.usefixtures("ck_smallseq_env") +class TestFusedAttnCkSmallseq: + """ + Small-sequence CK attention (1<=s_q<=16, 2<=s_kv<=16 THD self/cross and BSHD). + + On gfx942 with NVTE_FUSED_ATTN_CK_SMALLSEQ, exercises the dedicated MFMA + small-seq path. On gfx950 and newer, the same shapes run through normal CK/aiter + kernels, which now cover these cases at comparable performance. + """ + + @staticmethod + # fp16 is not supported on the CK small-seq path yet (the MFMA kernels are bf16-only); the + # backend guard rejects fp16 so it falls back to regular CK. Only bf16 is exercised here. + @pytest.mark.parametrize("dtype", [jnp.bfloat16], ids=["BF16"]) + @pytest.mark.parametrize("head_dim", [128, 256], ids=["d128", "d256"]) + @pytest.mark.parametrize("num_heads", [16, 32], ids=["h16", "h32"]) + @pytest.mark.parametrize( + "b, s_q, s_kv, qkv_layout", + [ + # cross-attention (s_q = 1, s_kv <= 16), THD + padding + pytest.param(4000, 1, 2, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-2"), + pytest.param(4000, 1, 3, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-3"), + pytest.param(4000, 1, 5, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-5"), + pytest.param(4000, 1, 6, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-6"), + pytest.param(4000, 1, 12, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-12"), + pytest.param(4000, 1, 13, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-13"), + pytest.param(4000, 1, 16, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-1-16"), + # cross-attention (s_q != s_kv), THD + padding + pytest.param(4000, 4, 8, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-4-8"), + pytest.param(4000, 8, 12, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-8-12"), + pytest.param(4000, 12, 16, QKVLayout.THD_THD_THD, id="cross-attn-THD_THD_THD-4000-12-16"), + # self-attention, THD + padding + pytest.param(4000, 2, 2, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-2-2"), + pytest.param(4000, 3, 3, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-3-3"), + pytest.param(4000, 5, 5, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-5-5"), + pytest.param(4000, 6, 6, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-6-6"), + pytest.param(4000, 8, 8, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-8-8"), + pytest.param(4000, 16, 16, QKVLayout.THD_THD_THD, id="self-attn-THD_THD_THD-4000-16-16"), + # self-attention, BSHD + pytest.param(4000, 2, 2, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-2-2"), + pytest.param(4000, 4, 4, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-4-4"), + pytest.param(4000, 8, 8, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-8-8"), + pytest.param(4000, 12, 12, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-12-12"), + pytest.param(4000, 16, 16, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-16-16"), + pytest.param(4000, 17, 17, QKVLayout.BSHD_BSHD_BSHD, id="self-attn-BSHD_BSHD_BSHD-4000-17-17"), + ], + ) + def test_smallseq( + dtype, + b, + s_q, + s_kv, + num_heads, + head_dim, + qkv_layout, + ): + """CK small-seq THD/BSHD: no bias; padding mask for THD, no mask for BSHD. + + """ + attn_mask_type = ( + AttnMaskType.NO_MASK + if qkv_layout == QKVLayout.BSHD_BSHD_BSHD + else AttnMaskType.PADDING_MASK + ) + runner = FusedAttnRunner( + batch_size=b, + max_seqlen_q=s_q, + max_seqlen_kv=s_kv, + num_heads_q=num_heads, + num_heads_kv=num_heads, + head_dim_qk=head_dim, + head_dim_v=head_dim, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask_type, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_prob=0.0, + use_old_rng=True, + dtype=dtype, + is_training=True, + qkv_layout=qkv_layout, + bias_shape=None, + window_size=None, + seq_desc_format=SeqDescFormat.Seqlens, + use_thd_smallseq_segments=True, + ) + runner.test_forward() + runner.test_backward() diff --git a/transformer_engine/common/ck_fused_attn/CMakeLists.txt b/transformer_engine/common/ck_fused_attn/CMakeLists.txt index 964ff6513..87fd6a5bd 100644 --- a/transformer_engine/common/ck_fused_attn/CMakeLists.txt +++ b/transformer_engine/common/ck_fused_attn/CMakeLists.txt @@ -211,7 +211,8 @@ set(ck_fused_attn_SOURCES) list(APPEND ck_fused_attn_SOURCES src/ck_fused_attn_fwd.cpp src/ck_fused_attn_bwd.cpp - src/ck_fused_attn_utils.cpp) + src/ck_fused_attn_utils.cpp + src/ck_fused_attn_smallseq.cpp) message(STATUS "Found the following fused attention files:") foreach(file ${ck_fused_attn_SOURCES}) @@ -235,6 +236,7 @@ target_include_directories(ck_fused_attn PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/inc target_include_directories(ck_fused_attn PRIVATE ${CK_INCLUDE_DIR} ${__CK_SOURCE_DIR}/example/ck_tile/01_fmha) target_include_directories(ck_fused_attn PRIVATE ${AITER_INCLUDE_DIR}) target_include_directories(ck_fused_attn PRIVATE ${__QOLA_INCLUDE_DIR}) +target_include_directories(ck_fused_attn PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src/small_seq_kernels") find_package(hip) target_link_directories(ck_fused_attn PUBLIC ${__AITER_MHA_PATH}) diff --git a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp index 91072f0ba..fe9528acc 100644 --- a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp +++ b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp @@ -170,6 +170,94 @@ int ck_attn_fwd_num_splits(const CKAttnFwdArgs& args); // Extra workspace in particular is needed for split-KV support. size_t ck_attn_fwd_workspace_size(const CKAttnFwdArgs& args); +// Gen the number of splits for split-KV support. +int ck_attn_fwd_num_splits(const CKAttnFwdArgs& args); + +// Gen the workspace size for fwd config. +// Extra workspace in particular is needed for split-KV support. +size_t ck_attn_fwd_workspace_size(const CKAttnFwdArgs& args); + +uint64_t get_runtime_max_seqlen(uint64_t b, + const void* cu_seqlen_ptr, + const void* cu_seqlen_padded_ptr, + void* workspace, + hipStream_t stream); + +// --------------------------------------------------------------------------- +// Small-sequence attention (gfx942/gfx950, NVTE_FUSED_ATTN_CK_SMALLSEQ=1) +// --------------------------------------------------------------------------- +size_t small_seq_thd_extra_workspace_bytes(); + +void ck_attn_smallseq_fwd_thd(size_t batch_size, + size_t num_heads, + size_t head_dim_qk, + size_t max_tokens_q, + size_t max_tokens_kv, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* o_ptr, + void* softmax_lse_ptr, + const void* cu_seqlens_q_ptr, + const void* cu_seqlens_q_padded_ptr, + const void* cu_seqlens_kv_ptr, + const void* cu_seqlens_kv_padded_ptr, + DType dtype, + hipStream_t stream); + +void ck_attn_smallseq_bwd_thd(size_t batch_size, + size_t num_heads, + size_t head_dim_qk, + size_t max_tokens_q, + size_t max_tokens_kv, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* do_ptr, + const void* softmax_lse_ptr, + void* dq_ptr, + void* dk_ptr, + void* dv_ptr, + const void* cu_seqlens_q_ptr, + const void* cu_seqlens_q_padded_ptr, + const void* cu_seqlens_kv_ptr, + const void* cu_seqlens_kv_padded_ptr, + DType dtype, + hipStream_t stream); + +void ck_attn_smallseq_fwd_bshd(size_t batch_size, + size_t num_heads, + size_t seqlen_q, + size_t seqlen_kv, + size_t head_dim_qk, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* o_ptr, + void* softmax_lse_ptr, + DType dtype, + hipStream_t stream); + +void ck_attn_smallseq_bwd_bshd(size_t batch_size, + size_t num_heads, + size_t seqlen_q, + size_t seqlen_kv, + size_t head_dim_qk, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* do_ptr, + const void* softmax_lse_ptr, + void* dq_ptr, + void* dk_ptr, + void* dv_ptr, + DType dtype, + hipStream_t stream); + }//namespace ck_fused_attn #endif // CK_FUSED_ATTN_H diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_smallseq.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_smallseq.cpp new file mode 100644 index 000000000..2eec76fb0 --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_smallseq.cpp @@ -0,0 +1,437 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include "ck_fused_attn/ck_fused_attn.hpp" + +#include +#include +#include + +#include "attn_bwd_mfma_16x16.h" +#include "attn_fwd_mfma_dispatch.h" + +#define CK_SMALL_SEQ_TYPE_SWITCH_16BIT(DTYPE, TYPE_NAME, ...) \ + do { \ + if((DTYPE) == DType::kBFloat16) { \ + using TYPE_NAME = hip_bfloat16; \ + __VA_ARGS__ \ + } else if((DTYPE) == DType::kFloat16) { \ + using TYPE_NAME = __half; \ + __VA_ARGS__ \ + } \ + } while(0) + +namespace ck_fused_attn { + +namespace { + +using namespace small_seq_kernels; + +template +void launch_fwd_thd_inst(size_t batch, + int total_padded_q, + float attn_scale, + const T* Q, + const T* K, + const T* V, + T* O, + float* softmax_lse, + const int* cu_q, + const int* cu_qp, + const int* cu_kv, + const int* cu_kvp, + hipStream_t stream) { + using Config = + FmhaKernelConfig; + using Launcher = AttnForwardMfmaDispatchLauncher; + Launcher::run_attn_fwd_kernel(Q, K, V, nullptr, 0.0f, attn_scale, O, softmax_lse, + 0, cu_q, cu_qp, cu_kv, cu_kvp, total_padded_q, + static_cast(batch), stream); + HIP_CHECK(hipGetLastError()); +} + +template +void launch_fwd_bshd_inst(size_t batch, + int seqlen_q, + int seqlen_kv, + float attn_scale, + const T* Q, + const T* K, + const T* V, + T* O, + float* softmax_lse, + hipStream_t stream) { + using Config = + FmhaKernelConfig; + using Launcher = AttnForwardMfmaDispatchLauncher; + const int total_tokens_q = static_cast(batch) * seqlen_q; + Launcher::run_attn_fwd_kernel(Q, K, V, nullptr, 0.0f, attn_scale, O, softmax_lse, seqlen_q, + nullptr, nullptr, nullptr, nullptr, total_tokens_q, + static_cast(batch), stream); + (void)seqlen_kv; + HIP_CHECK(hipGetLastError()); +} + +template +void launch_fwd_thd_dispatch(size_t batch, + size_t num_heads, + int head_dim, + int total_padded_q, + float attn_scale, + const T* Q, + const T* K, + const T* V, + T* O, + float* softmax_lse, + const int* cu_q, + const int* cu_qp, + const int* cu_kv, + const int* cu_kvp, + hipStream_t stream) { + if(num_heads == 16) { + if(head_dim == 128) { + launch_fwd_thd_inst(batch, total_padded_q, attn_scale, Q, K, V, O, softmax_lse, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + if(head_dim == 256) { + launch_fwd_thd_inst(batch, total_padded_q, attn_scale, Q, K, V, O, softmax_lse, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + } + if(num_heads == 32) { + if(head_dim == 128) { + launch_fwd_thd_inst(batch, total_padded_q, attn_scale, Q, K, V, O, softmax_lse, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + if(head_dim == 256) { + launch_fwd_thd_inst(batch, total_padded_q, attn_scale, Q, K, V, O, softmax_lse, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + } +} + +template +void launch_fwd_bshd_dispatch(size_t batch, + size_t num_heads, + int head_dim, + int seqlen_q, + int seqlen_kv, + float attn_scale, + const T* Q, + const T* K, + const T* V, + T* O, + float* softmax_lse, + hipStream_t stream) { + if(num_heads == 16) { + if(head_dim == 128) { + launch_fwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, O, + softmax_lse, stream); + return; + } + if(head_dim == 256) { + launch_fwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, O, + softmax_lse, stream); + return; + } + } + if(num_heads == 32) { + if(head_dim == 128) { + launch_fwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, O, + softmax_lse, stream); + return; + } + if(head_dim == 256) { + launch_fwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, O, + softmax_lse, stream); + return; + } + } +} + +template +void launch_bwd_thd_inst(size_t batch, + float attn_scale, + const T* Q, + const T* K, + const T* V, + const T* dO, + const float* softmax_lse, + T* dQ, + T* dK, + T* dV, + const int* cu_q, + const int* cu_qp, + const int* cu_kv, + const int* cu_kvp, + hipStream_t stream) { + using Config = + FmhaKernelConfig; + using Launcher = AttnBackwardMfma16x16KernelLauncher; + Launcher::run_attn_bwd_kernel(Q, K, V, dO, softmax_lse, dQ, dK, dV, attn_scale, 0, cu_q, cu_qp, + cu_kv, cu_kvp, static_cast(batch), stream); + HIP_CHECK(hipGetLastError()); +} + +template +void launch_bwd_bshd_inst(size_t batch, + int seqlen_q, + int seqlen_kv, + float attn_scale, + const T* Q, + const T* K, + const T* V, + const T* dO, + const float* softmax_lse, + T* dQ, + T* dK, + T* dV, + hipStream_t stream) { + using Config = + FmhaKernelConfig; + using Launcher = AttnBackwardMfma16x16KernelLauncher; + Launcher::run_attn_bwd_kernel(Q, K, V, dO, softmax_lse, dQ, dK, dV, attn_scale, seqlen_q, + nullptr, nullptr, nullptr, nullptr, static_cast(batch), + stream); + (void)seqlen_kv; + HIP_CHECK(hipGetLastError()); +} + +template +void launch_bwd_thd_dispatch(size_t batch, + size_t num_heads, + int head_dim, + float attn_scale, + const T* Q, + const T* K, + const T* V, + const T* dO, + const float* softmax_lse, + T* dQ, + T* dK, + T* dV, + const int* cu_q, + const int* cu_qp, + const int* cu_kv, + const int* cu_kvp, + hipStream_t stream) { + if(num_heads == 16) { + if(head_dim == 128) { + launch_bwd_thd_inst(batch, attn_scale, Q, K, V, dO, softmax_lse, dQ, dK, dV, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + if(head_dim == 256) { + launch_bwd_thd_inst(batch, attn_scale, Q, K, V, dO, softmax_lse, dQ, dK, dV, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + } + if(num_heads == 32) { + if(head_dim == 128) { + launch_bwd_thd_inst(batch, attn_scale, Q, K, V, dO, softmax_lse, dQ, dK, dV, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + if(head_dim == 256) { + launch_bwd_thd_inst(batch, attn_scale, Q, K, V, dO, softmax_lse, dQ, dK, dV, + cu_q, cu_qp, cu_kv, cu_kvp, stream); + return; + } + } +} + +template +void launch_bwd_bshd_dispatch(size_t batch, + size_t num_heads, + int head_dim, + int seqlen_q, + int seqlen_kv, + float attn_scale, + const T* Q, + const T* K, + const T* V, + const T* dO, + const float* softmax_lse, + T* dQ, + T* dK, + T* dV, + hipStream_t stream) { + if(num_heads == 16) { + if(head_dim == 128) { + launch_bwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, dO, + softmax_lse, dQ, dK, dV, stream); + return; + } + if(head_dim == 256) { + launch_bwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, dO, + softmax_lse, dQ, dK, dV, stream); + return; + } + } + if(num_heads == 32) { + if(head_dim == 128) { + launch_bwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, dO, + softmax_lse, dQ, dK, dV, stream); + return; + } + if(head_dim == 256) { + launch_bwd_bshd_inst(batch, seqlen_q, seqlen_kv, attn_scale, Q, K, V, dO, + softmax_lse, dQ, dK, dV, stream); + return; + } + } +} + +} // namespace + +size_t small_seq_thd_extra_workspace_bytes() { + // [max_seqlen_q probe][max_seqlen_kv probe] for THD runtime eligibility checks. + return 2 * sizeof(uint64_t); +} + +void ck_attn_smallseq_fwd_thd(size_t batch_size, + size_t num_heads, + size_t head_dim_qk, + size_t max_tokens_q, + size_t max_tokens_kv, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* o_ptr, + void* softmax_lse_ptr, + const void* cu_seqlens_q_ptr, + const void* cu_seqlens_q_padded_ptr, + const void* cu_seqlens_kv_ptr, + const void* cu_seqlens_kv_padded_ptr, + DType dtype, + hipStream_t stream) { + const int* cu_q = static_cast(cu_seqlens_q_ptr); + const int* cu_qp = static_cast(cu_seqlens_q_padded_ptr); + const int* cu_kv = static_cast(cu_seqlens_kv_ptr); + const int* cu_kvp = static_cast(cu_seqlens_kv_padded_ptr); + float* softmax_lse = static_cast(softmax_lse_ptr); + const int total_padded_q = static_cast(max_tokens_q); + const int hd = static_cast(head_dim_qk); + + CK_SMALL_SEQ_TYPE_SWITCH_16BIT(dtype, T, { + const T* Q = static_cast(q_ptr); + const T* K = static_cast(k_ptr); + const T* V = static_cast(v_ptr); + T* O = static_cast(o_ptr); + launch_fwd_thd_dispatch(batch_size, num_heads, hd, total_padded_q, attn_scale, Q, K, V, O, + softmax_lse, cu_q, cu_qp, cu_kv, cu_kvp, stream); + }); +} + +void ck_attn_smallseq_fwd_bshd(size_t batch_size, + size_t num_heads, + size_t seqlen_q, + size_t seqlen_kv, + size_t head_dim_qk, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* o_ptr, + void* softmax_lse_ptr, + DType dtype, + hipStream_t stream) { + float* softmax_lse = static_cast(softmax_lse_ptr); + const int hd = static_cast(head_dim_qk); + const int sq = static_cast(seqlen_q); + const int skv = static_cast(seqlen_kv); + + CK_SMALL_SEQ_TYPE_SWITCH_16BIT(dtype, T, { + const T* Q = static_cast(q_ptr); + const T* K = static_cast(k_ptr); + const T* V = static_cast(v_ptr); + T* O = static_cast(o_ptr); + launch_fwd_bshd_dispatch(batch_size, num_heads, hd, sq, skv, attn_scale, Q, K, V, O, + softmax_lse, stream); + }); +} + +void ck_attn_smallseq_bwd_thd(size_t batch_size, + size_t num_heads, + size_t head_dim_qk, + size_t max_tokens_q, + size_t max_tokens_kv, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* do_ptr, + const void* softmax_lse_ptr, + void* dq_ptr, + void* dk_ptr, + void* dv_ptr, + const void* cu_seqlens_q_ptr, + const void* cu_seqlens_q_padded_ptr, + const void* cu_seqlens_kv_ptr, + const void* cu_seqlens_kv_padded_ptr, + DType dtype, + hipStream_t stream) { + const int* cu_q = static_cast(cu_seqlens_q_ptr); + const int* cu_qp = static_cast(cu_seqlens_q_padded_ptr); + const int* cu_kv = static_cast(cu_seqlens_kv_ptr); + const int* cu_kvp = static_cast(cu_seqlens_kv_padded_ptr); + const float* softmax_lse = static_cast(softmax_lse_ptr); + const int hd = static_cast(head_dim_qk); + + CK_SMALL_SEQ_TYPE_SWITCH_16BIT(dtype, T, { + const T* Q = static_cast(q_ptr); + const T* K = static_cast(k_ptr); + const T* V = static_cast(v_ptr); + const T* dO = static_cast(do_ptr); + T* dQ = static_cast(dq_ptr); + T* dK = static_cast(dk_ptr); + T* dV = static_cast(dv_ptr); + launch_bwd_thd_dispatch(batch_size, num_heads, hd, attn_scale, Q, K, V, dO, softmax_lse, dQ, + dK, dV, cu_q, cu_qp, cu_kv, cu_kvp, stream); + }); +} + +void ck_attn_smallseq_bwd_bshd(size_t batch_size, + size_t num_heads, + size_t seqlen_q, + size_t seqlen_kv, + size_t head_dim_qk, + float attn_scale, + const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* do_ptr, + const void* softmax_lse_ptr, + void* dq_ptr, + void* dk_ptr, + void* dv_ptr, + DType dtype, + hipStream_t stream) { + const float* softmax_lse = static_cast(softmax_lse_ptr); + const int hd = static_cast(head_dim_qk); + const int sq = static_cast(seqlen_q); + const int skv = static_cast(seqlen_kv); + + CK_SMALL_SEQ_TYPE_SWITCH_16BIT(dtype, T, { + const T* Q = static_cast(q_ptr); + const T* K = static_cast(k_ptr); + const T* V = static_cast(v_ptr); + const T* dO = static_cast(do_ptr); + T* dQ = static_cast(dq_ptr); + T* dK = static_cast(dk_ptr); + T* dV = static_cast(dv_ptr); + launch_bwd_bshd_dispatch(batch_size, num_heads, hd, sq, skv, attn_scale, Q, K, V, dO, + softmax_lse, dQ, dK, dV, stream); + }); +} + +} // namespace ck_fused_attn diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.hpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.hpp index a926d230d..e166a8fec 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.hpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.hpp @@ -51,8 +51,6 @@ BiasShape get_bias_shape(uint64_t b, uint64_t h, uint64_t bias_b, uint64_t bias_ //get ck_tile bias_type and CK_FUSED_ATTN bias_shape from a fwd/bwd args struct std::pair get_ck_bias_type_shape(const CKAttnCommonArgs* args); -uint64_t get_runtime_max_seqlen(uint64_t b, const void* cu_seqlen_ptr, const void* cu_seqlen_padded_ptr, void* workspace, hipStream_t stream); - // This helper merely standardizes the logging to make it a bit easier to parse // through it at a glance while guaranteeing uniformity. template diff --git a/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_bwd_mfma_16x16.h b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_bwd_mfma_16x16.h new file mode 100644 index 000000000..79368a5cc --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_bwd_mfma_16x16.h @@ -0,0 +1,706 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once + +#include "attn_common.h" +#include +#include + +#ifndef ATTN_MFMA_TYPES_DEFINED +#define ATTN_MFMA_TYPES_DEFINED +using bhalf_t = __bf16; +using bf16x4 = __bf16 __attribute__((ext_vector_type(4))); +using bf16x8 = __bf16 __attribute__((ext_vector_type(8))); +using floatx4 = float __attribute__((ext_vector_type(4))); +#endif + +#ifndef CEIL_DIV +#define CEIL_DIV(a, b) (((a) + (b)-1) / (b)) +#endif + +namespace small_seq_kernels { + +template +__device__ __forceinline__ bf16x8 bwd_load_cvt_bf16x8(const T* src) +{ + if constexpr(sizeof(T) == 2) + { + return *(const bf16x8*)src; + } + else + { + bf16x8 r; + #pragma unroll + for(int i = 0; i < 8; i++) + r[i] = static_cast(src[i]); + return r; + } +} + +// --------------------------------------------------------------------------- +// grad_V kernel: grad_V = attn^T @ grad_O +// Grid: (1, head_num, bs), Block: 256 +// --------------------------------------------------------------------------- + +template +__launch_bounds__(256, 1) +__global__ void fmha_bwd_grad_v_mfma_16x16_kernel( + const T* Q, + const T* K, + const float* softmax_lse, + const T* grad_O, + T* grad_V, + float scale, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded) +{ + constexpr int head_dim = Config::head_dim; + constexpr int head_num = Config::head_num; + constexpr int max_seq_kv = Config::max_seq_kv; + constexpr int max_seq_q = Config::max_seq_q; + constexpr int hd_pad = head_dim + 4; + constexpr int q_tiles = CEIL_DIV(max_seq_q, 16); + constexpr int kv_tiles = CEIL_DIV(max_seq_kv, 16); + constexpr int lds_q_rows = q_tiles * 16; + constexpr int lds_kv_rows = kv_tiles * 16; + constexpr int attn_pad = lds_kv_rows + 4; + + const int batch_idx = blockIdx.z; + const int head_idx = blockIdx.y; + const int tid = threadIdx.x; + const int warp_id = tid / 64; + const int lane_id = tid % 64; + const int lane_row = lane_id / 16; + const int lane_col = lane_id % 16; + + const int actual_q = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_q); + if(actual_q == 0) + return; + + const int seq_kv = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_kv); + const int q_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_q_padded); + const int kv_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_kv_padded); + + __shared__ __attribute__((aligned(128))) float attn_lds[lds_q_rows * attn_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t Q_lds_bwd[lds_q_rows * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t K_lds_bwd[lds_kv_rows * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t dO_lds[lds_q_rows * hd_pad]; + + // Load Q → Q_lds_bwd + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_q_rows; r += (256 / threads_per_row)) + { + if(r < actual_q) + { + const T* q_src = Q + ((size_t)(q_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&Q_lds_bwd[r * hd_pad + col]) = bwd_load_cvt_bf16x8(q_src + col); + } + else + *(bf16x8*)(&Q_lds_bwd[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + + // Load K → K_lds_bwd + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_kv_rows; r += (256 / threads_per_row)) + { + if(r < seq_kv) + { + const T* k_src = K + ((size_t)(kv_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&K_lds_bwd[r * hd_pad + col]) = bwd_load_cvt_bf16x8(k_src + col); + } + else + *(bf16x8*)(&K_lds_bwd[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + + __syncthreads(); + + // QK^T (same MFMA tiling as forward) → exp(S - LSE) = P + float P_reg[q_tiles * kv_tiles * 4]; + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + floatx4 acc = {0, 0, 0, 0}; + constexpr int total_hd_tiles = CEIL_DIV(head_dim, 16); + + #pragma unroll + for(int k = 0; k < total_hd_tiles; ++k) + { + const int dim_off = k * 16; + bf16x4 a = *(const bf16x4*)(&Q_lds_bwd[(qt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + bf16x4 b = *(const bf16x4*)(&K_lds_bwd[(kvt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + int reg_base = (qt * kv_tiles + kvt) * 4; + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + int kv_pos = kvt * 16 + lane_col; + bool masked = (kv_pos >= seq_kv) || (q_row >= actual_q); + if constexpr(Config::mask_type == CausalMaskType::TOP_LEFT) + { + if(kv_pos > q_row) + masked = true; + } + float S = acc[i] * scale; + float lse = + softmax_lse[((size_t)(q_offset + q_row) * head_num + head_idx)]; + float pr = masked ? 0.0f : expf(S - lse); + P_reg[reg_base + i] = pr; + } + } + } + + // Scatter P_reg → attn_lds (same pattern as former workspace write) + if(warp_id == 0) + { + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + int kv_pos = kvt * 16 + lane_col; + if(q_row < actual_q && kv_pos < max_seq_kv) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + float w = (kv_pos < seq_kv) ? P_reg[reg_idx] : 0.0f; + attn_lds[q_row * attn_pad + kv_pos] = w; + } + } + } + } + } + + __syncthreads(); + + // Load grad_O → dO_lds + { + constexpr int threads_per_row = head_dim / 8; + const int do_row = tid / threads_per_row; + const int do_col = (tid % threads_per_row) * 8; + + for(int r = do_row; r < lds_q_rows; r += (256 / threads_per_row)) + { + if(r < actual_q) + { + const T* do_src = grad_O + ((size_t)(q_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&dO_lds[r * hd_pad + do_col]) = bwd_load_cvt_bf16x8(do_src + do_col); + } + else + { + *(bf16x8*)(&dO_lds[r * hd_pad + do_col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + } + + __syncthreads(); + + // MFMA: grad_V = attn^T @ grad_O (4 warps split head_dim) + constexpr int BK = 64; + + #pragma unroll + for(int kv_tile = 0; kv_tile < kv_tiles; kv_tile++) + { + constexpr int total_d_tiles = CEIL_DIV(head_dim, BK); + + #pragma unroll + for(int d = 0; d < total_d_tiles; d++) + { + const int dim_idx = d * BK + warp_id * 16; + + floatx4 acc = {0, 0, 0, 0}; + + #pragma unroll + for(int q_tile = 0; q_tile < q_tiles; q_tile++) + { + bf16x4 a; + #pragma unroll + for(int k = 0; k < 4; k++) + { + int q_row = q_tile * 16 + lane_row * 4 + k; + int kv_pos = kv_tile * 16 + lane_col; + float val = (q_row < actual_q && kv_pos < seq_kv) + ? attn_lds[q_row * attn_pad + kv_pos] : 0.0f; + a[k] = static_cast(val); + } + + // B: dO[q, d] + bf16x4 b; + #pragma unroll + for(int k = 0; k < 4; k++) + { + int q_row = q_tile * 16 + lane_row * 4 + k; + b[k] = dO_lds[q_row * hd_pad + dim_idx + lane_col]; + } + + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + // Write grad_V + #pragma unroll + for(int i = 0; i < 4; i++) + { + int kv_pos = kv_tile * 16 + lane_row * 4 + i; + if(kv_pos < seq_kv) + { + int gv_idx = (kv_offset + kv_pos) * head_num * head_dim + + head_idx * head_dim + dim_idx + lane_col; + grad_V[gv_idx] = static_cast(acc[i]); + } + } + } + } +} + +// --------------------------------------------------------------------------- +// Fused backward kernel: grad_attn → softmax_bwd → grad_Q + grad_K +// Grid: (1, head_num, bs), Block: 256 +// --------------------------------------------------------------------------- + +template +__launch_bounds__(256, 1) +__global__ void fmha_bwd_fused_mfma_16x16_kernel( + const T* Q, + const T* K, + const T* V, + const T* grad_O, + const float* softmax_lse, + T* grad_Q, + T* grad_K, + float scale, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded) +{ + constexpr int head_dim = Config::head_dim; + constexpr int head_num = Config::head_num; + constexpr int max_seq_kv = Config::max_seq_kv; + constexpr int max_seq_q = Config::max_seq_q; + constexpr int hd_pad = head_dim + 4; + constexpr int q_tiles = CEIL_DIV(max_seq_q, 16); + constexpr int kv_tiles = CEIL_DIV(max_seq_kv, 16); + constexpr int lds_q_rows = q_tiles * 16; + constexpr int lds_kv_rows = kv_tiles * 16; + constexpr int lds_sm_stride = lds_kv_rows + 4; + + const int batch_idx = blockIdx.z; + const int head_idx = blockIdx.y; + const int tid = threadIdx.x; + const int warp_id = tid / 64; + const int lane_id = tid % 64; + const int lane_row = lane_id / 16; + const int lane_col = lane_id % 16; + + const int actual_q = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_q); + if(actual_q == 0) + return; + + const int seq_kv = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_kv); + const int q_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_q_padded); + const int kv_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_kv_padded); + + __shared__ __attribute__((aligned(128))) bhalf_t Q_lds[lds_q_rows * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t dO_lds[lds_q_rows * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t KV_lds[lds_kv_rows * hd_pad]; + __shared__ float SM_lds[lds_q_rows * lds_sm_stride]; + + // Load Q → Q_lds + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_q_rows; r += (256 / threads_per_row)) + { + if(r < actual_q) + { + const T* q_src = Q + ((size_t)(q_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&Q_lds[r * hd_pad + col]) = bwd_load_cvt_bf16x8(q_src + col); + } + else + { + *(bf16x8*)(&Q_lds[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + } + + // Load dO → dO_lds + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_q_rows; r += (256 / threads_per_row)) + { + if(r < actual_q) + { + const T* do_src = grad_O + ((size_t)(q_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&dO_lds[r * hd_pad + col]) = bwd_load_cvt_bf16x8(do_src + col); + } + else + { + *(bf16x8*)(&dO_lds[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + } + + // Load V → KV_lds + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + const int clamped_max = max(seq_kv - 1, 0); + + for(int r = row; r < lds_kv_rows; r += (256 / threads_per_row)) + { + const int clamped_r = min(r, clamped_max); + const T* v_src = V + ((size_t)(kv_offset + clamped_r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&KV_lds[r * hd_pad + col]) = bwd_load_cvt_bf16x8(v_src + col); + } + } + + __syncthreads(); + + // grad_attn = dO @ V^T via MFMA (all 4 warps redundant) + float grad_attn[q_tiles * kv_tiles * 4]; + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + floatx4 acc = {0, 0, 0, 0}; + constexpr int total_d_tiles = CEIL_DIV(head_dim, 16); + + #pragma unroll + for(int dtile = 0; dtile < total_d_tiles; dtile++) + { + const int dim_off = dtile * 16; + // A: dO[q, d] + bf16x4 a = *(const bf16x4*)(&dO_lds[(qt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + // B: V[kv, d] + bf16x4 b = *(const bf16x4*)(&KV_lds[(kvt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + int reg_base = (qt * kv_tiles + kvt) * 4; + #pragma unroll + for(int i = 0; i < 4; i++) + grad_attn[reg_base + i] = acc[i]; + } + } + + // Reload K into KV_lds (overwrite V) and recompute P_ij = exp(S_ij - LSE_i) + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + const int clamped_max = max(seq_kv - 1, 0); + + for(int r = row; r < lds_kv_rows; r += (256 / threads_per_row)) + { + const int clamped_r = min(r, clamped_max); + const T* k_src = K + ((size_t)(kv_offset + clamped_r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&KV_lds[r * hd_pad + col]) = bwd_load_cvt_bf16x8(k_src + col); + } + } + + __syncthreads(); + + float attn_reg[q_tiles * kv_tiles * 4]; + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + floatx4 acc_s = {0, 0, 0, 0}; + constexpr int total_hd_tiles = CEIL_DIV(head_dim, 16); + + #pragma unroll + for(int k = 0; k < total_hd_tiles; ++k) + { + const int dim_off = k * 16; + bf16x4 a = *(const bf16x4*)(&Q_lds[(qt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + bf16x4 b = *(const bf16x4*)(&KV_lds[(kvt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + acc_s = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc_s, 0, 0, 0); + } + + int reg_base = (qt * kv_tiles + kvt) * 4; + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + int kv_pos = kvt * 16 + lane_col; + bool masked = (kv_pos >= seq_kv) || (q_row >= actual_q); + if constexpr(Config::mask_type == CausalMaskType::TOP_LEFT) + { + if(kv_pos > q_row) + masked = true; + } + float S = acc_s[i] * scale; + float lse = + softmax_lse[((size_t)(q_offset + q_row) * head_num + head_idx)]; + attn_reg[reg_base + i] = masked ? 0.0f : expf(S - lse); + } + } + } + + // Softmax backward: grad_score = attn * (grad_attn - dot_sum) + float grad_score[q_tiles * kv_tiles * 4]; + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + + // dot_sum = sum_kv(grad_attn * attn) + float dot_sum = 0.0f; + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + float partial = grad_attn[reg_idx] * attn_reg[reg_idx]; + + // Reduce across lane_col + #pragma unroll + for(int off = 8; off > 0; off /= 2) + partial += __shfl_xor(partial, off, 16); + + dot_sum += partial; + } + + // grad_score = attn * (grad_attn - dot_sum) + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + int kv_pos = kvt * 16 + lane_col; + float gs = attn_reg[reg_idx] * (grad_attn[reg_idx] - dot_sum); + + // Zero invalid + if(q_row >= actual_q || kv_pos >= seq_kv) + gs = 0.0f; + + grad_score[reg_idx] = gs; + } + } + } + + // Write grad_scores → SM_lds + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + int kv_pos = kvt * 16 + lane_col; + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + SM_lds[q_row * lds_sm_stride + kv_pos] = grad_score[reg_idx]; + } + } + } + + __syncthreads(); + + // K is already in KV_lds from P recomputation + + // grad_Q = grad_scores @ K * scale (4 warps split head_dim) + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + constexpr int BK = 64; + constexpr int total_d_tiles = CEIL_DIV(head_dim, BK); + + #pragma unroll + for(int d = 0; d < total_d_tiles; d++) + { + const int dim_idx = d * BK + warp_id * 16; + floatx4 acc = {0, 0, 0, 0}; + + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + // A: grad_scores (transposed SM_lds read) + bf16x4 a; + #pragma unroll + for(int k = 0; k < 4; k++) + { + int q_row = qt * 16 + lane_col; + int kv_pos = kvt * 16 + lane_row * 4 + k; + a[k] = static_cast(SM_lds[q_row * lds_sm_stride + kv_pos]); + } + + // B: K[kv, d] + bf16x4 b; + const int kv_base = kvt * 16; + #pragma unroll + for(int k = 0; k < 4; k++) + { + b[k] = KV_lds[(kv_base + lane_row * 4 + k) * hd_pad + dim_idx + lane_col]; + } + + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + // Write grad_Q + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + if(q_row < actual_q) + { + int gq_idx = ((size_t)(q_offset + q_row) * head_num + head_idx) * head_dim + + dim_idx + lane_col; + grad_Q[gq_idx] = static_cast(acc[i] * scale); + } + } + } + } + + // grad_K = grad_scores^T @ Q * scale (4 warps split head_dim) + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + constexpr int BK = 64; + constexpr int total_d_tiles = CEIL_DIV(head_dim, BK); + + #pragma unroll + for(int d = 0; d < total_d_tiles; d++) + { + const int dim_idx = d * BK + warp_id * 16; + floatx4 acc = {0, 0, 0, 0}; + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + // A: grad_scores^T (direct SM_lds read) + bf16x4 a; + #pragma unroll + for(int k = 0; k < 4; k++) + { + int q_row = qt * 16 + lane_row * 4 + k; + int kv_pos = kvt * 16 + lane_col; + a[k] = static_cast(SM_lds[q_row * lds_sm_stride + kv_pos]); + } + + // B: Q[q, d] + bf16x4 b; + const int q_base = qt * 16; + #pragma unroll + for(int k = 0; k < 4; k++) + { + b[k] = Q_lds[(q_base + lane_row * 4 + k) * hd_pad + dim_idx + lane_col]; + } + + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + // Write grad_K + #pragma unroll + for(int i = 0; i < 4; i++) + { + int kv_pos = kvt * 16 + lane_row * 4 + i; + if(kv_pos < seq_kv) + { + int gk_idx = (kv_offset + kv_pos) * head_num * head_dim + + head_idx * head_dim + dim_idx + lane_col; + grad_K[gk_idx] = static_cast(acc[i] * scale); + } + } + } + } +} + +// --------------------------------------------------------------------------- +// AttnBackwardMfma16x16KernelLauncher — Grid: (1, head_num, bs), Block: 256 +// --------------------------------------------------------------------------- + +template +struct AttnBackwardMfma16x16KernelLauncher +{ + using bwd_softmax_aux_scalar = float; + + /// Option A: backward recomputes P from Q, K, and softmax_lse — no P workspace. + static size_t calc_workspace_size(int total_padded_q) + { + (void)total_padded_q; + return 0; + } + + static void run_attn_bwd_kernel(const T* Q, + const T* K, + const T* V, + const T* grad_O, + const float* softmax_lse, + T* grad_Q, + T* grad_K, + T* grad_V, + float sqr_dk_scale, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded, + int batch, + hipStream_t stream = 0) + { + float scale = sqr_dk_scale; + + // Batch is a runtime argument mapped to the grid z-dimension. + dim3 grid(1, Config::head_num, batch); + dim3 block(256); + + // Kernel B: grad_V = P^T @ grad_O (P recomputed from Q, K, LSE) + fmha_bwd_grad_v_mfma_16x16_kernel<<>>( + Q, K, softmax_lse, grad_O, grad_V, scale, uniform_seq_len, + cu_seqlens_q, cu_seqlens_q_padded, + cu_seqlens_kv, cu_seqlens_kv_padded); + + // Kernel A: fused grad_attn / softmax_bwd / grad_Q / grad_K + fmha_bwd_fused_mfma_16x16_kernel<<>>( + Q, K, V, grad_O, softmax_lse, + grad_Q, grad_K, scale, uniform_seq_len, + cu_seqlens_q, cu_seqlens_q_padded, + cu_seqlens_kv, cu_seqlens_kv_padded); + } +}; + +} // namespace small_seq_kernels diff --git a/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_common.h b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_common.h new file mode 100644 index 000000000..dd140f0ab --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_common.h @@ -0,0 +1,101 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once + +#include +#include +#include +#include + +// --------------------------------------------------------------------------- +// Error checking macro +// --------------------------------------------------------------------------- + +#define HIP_CHECK(call) \ + do \ + { \ + hipError_t err = call; \ + if(err != hipSuccess) \ + { \ + printf("HIP error %s:%d: '%s'\n", __FILE__, __LINE__, hipGetErrorString(err)); \ + exit(1); \ + } \ + } while(0) + +// All named small-seq kernel entities live in this namespace to avoid clashing with +// similarly-named symbols in other fused-attn backends (e.g. FmhaKernelConfig). +namespace small_seq_kernels { + +// --------------------------------------------------------------------------- +// Causal mask type +// --------------------------------------------------------------------------- + +enum class CausalMaskType +{ + DISABLE = 0, + TOP_LEFT = 1, + BOTTOM_RIGHT = 2 +}; + +// inline to avoid ODR violation across multiple translation units (C++17) +inline std::map CausalMaskTypeName = { + {CausalMaskType::DISABLE, "DISABLE"}, + {CausalMaskType::TOP_LEFT, "TOP_LEFT"}, + {CausalMaskType::BOTTOM_RIGHT, "BOTTOM_RIGHT"}}; + +// --------------------------------------------------------------------------- +// Kernel configuration struct +// +// Template parameters encode the static layout dimensions used by all kernels. +// Runtime variability (actual batch size and actual Q/KV lengths per batch) is +// handled at runtime: batch is passed as a kernel-launch argument (mapped to the +// grid z-dimension) and per-batch sequence lengths come from cu_seqlens. +// --------------------------------------------------------------------------- + +template +struct FmhaKernelConfig +{ + static constexpr int head_num = HEAD_NUM; + static constexpr int max_seq_q = MAX_SEQ_Q; + // Backward compat alias for scalar fwd/bwd kernels (hardcoded seq_q=1) + static constexpr int seq_q = 1; + static constexpr int max_seq_kv = MAX_SEQ_KV; + static constexpr int head_dim = HEAD_DIM; + static constexpr int step2_block_size = STEP2_BLOCK_SIZE; + static constexpr bool enable_dropout_mask = ENABLE_DROPOUT_MASK; + static constexpr enum CausalMaskType mask_type = MAKS_TYPE; +}; + +// --------------------------------------------------------------------------- +// Device helpers: BSHD uniform layout vs THD varlen (cu_seqlens) +// +// When uniform_seq_len > 0 the kernel runs in dense BSHD mode: every batch +// has the same sequence length and offsets are batch_idx * uniform_seq_len. +// When uniform_seq_len == 0, lengths/offsets come from cu_seqlens_*. +// --------------------------------------------------------------------------- + +__device__ __forceinline__ int get_seq_len(int batch_idx, + int uniform_seq_len, + const int* cu_seqlens) { + if(uniform_seq_len > 0) return uniform_seq_len; + return cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]; +} + +__device__ __forceinline__ int get_token_offset(int batch_idx, + int uniform_seq_len, + const int* cu_seqlens_padded) { + if(uniform_seq_len > 0) return batch_idx * uniform_seq_len; + return cu_seqlens_padded[batch_idx]; +} + + +} // namespace small_seq_kernels diff --git a/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma.h b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma.h new file mode 100644 index 000000000..a3f2d639c --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma.h @@ -0,0 +1,408 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once + +#include "attn_common.h" +#include + +using bhalf_t = __bf16; +using bf16x4 = __bf16 __attribute__((ext_vector_type(4))); +using bf16x8 = __bf16 __attribute__((ext_vector_type(8))); +using floatx4 = float __attribute__((ext_vector_type(4))); + +#ifndef CEIL_DIV +#define CEIL_DIV(a, b) (((a) + (b)-1) / (b)) +#endif + +namespace small_seq_kernels { + +template +__device__ __forceinline__ bf16x8 load_cvt_bf16x8(const T* src) +{ + if constexpr(sizeof(T) == 2) + { + return *(const bf16x8*)src; + } + else + { + // T = float + bf16x8 r; + #pragma unroll + for(int i = 0; i < 8; i++) + { + r[i] = static_cast(src[i]); + } + return r; + } +} + +// --------------------------------------------------------------------------- +// MFMA 4x4x4 forward kernel (seq_q ≤ 4, online softmax, 16 heads/wave) +// +// Thread: warp[0-3], lane[0-63], mfma_block=lane/4 (head), mfma_tid=lane%4 (Q row) +// LDS: Q_lds[seq_q × 16 × hd_pad], KV_lds[4 × 16 × hd_pad] (reused K→V) +// Grid: (1, ceil(heads/16), bs), Block: 256 +// --------------------------------------------------------------------------- + +template +__launch_bounds__(256, (Config::head_dim == 128) ? 3 : 1) +__global__ void fmha_fwd_mfma_kernel( + const T* Q, + const T* K, + const T* V, + T* O, + T* workspace, + const T* dropout_mask, + float dropout_scale, + float scale, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded) +{ + // Compile-time constants + constexpr int head_dim = Config::head_dim; + constexpr int head_num = Config::head_num; + constexpr int max_seq_kv = Config::max_seq_kv; + constexpr int max_seq_q = Config::max_seq_q; + constexpr int hd_pad = head_dim + 4; + + static_assert(max_seq_q >= 1 && max_seq_q <= 4, "4x4x4 kernel supports max_seq_q 1..4"); + + // 4 warps split head_dim for Attn×V + constexpr int dims_per_warp = head_dim / 4; + constexpr int num_dim_groups = dims_per_warp / 4; + + // Thread mapping + const int batch_idx = blockIdx.z; + const int head_group = blockIdx.y; + const int tid = threadIdx.x; + const int warp_id = tid / 64; + const int lane_id = tid % 64; + const int mfma_block = lane_id / 4; // which head within group [0,16) + const int mfma_tid = lane_id % 4; // Q-row worker within MFMA block [0,4) + + const int head_base = head_group * 16; + const int head_idx = head_base + mfma_block; + const bool valid_head = (head_idx < head_num); + + const int actual_q = cu_seqlens_q[batch_idx + 1] - cu_seqlens_q[batch_idx]; + if(actual_q == 0) + return; + + const int seq_kv = cu_seqlens_kv[batch_idx + 1] - cu_seqlens_kv[batch_idx]; + const int kv_offset = cu_seqlens_kv_padded[batch_idx]; + const int q_offset = cu_seqlens_q_padded[batch_idx]; + + const int warp_dim_start = warp_id * dims_per_warp; + + // LDS + __shared__ __attribute__((aligned(128))) bhalf_t Q_lds[max_seq_q * 16 * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t KV_lds[4 * 16 * hd_pad]; + + // Cooperative load: each thread loads 8 bf16 values + const int load_idx = tid * 8; + const int load_head = load_idx / head_dim; + const int load_dim = load_idx % head_dim; + const int load_lds_off = load_head * hd_pad + load_dim; + + // MFMA LDS read offsets + const int q_lds_base = mfma_block * hd_pad; + const int k_lds_base = mfma_tid * 16 * hd_pad + mfma_block * hd_pad; + + // Load Q → Q_lds + #pragma unroll + for(int qr = 0; qr < max_seq_q; qr++) + { + const int q_lds_offset = qr * 16 * hd_pad; + + if(qr < actual_q && head_base + load_head < head_num) + { + const T* q_src = Q + ((size_t)(q_offset + qr) * head_num + head_base) * head_dim; + *(bf16x8*)(&Q_lds[q_lds_offset + load_lds_off]) = load_cvt_bf16x8(q_src + load_idx); + } + else + { + *(bf16x8*)(&Q_lds[q_lds_offset + load_lds_off]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + + // Online attention: fused QK^T → softmax → Attn×V per KV group of 4 + float running_max[max_seq_q]; + float running_sum[max_seq_q]; + float v_acc[max_seq_q][num_dim_groups]; + + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + { + running_max[m] = -INFINITY; + running_sum[m] = 0.0f; + #pragma unroll + for(int dg = 0; dg < num_dim_groups; dg++) + v_acc[m][dg] = 0.0f; + } + + const int num_kv_groups = CEIL_DIV(seq_kv, 4); + + for(int kv_grp = 0; kv_grp < num_kv_groups; kv_grp++) + { + const int kv_base = kv_grp * 4; + + // Load K[4 positions] → KV_lds + #pragma unroll + for(int kv = 0; kv < 4; kv++) + { + const int kv_pos = kv_base + kv; + const int clamped_kv = min(kv_pos, max(seq_kv - 1, 0)); + const T* k_src = K + ((size_t)(kv_offset + clamped_kv) * head_num + head_base) * head_dim; + const int kv_lds_offset = kv * 16 * hd_pad; + + if(head_base + load_head < head_num) + { + *(bf16x8*)(&KV_lds[kv_lds_offset + load_lds_off]) = load_cvt_bf16x8(k_src + load_idx); + } + else + { + *(bf16x8*)(&KV_lds[kv_lds_offset + load_lds_off]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + + __syncthreads(); + + // MFMA QK^T + floatx4 qk_acc = {0, 0, 0, 0}; + + #pragma unroll + for(int k = 0; k < head_dim; k += 4) + { + bf16x4 q_a, k_b; + + if(mfma_tid < actual_q) + { + q_a = *(const bf16x4*)(&Q_lds[mfma_tid * 16 * hd_pad + q_lds_base + k]); + } + else + { + q_a = bf16x4{0, 0, 0, 0}; + } + + k_b = *(const bf16x4*)(&KV_lds[k_lds_base + k]); + + qk_acc = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(q_a, k_b, qk_acc, 0, 0, 0); + } + + // Online softmax: extract scores, update running_max/sum, rescale v_acc + float my_weights[4]; + + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + { + float scores[4]; + #pragma unroll + for(int s = 0; s < 4; s++) + { + int kv_pos = kv_base + s; + bool masked = (kv_pos >= seq_kv) || (m >= actual_q); + if constexpr(Config::mask_type == CausalMaskType::TOP_LEFT) + { + if(kv_pos > m) + masked = true; + } + scores[s] = masked ? -INFINITY : (__shfl(qk_acc[m], s, 4) * scale); + } + + float tile_max = fmaxf(fmaxf(scores[0], scores[1]), fmaxf(scores[2], scores[3])); + float new_max = fmaxf(running_max[m], tile_max); + + // Rescale previous accumulations (guard -inf - (-inf) = NaN) + if(running_max[m] > -INFINITY) + { + float rescale = expf(running_max[m] - new_max); + running_sum[m] *= rescale; + #pragma unroll + for(int dg = 0; dg < num_dim_groups; dg++) + v_acc[m][dg] *= rescale; + } + running_max[m] = new_max; + + float weights[4]; + #pragma unroll + for(int s = 0; s < 4; s++) + { + weights[s] = (running_max[m] > -INFINITY) ? expf(scores[s] - running_max[m]) : 0.0f; + running_sum[m] += weights[s]; + } + + if(m == mfma_tid) + { + #pragma unroll + for(int s = 0; s < 4; s++) + my_weights[s] = weights[s]; + } + } + + // Apply dropout + if constexpr(Config::enable_dropout_mask) + { + if(valid_head && mfma_tid < actual_q) + { + const int ws_off = ((q_offset + mfma_tid) * head_num + head_idx) * max_seq_kv; + #pragma unroll + for(int s = 0; s < 4; s++) + { + int kv_pos = kv_base + s; + if(kv_pos < seq_kv) + { + my_weights[s] *= static_cast(dropout_mask[ws_off + kv_pos]) + * dropout_scale; + } + } + } + } + + // Convert weights to bf16 for V MFMA + bf16x4 weight_a; + if(mfma_tid < actual_q) + { + #pragma unroll + for(int i = 0; i < 4; i++) + weight_a[i] = static_cast(my_weights[i]); + } + else + { + weight_a = bf16x4{0, 0, 0, 0}; + } + + __syncthreads(); + + // Load V[4 positions] → KV_lds + #pragma unroll + for(int kv = 0; kv < 4; kv++) + { + const int kv_pos = kv_base + kv; + const int kv_lds_offset = kv * 16 * hd_pad; + + if(kv_pos < seq_kv && head_base + load_head < head_num) + { + const T* v_src = V + ((size_t)(kv_offset + kv_pos) * head_num + head_base) * head_dim; + *(bf16x8*)(&KV_lds[kv_lds_offset + load_lds_off]) = load_cvt_bf16x8(v_src + load_idx); + } + else + { + *(bf16x8*)(&KV_lds[kv_lds_offset + load_lds_off]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + + __syncthreads(); + + // MFMA weights × V → accumulate v_acc + #pragma unroll + for(int dg = 0; dg < num_dim_groups; dg++) + { + const int out_d = warp_dim_start + dg * 4 + mfma_tid; + + bf16x4 v_b; + #pragma unroll + for(int i = 0; i < 4; i++) + { + v_b[i] = KV_lds[i * 16 * hd_pad + mfma_block * hd_pad + out_d]; + } + + floatx4 mfma_acc; + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + mfma_acc[m] = v_acc[m][dg]; + #pragma unroll + for(int m = max_seq_q; m < 4; m++) + mfma_acc[m] = 0.0f; + + mfma_acc = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + weight_a, v_b, mfma_acc, 0, 0, 0); + + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + v_acc[m][dg] = mfma_acc[m]; + } + + __syncthreads(); + } + + // Normalize: v_acc /= running_sum + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + { + float inv_sum = (running_sum[m] > 0.0f) ? (1.0f / running_sum[m]) : 0.0f; + #pragma unroll + for(int dg = 0; dg < num_dim_groups; dg++) + v_acc[m][dg] *= inv_sum; + } + + // Write output O[total_padded_q, head_num, head_dim] + if(valid_head) + { + #pragma unroll + for(int m = 0; m < max_seq_q; m++) + { + if(m < actual_q) + { + #pragma unroll + for(int dg = 0; dg < num_dim_groups; dg++) + { + const int out_d = warp_dim_start + dg * 4 + mfma_tid; + O[((size_t)(q_offset + m) * head_num + head_idx) * head_dim + out_d] = + static_cast(v_acc[m][dg]); + } + } + } + } +} + +// --------------------------------------------------------------------------- +// AttnForwardMfmaKernelLauncher — Grid: (1, ceil(heads/16), bs), Block: 256 +// --------------------------------------------------------------------------- + +template +struct AttnForwardMfmaKernelLauncher +{ + using fwd_aux_buffer_scalar = T; + + static size_t calc_workspace_size(int total_padded_q) + { + return (size_t)total_padded_q * Config::head_num * Config::max_seq_kv * sizeof(T); + } + + static void run_attn_fwd_kernel(const T* Q, + const T* K, + const T* V, + const T* dropout_mask, + float dropout_p, + float sqr_dk_scale, + T* O, + T* workspace, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded, + const int* padded_q_to_batch, + int total_padded_q, + int batch, + hipStream_t stream = 0) + { + float dropout_scale = (dropout_p > 0.0f) ? (1.0f / (1.0f - dropout_p)) : 1.0f; + + // Batch is a runtime argument mapped to the grid z-dimension. + dim3 grid(1, CEIL_DIV(Config::head_num, 16), batch); + dim3 block(256); + + fmha_fwd_mfma_kernel<<>>( + Q, K, V, O, workspace, + dropout_mask, dropout_scale, sqr_dk_scale, + cu_seqlens_q, cu_seqlens_q_padded, + cu_seqlens_kv, cu_seqlens_kv_padded); + } +}; + +} // namespace small_seq_kernels diff --git a/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_16x16.h b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_16x16.h new file mode 100644 index 000000000..3047331ca --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_16x16.h @@ -0,0 +1,408 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once + +#include "attn_common.h" +#include + + +#ifndef ATTN_MFMA_TYPES_DEFINED +#define ATTN_MFMA_TYPES_DEFINED +using bhalf_t = __bf16; +using bf16x4 = __bf16 __attribute__((ext_vector_type(4))); +using bf16x8 = __bf16 __attribute__((ext_vector_type(8))); +using floatx4 = float __attribute__((ext_vector_type(4))); +#endif + +#ifndef CEIL_DIV +#define CEIL_DIV(a, b) (((a) + (b)-1) / (b)) +#endif + +namespace small_seq_kernels { + +template +__device__ __forceinline__ bf16x8 load_cvt_bf16x8_16(const T* src) +{ + if constexpr(sizeof(T) == 2) + { + return *(const bf16x8*)src; + } + else + { + // T = float + bf16x8 r; + #pragma unroll + for(int i = 0; i < 8; i++) + { + r[i] = static_cast(src[i]); + } + return r; + } +} + +// --------------------------------------------------------------------------- +// MFMA 16x16x16 forward kernel (tiled Q and KV, 1 head/block) +// +// Thread: warp[0-3], lane_row=lane/16 [0,4), lane_col=lane%16 [0,16) +// LDS: Q_lds[lds_q_rows × hd_pad], KV_lds[lds_kv_rows × hd_pad], +// SM_lds[lds_q_rows × lds_sm_stride] +// Grid: (1, head_num, bs), Block: 256 +// +// softmax_lse (Option A / FA2-style aux): one float per (padded Q row, head), +// index ((q_offset + q_row) * head_num + head_idx), +// value log(sum_j exp(scale * QK^T_{row,j})) = row_max + log(row_sum_exp). +// --------------------------------------------------------------------------- + +template +__launch_bounds__(256, 1) +__global__ void fmha_fwd_mfma_16x16_kernel( + const T* Q, + const T* K, + const T* V, + T* O, + float* softmax_lse, + const T* dropout_mask, + float dropout_scale, + float scale, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded) +{ + // Compile-time constants + constexpr int head_dim = Config::head_dim; + constexpr int head_num = Config::head_num; + constexpr int max_seq_kv = Config::max_seq_kv; + constexpr int max_seq_q = Config::max_seq_q; + constexpr int hd_pad = head_dim + 4; + constexpr int q_tiles = CEIL_DIV(max_seq_q, 16); + constexpr int kv_tiles = CEIL_DIV(max_seq_kv, 16); + constexpr int lds_q_rows = q_tiles * 16; + constexpr int lds_kv_rows = kv_tiles * 16; + constexpr int lds_sm_stride = lds_kv_rows + 4; + + static_assert(max_seq_q >= 1, "max_seq_q must be >= 1"); + + // Thread mapping + const int batch_idx = blockIdx.z; + const int head_idx = blockIdx.y; + const int tid = threadIdx.x; + const int warp_id = tid / 64; + const int lane_id = tid % 64; + const int lane_row = lane_id / 16; + const int lane_col = lane_id % 16; + + const int actual_q = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_q); + if(actual_q == 0) + return; + + const int seq_kv = get_seq_len(batch_idx, uniform_seq_len, cu_seqlens_kv); + const int kv_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_kv_padded); + const int q_offset = get_token_offset(batch_idx, uniform_seq_len, cu_seqlens_q_padded); + + // LDS + __shared__ __attribute__((aligned(128))) bhalf_t Q_lds[lds_q_rows * hd_pad]; + __shared__ __attribute__((aligned(128))) bhalf_t KV_lds[lds_kv_rows * hd_pad]; + __shared__ float SM_lds[lds_q_rows * lds_sm_stride]; + + // Load Q → Q_lds + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_q_rows; r += (256 / threads_per_row)) + { + if(r < actual_q) + { + const T* q_src = Q + ((size_t)(q_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&Q_lds[r * hd_pad + col]) = load_cvt_bf16x8_16(q_src + col); + } + else + { + *(bf16x8*)(&Q_lds[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + } + + // Load K → KV_lds + { + constexpr int threads_per_row = head_dim / 8; + const int row = tid / threads_per_row; + const int col = (tid % threads_per_row) * 8; + + for(int r = row; r < lds_kv_rows; r += (256 / threads_per_row)) + { + if(r < seq_kv) + { + const T* k_src = K + ((size_t)(kv_offset + r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&KV_lds[r * hd_pad + col]) = load_cvt_bf16x8_16(k_src + col); + } + else + { + *(bf16x8*)(&KV_lds[r * hd_pad + col]) = bf16x8{0, 0, 0, 0, 0, 0, 0, 0}; + } + } + } + + __syncthreads(); + + // QK^T via MFMA (all 4 warps redundant) + float attn_weight[q_tiles * kv_tiles * 4]; + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + floatx4 acc = {0, 0, 0, 0}; + constexpr int total_hd_tiles = CEIL_DIV(head_dim, 16); + + #pragma unroll + for(int k = 0; k < total_hd_tiles; ++k) + { + const int dim_off = k * 16; + bf16x4 a = *(const bf16x4*)(&Q_lds[(qt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + bf16x4 b = *(const bf16x4*)(&KV_lds[(kvt * 16 + lane_col) * hd_pad + dim_off + lane_row * 4]); + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + int reg_base = (qt * kv_tiles + kvt) * 4; + #pragma unroll + for(int i = 0; i < 4; i++) + attn_weight[reg_base + i] = acc[i] * scale; + } + } + + // Softmax: two-pass across KV tiles per Q row + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + + // Pass 1: find global row_max across all KV tiles + float row_max = -INFINITY; + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + int kv_pos = kvt * 16 + lane_col; + + bool masked = (kv_pos >= seq_kv) || (q_row >= actual_q); + if constexpr(Config::mask_type == CausalMaskType::TOP_LEFT) + { + if(kv_pos > q_row) + masked = true; + } + + float val = masked ? -INFINITY : attn_weight[reg_idx]; + + float tile_max = val; + #pragma unroll + for(int off = 8; off > 0; off /= 2) + tile_max = fmaxf(tile_max, __shfl_xor(tile_max, off, 16)); + + row_max = fmaxf(row_max, tile_max); + } + + // Pass 2: compute exp and sum across all KV tiles + float row_sum = 0.0f; + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + int kv_pos = kvt * 16 + lane_col; + + bool masked = (kv_pos >= seq_kv) || (q_row >= actual_q); + if constexpr(Config::mask_type == CausalMaskType::TOP_LEFT) + { + if(kv_pos > q_row) + masked = true; + } + + float exp_val = masked ? 0.0f : expf(attn_weight[reg_idx] - row_max); + attn_weight[reg_idx] = exp_val; + + float tile_sum = exp_val; + #pragma unroll + for(int off = 8; off > 0; off /= 2) + tile_sum += __shfl_xor(tile_sum, off, 16); + row_sum += tile_sum; + } + + // Log-sum-exp per row (matches FlashAttention-style LSE; pre-dropout) + float lse_row = (row_sum > 0.0f) ? (row_max + logf(row_sum)) : -INFINITY; + if(lane_col == 0 && q_row < actual_q) + { + softmax_lse[((size_t)(q_offset + q_row) * head_num + head_idx)] = lse_row; + } + + // Normalize and apply dropout + float inv_sum = __builtin_amdgcn_rcpf(row_sum); + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + attn_weight[reg_idx] *= inv_sum; + + if constexpr(Config::enable_dropout_mask) + { + int kv_pos = kvt * 16 + lane_col; + if(q_row < actual_q && kv_pos < seq_kv) + { + const int ws_offset = ((q_offset + q_row) * head_num + head_idx) * max_seq_kv; + attn_weight[reg_idx] *= static_cast(dropout_mask[ws_offset + kv_pos]) * dropout_scale; + } + } + } + } + } + + // Write weights to SM_lds for Attn×V + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + int kv_pos = kvt * 16 + lane_col; + int reg_idx = (qt * kv_tiles + kvt) * 4 + i; + SM_lds[q_row * lds_sm_stride + kv_pos] = attn_weight[reg_idx]; + } + } + } + + __syncthreads(); + + // Load V → KV_lds (clamped; invalid positions zeroed by softmax weights) + { + constexpr int threads_per_row = head_dim / 8; + const int v_row = tid / threads_per_row; + const int v_col = (tid % threads_per_row) * 8; + const int clamped_max = max(seq_kv - 1, 0); + + for(int r = v_row; r < lds_kv_rows; r += (256 / threads_per_row)) + { + const int clamped_r = min(r, clamped_max); + const T* v_src = V + ((size_t)(kv_offset + clamped_r) * head_num + head_idx) * head_dim; + *(bf16x8*)(&KV_lds[r * hd_pad + v_col]) = load_cvt_bf16x8_16(v_src + v_col); + } + } + + __syncthreads(); + + // Attn×V via MFMA (4 warps split head_dim, tiled over Q and KV) + { + constexpr int BK = 64; + constexpr int total_d_tiles = CEIL_DIV(head_dim, BK); + + #pragma unroll + for(int qt = 0; qt < q_tiles; qt++) + { + #pragma unroll + for(int d = 0; d < total_d_tiles; d++) + { + const int dim_idx = d * BK + warp_id * 16; + floatx4 acc = {0, 0, 0, 0}; + + #pragma unroll + for(int kvt = 0; kvt < kv_tiles; kvt++) + { + // A: softmax weights (transposed read from SM_lds) + bf16x4 a; + #pragma unroll + for(int k = 0; k < 4; k++) + { + int q_idx = qt * 16 + lane_col; + int kv_pos = kvt * 16 + lane_row * 4 + k; + a[k] = static_cast(SM_lds[q_idx * lds_sm_stride + kv_pos]); + } + + // B: V[kv, d] + bf16x4 b; + const int kv_base = kvt * 16; + #pragma unroll + for(int k = 0; k < 4; k++) + { + b[k] = KV_lds[(kv_base + lane_row * 4 + k) * hd_pad + dim_idx + lane_col]; + } + + acc = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a, b, acc, 0, 0, 0); + } + + // Write output + #pragma unroll + for(int i = 0; i < 4; i++) + { + int q_row = qt * 16 + lane_row * 4 + i; + if(q_row < actual_q) + { + O[((size_t)(q_offset + q_row) * head_num + head_idx) * head_dim + dim_idx + lane_col] = + static_cast(acc[i]); + } + } + } + } + } +} + +// --------------------------------------------------------------------------- +// AttnForwardMfma16x16KernelLauncher — Grid: (1, head_num, bs), Block: 256 +// --------------------------------------------------------------------------- + +template +struct AttnForwardMfma16x16KernelLauncher +{ + using fwd_aux_buffer_scalar = float; + + /// Per-(padded Q row, head) softmax log-sum-exp (float), FA2-compatible aux. + static size_t calc_workspace_size(int total_padded_q) + { + return (size_t)total_padded_q * Config::head_num * sizeof(float); + } + + static void run_attn_fwd_kernel(const T* Q, + const T* K, + const T* V, + const T* dropout_mask, + float dropout_p, + float sqr_dk_scale, + T* O, + float* softmax_lse, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded, + int total_padded_q, + int batch, + hipStream_t stream = 0) + { + float dropout_scale = (dropout_p > 0.0f) ? (1.0f / (1.0f - dropout_p)) : 1.0f; + float scale = sqr_dk_scale; + + // Batch is a runtime argument mapped to the grid z-dimension. + dim3 grid(1, Config::head_num, batch); + dim3 block(256); + + fmha_fwd_mfma_16x16_kernel<<>>( + Q, K, V, O, softmax_lse, + dropout_mask, dropout_scale, scale, uniform_seq_len, + cu_seqlens_q, cu_seqlens_q_padded, + cu_seqlens_kv, cu_seqlens_kv_padded); + } +}; + +} // namespace small_seq_kernels diff --git a/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_dispatch.h b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_dispatch.h new file mode 100644 index 000000000..994a36d39 --- /dev/null +++ b/transformer_engine/common/ck_fused_attn/src/small_seq_kernels/attn_fwd_mfma_dispatch.h @@ -0,0 +1,67 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once + +#include "attn_fwd_mfma.h" +#include "attn_fwd_mfma_16x16.h" + +namespace small_seq_kernels { + +// --------------------------------------------------------------------------- +// Dispatch: seq_q ≤ 4 → 4x4x4 (16 heads/wave), seq_q > 4 → 16x16x16 +// --------------------------------------------------------------------------- + +template +struct AttnForwardMfmaDispatchLauncher +{ + static_assert(Config::max_seq_q >= 1, + "max_seq_q must be >= 1"); + + static size_t calc_workspace_size(int total_padded_q) + { + if constexpr(Config::max_seq_q <= 4) + return AttnForwardMfmaKernelLauncher::calc_workspace_size(total_padded_q); + else + return AttnForwardMfma16x16KernelLauncher::calc_workspace_size(total_padded_q); + } + + /// `aux`: 4x4 path = `T*` attention workspace; 16x16 path = `float*` softmax LSE (see + /// AttnForwardMfma16x16KernelLauncher::calc_workspace_size). + static void run_attn_fwd_kernel(const T* Q, + const T* K, + const T* V, + const T* dropout_mask, + float dropout_p, + float sqr_dk_scale, + T* O, + void* aux, + int uniform_seq_len, + const int* cu_seqlens_q, + const int* cu_seqlens_q_padded, + const int* cu_seqlens_kv, + const int* cu_seqlens_kv_padded, + int total_padded_q, + int batch, + hipStream_t stream = 0) + { + if constexpr(Config::max_seq_q <= 4) + { + AttnForwardMfmaKernelLauncher::run_attn_fwd_kernel( + Q, K, V, dropout_mask, dropout_p, sqr_dk_scale, O, static_cast(aux), + uniform_seq_len, cu_seqlens_q, cu_seqlens_q_padded, cu_seqlens_kv, + cu_seqlens_kv_padded, total_padded_q, batch, stream); + } + else + { + AttnForwardMfma16x16KernelLauncher::run_attn_fwd_kernel( + Q, K, V, dropout_mask, dropout_p, sqr_dk_scale, O, static_cast(aux), + uniform_seq_len, cu_seqlens_q, cu_seqlens_q_padded, cu_seqlens_kv, + cu_seqlens_kv_padded, total_padded_q, batch, stream); + } + } +}; + +} // namespace small_seq_kernels diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp index 21beff7ca..578a02c07 100644 --- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp +++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp @@ -151,6 +151,32 @@ bool is_ck_backend_supported( #endif // USE_FUSED_ATTN_CK } +bool is_small_seq_supported_static(DType dtype, + NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, + float dropout, + size_t head_dim_qk, + size_t head_dim_v, + size_t num_attn_heads, + size_t num_gqa_groups) { + if(dropout != 0.0f) return false; + if(bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) return false; + if(dtype != DType::kBFloat16) return false; + if(head_dim_qk != head_dim_v) return false; + if(head_dim_qk != 128 && head_dim_qk != 256) return false; + if(num_gqa_groups == 0 || num_attn_heads % num_gqa_groups != 0) return false; + if(num_attn_heads != num_gqa_groups) return false; + if(num_attn_heads != 16 && num_attn_heads != 32) return false; + if(!is_padding_mask(mask_type) && mask_type != NVTE_Mask_Type::NVTE_NO_MASK) return false; + return true; +} + +bool is_small_seq_supported_runtime(size_t runtime_max_seqlen_q, + size_t runtime_max_seqlen_kv) { + return runtime_max_seqlen_q > 0 && runtime_max_seqlen_q <= kSmallSeqMaxSeqlen && + runtime_max_seqlen_kv > 0 && runtime_max_seqlen_kv <= kSmallSeqMaxSeqlen; +} + #ifdef USE_FUSED_ATTN_CK ck_fused_attn::DType nvte_to_ck_dtype(DType t_dtype){ @@ -487,6 +513,11 @@ void fused_attn_ck_fwd_impl( bool is_padding = is_padding_mask(mask_type); bool bshd_to_thd = is_BSHD && is_padding; + const bool ck_small_seq_env_enabled = + getenv("NVTE_FUSED_ATTN_CK_SMALLSEQ") == "1"; + bool ck_small_seq_enabled = false; + void* ck_smallseq_workspace_prefix = nullptr; + // extract the qkv and o storage bytes to allocate buffer for padding removing // b from cu_seqlen is not the actual storage batch for pad_between_seqs case size_t q_storage_bytes = max_tokens_q*h*d_qk*nvte_dtype_size(dtype); @@ -498,6 +529,22 @@ void fused_attn_ck_fwd_impl( // (planner returns nullptr, accumulates total) and execution mode. WorkspacePlanner planner(workspace); + if(ck_small_seq_env_enabled) { + if(cuda::sm_arch() == 94) { + if(is_small_seq_supported_static(dtype, bias_type, mask_type, dropout_probability, d_qk, d_v, + h, hg)) { + if(is_ragged) { + ck_small_seq_enabled = true; + ck_smallseq_workspace_prefix = + planner.allocate(ck_fused_attn::small_seq_thd_extra_workspace_bytes()); + } else { + ck_small_seq_enabled = + is_BSHD && s_q == s_kv && s_q >= 2 && s_q <= kSmallSeqMaxSeqlen; + } + } + } + } + void* devPtrAlibiSlope = nullptr; if(bias_type == NVTE_Bias_Type::NVTE_ALIBI){ // ck requires an alibi slope array even if in standard (vanilla) mode @@ -505,7 +552,7 @@ void fused_attn_ck_fwd_impl( } void* devPtrSoftmaxLSEWithoutPadding = nullptr; - if((is_SBHD && is_padding) || bshd_to_thd || is_ragged){ + if((is_SBHD && is_padding) || bshd_to_thd || is_ragged || ck_small_seq_enabled){ devPtrSoftmaxLSEWithoutPadding = planner.allocate(h*max_tokens_q*sizeof(float)); } @@ -660,6 +707,59 @@ void fused_attn_ck_fwd_impl( std::cout<<"num_splits: "< pad_remap_lse -> devPtrSoftmaxAux [b,h,s]. + // + // THD (is_ragged): probe runtime max seqlen via cu_seqlens; if <= 17, + // ck_attn_smallseq_fwd_thd with cu_seqlens; LSE remap to THD layout on O. + // + // --------------------------------------------------------------------------- + if(ck_small_seq_enabled) { + if(is_BSHD) { + if(nvte_log_ck_config) { + std::cout << std::endl << "attn_fwd(ck small-seq, BSHD self-attn): b: " << b + << ", s: " << s_q << ", flow: ck-smallseq" << std::endl; + } + ck_fused_attn::ck_attn_smallseq_fwd_bshd( + b, h, s_q, s_kv, d_qk, scaling_factor, devPtrQ, devPtrK, devPtrV, devPtrO, + devPtrSoftmaxLSEWithoutPadding, nvte_to_ck_dtype(dtype), stream); + pad_remap_lse(b, h, s_q, max_tokens_q, false, devPtrSoftmaxAux, + devPtrCuSeqlensQ, devPtrSeqOffsetsQ, + devPtrSoftmaxLSEWithoutPadding, stream); + return; + } else { + void* max_seqlen_workspace_q = ck_smallseq_workspace_prefix; + void* max_seqlen_workspace_kv = + static_cast(static_cast(ck_smallseq_workspace_prefix) + sizeof(uint64_t)); + const size_t runtime_max_seqlen_q = static_cast(ck_fused_attn::get_runtime_max_seqlen( + b, devPtrCuSeqlensQ, devPtrCuSeqlenPaddedQ, max_seqlen_workspace_q, stream)); + const size_t runtime_max_seqlen_kv = static_cast(ck_fused_attn::get_runtime_max_seqlen( + b, devPtrCuSeqlensKV, devPtrCuSeqlenPaddedKV, max_seqlen_workspace_kv, stream)); + const bool run_smallseq = + is_small_seq_supported_runtime(runtime_max_seqlen_q, runtime_max_seqlen_kv); + if(nvte_log_ck_config) { + std::cout << std::endl << "attn_fwd(ck small-seq, THD): b: " << b + << ", runtime_max_seqlen_q: " << runtime_max_seqlen_q + << ", runtime_max_seqlen_kv: " << runtime_max_seqlen_kv + << ", flow: " << (run_smallseq ? "ck-smallseq" : "regular ck/aiter") << std::endl; + } + + if(run_smallseq) { + ck_fused_attn::ck_attn_smallseq_fwd_thd( + b, h, d_qk, max_tokens_q, max_tokens_kv, scaling_factor, devPtrQ, devPtrK, devPtrV, + devPtrO, devPtrSoftmaxLSEWithoutPadding, devPtrCuSeqlensQ, devPtrCuSeqlenPaddedQ, + devPtrCuSeqlensKV, devPtrCuSeqlenPaddedKV, nvte_to_ck_dtype(dtype), stream); + pad_remap_lse(b, h, s_q, max_tokens_q, true, devPtrSoftmaxAux, + devPtrCuSeqlenPaddedQ, devPtrCuSeqlenPaddedQ, + devPtrSoftmaxLSEWithoutPadding, stream); + return; + } + } + } + if(is_SBHD && is_padding){ // remove padding for q, k, v pad_remap(dtype, b, h, s_q, d_qk, max_tokens_q, false, q_stride[0], q_stride[1], q_stride[2], devPtrQ, devPtrCuSeqlensQ, devPtrCuSeqlenPaddedQ, devPtrQWithoutPadding, stream); @@ -758,6 +858,11 @@ void fused_attn_ck_bwd_impl( bool bshd_to_thd = is_BSHD && is_padding; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(layout); + const bool ck_small_seq_env_enabled = + getenv("NVTE_FUSED_ATTN_CK_SMALLSEQ") == "1"; + bool ck_small_seq_enabled = false; + void* ck_smallseq_workspace_prefix = nullptr; + // extract the qkv and o storage bytes to allocate buffer for padding removing // b from cu_seqlen is not the actual storage batch for pad_between_seqs case size_t q_storage_bytes = max_tokens_q*h*d_qk*nvte_dtype_size(dtype); @@ -769,6 +874,22 @@ void fused_attn_ck_bwd_impl( // (planner returns nullptr, accumulates total) and execution mode. WorkspacePlanner planner(workspace); + if(ck_small_seq_env_enabled) { + if(cuda::sm_arch() == 94) { + if(is_small_seq_supported_static(dtype, bias_type, mask_type, dropout_probability, d_qk, d_v, + h, hg)) { + if(is_ragged) { + ck_small_seq_enabled = true; + ck_smallseq_workspace_prefix = + planner.allocate(ck_fused_attn::small_seq_thd_extra_workspace_bytes()); + } else { + ck_small_seq_enabled = + is_BSHD && s_q == s_kv && s_q >= 2 && s_q <= kSmallSeqMaxSeqlen; + } + } + } + } + // First h*max_tokens_q*sizeof(float) is the lse-d buffer (passed as softmax_lsed) void* lse_workspace = planner.allocate(h*max_tokens_q*sizeof(float)); @@ -850,7 +971,7 @@ void fused_attn_ck_bwd_impl( void* devPtrCuSeqlenPaddedQ = devPtrSeqOffsetsQ; void* devPtrCuSeqlenPaddedKV = devPtrSeqOffsetsKV; - if((is_SBHD && is_padding) || bshd_to_thd || is_ragged){ + if((is_SBHD && is_padding) || bshd_to_thd || is_ragged || ck_small_seq_enabled){ devPtrSoftmaxLSEWithoutPadding = planner.allocate(h*max_tokens_q*sizeof(float)); } if(is_SBHD && is_padding){ @@ -1034,6 +1155,54 @@ void fused_attn_ck_bwd_impl( ck_args.is_v3_atomic_fp32 = nvte_ck_is_v3_atomic_fp32; ck_args.how_v3_bf16_cvt = nvte_ck_how_v3_bf16_cvt; + // --------------------------------------------------------------------------- + // CK small-seq backward (mirrors forward). + // --------------------------------------------------------------------------- + if(ck_small_seq_enabled) { + if(is_BSHD) { + if(nvte_log_ck_config) { + std::cout << std::endl << "attn_bwd(ck small-seq, BSHD self-attn): b: " << b + << ", s: " << s_q << ", flow: ck-smallseq" << std::endl; + } + pad_remap_lse(b, h, s_q, max_tokens_q, false, devPtrSoftmaxAux, + devPtrCuSeqlensQ, devPtrSeqOffsetsQ, + devPtrSoftmaxLSEWithoutPadding, stream); + ck_fused_attn::ck_attn_smallseq_bwd_bshd( + b, h, s_q, s_kv, d_qk, scaling_factor, devPtrQ, devPtrK, devPtrV, devPtrdO, + devPtrSoftmaxLSEWithoutPadding, devPtrdQ, devPtrdK, devPtrdV, nvte_to_ck_dtype(dtype), + stream); + return; + } else { + void* max_seqlen_workspace_q = ck_smallseq_workspace_prefix; + void* max_seqlen_workspace_kv = + static_cast(static_cast(ck_smallseq_workspace_prefix) + sizeof(uint64_t)); + const size_t runtime_max_seqlen_q = static_cast(ck_fused_attn::get_runtime_max_seqlen( + b, devPtrCuSeqlensQ, devPtrCuSeqlenPaddedQ, max_seqlen_workspace_q, stream)); + const size_t runtime_max_seqlen_kv = static_cast(ck_fused_attn::get_runtime_max_seqlen( + b, devPtrCuSeqlensKV, devPtrCuSeqlenPaddedKV, max_seqlen_workspace_kv, stream)); + const bool run_smallseq = + is_small_seq_supported_runtime(runtime_max_seqlen_q, runtime_max_seqlen_kv); + if(nvte_log_ck_config) { + std::cout << std::endl << "attn_bwd(ck small-seq, THD): b: " << b + << ", runtime_max_seqlen_q: " << runtime_max_seqlen_q + << ", runtime_max_seqlen_kv: " << runtime_max_seqlen_kv + << ", flow: " << (run_smallseq ? "ck-smallseq" : "regular ck/aiter") << std::endl; + } + + if(run_smallseq) { + pad_remap_lse(b, h, s_q, max_tokens_q, true, devPtrSoftmaxAux, + devPtrCuSeqlenPaddedQ, devPtrCuSeqlenPaddedQ, + devPtrSoftmaxLSEWithoutPadding, stream); + ck_fused_attn::ck_attn_smallseq_bwd_thd( + b, h, d_qk, max_tokens_q, max_tokens_kv, scaling_factor, devPtrQ, devPtrK, devPtrV, + devPtrdO, devPtrSoftmaxLSEWithoutPadding, devPtrdQ, devPtrdK, devPtrdV, + devPtrCuSeqlensQ, devPtrCuSeqlenPaddedQ, devPtrCuSeqlensKV, devPtrCuSeqlenPaddedKV, + nvte_to_ck_dtype(dtype), stream); + return; + } + } + } + if(is_SBHD && is_padding){ // remove padding for q, k, v, o, do pad_remap(dtype, b, h, s_q, d_qk, max_tokens_q, false, q_stride[0], q_stride[1], q_stride[2], devPtrQ, devPtrCuSeqlensQ, devPtrSeqOffsetsQ, devPtrQWithoutPadding, stream); diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h index e98281d1a..8bb807139 100644 --- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h +++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h @@ -30,6 +30,27 @@ bool is_ck_backend_supported( size_t head_dim_v, int64_t window_size_left, int64_t window_size_right); + +constexpr size_t kSmallSeqMaxSeqlen = 17; + +// Small-seq eligibility is split into static and runtime checks: +// - Static: config known at call time (dtype, head dims, heads, bias, dropout, mask). +// Used during workspace sizing to reserve THD probe buffers and to gate the small-seq path +// before any device data is read. +// - Runtime: actual per-batch max seqlen for THD/ragged inputs (from cu_seqlens on device). +// Even when static config matches, individual batches may exceed kSmallSeqMaxSeqlen; probe +// at execute time and fall back to regular CK when out of range. +bool is_small_seq_supported_static(DType dtype, + NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, + float dropout, + size_t head_dim_qk, + size_t head_dim_v, + size_t num_attn_heads, + size_t num_gqa_groups); + +bool is_small_seq_supported_runtime(size_t runtime_max_seqlen_q, + size_t runtime_max_seqlen_kv); } // namespace fused_attn_rocm void fused_attn_ck_fwd(