# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import torch
import torch.nn.functional as F
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, tl, triton


@triton.jit
def _layernorm_select01_kernel(
    output_ptr,
    gate_out_ptr,
    x_ptr,
    weight_ptr,
    bias_ptr,
    scale0_ptr,
    shift0_ptr,
    gate0_ptr,
    scale1_ptr,
    shift1_ptr,
    gate1_ptr,
    index_ptr,
    inner_dim,
    seq_len,
    stride_x_row,
    stride_out_row,
    stride_gate_out_row,
    stride_w,
    stride_b,
    stride_s0_b,
    stride_s0_c,
    stride_sh0_b,
    stride_sh0_c,
    stride_g0_b,
    stride_g0_c,
    stride_s1_b,
    stride_s1_c,
    stride_sh1_b,
    stride_sh1_c,
    stride_g1_b,
    stride_g1_c,
    stride_i_b,
    stride_i_l,
    eps,
    has_weight: tl.constexpr,
    has_bias: tl.constexpr,
    block_n: tl.constexpr,
):
    row = tl.program_id(0)
    cols = tl.arange(0, block_n)
    mask = cols < inner_dim

    x_row_ptr = x_ptr + row * stride_x_row
    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)

    mean = tl.sum(x, axis=0) / inner_dim
    xbar = tl.where(mask, x - mean, 0.0)
    var = tl.sum(xbar * xbar, axis=0) / inner_dim
    x_hat = (x - mean) * tl.rsqrt(var + eps)

    if has_weight:
        weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
        x_hat = x_hat * weight
    if has_bias:
        bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
        x_hat = x_hat + bias

    batch_idx = row // seq_len
    seq_idx = row % seq_len
    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)

    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c

    scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs)
    shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs)
    gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs)

    scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32)
    shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32)
    gate = tl.load(gate_ptrs, mask=mask, other=0.0)

    y = x_hat * (1.0 + scale) + shift
    tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask)
    tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask)


@triton.jit
def _residual_layernorm_select01_kernel(
    output_ptr,
    residual_out_ptr,
    gate_out_ptr,
    x_ptr,
    residual_ptr,
    residual_gate_ptr,
    weight_ptr,
    bias_ptr,
    scale0_ptr,
    shift0_ptr,
    gate0_ptr,
    scale1_ptr,
    shift1_ptr,
    gate1_ptr,
    index_ptr,
    inner_dim,
    seq_len,
    stride_x_row,
    stride_res_row,
    stride_res_gate_row,
    stride_out_row,
    stride_res_out_row,
    stride_gate_out_row,
    stride_w,
    stride_b,
    stride_s0_b,
    stride_s0_c,
    stride_sh0_b,
    stride_sh0_c,
    stride_g0_b,
    stride_g0_c,
    stride_s1_b,
    stride_s1_c,
    stride_sh1_b,
    stride_sh1_c,
    stride_g1_b,
    stride_g1_c,
    stride_i_b,
    stride_i_l,
    eps,
    has_weight: tl.constexpr,
    has_bias: tl.constexpr,
    block_n: tl.constexpr,
):
    row = tl.program_id(0)
    cols = tl.arange(0, block_n)
    mask = cols < inner_dim

    x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(tl.float32)
    residual = tl.load(residual_ptr + row * stride_res_row + cols, mask=mask, other=0.0).to(tl.float32)
    residual_gate = tl.load(
        residual_gate_ptr + row * stride_res_gate_row + cols,
        mask=mask,
        other=0.0,
    ).to(tl.float32)

    residual_out = residual + residual_gate * x
    tl.store(residual_out_ptr + row * stride_res_out_row + cols, residual_out, mask=mask)

    mean = tl.sum(residual_out, axis=0) / inner_dim
    xbar = tl.where(mask, residual_out - mean, 0.0)
    var = tl.sum(xbar * xbar, axis=0) / inner_dim
    x_hat = (residual_out - mean) * tl.rsqrt(var + eps)

    if has_weight:
        weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
        x_hat = x_hat * weight
    if has_bias:
        bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
        x_hat = x_hat + bias

    batch_idx = row // seq_len
    seq_idx = row % seq_len
    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)

    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c

    scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs)
    shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs)
    gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs)

    scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32)
    shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32)
    gate = tl.load(gate_ptrs, mask=mask, other=0.0)

    y = x_hat * (1.0 + scale) + shift
    tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask)
    tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask)


