# SPDX-License-Identifier: Apache-2.0
"""Shared MoE activation helpers and Triton kernels.

Terminology:
- "gated" activation means gate-up fusion: input is interpreted as [gate, up]
  (typically shape [..., 2D]) and output is activation(gate) * up (shape [..., D]).
- "non-gated" activation means no gate multiplication: input/output keep the
  same last-dimension size (typically [..., D] -> [..., D]).
"""

from typing import Optional, Tuple

import torch
import triton
import triton.language as tl


@triton.heuristics(
    {
        "N_DIV": lambda args: (args["N"] % args["BLOCK_SIZE_N"]) == 0,
    }
)
@triton.jit
def activation_and_mul_kernel(
    out_ptr,  # [M, N]
    in_ptr,  # [M, 2N], chunked layout [gate(0:N), up(N:2N)]
    M,
    N: tl.constexpr,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    ACT: tl.constexpr,
    N_DIV: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    # 1D launch over [M, N] tiles:
    # pid -> (pid_m, pid_n), each program computes one [1, BLOCK_SIZE_N] slice.
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(stride_in0 > 0)
    tl.assume(stride_in1 > 0)
    tl.assume(stride_out0 > 0)
    tl.assume(stride_out1 > 0)
    tl.assume(in_ptr.to(tl.int64) >= 0)
    tl.assume(out_ptr.to(tl.int64) >= 0)

    pid = tl.program_id(axis=0)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    mask_n = offs_n < N
    row_in = in_ptr + pid_m.to(tl.int64) * stride_in0
    gate_ptrs = row_in + offs_n * stride_in1
    up_ptrs = row_in + (offs_n + N) * stride_in1
    if N_DIV:
        gate = tl.load(gate_ptrs)
        up = tl.load(up_ptrs)
    else:
        gate = tl.load(gate_ptrs, mask=mask_n, other=0.0)
        up = tl.load(up_ptrs, mask=mask_n, other=0.0)

    gate_f = gate.to(tl.float32)
    up_f = up.to(tl.float32)
    # ACT=0: SiLU(x) * y, where SiLU(x)=x*sigmoid(x)
    # ACT=1: GELU(x) * y, exact erf form:
    #        GELU(x)=0.5*x*(1+erf(x/sqrt(2)))
    if ACT == 0:
        act = gate_f * (1.0 / (1.0 + tl.exp(-gate_f)))
    else:
        act = gate_f * 0.5 * (1.0 + tl.erf(gate_f * 0.7071067811865476))
    y = act * up_f

    out_ptrs = out_ptr + pid_m.to(tl.int64) * stride_out0 + offs_n * stride_out1
    if N_DIV:
        tl.store(out_ptrs, y.to(gate.dtype))
    else:
        tl.store(out_ptrs, y.to(gate.dtype), mask=mask_n)


def get_triton_activation_and_mul_config(M, _N):
    # 2 configs tuned on gfx930 (BW1000B), bf16:
    # small-M: BS128_W1 — stable, low overhead for decode
    # large-M: BS1024_W1 — ~1.4-1.7x faster than old BS512_W4 at M>=4K
    if M <= 16:
        return {"BLOCK_SIZE_N": 128, "num_warps": 1}
    return {"BLOCK_SIZE_N": 1024, "num_warps": 1}


def get_triton_activation_no_mul_config(M, N):
    # 2 configs tuned on gfx930 (BW1000B), bf16:
    # small-M: BS128_W1 — stable, low overhead for decode
    # large-M: BS1024_W2 — ~1.4-2.6x faster than old BS256_W2 at M>=4K
    if M <= 16:
        return {"BLOCK_SIZE_N": 128, "num_warps": 1}
    return {"BLOCK_SIZE_N": 1024, "num_warps": 2}


@triton.heuristics(
    {
        "N_DIV": lambda args: (args["N"] % args["BLOCK_SIZE_N"]) == 0,
    }
)
@triton.jit
def activation_no_mul_1d_kernel(
    out_ptr,  # [M, N]
    in_ptr,  # [M, N]
    M,
    N: tl.constexpr,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    ACT: tl.constexpr,  # 0:silu 1:gelu 2:relu2
    BLOCK_SIZE_N: tl.constexpr,
    N_DIV: tl.constexpr,
):
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(stride_in0 > 0)
    tl.assume(stride_in1 > 0)
    tl.assume(stride_out0 > 0)
    tl.assume(stride_out1 > 0)
    tl.assume(in_ptr.to(tl.int64) >= 0)
    tl.assume(out_ptr.to(tl.int64) >= 0)

    pid = tl.program_id(axis=0)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    mask_n = offs_n < N

    in_ptrs = in_ptr + pid_m.to(tl.int64) * stride_in0 + offs_n * stride_in1
    if N_DIV:
        x = tl.load(in_ptrs)
    else:
        x = tl.load(in_ptrs, mask=mask_n, other=0.0)

    xf = x.to(tl.float32)
    # ACT=0: SiLU(x)=x*sigmoid(x)
    # ACT=1: GELU(x)=0.5*x*(1+erf(x/sqrt(2)))
    # ACT=2: ReLU^2(x)=max(x,0)^2
    if ACT == 0:
        y = xf * (1.0 / (1.0 + tl.exp(-xf)))
    elif ACT == 1:
        y = xf * 0.5 * (1.0 + tl.erf(xf * 0.7071067811865476))
    else:
        y = tl.where(xf > 0.0, xf * xf, 0.0)

    out_ptrs = out_ptr + pid_m.to(tl.int64) * stride_out0 + offs_n * stride_out1
    if N_DIV:
        tl.store(out_ptrs, y.to(x.dtype))
    else:
        tl.store(out_ptrs, y.to(x.dtype), mask=mask_n)


def triton_relu2(out: torch.Tensor, inp: torch.Tensor) -> None:
    # ReLU^2(x) = max(x, 0)^2
    assert inp.shape == out.shape
    assert inp.is_contiguous() and out.is_contiguous()
    M, N = inp.shape
    config = get_triton_activation_no_mul_config(M, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
    activation_no_mul_1d_kernel[grid](
        out,
        inp,
        M,
        N,
        inp.stride(0),
        inp.stride(1),
        out.stride(0),
        out.stride(1),
        ACT=2,
        **config,
    )


def triton_silu_no_mul(out: torch.Tensor, inp: torch.Tensor) -> None:
    # SiLU(x) = x * sigmoid(x)
    assert inp.shape == out.shape
    assert inp.is_contiguous() and out.is_contiguous()
    M, N = inp.shape
    config = get_triton_activation_no_mul_config(M, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
    activation_no_mul_1d_kernel[grid](
        out,
        inp,
        M,
        N,
        inp.stride(0),
        inp.stride(1),
        out.stride(0),
        out.stride(1),
        ACT=0,
        **config,
    )


def triton_gelu_no_mul(out: torch.Tensor, inp: torch.Tensor) -> None:
    # GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
    assert inp.shape == out.shape
    assert inp.is_contiguous() and out.is_contiguous()
    M, N = inp.shape
    config = get_triton_activation_no_mul_config(M, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
    activation_no_mul_1d_kernel[grid](
        out,
        inp,
        M,
        N,
        inp.stride(0),
        inp.stride(1),
        out.stride(0),
        out.stride(1),
        ACT=1,
        **config,
    )


def triton_silu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None:
    # Split input into [x, y] along last dim and compute:
    # out = SiLU(x) * y
    # Shapes: input=[M, 2N], out=[M, N]
    assert input.shape[-1] % 2 == 0
    assert input.is_contiguous()
    assert out.is_contiguous()
    M = input.numel() // input.shape[-1]
    N = input.shape[-1] // 2
    input_2d = input.view(M, input.shape[-1])
    out_2d = out.view(M, N)
    config = get_triton_activation_and_mul_config(M, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
    activation_and_mul_kernel[grid](
        out_2d,
        input_2d,
        M,
        N,
        input_2d.stride(0),
        input_2d.stride(1),
        out_2d.stride(0),
        out_2d.stride(1),
        ACT=0,
        **config,
    )


def triton_gelu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None:
    # Split input into [x, y] along last dim and compute:
    # out = GELU(x) * y
    # Shapes: input=[M, 2N], out=[M, N]
    assert input.shape[-1] % 2 == 0
    assert input.is_contiguous()
    assert out.is_contiguous()
    M = input.numel() // input.shape[-1]
    N = input.shape[-1] // 2
    input_2d = input.view(M, input.shape[-1])
    out_2d = out.view(M, N)
    config = get_triton_activation_and_mul_config(M, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
    activation_and_mul_kernel[grid](
        out_2d,
        input_2d,
        M,
        N,
        input_2d.stride(0),
        input_2d.stride(1),
        out_2d.stride(0),
        out_2d.stride(1),
        ACT=1,
        **config,
    )


def get_triton_swiglu_variant_config(M, D):
    # 2 configs tuned on gfx930 (BW1000B), bf16:
    # small-M: BS128_W1 — stable, low overhead for decode
    # large-M: BS1024_W1 — ~1.3-2.1x faster than old BS256_W2 at M>=4K
    if M <= 16:
        return {"BLOCK_SIZE_D": 128, "num_warps": 1}
    return {"BLOCK_SIZE_D": 1024, "num_warps": 1}


def get_triton_swiglu_oai_interleaved_config(M, D):
    # BLOCK_SIZE_D >= 128 produces incorrect results due to interleaved
    # address pattern (offs_d*2) failing to generate correct buffer_load.
    # Must keep BS64. small-M: W1; large-M: W4 for better occupancy.
    if M <= 16:
        return {"BLOCK_SIZE_D": 64, "num_warps": 1}
    return {"BLOCK_SIZE_D": 64, "num_warps": 4}


@triton.heuristics(
    {
        "D_DIV": lambda args: (args["D"] % args["BLOCK_SIZE_D"]) == 0,
    }
)
@triton.jit
def swiglu_variant_1d_kernel(
    out_ptr,  # [M, D]
    in_ptr,  # [M, 2*D]
    M,
    D: tl.constexpr,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    ALPHA,
    LIMIT,
    MODE: tl.constexpr,  # 0:gpt_oss 1:silu_clamp_mul 2:step
    BLOCK_SIZE_D: tl.constexpr,
    D_DIV: tl.constexpr,
):
    tl.assume(M > 0)
    tl.assume(D > 0)
    tl.assume(stride_in0 > 0)
    tl.assume(stride_in1 > 0)
    tl.assume(stride_out0 > 0)
    tl.assume(stride_out1 > 0)
    tl.assume(in_ptr.to(tl.int64) >= 0)
    tl.assume(out_ptr.to(tl.int64) >= 0)

    pid = tl.program_id(axis=0)
    num_pid_d = tl.cdiv(D, BLOCK_SIZE_D)
    pid_m = pid // num_pid_d
    pid_d = pid % num_pid_d

    offs_d = pid_d * BLOCK_SIZE_D + tl.arange(0, BLOCK_SIZE_D)
    mask_d = offs_d < D

    row_in = in_ptr + pid_m.to(tl.int64) * stride_in0
    gate_ptrs = row_in + offs_d * stride_in1
    up_ptrs = row_in + (offs_d + D) * stride_in1

    if D_DIV:
        gate = tl.load(gate_ptrs)
        up = tl.load(up_ptrs)
    else:
        gate = tl.load(gate_ptrs, mask=mask_d, other=0.0)
        up = tl.load(up_ptrs, mask=mask_d, other=0.0)

    gate_f = gate.to(tl.float32)
    up_f = up.to(tl.float32)
    gate_clamp = tl.minimum(gate_f, LIMIT)
    up_clamp = tl.minimum(tl.maximum(up_f, -LIMIT), LIMIT)

    # MODE=0 (gpt_oss/oai):
    # out = gate' * sigmoid(alpha*gate') * (up' + 1)
    # MODE=1 (silu_clamp_mul):
    # out = SiLU(gate') * up'
    # MODE=2 (step):
    # out = clamp(SiLU(gate), max=LIMIT) * up'
    if MODE == 0:
        y = gate_clamp * (1.0 / (1.0 + tl.exp(-(gate_clamp * ALPHA)))) * (up_clamp + 1.0)
    elif MODE == 1:
        y = gate_clamp * (1.0 / (1.0 + tl.exp(-gate_clamp))) * up_clamp
    else:
        silu_gate = gate_f * (1.0 / (1.0 + tl.exp(-gate_f)))
        silu_clamp = tl.minimum(silu_gate, LIMIT)
        y = silu_clamp * up_clamp

    out_ptrs = out_ptr + pid_m.to(tl.int64) * stride_out0 + offs_d * stride_out1
    if D_DIV:
        tl.store(out_ptrs, y.to(gate.dtype))
    else:
        tl.store(out_ptrs, y.to(gate.dtype), mask=mask_d)


def _triton_chunked_swiglu_variant(
    out: torch.Tensor,
    inp: torch.Tensor,
    *,
    alpha: float,
    limit: float,
    mode: int,
) -> None:
    assert inp.shape[-1] % 2 == 0
    assert inp.is_contiguous() and out.is_contiguous()
    M = inp.numel() // inp.shape[-1]
    D = inp.shape[-1] // 2
    out_rows = out.view(M, D)
    inp_rows = inp.view(M, inp.shape[-1])
    config = get_triton_swiglu_variant_config(M, D)
    # 1D launch: a single pid is mapped to (row_id, d_tile_id) inside kernel.
    grid = (M * triton.cdiv(D, config["BLOCK_SIZE_D"]),)
    swiglu_variant_1d_kernel[grid](
        out_rows,
        inp_rows,
        M,
        D,
        inp_rows.stride(0),
        inp_rows.stride(1),
        out_rows.stride(0),
        out_rows.stride(1),
        alpha,
        limit,
        MODE=mode,
        **config,
    )


def triton_swiglu_silu_clamp_mul(out: torch.Tensor, inp: torch.Tensor, limit: float) -> None:
    # out = SiLU(clamp(gate,max=limit)) * clamp(up,-limit,limit)
    _triton_chunked_swiglu_variant(out, inp, alpha=1.0, limit=limit, mode=1)


def triton_swiglu_gpt_oss_sigmoid_alpha(
    out: torch.Tensor, inp: torch.Tensor, alpha: float, limit: float
) -> None:
    # out = gate' * sigmoid(alpha * gate') * (up' + 1), chunked layout
    _triton_chunked_swiglu_variant(out, inp, alpha=alpha, limit=limit, mode=0)


def triton_swiglu_step_and_mul(out: torch.Tensor, inp: torch.Tensor, limit: float) -> None:
    # out = clamp(SiLU(gate),max=limit) * clamp(up,-limit,limit)
    _triton_chunked_swiglu_variant(out, inp, alpha=1.0, limit=limit, mode=2)


@triton.heuristics(
    {
        "D_DIV": lambda args: (args["D"] % args["BLOCK_SIZE_D"]) == 0,
    }
)
@triton.jit
def swiglu_oai_interleaved_1d_kernel(
    out_ptr,  # [M, D]
    in_ptr,  # [M, 2*D], interleaved layout [g0, u0, g1, u1, ...]
    M,
    D,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    ALPHA,
    LIMIT,
    BLOCK_SIZE_D: tl.constexpr,
    D_DIV: tl.constexpr,
):
    tl.assume(M > 0)
    tl.assume(D > 0)
    tl.assume(stride_in0 > 0)
    tl.assume(stride_in1 > 0)
    tl.assume(stride_out0 > 0)
    tl.assume(stride_out1 > 0)
    tl.assume(in_ptr.to(tl.int64) >= 0)
    tl.assume(out_ptr.to(tl.int64) >= 0)

    pid = tl.program_id(axis=0)
    num_pid_d = tl.cdiv(D, BLOCK_SIZE_D)
    pid_m = pid // num_pid_d
    pid_d = pid % num_pid_d

    offs_d = pid_d * BLOCK_SIZE_D + tl.arange(0, BLOCK_SIZE_D)
    mask_d = offs_d < D

    row_in = in_ptr + pid_m.to(tl.int64) * stride_in0
    gate_ptrs = row_in + (offs_d * 2) * stride_in1
    up_ptrs = row_in + (offs_d * 2 + 1) * stride_in1

    if D_DIV:
        gate = tl.load(gate_ptrs)
        up = tl.load(up_ptrs)
    else:
        gate = tl.load(gate_ptrs, mask=mask_d, other=0.0)
        up = tl.load(up_ptrs, mask=mask_d, other=0.0)

    gate_f = gate.to(tl.float32)
    up_f = up.to(tl.float32)
    gate_clamp = tl.minimum(gate_f, LIMIT)
    up_clamp = tl.minimum(tl.maximum(up_f, -LIMIT), LIMIT)

    # OpenAI interleaved SwiGLU variant:
    # out = gate' * sigmoid(alpha * gate') * (up' + 1)
    y = gate_clamp * (1.0 / (1.0 + tl.exp(-(gate_clamp * ALPHA)))) * (up_clamp + 1.0)

    out_ptrs = out_ptr + pid_m.to(tl.int64) * stride_out0 + offs_d * stride_out1
    if D_DIV:
        tl.store(out_ptrs, y.to(gate.dtype))
    else:
        tl.store(out_ptrs, y.to(gate.dtype), mask=mask_d)


def triton_swiglu_oai_and_mul(out: torch.Tensor, inp: torch.Tensor, alpha: float, limit: float) -> None:
    # out = gate' * sigmoid(alpha * gate') * (up' + 1), interleaved layout
    assert inp.shape[-1] % 2 == 0
    assert inp.is_contiguous() and out.is_contiguous()
    M = inp.numel() // inp.shape[-1]
    D = inp.shape[-1] // 2
    out_2d = out.view(M, D)
    inp_2d = inp.view(M, inp.shape[-1])
    config = get_triton_swiglu_oai_interleaved_config(M, D)
    grid = (M * triton.cdiv(D, config["BLOCK_SIZE_D"]),)
    swiglu_oai_interleaved_1d_kernel[grid](
        out_2d,
        inp_2d,
        M,
        D,
        inp_2d.stride(0),
        inp_2d.stride(1),
        out_2d.stride(0),
        out_2d.stride(1),
        alpha,
        limit,
        **config,
    )


_NO_MUL_ACTIVATIONS = {"silu_no_mul", "gelu_no_mul", "relu2_no_mul", "relu2"}
_GATED_ONLY_ACTIVATIONS = {"swigluoai", "swiglustep"}
_SUPPORTED_ACTIVATIONS = {
    "silu",
    "gelu",
    "swigluoai",
    "swiglustep",
    "silu_no_mul",
    "gelu_no_mul",
    "relu2_no_mul",
    "relu2",
}


def _normalize_activation_and_gate(
    activation: str,
    is_gated: Optional[bool],
) -> Tuple[str, bool]:
    # Match upstream style: keep (activation, is_gated) as a direct pair.
    # 1) validate string, 2) infer is_gated when not explicitly given,
    # 3) validate gated/non-gated constraints.
    activation_name = activation.lower()
    if activation_name not in _SUPPORTED_ACTIVATIONS:
        raise ValueError(f"Unsupported activation: {activation}")

    if is_gated is None:
        is_gated = activation_name not in _NO_MUL_ACTIVATIONS

    if is_gated and activation_name in _NO_MUL_ACTIVATIONS:
        raise ValueError(f"Activation '{activation}' is non-gated but is_gated=True")
    if (not is_gated) and activation_name in _GATED_ONLY_ACTIVATIONS:
        raise ValueError(f"Activation '{activation}' requires gated mode")

    return activation_name, is_gated


def adjust_N_for_activation(n: int, is_gated: bool) -> int:
    """Align with vLLM naming: output dim after expert activation."""
    return n // 2 if is_gated else n


def _apply_activation(
    activation: str,
    is_gated: bool,
    activated_out: torch.Tensor,
    ffn1_out_2d: torch.Tensor,
    gemm1_alpha: Optional[float],
    gemm1_limit: Optional[float],
) -> None:
    """Apply MoE expert activation.

    Shapes:
    - gated path: `ffn1_out_2d` is [rows, 2D], `activated_out` is [rows, D]
    - non-gated path: `ffn1_out_2d` and `activated_out` are both [rows, D]

    Meaning of `is_gated`:
    - True: run gate-up fusion (activation(gate) * up).
    - False: run plain activation without gate multiplication.
    """
    if is_gated:
        if activation == "silu":
            # Gated SiLU family:
            # base: out = SiLU(gate) * up
            # gpt-oss(alpha,limit): out = gate'*sigmoid(alpha*gate')*(up'+1)
            # limit-only: out = SiLU(gate')*up'
            if gemm1_alpha is not None:
                if gemm1_limit is None:
                    raise ValueError("gemm1_limit must be set when gemm1_alpha is set")
                triton_swiglu_gpt_oss_sigmoid_alpha(
                    activated_out,
                    ffn1_out_2d,
                    gemm1_alpha,
                    gemm1_limit,
                )
            elif gemm1_limit is not None:
                triton_swiglu_silu_clamp_mul(
                    activated_out,
                    ffn1_out_2d,
                    gemm1_limit,
                )
            else:
                triton_silu_and_mul(activated_out, ffn1_out_2d)
        elif activation == "gelu":
            # Gated GELU: out = GELU(gate) * up
            if gemm1_alpha is not None or gemm1_limit is not None:
                raise ValueError("gemm1_alpha/gemm1_limit are not supported for gated GELU")
            triton_gelu_and_mul(activated_out, ffn1_out_2d)
        elif activation == "swigluoai":
            # OpenAI interleaved SwiGLU variant:
            # out = gate'*sigmoid(alpha*gate')*(up'+1)
            alpha = 1.702 if gemm1_alpha is None else gemm1_alpha
            limit = 7.0 if gemm1_limit is None else gemm1_limit
            triton_swiglu_oai_and_mul(activated_out, ffn1_out_2d, alpha, limit)
        elif activation == "swiglustep":
            # Step variant:
            # out = clamp(SiLU(gate),max=limit) * clamp(up,-limit,limit)
            limit = 7.0 if gemm1_limit is None else gemm1_limit
            if gemm1_alpha is not None:
                raise ValueError("gemm1_alpha is not supported for swiglustep")
            triton_swiglu_step_and_mul(activated_out, ffn1_out_2d, limit)
        else:
            raise ValueError(f"Unsupported gated activation: {activation}")
        return

    if activation in {"silu", "silu_no_mul"}:
        # Non-gated: out = SiLU(x)
        triton_silu_no_mul(activated_out, ffn1_out_2d)
    elif activation in {"gelu", "gelu_no_mul"}:
        # Non-gated: out = GELU(x)
        triton_gelu_no_mul(activated_out, ffn1_out_2d)
    elif activation in {"relu2", "relu2_no_mul"}:
        # Non-gated: out = ReLU(x)^2
        triton_relu2(activated_out, ffn1_out_2d)
    else:
        raise ValueError(f"Unsupported non-gated activation: {activation}")
