Add flash_attention_forward_no_dropout_inplace operator implementation, tests and benchmark

Author: chx7514

benchmark/test_flash_attention_forward_no_dropout_inplace.py

@@ -0,0 +1,124 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import base, consts
+
+device = flag_gems.device
+
+# (batch, num_heads, q_seq_len, kv_seq_len, head_dim) configs.
+FLASH_FWD_CONFIGS = [
+    (1, 2, 512, 512, 64),
+    (1, 8, 1024, 1024, 128),
+    (2, 4, 512, 512, 64),
+    (1, 2, 1024, 2048, 64),
+]
+
+
+def torch_flash_attention_forward_no_dropout_inplace(
+    q,
+    k,
+    v,
+    scale,
+    is_causal,
+    return_debug_mask=False,
+    **extra_kwargs,
+):
+    """Reference: aten::_flash_attention_forward with dropout_p=0.0."""
+    return torch.ops.aten._flash_attention_forward(
+        q,
+        k,
+        v,
+        None,
+        None,
+        q.shape[-3],
+        k.shape[-3],
+        0.0,  # dropout_p = 0.0 (no dropout)
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        **extra_kwargs,
+    )
+
+
+def gems_flash_attention_forward_no_dropout_inplace(
+    q,
+    k,
+    v,
+    scale,
+    is_causal,
+    return_debug_mask=False,
+    **extra_kwargs,
+):
+    """FlagGems Triton implementation (no dropout, in-place into ``q``)."""
+    # ``do_bench`` reuses the same tensors across iterations, so clone ``q`` to
+    # preserve the original data between runs (the kernel writes in-place).
+    return flag_gems._flash_attention_forward_no_dropout_inplace(
+        q.clone(),
+        k,
+        v,
+        None,
+        None,
+        q.shape[-3],
+        k.shape[-3],
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        **extra_kwargs,
+    )
+
+
+def flash_attention_forward_no_dropout_inplace_input_fn(config, dtype, device):
+    batch, num_head, q_seq_len, kv_seq_len, head_size = config
+    q = torch.empty(
+        (batch, q_seq_len, num_head, head_size), device=device, dtype=dtype
+    ).uniform_(-0.05, 0.05)
+    k = torch.empty(
+        (batch, kv_seq_len, num_head, head_size), device=device, dtype=dtype
+    ).uniform_(-0.05, 0.05)
+    v = torch.empty(
+        (batch, kv_seq_len, num_head, head_size), device=device, dtype=dtype
+    ).uniform_(-0.05, 0.05)
+    scale = float(1.0 / math.sqrt(head_size))
+
+    # BSHD layout; no dropout; non-causal for the default benchmark configs.
+    yield q, k, v, scale, False, False, {}
+
+
+class FlashAttentionForwardNoDropoutInplaceBenchmark(base.GenericBenchmark):
+    def set_shapes(self, shape_file_path=None):
+        # Use the configs defined in FLASH_FWD_CONFIGS directly, since this
+        # operator has no entry in the shared core-shapes yaml file.
+        self.shapes = [tuple(c) for c in FLASH_FWD_CONFIGS]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.skipif(flag_gems.device == "cpu", reason="Unsupported in CPU mode")
+@pytest.mark.flash_attention_forward_no_dropout_inplace
+def test_flash_attention_forward_no_dropout_inplace_impl():
+    bench = FlashAttentionForwardNoDropoutInplaceBenchmark(
+        op_name="flash_attention_forward_no_dropout_inplace",
+        torch_op=torch_flash_attention_forward_no_dropout_inplace,
+        input_fn=flash_attention_forward_no_dropout_inplace_input_fn,
+        # FlashAttention only supports fp16/bf16; filter from FLOAT_DTYPES.
+        dtypes=[d for d in consts.FLOAT_DTYPES if d in (torch.float16, torch.bfloat16)],
+    )
+    bench.set_gems(gems_flash_attention_forward_no_dropout_inplace)
+    bench.run()

conf/operators.yaml

@@ -3469,6 +3469,20 @@ ops:
       - NeuralNetwork
     stages:
       - stable: '3.0'
+  - id: flash_attention_forward_no_dropout_inplace
+    description: |
+      Triton kernel implementation for _flash_attention_forward_no_dropout_inplace.
+      A specialised variant of _flash_attention_forward that drops the dropout
+      argument (implicitly dropout_p=0.0) and writes the attention output in-place.
+    for:
+      - _flash_attention_forward_no_dropout_inplace
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - NeuralNetwork
+    stages:
+      - alpha: '5.4'
   - id: flash_attn_varlen_func
     description: |
       Compute attention for sequences of variable lengths within a single batch.