def split_select01_mod_params(mod_params: torch.Tensor) -> tuple[torch.Tensor, ...]:
    shift, scale, gate = mod_params.chunk(3, dim=-1)
    actual_batch = shift.size(0) // 2
    if actual_batch == 0 or shift.size(0) != actual_batch * 2:
        raise ValueError("select01 modulation expects mod_params batch dim to be 2 * batch")
    return (
        scale[:actual_batch],
        shift[:actual_batch],
        gate[:actual_batch],
        scale[actual_batch:],
        shift[actual_batch:],
        gate[actual_batch:],
    )


def select01_modulation_native(
    mod_params: torch.Tensor,
    index: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
    idx = index.to(dtype=torch.bool).unsqueeze(-1)
    scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
    shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
    gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
    return scale, shift, gate


def _validate_select01_inputs(
    x: torch.Tensor,
    index: torch.Tensor,
    tensors: tuple[torch.Tensor, ...],
) -> None:
    if x.dim() != 3:
        raise ValueError("x must be 3D [B, L, C]")
    if index.dim() != 2 or index.shape[1] != x.shape[1] or index.shape[0] not in (1, x.shape[0]):
        raise ValueError("index must be 2D [B or 1, L] and match x batch/sequence")
    batch_size, _, hidden_size = x.shape
    for tensor in tensors:
        if tensor.dim() != 2 or tensor.shape != (batch_size, hidden_size):
            raise ValueError("scale/shift/gate tensors must be 2D [B, C]")


def _maybe_contiguous(x: torch.Tensor) -> torch.Tensor:
    return x if x.is_contiguous() else x.contiguous()


def _launch_layernorm_select01(
    x: torch.Tensor,
    weight: torch.Tensor | None,
    bias: torch.Tensor | None,
    scale0: torch.Tensor,
    shift0: torch.Tensor,
    gate0: torch.Tensor,
    scale1: torch.Tensor,
    shift1: torch.Tensor,
    gate1: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1))
    batch_size, seq_len, hidden_size = x.shape
    output = torch.empty_like(x)
    gate_out = torch.empty_like(x)

    x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size)
    output_2d = output.view(batch_size * seq_len, hidden_size)
    gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size)
    index = _maybe_contiguous(index)
    scale0 = _maybe_contiguous(scale0)
    shift0 = _maybe_contiguous(shift0)
    gate0 = _maybe_contiguous(gate0)
    scale1 = _maybe_contiguous(scale1)
    shift1 = _maybe_contiguous(shift1)
    gate1 = _maybe_contiguous(gate1)

    weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d
    bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d
    block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size))
    if hidden_size > block_n:
        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")

    _layernorm_select01_kernel[(batch_size * seq_len,)](
        output_2d,
        gate_out_2d,
        x_2d,
        weight_arg,
        bias_arg,
        scale0,
        shift0,
        gate0,
        scale1,
        shift1,
        gate1,
        index,
        hidden_size,
        seq_len,
        x_2d.stride(0),
        output_2d.stride(0),
        gate_out_2d.stride(0),
        weight_arg.stride(0) if weight is not None else 0,
        bias_arg.stride(0) if bias is not None else 0,
        scale0.stride(0),
        scale0.stride(1),
        shift0.stride(0),
        shift0.stride(1),
        gate0.stride(0),
        gate0.stride(1),
        scale1.stride(0),
        scale1.stride(1),
        shift1.stride(0),
        shift1.stride(1),
        gate1.stride(0),
        gate1.stride(1),
        index.stride(0) if index.shape[0] > 1 else 0,
        index.stride(1),
        eps,
        has_weight=weight is not None,
        has_bias=bias is not None,
        block_n=block_n,
        num_warps=4,
        num_stages=4,
    )
    return output, gate_out


