Add ROCm HIP small-seq fused attention via crossattn_hip_kernel

README.rst

@@ -3,7 +3,7 @@
     Copyright (c) 2023-2026, Advanced Micro Devices, Inc. All rights reserved.
     Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-    See LICENSE for license information.
+    See LICENSE for license information.Please also add one subsection in readme to tell our customers who to use small_seq attn
 
 |License|
 
@@ -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.

ci/jax.sh

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

tests/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()