"""DeepGEMM-backed helpers for SigmaMoE forward pass.

This module provides utilities to run SigmaMoE's expert FFN using DeepGEMM's
highly-optimized FP8 kernels. It focuses on the forward pass (prefill/inference),
building an m-grouped contiguous layout per expert, casting activations/weights
to FP8, invoking DeepGEMM kernels, and scattering results back to the original
token order with top-k gating weights applied.

Usage sketch (inside SigmaMoE.forward):

    from mt_dao.models.ffns.deepgemm_moe import (
        is_deepgemm_available,
        moe_forward_deepgemm,
    )

    if is_deepgemm_available() and self.n_active_experts > 1:
        # sel_index: [B, T, k], sel_val: [B, T, k]
        y = moe_forward_deepgemm(
            x, sel_index, sel_val,
            self.up_proj, self.down_proj,
            activation=self.activation,
        )
        return y

Notes
-----
- Requires NVIDIA SM90/SM100 device and the DeepGEMM package built/installed.
- We assume the MoE top-k router has already produced sel_index and sel_val.
- Output dtype is bfloat16 by default to match common training/inference setups.

"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, NamedTuple

import torch
import torch.nn.functional as F  # noqa: N812

if TYPE_CHECKING:
    from collections.abc import Callable

try:
    from deep_gemm import m_grouped_fp8_gemm_nt_contiguous
    from deep_gemm.utils import (
        ceil_div,
        get_mk_alignment_for_contiguous_layout,
        per_block_cast_to_fp8,
        per_token_cast_to_fp8,
    )

    _DG_AVAILABLE = True
except (ImportError, ModuleNotFoundError):  # pragma: no cover - best effort soft import
    # Define minimal stubs to keep module importable; real calls are guarded
    def m_grouped_fp8_gemm_nt_contiguous(*_args: object, **_kwargs: object) -> None:
        """Stub: DeepGEMM not available.

        Parameters
        ----------
        *_args : object
            Positional arguments for the DeepGEMM kernel.
        **_kwargs : object
            Keyword arguments for the DeepGEMM kernel.

        Raises
        ------
        RuntimeError
            If DeepGEMM is not available.

        """
        msg = "DeepGEMM not available"
        raise RuntimeError(msg)

    def get_mk_alignment_for_contiguous_layout() -> int:
        """Fallback alignment when DeepGEMM is missing.

        Returns
        -------
        int
            The fallback alignment value: 128.

        """
        return 128

    def per_block_cast_to_fp8(
        x: torch.Tensor,  # noqa: ARG001
        *,
        use_ue8m0: bool,  # noqa: ARG001
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Stub: DeepGEMM not available.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor to cast.
        use_ue8m0 : bool
            Whether to use packed UE8M0 format.

        Raises
        ------
        RuntimeError
            If DeepGEMM is not available.

        """
        msg = "DeepGEMM not available"
        raise RuntimeError(msg)

    def per_token_cast_to_fp8(
        x: torch.Tensor,  # noqa: ARG001
        *,
        use_ue8m0: bool,  # noqa: ARG001
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Stub: DeepGEMM not available.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor to cast.
        use_ue8m0 : bool
            Whether to use packed UE8M0 format.

        Raises
        ------
        RuntimeError
            If DeepGEMM is not available.

        """
        msg = "DeepGEMM not available"
        raise RuntimeError(msg)

    def ceil_div(x: int, y: int) -> int:
        """Ceiling division fallback.

        Returns
        -------
        int
            The ceiling division result.

        """
        return (x + y - 1) // y

    _DG_AVAILABLE = False


_SM100_MAJOR = 10
_RANK3 = 3
_RANK2 = 2


def _use_ue8m0() -> bool:
    """Return True if we should use packed UE8M0 scales (SM100+).

    Returns
    -------
    bool
        True on SM100+, False otherwise.

    """
    if not torch.cuda.is_available():
        return False
    major, _ = torch.cuda.get_device_capability()
    return major >= _SM100_MAJOR


def is_deepgemm_available() -> bool:
    """Return True if DeepGEMM can be imported and CUDA is available.

    Returns
    -------
    bool
        True when DeepGEMM is importable and CUDA is available.

    """
    return bool(_DG_AVAILABLE and torch.cuda.is_available())


# -----------------------------------------------------------------------------
# Lightweight buffer pool to trade spare memory for lower allocator overhead
# Enabled only when gradients are disabled (e.g., eval/inference) by default.
# Set mt_dao_MOE_DG_BUFFER_POOL=0 to disable, or =1 to force-enable.
# -----------------------------------------------------------------------------
_USE_BUFFER_POOL_ENV = os.environ.get("mt_dao_MOE_DG_BUFFER_POOL", "auto").lower()


def _buffer_pool_enabled() -> bool:
    if _USE_BUFFER_POOL_ENV in {"1", "true", "yes", "on"}:
        return True
    if _USE_BUFFER_POOL_ENV in {"0", "false", "no", "off"}:
        return False
    # auto: enable when not tracking gradients
    return not torch.is_grad_enabled()


_BUFFER_POOL: dict[tuple[torch.device, torch.dtype, tuple[int, ...]], torch.Tensor] = {}


def _get_buffer(
    shape: tuple[int, ...],
    dtype: torch.dtype,
    device: torch.device,
    *,
    zero: bool = False,
) -> torch.Tensor:
    key = (device, dtype, shape)
    buf = _BUFFER_POOL.get(key)
    if (
        buf is None
        or buf.device != device
        or buf.dtype != dtype
        or tuple(buf.shape) != shape
    ):
        buf = (
            torch.zeros(shape, device=device, dtype=dtype)
            if zero
            else torch.empty(shape, device=device, dtype=dtype)
        )
        _BUFFER_POOL[key] = buf
    elif zero:
        buf.zero_()
    return buf


@dataclass
class GroupedLayout:
    """Holds the m-grouped contiguous layout for DeepGEMM calls.

    Parameters
    ----------
    a_index : torch.Tensor
        Indices into x_flat selecting rows in grouped, expert-major order.
        Shape: [M_aligned], dtype: long.
    m_indices : torch.Tensor
        Per-row expert id or -1 for padded rows. Shape: [M_aligned], int32.
    row_weight : torch.Tensor
        Per-row gating weight (0 for padding). Shape: [M_aligned].
    row_token : torch.Tensor
        Original token index for each row (ignored if m_indices == -1).
        Shape: [M_aligned], dtype: long.
    counts : torch.Tensor
        Valid (unpadded) row counts per expert. Shape: [E], dtype: long.
    aligned_counts : torch.Tensor
        Aligned row counts per expert. Shape: [E], dtype: long.
    total_aligned : int
        Sum of aligned_counts.

    """

    a_index: torch.Tensor
    m_indices: torch.Tensor
    row_weight: torch.Tensor
    row_token: torch.Tensor
    counts: torch.Tensor
    aligned_counts: torch.Tensor
    total_aligned: int


def _build_grouped_layout(  # noqa: PLR0914
    sel_index: torch.Tensor,
    sel_weight: torch.Tensor,
    n_experts: int,
) -> GroupedLayout:
    """Create expert-major contiguous layout with per-expert alignment.

    Parameters
    ----------
    sel_index : LongTensor [B, T, K]
        Top-k expert indices per token.
    sel_weight : Tensor [B, T, K]
        Associated gating weights per token (already activated, e.g., sigmoid/softmax).
    n_experts : int
        Total number of experts.

    Returns
    -------
    GroupedLayout
        Indices/mapping arrays to materialize (and later scatter) grouped rows.

    """
    device = sel_index.device
    bsz, seqlen, ksel = sel_index.shape
    n_tokens = bsz * seqlen

    sel_index_f = sel_index.reshape(n_tokens, ksel)
    sel_weight_f = sel_weight.reshape(n_tokens, ksel)

    # Replicate tokens by selected experts: build (row_token, row_expert, row_weight)
    row_token = torch.arange(n_tokens, device=device).repeat_interleave(ksel)
    row_expert = sel_index_f.reshape(-1)
    row_weight = sel_weight_f.reshape(-1)

    # Group by expert id via sort; compute counts per expert
    order = torch.argsort(row_expert)
    row_token = row_token[order]
    row_expert = row_expert[order]
    row_weight = row_weight[order]

    counts = torch.bincount(row_expert, minlength=n_experts)

    # Per-expert alignment required by DeepGEMM grouped-contiguous API
    alignment = get_mk_alignment_for_contiguous_layout()
    aligned_counts = ((counts + alignment - 1) // alignment) * alignment
    offsets = torch.nn.functional.pad(aligned_counts.cumsum(0), (1, 0))[:-1]
    total_aligned = int(aligned_counts.sum().item())

    # Allocate outputs (initialize paddings to zeros)
    a_index = torch.zeros(total_aligned, dtype=torch.long, device=device)
    m_indices = torch.full((total_aligned,), -1, dtype=torch.int32, device=device)
    row_weight_aligned = torch.zeros(
        total_aligned,
        dtype=row_weight.dtype,
        device=device,
    )
    row_token_aligned = torch.zeros(total_aligned, dtype=torch.long, device=device)

    # Fill valid region for each expert; pad tail with -1/0s
    ptr = 0
    for e in range(n_experts):
        cnt = int(counts[e].item())
        if cnt == 0:
            continue
        start = int(offsets[e].item())
        end_valid = start + cnt
        sl = slice(ptr, ptr + cnt)
        a_index[start:end_valid] = row_token[sl]
        m_indices[start:end_valid] = int(e)
        row_weight_aligned[start:end_valid] = row_weight[sl]
        row_token_aligned[start:end_valid] = row_token[sl]
        ptr += cnt

    return GroupedLayout(
        a_index=a_index,
        m_indices=m_indices,
        row_weight=row_weight_aligned,
        row_token=row_token_aligned,
        counts=counts,
        aligned_counts=aligned_counts,
        total_aligned=total_aligned,
    )


def _cast_grouped_weights_fp8(
    w_bf16: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Cast per-expert weight matrix to FP8-block layout expected by DeepGEMM.

    Parameters
    ----------
    w_bf16 : torch.Tensor
        Tensor of shape [E, N, K] (K-major memory is fine; DeepGEMM handles aliasing)

    Returns
    -------
    (w_fp8, w_scales)
        - w_fp8: float8_e4m3fn tensor with same leading dims as w_bf16
        - w_scales: float32 tensor [E, ceil(N/128), ceil(K/128)]

    """
    assert w_bf16.dim() == _RANK3, "Expect [E, N, K]"
    e_count, n_dim, k_dim = w_bf16.shape
    device = w_bf16.device
    w_fp8 = torch.empty(
        (e_count, n_dim, k_dim),
        device=device,
        dtype=torch.float8_e4m3fn,
    )
    w_sf = torch.empty(
        (e_count, ceil_div(n_dim, 128), ceil_div(k_dim, 128)),
        device=device,
        dtype=torch.float,
    )
    use_ue8m0 = _use_ue8m0()
    for e in range(e_count):
        w_fp8[e], w_sf[e] = per_block_cast_to_fp8(w_bf16[e], use_ue8m0=use_ue8m0)
    return w_fp8, w_sf


class _FP8CacheEntry(NamedTuple):
    fp8: torch.Tensor
    scales: torch.Tensor
    version: int | None
    device: torch.device


# Simple module-level caches keyed by storage pointer and role
_FP8_UP_CACHE: dict[tuple[int, str], _FP8CacheEntry] = {}
_FP8_DOWN_CACHE: dict[tuple[int, str], _FP8CacheEntry] = {}


def clear_fp8_weight_cache() -> None:
    """Clear the small FP8 weight caches."""
    _FP8_UP_CACHE.clear()
    _FP8_DOWN_CACHE.clear()


def _get_up_weights_fp8_cached(
    up_bf16: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Get or compute FP8 up weights with caching per tensor version.

    The cache key uses the storage pointer and a role tag to avoid collisions.
    It refreshes when the tensor's internal version counter changes.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The cached or freshly cast FP8 weights and their scales.

    """
    key = (up_bf16.storage().data_ptr(), "up")
    ver = getattr(up_bf16, "_version", None)
    entry = _FP8_UP_CACHE.get(key)
    if (
        entry is not None
        and entry.version == ver
        and entry.device == up_bf16.device
        and tuple(entry.fp8.shape) == tuple(up_bf16.shape)
    ):
        return entry.fp8, entry.scales
    w_fp8, w_sf = _cast_grouped_weights_fp8(up_bf16)
    _FP8_UP_CACHE[key] = _FP8CacheEntry(w_fp8, w_sf, ver, up_bf16.device)
    return w_fp8, w_sf


def _get_down_weights_fp8_cached(
    down_bf16: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Get or compute FP8 down weights (transposed) with caching.

    Expects input in [E, S, V] BF16 and internally transposes to [E, V, S]
    before casting.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The cached or freshly cast FP8 weights (transposed) and their scales.

    """
    key = (down_bf16.storage().data_ptr(), "down_T")
    ver = getattr(down_bf16, "_version", None)
    entry = _FP8_DOWN_CACHE.get(key)
    # Cached entry stores FP8 for transposed shape; verify device matches version
    if (
        entry is not None
        and entry.version == ver
        and entry.device == down_bf16.device
    ):
        return entry.fp8, entry.scales
    # Create transposed contiguous view [E, V, S] for casting once per version
    w_down_bf16_t = down_bf16.permute(0, 2, 1).contiguous()
    w_fp8, w_sf = _cast_grouped_weights_fp8(w_down_bf16_t)
    _FP8_DOWN_CACHE[key] = _FP8CacheEntry(w_fp8, w_sf, ver, down_bf16.device)
    return w_fp8, w_sf


def _cast_tokens_fp8(a_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Cast tokens (activations) to FP8 per-token layout.

    Parameters
    ----------
    a_bf16 : torch.Tensor
        Tensor of tokens to be casted to FP8.

    Returns
    -------
    torch.Tensor
        Casted tensor in FP8 format.

    """
    use_ue8m0 = _use_ue8m0()
    return per_token_cast_to_fp8(a_bf16, use_ue8m0=use_ue8m0)


def _build_all_experts_layout(sel_weight: torch.Tensor) -> GroupedLayout:  # noqa: PLR0914
    # pylint: disable=too-many-locals
    """Build grouped-contiguous layout for the all-experts-active case.

    Parameters
    ----------
    sel_weight : torch.Tensor
        Tensor of shape [B, T, E] with gating weights for all experts.

    Returns
    -------
    GroupedLayout
        Layout with per-expert aligned rows, no sorting or indices needed.

    """
    device = sel_weight.device
    bsz, seqlen, n_experts = sel_weight.shape
    n_tokens = bsz * seqlen

    alignment = get_mk_alignment_for_contiguous_layout()
    aligned_count = ((n_tokens + alignment - 1) // alignment) * alignment

    counts = torch.full((n_experts,), n_tokens, device=device, dtype=torch.long)
    aligned_counts = torch.full(
        (n_experts,), aligned_count, device=device, dtype=torch.long,
    )
    offsets = torch.nn.functional.pad(aligned_counts.cumsum(0), (1, 0))[:-1]
    total_aligned = int(aligned_counts.sum().item())

    a_index = torch.zeros(total_aligned, dtype=torch.long, device=device)
    m_indices = torch.full((total_aligned,), -1, dtype=torch.int32, device=device)
    row_weight = torch.zeros(total_aligned, dtype=sel_weight.dtype, device=device)
    row_token = torch.zeros(total_aligned, dtype=torch.long, device=device)

    base_tokens = torch.arange(n_tokens, device=device, dtype=torch.long)
    sel_w_flat = sel_weight.reshape(n_tokens, n_experts)

    for e in range(n_experts):
        start = int(offsets[e].item())
        end_valid = start + n_tokens

        a_index[start:end_valid] = base_tokens
        m_indices[start:end_valid] = e
        row_token[start:end_valid] = base_tokens
        row_weight[start:end_valid] = sel_w_flat[:, e]
        # Tail remains padding with defaults

    return GroupedLayout(
        a_index=a_index,
        m_indices=m_indices,
        row_weight=row_weight,
        row_token=row_token,
        counts=counts,
        aligned_counts=aligned_counts,
        total_aligned=total_aligned,
    )


def moe_forward_deepgemm(  # noqa: PLR0912, PLR0913, PLR0914, PLR0917
    x: torch.Tensor,
    sel_index: torch.Tensor,
    sel_weight: torch.Tensor,
    up_proj: torch.Tensor,
    down_proj: torch.Tensor,
    activation: Callable[[torch.Tensor], torch.Tensor] = F.silu,
    *,
    out_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """Run SigmaMoE FFN via DeepGEMM grouped FP8 kernels.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor of shape [B, T, D].
    sel_index : torch.Tensor
        Selection index tensor of shape [B, T, K].
    sel_weight : torch.Tensor
        Selection weight tensor of shape [B, T, K].
    up_proj : torch.Tensor
        Up-projection weight tensor of shape [E, S, D].
    down_proj : torch.Tensor
        Down-projection weight tensor of shape [E, S, V].
    activation : Callable[[torch.Tensor], torch.Tensor]
        Activation function to apply.
    out_dtype : torch.dtype
        Output data type.

    Returns
    -------
    torch.Tensor
        Output tensor of shape [B, T, V].

    Raises
    ------
    ValueError
        If input tensor shapes are incompatible.
    RuntimeError
        If DeepGEMM is not available.

    """
    if not is_deepgemm_available():
        msg = (
            "DeepGEMM is not available. Build/install it to use "
            "moe_forward_deepgemm."
        )
        raise RuntimeError(msg)

    bsz, seqlen, d_model = x.shape
    e_count, s_dim, d_w = up_proj.shape
    assert d_w == d_model, "up_proj last dim must equal x hidden size"
    assert down_proj.shape[:2] == (e_count, s_dim), "down_proj must be [E, S, V]"
    v_dim = int(down_proj.shape[2])

    # Check per-token casting alignment requirement
    if (d_model % 128) != 0 or (s_dim % 128) != 0:
        msg = "DeepGEMM per-token FP8 casting requires D and S to be multiples of 128."
        raise ValueError(msg)

    # Prepare grouped contiguous layout
    layout = _build_grouped_layout(
        sel_index=sel_index,
        sel_weight=sel_weight,
        n_experts=e_count,
    )

    # Gather tokens into expert-major contiguous buffer (BF16 for casting)
    x_flat = x.reshape(-1, d_model)
    a_bf16 = x_flat.index_select(0, layout.a_index).to(torch.bfloat16)

    # Up-projection weights: [E, N=S, K=D]
    w_up_bf16 = up_proj.to(torch.bfloat16)
    b_up_fp8 = _get_up_weights_fp8_cached(w_up_bf16)
    a_up_fp8 = _cast_tokens_fp8(a_bf16)

    # Output buffer for up-proj: [M_aligned, S]
    if _buffer_pool_enabled():
        d_up = _get_buffer((layout.total_aligned, s_dim), out_dtype, x.device)
    else:
        d_up = torch.empty(
            (layout.total_aligned, s_dim),
            device=x.device,
            dtype=out_dtype,
        )
    m_grouped_fp8_gemm_nt_contiguous(a_up_fp8, b_up_fp8, d_up, layout.m_indices)

    # Zero-out padded rows (m_indices == -1) in-place to avoid extra allocation
    if (layout.m_indices == -1).any():
        d_up.masked_fill_(layout.m_indices.eq(-1).unsqueeze(1), 0)

    # Apply activation in-place and pre-scale rows by gating weights before down GEMM
    # Prefer in-place SiLU if available
    if activation is F.silu or activation == F.silu:
        F.silu(d_up, inplace=True)
    else:
        # Fallback: apply provided activation out-of-place into same tensor
        d_up.copy_(activation(d_up))
    d_up.mul_(layout.row_weight.unsqueeze(1))

    # Down-projection weights: need [E, N=V, K=S] (cache transposed FP8)
    w_down_bf16 = down_proj.to(torch.bfloat16)
    b_down_fp8 = _get_down_weights_fp8_cached(w_down_bf16)

    a_down_fp8 = _cast_tokens_fp8(d_up.to(torch.bfloat16))
    if _buffer_pool_enabled():
        y_concat = _get_buffer((layout.total_aligned, v_dim), out_dtype, x.device)
    else:
        y_concat = torch.empty(
            (layout.total_aligned, v_dim),
            device=x.device,
            dtype=out_dtype,
        )
    m_grouped_fp8_gemm_nt_contiguous(a_down_fp8, b_down_fp8, y_concat, layout.m_indices)

    # Zero padded rows (already pre-scaled above) in-place
    if (layout.m_indices == -1).any():
        y_concat.masked_fill_(layout.m_indices.eq(-1).unsqueeze(1), 0)

    # Scatter-add back to [N, V], then reshape to [B, T, V]
    if _buffer_pool_enabled():
        y_flat = _get_buffer((bsz * seqlen, v_dim), out_dtype, x.device, zero=True)
    else:
        y_flat = torch.zeros((bsz * seqlen, v_dim), device=x.device, dtype=out_dtype)
    valid = layout.m_indices.ne(-1)
    if valid.any():
        y_flat.index_add_(0, layout.row_token[valid], y_concat[valid])
    return y_flat.view(bsz, seqlen, v_dim)


def _dense_single_expert_forward(
    x: torch.Tensor,
    up_weight: torch.Tensor,
    down_weight: torch.Tensor,
    activation: Callable[[torch.Tensor], torch.Tensor],
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """Single-expert fast path using non-grouped DeepGEMM with FP8 activations.

    Expects up_weight shape [S, D] and down_weight shape [S, V].

    Returns
    -------
    torch.Tensor
        Output tensor with shape [B, T, V].

    Raises
    ------
    ValueError
        If D or S are not multiples of 128 required by per-token FP8 casting.

    """
    bsz, seqlen, d_model = x.shape
    s_dim, d_w = up_weight.shape
    assert d_w == d_model
    s2_dim, v_dim = down_weight.shape
    assert s2_dim == s_dim

    if (d_model % 128) != 0 or (s_dim % 128) != 0:
        msg = (
            "DeepGEMM per-token FP8 casting requires D and S to be multiples of 128."
        )
        raise ValueError(msg)

    x2d = x.reshape(-1, d_model).to(torch.bfloat16)
    a_fp8 = _cast_tokens_fp8(x2d)

    # Up: treat as 1-group weights of shape [1, N=S, K=D] with caching
    b_up_fp8 = _get_up_weights_fp8_cached(up_weight.unsqueeze(0).to(torch.bfloat16))
    if _buffer_pool_enabled():
        up = _get_buffer((x2d.shape[0], s_dim), out_dtype, x.device)
    else:
        up = torch.empty((x2d.shape[0], s_dim), device=x.device, dtype=out_dtype)
    m_indices = torch.zeros(x2d.shape[0], dtype=torch.int32, device=x.device)
    m_grouped_fp8_gemm_nt_contiguous(a_fp8, b_up_fp8, up, m_indices)
    # In-place activation to avoid an extra allocation
    if activation is F.silu or activation == F.silu:
        F.silu(up, inplace=True)
    else:
        up.copy_(activation(up))

    # Down: weights need [1, N=V, K=S] with caching (cache does transpose)
    b_down_fp8 = _get_down_weights_fp8_cached(
        down_weight.unsqueeze(0).to(torch.bfloat16),
    )
    if _buffer_pool_enabled():
        y2d = _get_buffer((x2d.shape[0], v_dim), out_dtype, x.device)
    else:
        y2d = torch.empty((x2d.shape[0], v_dim), device=x.device, dtype=out_dtype)
    m_grouped_fp8_gemm_nt_contiguous(
        _cast_tokens_fp8(up.to(torch.bfloat16)),
        b_down_fp8,
        y2d,
        m_indices,
    )
    return y2d.view(bsz, seqlen, v_dim)


def _dense_all_experts_forward(  # noqa: PLR0914
    x: torch.Tensor,
    sel_weight: torch.Tensor,
    weights: tuple[torch.Tensor, torch.Tensor],
    activation: Callable[[torch.Tensor], torch.Tensor],
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """All-experts-active path via grouped MoE kernel without sel_index.

    Returns
    -------
    torch.Tensor
        Output tensor with shape [B, T, V].

    Raises
    ------
    ValueError
        If D or S are not multiples of 128 required by per-token FP8 casting.

    """
    up_proj, down_proj = weights
    bsz, seqlen, d_model = x.shape
    assert up_proj.dim() == _RANK3
    assert down_proj.dim() == _RANK3
    e_count, s_dim, d_w = up_proj.shape
    assert d_w == d_model, "up_proj last dim must equal x hidden size"
    assert down_proj.shape[:2] == (e_count, s_dim), "down_proj must be [E,S,V]"
    v_dim = int(down_proj.shape[2])
    if (d_model % 128) != 0 or (s_dim % 128) != 0:
        msg = (
            "DeepGEMM per-token FP8 casting requires D and S to be multiples of 128."
        )
        raise ValueError(msg)

    layout = _build_all_experts_layout(sel_weight)

    # Gather tokens into expert-major contiguous buffer (BF16 for casting)
    x_flat = x.reshape(-1, d_model)
    a_bf16 = x_flat.index_select(0, layout.a_index).to(torch.bfloat16)

    # Up-projection weights: [E, N=S, K=D]
    w_up_bf16 = up_proj.to(torch.bfloat16)
    b_up_fp8 = _get_up_weights_fp8_cached(w_up_bf16)
    a_up_fp8 = _cast_tokens_fp8(a_bf16)

    # Output buffer for up-proj: [M_aligned, S]
    if _buffer_pool_enabled():
        d_up = _get_buffer((layout.total_aligned, s_dim), out_dtype, x.device)
    else:
        d_up = torch.empty(
            (layout.total_aligned, s_dim),
            device=x.device,
            dtype=out_dtype,
        )
    m_grouped_fp8_gemm_nt_contiguous(a_up_fp8, b_up_fp8, d_up, layout.m_indices)

    # Zero-out padded rows (m_indices == -1) in-place to avoid extra allocation
    pad_mask = layout.m_indices.eq(-1).unsqueeze(1)
    if pad_mask.any():
        d_up.masked_fill_(pad_mask, 0)

    # Apply activation and pre-scale rows by gating weights before down GEMM
    if activation is F.silu or activation == F.silu:
        F.silu(d_up, inplace=True)
    else:
        d_up.copy_(activation(d_up))
    d_up.mul_(layout.row_weight.unsqueeze(1))

    # Down-projection weights: need [E, N=V, K=S] (cache transposed FP8)
    w_down_bf16 = down_proj.to(torch.bfloat16)
    b_down_fp8 = _get_down_weights_fp8_cached(w_down_bf16)

    a_down_fp8 = _cast_tokens_fp8(d_up.to(torch.bfloat16))
    if _buffer_pool_enabled():
        y_concat = _get_buffer((layout.total_aligned, v_dim), out_dtype, x.device)
    else:
        y_concat = torch.empty(
            (layout.total_aligned, v_dim),
            device=x.device,
            dtype=out_dtype,
        )
    m_grouped_fp8_gemm_nt_contiguous(a_down_fp8, b_down_fp8, y_concat, layout.m_indices)

    # Zero-out padded rows in-place
    if pad_mask.any():
        y_concat.masked_fill_(pad_mask, 0)

    # Scatter-add back to [N, V], then reshape to [B, T, V]
    if _buffer_pool_enabled():
        y_flat = _get_buffer((bsz * seqlen, v_dim), out_dtype, x.device, zero=True)
    else:
        y_flat = torch.zeros((bsz * seqlen, v_dim), device=x.device, dtype=out_dtype)
    valid = layout.m_indices.ne(-1)
    if valid.any():
        y_flat.index_add_(0, layout.row_token[valid], y_concat[valid])
    return y_flat.view(bsz, seqlen, v_dim)


def dense_forward_deepgemm(  # noqa: PLR0913
    x: torch.Tensor,
    sel_weight: torch.Tensor | None,
    up_proj: torch.Tensor,
    down_proj: torch.Tensor,
    activation: Callable[[torch.Tensor], torch.Tensor] = F.silu,
    *,
    out_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    # pylint: disable=too-many-arguments, too-many-branches, too-many-statements
    """Fast path using DeepGEMM for either single expert or all-experts-active.

    Modes
    -----
    1) Single-expert: up_proj.shape == [S, D], down_proj.shape == [S, V]
       Computes y = down(activation(up(x))).
    2) All experts active: up_proj.shape == [E, S, D], down_proj.shape == [E, S, V],
       sel_weight.shape == [B, T, E]. Delegates to grouped MoE path with
       sel_index covering all experts per token.

    Returns
    -------
    torch.Tensor
        Output tensor of shape [B, T, V].

    Raises
    ------
    RuntimeError
        If DeepGEMM is not available.
    ValueError
        If tensor shapes are incompatible or required multiples-of-128
        constraints are not met, or if sel_weight is missing in the
        all-experts-active path.

    """
    if not is_deepgemm_available():
        msg = (
            "DeepGEMM is not available. Build/install it to use dense_forward_deepgemm."
        )
        raise RuntimeError(msg)

    # Branch on weight dimensionality
    if up_proj.dim() == _RANK2 and down_proj.dim() == _RANK2:
        return _dense_single_expert_forward(
            x=x,
            up_weight=up_proj,
            down_weight=down_proj,
            activation=activation,
            out_dtype=out_dtype,
        )

    # All-experts-active dense path via grouped MoE kernel without sel_index
    if sel_weight is None:
        msg = "sel_weight must be provided when all experts are active"
        raise ValueError(msg)
    return _dense_all_experts_forward(
        x=x,
        sel_weight=sel_weight,
        weights=(up_proj, down_proj),
        activation=activation,
        out_dtype=out_dtype,
    )