def _launch_residual_layernorm_select01(
    x: torch.Tensor,
    residual: torch.Tensor,
    residual_gate: torch.Tensor,
    weight: torch.Tensor | None,
    bias: torch.Tensor | None,
    scale0: torch.Tensor,
    shift0: torch.Tensor,
    gate0: torch.Tensor,
    scale1: torch.Tensor,
    shift1: torch.Tensor,
    gate1: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1))
    if residual.shape != x.shape:
        raise ValueError("residual must have the same shape as x")
    if residual_gate.shape != x.shape:
        raise ValueError("residual_gate must have the same shape as x")

    batch_size, seq_len, hidden_size = x.shape
    output = torch.empty_like(x)
    residual_out = torch.empty_like(x)
    gate_out = torch.empty_like(x)

    x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size)
    residual_2d = _maybe_contiguous(residual).view(batch_size * seq_len, hidden_size)
    residual_gate_2d = _maybe_contiguous(residual_gate).view(batch_size * seq_len, hidden_size)
    output_2d = output.view(batch_size * seq_len, hidden_size)
    residual_out_2d = residual_out.view(batch_size * seq_len, hidden_size)
    gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size)
    index = _maybe_contiguous(index)
    scale0 = _maybe_contiguous(scale0)
    shift0 = _maybe_contiguous(shift0)
    gate0 = _maybe_contiguous(gate0)
    scale1 = _maybe_contiguous(scale1)
    shift1 = _maybe_contiguous(shift1)
    gate1 = _maybe_contiguous(gate1)

    weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d
    bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d
    block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size))
    if hidden_size > block_n:
        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")

    _residual_layernorm_select01_kernel[(batch_size * seq_len,)](
        output_2d,
        residual_out_2d,
        gate_out_2d,
        x_2d,
        residual_2d,
        residual_gate_2d,
        weight_arg,
        bias_arg,
        scale0,
        shift0,
        gate0,
        scale1,
        shift1,
        gate1,
        index,
        hidden_size,
        seq_len,
        x_2d.stride(0),
        residual_2d.stride(0),
        residual_gate_2d.stride(0),
        output_2d.stride(0),
        residual_out_2d.stride(0),
        gate_out_2d.stride(0),
        weight_arg.stride(0) if weight is not None else 0,
        bias_arg.stride(0) if bias is not None else 0,
        scale0.stride(0),
        scale0.stride(1),
        shift0.stride(0),
        shift0.stride(1),
        gate0.stride(0),
        gate0.stride(1),
        scale1.stride(0),
        scale1.stride(1),
        shift1.stride(0),
        shift1.stride(1),
        gate1.stride(0),
        gate1.stride(1),
        index.stride(0) if index.shape[0] > 1 else 0,
        index.stride(1),
        eps,
        has_weight=weight is not None,
        has_bias=bias is not None,
        block_n=block_n,
        num_warps=4,
        num_stages=4,
    )
    return output, residual_out, gate_out


def can_use_qwen_select01_triton(x: torch.Tensor) -> bool:
    return HAS_TRITON and current_platform.is_cuda() and x.is_cuda


def fused_layernorm_select01(
    x: torch.Tensor,
    mod_params: torch.Tensor,
    index: torch.Tensor,
    eps: float,
    weight: torch.Tensor | None = None,
    bias: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
    is_compiling = torch.compiler.is_compiling()
    if can_use_qwen_select01_triton(x) and not is_compiling:
        return _launch_layernorm_select01(
            x,
            weight,
            bias,
            scale0,
            shift0,
            gate0,
            scale1,
            shift1,
            gate1,
            index,
            eps,
        )

    scale, shift, gate = select01_modulation_native(mod_params, index)
    out = F.layer_norm(
        x.float(),
        (x.shape[-1],),
        weight=weight.float() if weight is not None else None,
        bias=bias.float() if bias is not None else None,
        eps=eps,
    ).to(x.dtype)
    return out * (1 + scale) + shift, gate


def fused_residual_layernorm_select01(
    x: torch.Tensor,
    residual: torch.Tensor,
    residual_gate: torch.Tensor,
    mod_params: torch.Tensor,
    index: torch.Tensor,
    eps: float,
    weight: torch.Tensor | None = None,
    bias: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
    is_compiling = torch.compiler.is_compiling()
    if can_use_qwen_select01_triton(x) and not is_compiling:
        return _launch_residual_layernorm_select01(
            x,
            residual,
            residual_gate,
            weight,
            bias,
            scale0,
            shift0,
            gate0,
            scale1,
            shift1,
            gate1,
            index,
            eps,
        )

    scale, shift, gate = select01_modulation_native(mod_params, index)
    residual_out = residual + residual_gate * x
    out = F.layer_norm(
        residual_out.float(),
        (residual_out.shape[-1],),
        weight=weight.float() if weight is not None else None,
        bias=bias.float() if bias is not None else None,
        eps=eps,
    ).to(residual_out.dtype)
    return out * (1 + scale) + shift, residual_out, gate