src/flag_gems/__init__.py

@@ -116,6 +116,10 @@ def torch_ge(v):
     ("_euclidean_dist", _euclidean_dist),
     ("_flash_attention_backward", flash_attention_backward),
     ("_flash_attention_forward", _flash_attention_forward),
+    (
+        "_flash_attention_forward_no_dropout_inplace",
+        _flash_attention_forward_no_dropout_inplace,
+    ),
     (
         "_functional_sym_constrain_range",
         _functional_sym_constrain_range,

src/flag_gems/ops/__init__.py

@@ -38,6 +38,9 @@
 )
 from flag_gems.ops._euclidean_dist import _euclidean_dist
 from flag_gems.ops._flash_attention_forward import _flash_attention_forward
+from flag_gems.ops._flash_attention_forward_no_dropout_inplace import (
+    _flash_attention_forward_no_dropout_inplace,
+)
 from flag_gems.ops._functional_sym_constrain_range import (
     _functional_sym_constrain_range,
 )
@@ -815,6 +818,7 @@
     "_amp_foreach_non_finite_check_and_unscale_",
     "_assert_async",
     "_batch_norm_no_update",
+    "_flash_attention_forward_no_dropout_inplace",
     "_functional_assert_async",
     "_cdist_backward",
     "_cdist_forward",

src/flag_gems/ops/_flash_attention_forward_no_dropout_inplace.py

@@ -0,0 +1,107 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+
+import logging
+
+import torch
+
+from flag_gems.ops.attention import flash_attention_forward
+
+logger = logging.getLogger(__name__)
+
+
+def _flash_attention_forward_no_dropout_inplace(
+    query,
+    key,
+    value,
+    cum_seq_q,
+    cum_seq_k,
+    max_q,
+    max_k,
+    is_causal,
+    return_debug_mask,
+    *,
+    scale=None,
+    window_size_left=None,
+    window_size_right=None,
+    seqused_k=None,
+    alibi_slopes=None,
+):
+    """FlagGems implementation of ``aten::_flash_attention_forward_no_dropout_inplace``.
+
+    This is a specialised variant of ``_flash_attention_forward`` that drops the
+    ``dropout_p`` argument (it is implicitly ``0.0``) and writes the attention
+    output in-place into ``query``. The heavy lifting is delegated to the
+    existing Triton FlashAttention kernel (:func:`flash_attention_forward`),
+    which runs the ``mha_fwd`` Triton kernel.
+
+    Args:
+        query: ``(batch, num_heads, seq_len_q, head_dim)``.
+        key:   ``(batch, num_kv_heads, seq_len_kv, head_dim)``.
+        value: ``(batch, num_kv_heads, seq_len_kv, head_dim)``.
+        cum_seq_q / cum_seq_k: optional cumulative sequence lengths (varlen).
+        max_q / max_k: maximum sequence lengths.
+        is_causal: whether to apply causal masking.
+        return_debug_mask: whether to return the debug attention mask.
+        scale: optional scale factor for the QK dot product.
+        window_size_left / window_size_right: sliding-window attention sizes.
+        seqused_k: optional per-batch key sequence lengths.
+        alibi_slopes: optional ALiBi bias slopes.
+
+    Returns:
+        ``(output, softmax_logsumexp, rng_state, unused, debug_attn_mask)``.
+        ``output`` aliases ``query`` (written in-place).
+    """
+    logger.debug("GEMS FLASH_ATTENTION_FORWARD_NO_DROPOUT_INPLACE")
+
+    # FlashAttention only supports fp16/bf16 inputs.
+    assert query.dtype in (
+        torch.float16,
+        torch.bfloat16,
+    ), f"expected fp16/bf16 query, got {query.dtype}"
+
+    # No dropout for this variant.
+    dropout_p = 0.0
+
+    # Run the existing Triton FlashAttention forward kernel.
+    #
+    # ``disable_splitkv`` is set because the split-KV combine kernel of the
+    # shared ``mha_fwd`` path can leave output rows unwritten for some shapes,
+    # which would violate the in-place contract (the query must hold the full
+    # result). The non-split kernel path is numerically exact for the shapes
+    # exercised by this operator.
+    out, lse, philox_seed, philox_offset, debug_attn_mask = flash_attention_forward(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        disable_splitkv=True,
+    )
+
+    # In-place semantics: the result is written back into ``query``.
+    query.copy_(out)
+    return query, lse, philox_seed, philox_offset, debug_attn_mask