# ********************************************************************************
# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
# ********************************************************************************

import os

import torch
import torch.nn.functional as F
from quack.gemm_interface import gemm

from ..count_cumsum import count_cumsum
from ..enums import ActivationType, is_glu
from ..quack_utils import gemm_dgated, gemm_gated
from .backward import (
    _down_projection_backward_act,
    _down_projection_backward_weight,
    _softmax_topk_bwd,
    _token_broadcast_backward,
    _up_projection_backward_act,
    _up_projection_backward_weight,
)
from .forward import _down_projection_forward, _router_forward, _softmax_topk_fwd, _up_projection_forward
from .triton_kernels import TC_topk_router_metadata_triton
from .utils import enable_quack_gemm, is_using_quack_gemm
from triton.testing import do_bench


def general_routing_router_metadata(
    router_scores_selected: torch.Tensor, sorted_selected_T: torch.Tensor, selected_E: torch.Tensor, T: int, E: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

    device = router_scores_selected.device

    expert_frequency, expert_frequency_offset = count_cumsum(selected_E, E, do_cumsum=True)
    expert_frequency_offset = torch.cat([torch.zeros(1, dtype=torch.int32, device=device), expert_frequency_offset])

    s_scatter_idx = selected_E.argsort().int()
    s_reverse_scatter_idx = torch.empty_like(s_scatter_idx)
    s_reverse_scatter_idx[s_scatter_idx] = torch.arange(
        s_scatter_idx.size(0), device=s_scatter_idx.device, dtype=s_scatter_idx.dtype
    )

    x_gather_idx = sorted_selected_T[s_scatter_idx]

    if T % 4 == 0 and T <= 50000:
        _, num_activated_expert_per_token_offset = count_cumsum(sorted_selected_T, T, do_cumsum=True)
    else:
        num_activated_expert_per_token_offset = torch.bincount(sorted_selected_T, minlength=T).cumsum(0).int()

    num_activated_expert_per_token_offset = torch.cat(
        [torch.zeros(1, dtype=torch.int32, device=device), num_activated_expert_per_token_offset]
    )

    return (
        expert_frequency,
        expert_frequency_offset,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        num_activated_expert_per_token_offset,
    )


class TC_Softmax_Topk_Router_Function(torch.autograd.Function):
    @staticmethod
    def forward(ctx, router_logits: torch.Tensor, E: int, K: int) -> tuple[torch.Tensor, torch.Tensor]:
        T = router_logits.size(0)

        # change this to router_logits.dtype (bfloat16) increase another 5 tflops at fwd at the cost of numerical accuracy
        topk_router_score = torch.empty(T, K, dtype=torch.float32, device=router_logits.device)
        topk_router_indices = torch.empty(T, K, dtype=torch.int32, device=router_logits.device)

        _softmax_topk_fwd(router_logits, topk_router_score, topk_router_indices, E, K)

        ctx.save_for_backward(topk_router_score, topk_router_indices)
        ctx.E = E
        ctx.dtype = router_logits.dtype

        return topk_router_score, topk_router_indices

    @staticmethod
    def backward(ctx, dtopk_score: torch.Tensor, _: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        T, K = dtopk_score.size()

        topk_router_score, topk_router_indices = ctx.saved_tensors
        dlogits = torch.zeros(T, ctx.E, dtype=ctx.dtype, device=topk_router_score.device)

        _softmax_topk_bwd(dlogits, None, dtopk_score, topk_router_score, topk_router_indices, K)

        return dlogits, None, None


class _UpProjection(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        x: torch.Tensor,
        w1: torch.Tensor,
        b1: torch.Tensor | None,
        expert_frequency_offset: torch.Tensor,
        total_expert_freq: int,
        K: int,
        stream_id: int,
        x_gather_idx: torch.Tensor,
        s_scatter_idx: torch.Tensor,
        s_reverse_scatter_idx: torch.Tensor,
        num_activated_expert_per_token_offset: torch.Tensor,
        is_varlen_K: bool,
        activation_type: ActivationType,
        is_inference_mode_enabled: bool,
    ) -> torch.Tensor:
        T, H = x.shape
        I, H, E = w1.shape
        is_glu_activation = is_glu(activation_type)
        if is_glu_activation:
            I //= 2
        TK = total_expert_freq

        if is_using_quack_gemm():
            assert not torch.compiler.is_compiling()
            assert is_glu_activation, "QuACK GEMM does not support non GLU activation yet"
            z, y1 = gemm_gated(
                x,
                w1.permute(2, 1, 0),
                activation="swiglu",
                cu_seqlens_m=expert_frequency_offset,
                A_idx=x_gather_idx,
                dynamic_scheduler=False,
            )
            timing = do_bench((lambda: gemm_gated(
                x,
                w1.permute(2, 1, 0),
                activation="swiglu",
                cu_seqlens_m=expert_frequency_offset,
                A_idx=x_gather_idx,
                dynamic_scheduler=False,
            )), 50, 500)
            print(f"up-proj-fwd Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
        else:
            z = torch.empty(TK, (2 * I if is_glu_activation else I), dtype=x.dtype, device=x.device)
            y1 = torch.empty(TK, I, dtype=x.dtype, device=x.device)
            _up_projection_forward(
                x=x,
                w1=w1,
                z=z,
                y1=y1,
                b1=b1,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                stream_id=stream_id,
                activation_type=activation_type.value,
                is_glu_activation=is_glu_activation,
                is_inference_mode_enabled=is_inference_mode_enabled,
            )
            timing = do_bench((lambda: _up_projection_forward(
                x=x,
                w1=w1,
                z=z,
                y1=y1,
                b1=b1,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                stream_id=stream_id,
                activation_type=activation_type.value,
                is_glu_activation=is_glu_activation,
                is_inference_mode_enabled=is_inference_mode_enabled,
            )), 50, 500)
            print(f"up-proj-fwd Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            
        ctx.T = T
        ctx.TK = TK
        ctx.E = E
        ctx.K = K
        ctx.H = H
        ctx.I = I
        ctx.is_varlen_K = is_varlen_K
        ctx.is_glu_activation = is_glu_activation
        ctx.stream_id = stream_id

        ctx.save_for_backward(
            x,
            w1,
            b1,
            expert_frequency_offset,
            x_gather_idx,
            s_scatter_idx,
            s_reverse_scatter_idx,
            num_activated_expert_per_token_offset,
        )

        ctx.mark_non_differentiable(y1)
        ctx.set_materialize_grads(False)

        return y1, z

    @staticmethod
    def backward(ctx, _: None, dz: torch.Tensor):
        is_compiling = torch.compiler.is_compiling()

        if not is_compiling:
            assert _ is None

        T = ctx.T
        TK = ctx.TK
        E = ctx.E
        K = ctx.K
        I = ctx.I
        H = ctx.H
        is_glu_activation = ctx.is_glu_activation
        is_varlen_K = ctx.is_varlen_K
        stream_id = ctx.stream_id

        (
            x,
            w1,
            b1,
            expert_frequency_offset,
            x_gather_idx,
            s_scatter_idx,
            s_reverse_scatter_idx,
            num_activated_expert_per_token_offset,
        ) = ctx.saved_tensors

        dw1 = torch.empty_like(w1)
        db1 = None if b1 is None else torch.empty_like(b1)

        if is_using_quack_gemm():
            assert not is_compiling

            gemm(
                x.T,
                dz,
                out=dw1.permute(2, 1, 0),
                cu_seqlens_k=expert_frequency_offset,
                A_idx=x_gather_idx,
                batch_idx_permute=None,
                dynamic_scheduler=False,
            )
            # timing = do_bench((lambda: gemm(
            #     x.T,
            #     dz,
            #     out=dw1.permute(2, 1, 0),
            #     cu_seqlens_k=expert_frequency_offset,
            #     A_idx=x_gather_idx,
            #     batch_idx_permute=None,
            #     dynamic_scheduler=False,
            # )), 50, 500)
            # print(f"dw1 Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            # dx_expanded = gemm(dz, w1.permute(2, 0, 1), cu_seqlens_m=expert_frequency_offset, dynamic_scheduler=False)
            
            # timing = do_bench((lambda: gemm(dz, w1.permute(2, 0, 1), cu_seqlens_m=expert_frequency_offset, dynamic_scheduler=False)), 50, 500)
            # print(f"dx Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
        else:
            dx_expanded = torch.empty(TK, H, dtype=dz.dtype, device=dz.device)

            _up_projection_backward_act(
                w1=w1,
                dx_expanded=dx_expanded,
                dz=dz,
                db1=db1,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                s_scatter_idx=s_scatter_idx,
                is_glu_activation=is_glu_activation,
                stream_id=stream_id,
            )

            _up_projection_backward_weight(
                x=x,
                dw1=dw1,
                dz=dz,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                is_glu_activation=is_glu_activation,
                stream_id=stream_id,
            )
            # timing = do_bench((lambda: _up_projection_backward_weight(
            #     x=x,
            #     dw1=dw1,
            #     dz=dz,
            #     expert_frequency_offset=expert_frequency_offset,
            #     expert_schedule_order=None,
            #     x_gather_idx=x_gather_idx,
            #     is_glu_activation=is_glu_activation,
            #     stream_id=stream_id,
            # )), 50, 500)
            # print(f"dw1 Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            
            # timing = do_bench((lambda: _up_projection_backward_act(
            #     w1=w1,
            #     dx_expanded=dx_expanded,
            #     dz=dz,
            #     db1=db1,
            #     expert_frequency_offset=expert_frequency_offset,
            #     expert_schedule_order=None,
            #     x_gather_idx=x_gather_idx,
            #     s_scatter_idx=s_scatter_idx,
            #     is_glu_activation=is_glu_activation,
            #     stream_id=stream_id,
            # )), 50, 500)
            # print(f"dx Average time: {timing:.3f} ms, TFLOPS: {(4 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            
        dx_reduced = torch.empty(T, H, dtype=dz.dtype, device=dz.device)

        _token_broadcast_backward(
            dx_reduced=dx_reduced,
            dx_expanded=dx_expanded,
            s_reverse_scatter_idx=s_reverse_scatter_idx,
            num_activated_expert_per_token_offset=num_activated_expert_per_token_offset,
            varlen_K_max=(E if is_varlen_K else K),
            H=H,
            is_varlen_K=is_varlen_K,
        )
        # timing = do_bench((lambda: _token_broadcast_backward(
        #     dx_reduced=dx_reduced,
        #     dx_expanded=dx_expanded,
        #     s_reverse_scatter_idx=s_reverse_scatter_idx,
        #     num_activated_expert_per_token_offset=num_activated_expert_per_token_offset,
        #     varlen_K_max=(E if is_varlen_K else K),
        #     H=H,
        #     is_varlen_K=is_varlen_K,
        # )), 50, 500)
        # print(f"expert-agg-bwd Average time: {timing:.3f} ms, TB/s: {(2 * T * H * K + 2 * T * H) / (timing / 1e3) / 1e12:.2f}")
        return dx_reduced, dw1, db1, *[None] * 12


class _DownProjection(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        y1: torch.Tensor,
        z: torch.Tensor,
        w2: torch.Tensor,
        b2: torch.Tensor | None,
        topk_scores: torch.Tensor,
        expert_frequency_offset: torch.Tensor,
        T: int,
        K: int,
        stream_id: int,
        x_gather_idx: torch.Tensor,
        s_scatter_idx: torch.Tensor,
        s_reverse_scatter_idx: torch.Tensor,
        num_activated_expert_per_token_offset: torch.Tensor,
        is_varlen_K: bool,
        activation_type: ActivationType,
    ) -> torch.Tensor:
        TK = y1.size(0)
        H, I, E = w2.shape

        y2 = torch.empty(TK, H, dtype=y1.dtype, device=y1.device)
        if is_using_quack_gemm():
            assert not torch.compiler.is_compiling()

            assert b2 is None
            
            y2 = gemm(y1, w2.permute(2, 1, 0), out=y2, cu_seqlens_m=expert_frequency_offset)
            # timing = do_bench((lambda: gemm(y1, w2.permute(2, 1, 0), cu_seqlens_m=expert_frequency_offset)), 50, 500)
            # print(f"down-proj-fwd Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
        else:
            _down_projection_forward(
                w2=w2,
                y1=y1,
                y2=y2,
                b2=b2,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                stream_id=stream_id,
            )
            # timing = do_bench((lambda: _down_projection_forward(
            #     w2=w2,
            #     y1=y1,
            #     y2=y2,
            #     b2=b2,
            #     expert_frequency_offset=expert_frequency_offset,
            #     expert_schedule_order=None,
            #     x_gather_idx=x_gather_idx,
            #     stream_id=stream_id,
            # )), 50, 500)
            # print(f"down-proj-fwd Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
        o = torch.empty(T, H, device=z.device, dtype=z.dtype)
        topk_scores = topk_scores.flatten()

        _router_forward(
            y2=y2,
            o=o,
            topk_scores=topk_scores,
            s_reverse_scatter_idx=s_reverse_scatter_idx,
            num_activated_expert_per_token_offset=num_activated_expert_per_token_offset,
            varlen_K_max=(E if is_varlen_K else K),
            H=H,
            is_varlen_K=is_varlen_K,
        )
        # timing = do_bench((lambda: _router_forward(
        #     y2=y2,
        #     o=o,
        #     topk_scores=topk_scores,
        #     s_reverse_scatter_idx=s_reverse_scatter_idx,
        #     num_activated_expert_per_token_offset=num_activated_expert_per_token_offset,
        #     varlen_K_max=(E if is_varlen_K else K),
        #     H=H,
        #     is_varlen_K=is_varlen_K,
        # )), 50, 500)
        # print(f"expert-agg-fwd Average time: {timing:.3f} ms, TB/s: {(2 * T * H * K + 4 * T * K + 2 * T * H) / (timing / 1e3) / 1e12:.2f}")
        ctx.T = T
        ctx.K = K
        ctx.H = H
        ctx.I = I
        ctx.is_varlen_K = is_varlen_K
        ctx.activation_type = activation_type
        ctx.stream_id = stream_id

        ctx.save_for_backward(
            z,
            w2,
            b2,
            topk_scores,
            expert_frequency_offset,
            x_gather_idx,
            s_scatter_idx,
            s_reverse_scatter_idx,
        )

        return o

    @staticmethod
    def backward(ctx, dout: torch.Tensor):
        T = ctx.T
        K = ctx.K
        H = ctx.H
        I = ctx.I
        stream_id = ctx.stream_id
        is_varlen_K = ctx.is_varlen_K
        activation_type = ctx.activation_type

        (
            z,
            w2,
            b2,
            topk_scores,
            expert_frequency_offset,
            x_gather_idx,
            s_scatter_idx,
            s_reverse_scatter_idx,
        ) = ctx.saved_tensors

        dw2 = torch.empty_like(w2)
        db2 = None if b2 is None else torch.empty_like(b2)
        dz = torch.empty_like(z)

        if is_using_quack_gemm():
            assert not torch.compiler.is_compiling()
            assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet"

            s = topk_scores[s_scatter_idx]
            _, y1s, ds = gemm_dgated(
                dout,
                w2.permute(2, 0, 1),
                PreAct=z,
                activation="swiglu",
                dx_out=dz,
                colvec_scale=s,
                colvec_reduce=True,
                cu_seqlens_m=expert_frequency_offset,
                A_idx=x_gather_idx,
                dynamic_scheduler=False,
            )
            # timing = do_bench((lambda: gemm(
            #     dout.T,
            #     y1s,
            #     out=dw2.permute(2, 0, 1),
            #     cu_seqlens_k=expert_frequency_offset,
            #     A_idx=x_gather_idx,
            #     batch_idx_permute=None,
            #     dynamic_scheduler=False,
            # )), 50, 500)
            # print(f"dw2 Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            # timing = do_bench((lambda: gemm_dgated(
            #     dout,
            #     w2.permute(2, 0, 1),
            #     PreAct=z,
            #     activation="swiglu",
            #     dx_out=dz,
            #     colvec_scale=s,
            #     colvec_reduce=True,
            #     cu_seqlens_m=expert_frequency_offset,
            #     A_idx=x_gather_idx,
            #     dynamic_scheduler=False,
            # )), 50, 500)
            # print(f"dz Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            gemm(
                dout.T,
                y1s,
                out=dw2.permute(2, 0, 1),
                cu_seqlens_k=expert_frequency_offset,
                A_idx=x_gather_idx,
                batch_idx_permute=None,
                dynamic_scheduler=False,
            )

            ds = ds[s_reverse_scatter_idx]
        else:
            ds = torch.empty_like(topk_scores)

            I = w2.size(1)
            TK = x_gather_idx.size(0)

            y1s = torch.empty(TK, I, dtype=z.dtype, device=z.device)
            is_glu_activation = is_glu(activation_type)


            _down_projection_backward_act(
                dout=dout,
                z=z,
                w2=w2,
                dz=dz,
                ds=ds,
                b2=b2,
                db2=db2,
                y1s=y1s,
                topk_scores=topk_scores,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                s_scatter_idx=s_scatter_idx,
                is_glu_activation=is_glu_activation,
                activation_type=activation_type.value,
                stream_id=stream_id,
            )

            _down_projection_backward_weight(
                dout=dout,
                y1s=y1s,
                dw2=dw2,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                stream_id=stream_id,
            )
            timing = do_bench((lambda: _down_projection_backward_weight(
                dout=dout,
                y1s=y1s,
                dw2=dw2,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                stream_id=stream_id,
            )), 50, 500)
            print(f"dw2 Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
            timing = do_bench((lambda: _down_projection_backward_act(
                dout=dout,
                z=z,
                w2=w2,
                dz=dz,
                ds=ds,
                b2=b2,
                db2=db2,
                y1s=y1s,
                topk_scores=topk_scores,
                expert_frequency_offset=expert_frequency_offset,
                expert_schedule_order=None,
                x_gather_idx=x_gather_idx,
                s_scatter_idx=s_scatter_idx,
                is_glu_activation=is_glu_activation,
                activation_type=activation_type.value,
                stream_id=stream_id,
            )), 50, 500)
            print(f"dz Average time: {timing:.3f} ms, TFLOPS: {(2 * T * I * H * K) / (timing / 1e3) / 1e12:.1f}")
        # TC top-K routing
        if not is_varlen_K:
            ds = ds.view(T, K)

        return None, dz, dw2, db2, ds, *[None] * 10


def moe_TC_softmax_topk_layer(
    x: torch.Tensor,
    router_w: torch.Tensor,
    w1: torch.Tensor,
    b1: torch.Tensor | None,
    w2: torch.Tensor,
    b2: torch.Tensor | None,
    K: int,
    stream_id: int,
    activation_type: ActivationType | str = ActivationType.SWIGLU,
    is_inference_mode_enabled: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    assert ((b1 is None) and (b2 is None)) or (
        (b1 is not None) and (b2 is not None)
    ), "b1 and b2 has to be None or not None at the same time!"
    E = router_w.size(0)
    router_logits = F.linear(x, router_w)
    topk_scores, topk_indices = TC_Softmax_Topk_Router_Function.apply(router_logits, E, K)

    (expert_frequency, expert_frequency_offset, x_gather_idx, s_scatter_idx, s_reverse_scatter_idx) = (
        TC_topk_router_metadata_triton(topk_indices, E)
    )

    T = x.size(0)

    if type(activation_type) == str:
        activation_type = ActivationType(activation_type)

    y1, z = _UpProjection.apply(
        x,
        w1,
        b1,
        expert_frequency_offset,
        T * K,
        K,
        stream_id,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        None,
        False,  # is_varlen_K
        activation_type,
        is_inference_mode_enabled,
    )

    o = _DownProjection.apply(
        y1,
        z,
        w2,
        b2,
        topk_scores,
        expert_frequency_offset,
        T,
        K,
        stream_id,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        None,
        False,  # is_varlen_K
        activation_type,
    )

    return o, router_logits, expert_frequency


# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# Weight format requirements:
# - w1_weight: Shape (2*I, H, E), stride order (2, 0, 1), must be interleaved [gate_row0, up_row0, gate_row1, up_row1, ...]
# - w2_weight: Shape (H, I, E), stride order (2, 0, 1)


# We assume token_indices is already SORTED ascendingly !!!
#   and len(token_indices) = len(expert_indices) = len(router_scores)
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
def moe_general_routing_inputs(
    x: torch.Tensor,
    router_scores: torch.Tensor,
    token_indices: torch.Tensor,
    expert_indices: torch.Tensor,
    w1: torch.Tensor,
    b1: torch.Tensor | None,
    w2: torch.Tensor,
    b2: torch.Tensor | None,
    E: int,
    stream_id: int,
    activation_type: ActivationType,
    is_inference_mode_enabled: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    assert ((b1 is None) and (b2 is None)) or (
        (b1 is not None) and (b2 is not None)
    ), "b1 and b2 has to be None or not None at the same time!"

    T = x.size(0)
    TK = router_scores.size(0)
    E = w2.size(-1)
    (
        expert_frequency,
        expert_frequency_offset,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        num_activated_expert_per_token_offset,
    ) = general_routing_router_metadata(router_scores, token_indices, expert_indices, T, E)

    y1, z = _UpProjection.apply(
        x,
        w1,
        b1,
        expert_frequency_offset,
        TK,
        None,  # K, not needed
        stream_id,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        num_activated_expert_per_token_offset,
        True,  # is_varlen_K
        activation_type,
        is_inference_mode_enabled,
    )

    o = _DownProjection.apply(
        y1,
        z,
        w2,
        b2,
        router_scores,
        expert_frequency_offset,
        T,
        None,  # K, not needed
        stream_id,
        x_gather_idx,
        s_scatter_idx,
        s_reverse_scatter_idx,
        num_activated_expert_per_token_offset,
        True,  # is_varlen_K
        activation_type,
    )

    return o, expert_frequency
