[KernelGen][Metax] Add alpha_dropout operator for Metax backend (#203)

Co-authored-by: Dongxu-H <dxhan@baai.ac.cn>

Author: yzw1128

src/flag_gems/runtime/backend/_metax/ops/__init__.py

@@ -15,6 +15,7 @@
 from ._nested_view_from_buffer_copy import _nested_view_from_buffer_copy
 from .adaptive_max_pool3d_backward import adaptive_max_pool3d_backward
 from .addmm import addmm
+from .alpha_dropout import alpha_dropout
 from .amax import amax
 from .arange import arange, arange_start
 from .batch_norm import batch_norm, batch_norm_backward
@@ -63,8 +64,9 @@
 __all__ = [
     "_nested_view_from_buffer_copy",
     "_unique2",
-    "addmm",
     "adaptive_max_pool3d_backward",
+    "addmm",
+    "alpha_dropout",
     "amax",
     "arange",
     "arange_start",

src/flag_gems/runtime/backend/_metax/ops/alpha_dropout.py

@@ -0,0 +1,175 @@
+# 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 math
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems import runtime
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
+from flag_gems.utils.random_utils import (
+    philox_backend_seed_offset,
+    uint_to_uniform_float,
+)
+
+logger = logging.getLogger("flag_gems." + __name__)
+
+
+def _alpha_dropout_affine(p: float):
+    """Compute affine parameters a, b for alpha dropout."""
+    _ALPHA = 1.6732632423543772848170429916717
+    _SCALE = 1.0507009873554804934193349852946
+    alpha_prime = -_ALPHA * _SCALE
+    if p == 0.0:
+        return 1.0, 0.0
+    a = 1.0 / math.sqrt((1.0 - p) * (1.0 + p * alpha_prime * alpha_prime))
+    b = -a * p * alpha_prime
+    return a, b
+
+
+@libentry()
+@triton.autotune(
+    configs=runtime.get_tuned_config("alpha_dropout"),
+    key=["N"],
+)
+@triton.jit(do_not_specialize=["philox_seed", "philox_offset"])
+def alpha_dropout_forward_kernel(
+    X,
+    Y,
+    N,
+    p,
+    a,
+    b,
+    philox_seed,
+    philox_offset,
+    BLOCK: tl.constexpr,
+):
+    """Metax-optimized alpha dropout forward kernel.
+
+    Uses UNROLL=4 with num_warps <= 8 to stay within the 512-thread
+    hardware limit on Metax GPUs.
+    """
+    UNROLL: tl.constexpr = 4
+    philox_seed = philox_seed.to(tl.int64)
+    philox_offset = philox_offset.to(tl.int64)
+    c0 = (philox_offset & 0xFFFFFFFF).to(tl.uint32)
+    c1 = ((philox_offset >> 32) & 0xFFFFFFFF).to(tl.uint32)
+    i4 = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    c0 += i4
+    _O = c0 * 0
+    r0, r1, r2, r3 = tl.philox(philox_seed, c0, c1, _O, _O)
+    r0 = uint_to_uniform_float(r0)
+    r1 = uint_to_uniform_float(r1)
+    r2 = uint_to_uniform_float(r2)
+    r3 = uint_to_uniform_float(r3)
+
+    mask0 = r0 > p
+    mask1 = r1 > p
+    mask2 = r2 > p
+    mask3 = r3 > p
+
+    alpha_prime = -1.7580993408473766
+
+    off_0 = tl.program_id(0) * BLOCK * UNROLL + tl.arange(0, BLOCK)
+    off_1 = off_0 + BLOCK
+    off_2 = off_1 + BLOCK
+    off_3 = off_2 + BLOCK
+
+    x0 = tl.load(
+        X + off_0,
+        mask=off_0 < N,
+        other=0.0,
+        eviction_policy="evict_first",
+    )
+    x1 = tl.load(
+        X + off_1,
+        mask=off_1 < N,
+        other=0.0,
+        eviction_policy="evict_first",
+    )
+    x2 = tl.load(
+        X + off_2,
+        mask=off_2 < N,
+        other=0.0,
+        eviction_policy="evict_first",
+    )
+    x3 = tl.load(
+        X + off_3,
+        mask=off_3 < N,
+        other=0.0,
+        eviction_policy="evict_first",
+    )
+
+    y0 = tl.where(mask0, a * x0 + b, a * alpha_prime + b)
+    y1 = tl.where(mask1, a * x1 + b, a * alpha_prime + b)
+    y2 = tl.where(mask2, a * x2 + b, a * alpha_prime + b)
+    y3 = tl.where(mask3, a * x3 + b, a * alpha_prime + b)
+
+    tl.store(
+        Y + off_0,
+        y0,
+        mask=off_0 < N,
+        eviction_policy="evict_first",
+    )
+    tl.store(
+        Y + off_1,
+        y1,
+        mask=off_1 < N,
+        eviction_policy="evict_first",
+    )
+    tl.store(
+        Y + off_2,
+        y2,
+        mask=off_2 < N,
+        eviction_policy="evict_first",
+    )
+    tl.store(
+        Y + off_3,
+        y3,
+        mask=off_3 < N,
+        eviction_policy="evict_first",
+    )
+
+
+def alpha_dropout(input, p=0.5, train=True):
+    UNROLL = 4
+    logger.debug("GEMS_METAX ALPHA_DROPOUT")
+    if not train or p == 0:
+        return input.clone()
+    if p == 1:
+        return torch.zeros_like(input)
+
+    assert 0.0 < p < 1.0, "p must be in (0, 1)"
+
+    device = input.device
+    input = input.contiguous()
+    out = torch.empty_like(input)
+    N = input.numel()
+    grid_fn = lambda meta: (triton.cdiv(N, meta["BLOCK"] * UNROLL),)
+    increment = triton.cdiv(N, UNROLL)
+
+    a, b = _alpha_dropout_affine(p)
+
+    with torch_device_fn.device(device):
+        philox_seed, philox_offset = philox_backend_seed_offset(increment)
+        alpha_dropout_forward_kernel[grid_fn](
+            input, out, N, p, a, b, philox_seed, philox_offset
+        )
+    return out

src/flag_gems/runtime/backend/_metax/tune_configs.yaml

@@ -1468,3 +1468,37 @@ diff:
   - 16
   - 256
   - 1024
+
+alpha_dropout:
+- META:
+    BLOCK: 256
+  num_warps: 4
+  num_stages: 2
+- META:
+    BLOCK: 256
+  num_warps: 8
+  num_stages: 2
+- META:
+    BLOCK: 512
+  num_warps: 4
+  num_stages: 2
+- META:
+    BLOCK: 512
+  num_warps: 8
+  num_stages: 2
+- META:
+    BLOCK: 256
+  num_warps: 4
+  num_stages: 3
+- META:
+    BLOCK: 256
+  num_warps: 8
+  num_stages: 3
+- META:
+    BLOCK: 512
+  num_warps: 4
+  num_stages: 3
+- META:
+    BLOCK: 512
+  num_warps: 8
+  num_stages: 3