Blame for build/torch-cuda/functional/backward.py

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

from typing import Optional

import cuda.bindings.driver as cuda
import cutlass.cute as cute
import torch
import triton
import triton.language as tl
from ..quack.gemm_interface import gemm, gemm_dgated

from .._ops_compat import add_op_namespace_prefix
from ..utils import get_powers_of_2
from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton


def _get_autotune_configs_for_db2_and_ds() -> list[triton.Config]:
    configs = []
    for BLOCK_TK in get_powers_of_2(4, 32):
        configs.append(triton.Config({"BLOCK_TK": BLOCK_TK}, num_warps=8, num_stages=4))
    return configs


@triton.autotune(
    configs=_get_autotune_configs_for_db2_and_ds(),
    key=["H", "E"],
)
@triton.jit
def db2_and_ds_kernel(
    dout_ptr,  # (T, H)
    s_ptr,  # (TK,)
    new_ds_partial_ptr,  # (TK, n_h_blocks)
    old_ds_partial_ptr,  # (TK, OLD_DS_PARTIAL_N)
    b2_ptr,  # (E, H),
    db2_ptr,  # (E, H),
    x_gather_idx_ptr,  # (TK,), maps grouped -> token index
    s_scatter_idx_ptr,  # (TK,), maps grouped -> scatter index
    expert_offset_ptr,  # (E+1,), offsets in grouped layout
    H: tl.constexpr,
    E: tl.constexpr,
    OLD_DS_PARTIAL_N: tl.constexpr,
    BLOCK_H: tl.constexpr,  # Block size for H dimension
    BLOCK_TK: tl.constexpr,  # Block size for token dimension
    BLOCK_OLD_DS_PARTIAL_N: tl.constexpr,
):
    Eidx = tl.program_id(0)  # expert id
    Hidx = tl.program_id(1)  # h-block id
    NUM_H_BLOCKS: tl.constexpr = tl.num_programs(1)

    # Hidden dimension indices for this block
    h_offsets = Hidx * BLOCK_H + tl.arange(0, BLOCK_H)
    h_mask = h_offsets < H

    E_count_start = tl.load(expert_offset_ptr + Eidx)
    E_count_end = tl.load(expert_offset_ptr + Eidx + 1)
    n_tokens = E_count_end - E_count_start

    b2 = tl.load(b2_ptr + Eidx * H + h_offsets, mask=h_mask, other=0.0).to(tl.float32)

    db2_acc = tl.zeros([BLOCK_H], dtype=tl.float32)

    # Process tokens in blocks of BLOCK_TK
    for block_start in tl.range(0, n_tokens, BLOCK_TK):
        # Token offsets within this block
        tk_offsets = block_start + tl.arange(0, BLOCK_TK)
        tk_mask = tk_offsets < n_tokens
        tk_grouped = E_count_start + tk_offsets

        # Gather token indices: [BLOCK_TK]
        token_indices = tl.load(x_gather_idx_ptr + tk_grouped, mask=tk_mask, other=0).to(tl.uint32)

        # Get scatter indices: [BLOCK_TK]
        scatter_indices = tl.load(s_scatter_idx_ptr + tk_grouped, mask=tk_mask, other=0).to(tl.uint32)

        s = tl.load(s_ptr + scatter_indices, mask=tk_mask, other=0.0).to(tl.float32)

        # Gather dout: [BLOCK_TK, BLOCK_H]
        dout_offsets = token_indices[:, None] * H + h_offsets[None, :]
        dout_mask = tk_mask[:, None] & h_mask[None, :]
        dout = tl.load(dout_ptr + dout_offsets, mask=dout_mask, other=0.0).to(tl.float32)

        # Accumulate db2: sum over tokens of (dout * s)
        db2_acc += tl.sum(dout * s[:, None], axis=0)  # Sum over BLOCK_TK dimension

        # Compute ds: dot(dout, b2) for this H-block
        ds_partial = tl.sum(dout * b2[None, :], axis=1)  # [BLOCK_TK]

        # On first H-block, add old_ds_partial.sum(dim=1)
        if Hidx == 0:
            n_offsets = tl.arange(0, BLOCK_OLD_DS_PARTIAL_N)
            old_ds_partial_offsets = scatter_indices[:, None] * OLD_DS_PARTIAL_N + n_offsets[None, :]
            old_ds_partial_mask = tk_mask[:, None] & (n_offsets[None, :] < OLD_DS_PARTIAL_N)
            old_ds_partial_vals = tl.load(
                old_ds_partial_ptr + old_ds_partial_offsets, mask=old_ds_partial_mask, other=0.0
            ).to(tl.float32)
            ds_partial += tl.sum(old_ds_partial_vals, axis=1)

        tl.store(new_ds_partial_ptr + scatter_indices * NUM_H_BLOCKS + Hidx, ds_partial, mask=tk_mask)

    tl.store(db2_ptr + Eidx * H + h_offsets, db2_acc, mask=h_mask)


def _get_autotune_configs_for_db1() -> list[triton.Config]:
    configs = []
    for BLOCK_TK in get_powers_of_2(4, 128):
        for BLOCK_I in get_powers_of_2(64, 4096):
            if 4096 <= BLOCK_I * BLOCK_TK <= 16384:
                configs.append(triton.Config({"BLOCK_I": BLOCK_I, "BLOCK_TK": BLOCK_TK}, num_warps=8, num_stages=4))
    return configs


def _prune_triton_autotune_config(configs, nargs, **kw):
    pruned_configs = []
    for c in configs:
        if c.kwargs["BLOCK_I"] <= triton.next_power_of_2(nargs["I"]):
            pruned_configs.append(c)
    return pruned_configs


@triton.autotune(
    configs=_get_autotune_configs_for_db1(),
    key=["I", "E"],
    prune_configs_by={"early_config_prune": _prune_triton_autotune_config},
)
@triton.jit
def db1_kernel(
    dh_ptr,  # (TK, I)  — always interleaved
    db1_ptr,  # (E, I)
    expert_offset_ptr,  # (E+1,)
    I: tl.constexpr,
    E: tl.constexpr,
    BLOCK_I: tl.constexpr,
    BLOCK_TK: tl.constexpr,
    CONCAT_LAYOUT: tl.constexpr = False,
):
    Eidx = tl.program_id(0)

    E_count_start = tl.load(expert_offset_ptr + Eidx).to(tl.int64)
    E_count_end = tl.load(expert_offset_ptr + Eidx + 1).to(tl.int64)
    n_tokens = E_count_end - E_count_start

    NUM_I_BLOCKS: tl.constexpr = triton.cdiv(I, BLOCK_I)
    I_HALF: tl.constexpr = I // 2
    for Iidx in tl.static_range(0, NUM_I_BLOCKS, 1):
        i_offsets = Iidx * BLOCK_I + tl.arange(0, BLOCK_I)
        i_mask = i_offsets < I

        db1_acc = tl.zeros([BLOCK_I], dtype=tl.float32)

        for block_start in tl.range(0, n_tokens, BLOCK_TK):
            # Token offsets within this block
            tk_offsets = block_start + tl.arange(0, BLOCK_TK)
            tk_mask = tk_offsets < n_tokens
            tk_grouped = E_count_start + tk_offsets

            dz_offsets = tk_grouped[:, None] * I + i_offsets[None, :]
            dz_mask = tk_mask[:, None] & i_mask[None, :]
            dz = tl.load(dh_ptr + dz_offsets, mask=dz_mask, other=0.0).to(tl.float32)

            db1_acc += tl.sum(dz, axis=0)

        # Write: remap interleaved → concat if needed
        if CONCAT_LAYOUT:
            out_offsets = i_offsets // 2 + (i_offsets % 2) * I_HALF
        else:
            out_offsets = i_offsets
        db1_offsets = Eidx.to(tl.int64) * I + out_offsets
        tl.store(db1_ptr + db1_offsets, db1_acc, mask=i_mask)


@torch.library.custom_op(add_op_namespace_prefix("_up_projection_backward_act"), mutates_args={"dx_expanded", "db1"})
def _up_projection_backward_act(
    w1: torch.Tensor,
    dx_expanded: torch.Tensor,
    dh: torch.Tensor,
    db1: torch.Tensor | None,
    expert_frequency_offset: torch.Tensor,
    is_glu_activation: bool,
    concat_layout: bool = False,
) -> None:
    I, H, E = w1.size()
    if is_glu_activation:
        I //= 2

    gemm(
        dh,
        w1.permute(2, 0, 1),
        cu_seqlens_m=expert_frequency_offset,
        dynamic_scheduler=False,
        out=dx_expanded,
        concat_layout=(("B",) if concat_layout else None),
    )

    # db1 computation
    if db1 is not None:
        db1_kernel[(E,)](
            dh,
            db1,
            expert_frequency_offset,
            (2 * I if is_glu_activation else I),
            E,
            CONCAT_LAYOUT=concat_layout and is_glu_activation,
        )


_up_projection_backward_act.compile_cache = {}


@torch.library.custom_op(add_op_namespace_prefix("_up_projection_backward_weight"), mutates_args={"dw1"})
def _up_projection_backward_weight(
    x: torch.Tensor,
    dw1: torch.Tensor,
    dh: torch.Tensor,
    expert_frequency_offset: torch.Tensor,
    x_gather_idx: torch.Tensor,
    is_glu_activation: bool,
    concat_layout: bool = False,
) -> None:
    I, H, E = dw1.size()
    if is_glu_activation:
        I //= 2

    gemm(
        x.T,
        dh,
        out=dw1.permute(2, 1, 0),
        cu_seqlens_k=expert_frequency_offset,
        A_idx=x_gather_idx,
        batch_idx_permute=None,
        dynamic_scheduler=False,
        concat_layout=(("out",) if concat_layout else None),
    )


_up_projection_backward_weight.compile_cache = {}


@torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_act"), mutates_args={"dh", "ds", "db2", "a_prime"})
def _down_projection_backward_act(
    dout: torch.Tensor,
    h: torch.Tensor,
    w2: torch.Tensor,
    dh: torch.Tensor,
    ds: torch.Tensor,
    b2: torch.Tensor | None,
    db2: torch.Tensor | None,  # add impl later
    a_prime: torch.Tensor,
    topk_scores: torch.Tensor,
    expert_frequency_offset: torch.Tensor,
    x_gather_idx: torch.Tensor,
    s_scatter_idx: torch.Tensor,
    activation_type: str,
) -> None:
    assert activation_type in (
        "swiglu",
        "geglu",
    ), f"QuACK gemm_gated only supports glu activations, got {activation_type}"

    s = topk_scores[s_scatter_idx]
    _, _, ds_scattered = gemm_dgated(
        dout,
        w2.permute(2, 0, 1),
        PreAct=h,
        activation=activation_type,
        dx_out=dh,
        postact_out=a_prime,
        colvec_scale=s,
        colvec_reduce=True,
        cu_seqlens_m=expert_frequency_offset,
        A_idx=x_gather_idx,
        dynamic_scheduler=False,
    )
    ds[s_scatter_idx] = ds_scattered

    if db2 is None:
        ds[s_scatter_idx] = ds_scattered
    else:
        H = w2.size(0)
        E = expert_frequency_offset.size(0) - 1
        TK = x_gather_idx.size(0)

        old_ds_partial = torch.empty(TK, 1, device=ds_scattered.device, dtype=ds_scattered.dtype)
        old_ds_partial[s_scatter_idx, 0] = ds_scattered

        BLOCK_H = min(triton.next_power_of_2(H), 2048)
        NUM_H_BLOCKS = triton.cdiv(H, BLOCK_H)
        new_ds_partial = torch.empty(TK, NUM_H_BLOCKS, dtype=torch.float32, device=ds.device)

        db2_and_ds_kernel[(E, NUM_H_BLOCKS)](
            dout,
            topk_scores,
            new_ds_partial,
            old_ds_partial,
            b2,
            db2,
            x_gather_idx,
            s_scatter_idx,
            expert_frequency_offset,
            H,
            E,
            1,  # OLD_DS_PARTIAL_N = 1
            BLOCK_H=BLOCK_H,
            BLOCK_OLD_DS_PARTIAL_N=1,
        )

        if NUM_H_BLOCKS == 1:
            ds.copy_(new_ds_partial.view(-1).to(dtype=ds.dtype))
        else:
            ds.copy_(new_ds_partial.sum(dim=-1, dtype=ds.dtype))


_down_projection_backward_act.compile_cache = {}


@torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_weight"), mutates_args={"dw2"})
def _down_projection_backward_weight(
    dout: torch.Tensor,
    a_prime: torch.Tensor,
    dw2: torch.Tensor,
    expert_frequency_offset: torch.Tensor,
    x_gather_idx: torch.Tensor,
) -> None:
    gemm(
        dout.T,
        a_prime,
        out=dw2.permute(2, 0, 1),
        cu_seqlens_k=expert_frequency_offset,
        A_idx=x_gather_idx,
        batch_idx_permute=None,
        dynamic_scheduler=False,
    )


_down_projection_backward_weight.compile_cache = {}


@torch.library.custom_op(add_op_namespace_prefix("_token_broadcast_backward"), mutates_args={"dx_reduced"})
def _token_broadcast_backward(
    dx_reduced: torch.Tensor,
    dx_expanded: torch.Tensor,
    s_reverse_scatter_idx: torch.Tensor,
    num_activated_expert_per_token_offset: Optional[torch.Tensor],
    varlen_K_max: int,
    H: int,
    is_varlen_K: bool,
) -> None:
    if num_activated_expert_per_token_offset is None:
        assert not is_varlen_K, "`num_activated_expert_per_token_offset` as None requires fixed top-K routing"
    token_gather_and_sum_varlen_K_triton(
        dx_expanded,
        None,
        dx_reduced,
        s_reverse_scatter_idx,
        num_activated_expert_per_token_offset,
        dx_reduced.size(0),
        varlen_K_max,
        H,
        is_varlen_K,
    )


@triton.jit
def _softmax_over_topk_bwd_kernel(
    dlogits_ptr,
    dlogits_full_ptr,
    score_ptr,
    dscore_ptr,
    idx_ptr,
    stride_dm: tl.constexpr,
    stride_dn: tl.constexpr,
    stride_sm: tl.constexpr,
    stride_sn: tl.constexpr,
    stride_gm: tl.constexpr,
    stride_gk: tl.constexpr,
    stride_im: tl.constexpr,
    stride_ik: tl.constexpr,
    K: tl.constexpr,
    BLOCK_K: tl.constexpr,
    dlogits_is_none: tl.constexpr,
):
    row = tl.program_id(axis=0)

    # tl.assume(K <= BLOCK_K)
    k_offs = tl.arange(0, BLOCK_K)
    k_mask = k_offs < K

    idx = tl.load(idx_ptr + row * stride_im + k_offs * stride_ik, mask=k_mask, other=0).to(tl.int32)
    s_sel = tl.load(score_ptr + row * stride_sm + k_offs * stride_sn, mask=k_mask, other=0).to(tl.float32)
    g_sel = tl.load(dscore_ptr + row * stride_gm + k_offs * stride_gk, mask=k_mask, other=0).to(tl.float32)

    # dot = sum_j g_j * y_j over selected columns
    dot = tl.sum(g_sel * s_sel, axis=0)

    # scatter-only: dx[idx] += y_sel * (g_sel - dot)
    add_vals = s_sel * (g_sel - dot)

    indices = row * stride_dm + idx * stride_dn
    if not dlogits_is_none:
        add_vals += tl.load(dlogits_ptr + indices, mask=k_mask)
    tl.store(dlogits_full_ptr + indices, add_vals, mask=k_mask)


@triton.jit
def _topk_over_softmax_bwd_kernel(
    logits_ptr,  # (T, N) saved router logits
    dlogits_ptr,  # (T, N) output gradient
    dscore_ptr,  # (T, K) upstream gradient
    idx_ptr,  # (T, K) selected indices (int32)
    score_ptr,  # (T, K) forward scores (only used for renorm)
    stride_lm: tl.constexpr,
    stride_le: tl.constexpr,
    stride_dm: tl.constexpr,
    stride_dn: tl.constexpr,
    stride_sm: tl.constexpr,
    stride_sn: tl.constexpr,
    stride_im: tl.constexpr,
    stride_ik: tl.constexpr,
    stride_scm: tl.constexpr,
    stride_scn: tl.constexpr,
    E: tl.constexpr,
    K: tl.constexpr,
    BLOCK_E: tl.constexpr,
    BLOCK_K: tl.constexpr,
    norm_topk_probs: tl.constexpr,
):
    """
    Full topk(softmax()) backward over ALL E indices.

    Forward: logits → p = softmax(logits) → [raw, idx] = topk(p, K)
             → scores = raw / sum(raw)  (if norm_topk_probs)

    Backward:
      1. Recompute p = softmax(logits) over all E
      2. If renorm: dp_sel = (dscore - dot_s) / S
         Else:      dp_sel = dscore
      3. dot = Σ dp_sel_j * p_sel_j
      4. Scatter dp_sel into E-wide dp (zero at non-selected)
      5. dlogits = p * (dp - dot)  for all E
    """
    row = tl.program_id(axis=0)

    e_offs = tl.arange(0, BLOCK_E)
    e_mask = e_offs < E
    logits = tl.load(logits_ptr + row * stride_lm + e_offs * stride_le, mask=e_mask, other=-float("inf")).to(
        tl.float32
    )
    row_max = tl.max(logits, axis=0)
    exp_vals = tl.exp(logits - row_max)
    row_sum = tl.sum(exp_vals, axis=0)
    p = exp_vals / row_sum  # (BLOCK_E,)

    # --- Load K selected indices and upstream gradient ---
    k_offs = tl.arange(0, BLOCK_K)
    k_mask = k_offs < K
    idx = tl.load(
        idx_ptr + row * stride_im + k_offs * stride_ik,
        mask=k_mask,
        other=0,
    ).to(tl.int32)
    g_sel = tl.load(
        dscore_ptr + row * stride_sm + k_offs * stride_sn,
        mask=k_mask,
        other=0,
    ).to(tl.float32)

    # p at selected indices (gather from global mem; can't index register tensor)
    sel_logits = tl.load(
        logits_ptr + row * stride_lm + idx * stride_le,
        mask=k_mask,
        other=-float("inf"),
    ).to(tl.float32)
    p_sel = tl.exp(sel_logits - row_max) / row_sum  # (BLOCK_K,)

    # --- Backward through optional renormalization ---
    if norm_topk_probs:
        scores = tl.load(
            score_ptr + row * stride_scm + k_offs * stride_scn,
            mask=k_mask,
            other=0,
        ).to(tl.float32)
        dot_s = tl.sum(g_sel * scores, axis=0)
        S = tl.sum(p_sel, axis=0)
        dp_sel = (g_sel - dot_s) / S
    else:
        dp_sel = g_sel

    # dot = Σ dp_sel_j * p_sel_j
    dot = tl.sum(dp_sel * p_sel, axis=0)

    # --- Scatter dp_sel into N-wide dp ---
    # dp[i] = dp_sel[k] if i == idx[k], else 0
    # Loop over K (unrolled at compile time since K is constexpr)
    dp = tl.zeros([BLOCK_E], dtype=tl.float32)
    for k_iter in tl.static_range(K):
        cur_dp = tl.sum(tl.where(k_offs == k_iter, dp_sel, 0.0))
        cur_idx = tl.sum(tl.where(k_offs == k_iter, idx, 0))
        dp = tl.where(e_offs == cur_idx, cur_dp, dp)

    # --- dlogits = p * (dp - dot) for all E ---
    dlogits = p * (dp - dot)
    tl.store(
        dlogits_ptr + row * stride_dm + e_offs * stride_dn,
        dlogits,
        mask=e_mask,
    )


@torch.library.custom_op(add_op_namespace_prefix("_topk_softmax_bwd"), mutates_args={"dlogits_full"})
def _topk_softmax_bwd(
    router_logits: torch.Tensor,
    dlogits_full: torch.Tensor,
    dlogits: Optional[torch.Tensor],
    dtopk_score: torch.Tensor,
    topk_router_score: torch.Tensor,
    topk_router_indices: torch.Tensor,
    E: int,
    K: int,
    is_softmax_over_topk: bool = True,
    norm_topk_probs: bool = False,
) -> None:
    T = dtopk_score.shape[0]

    if is_softmax_over_topk:
        # non-selected gradient is zero.
        _softmax_over_topk_bwd_kernel[T,](
            dlogits,
            dlogits_full,
            topk_router_score,
            dtopk_score,
            topk_router_indices,
            dlogits_full.stride(0),
            dlogits_full.stride(1),
            topk_router_score.stride(0),
            topk_router_score.stride(1),
            dtopk_score.stride(0),
            dtopk_score.stride(1),
            topk_router_indices.stride(0),
            topk_router_indices.stride(1),
            K,
            triton.next_power_of_2(K),
            (dlogits is None),
        )
    else:
        # topk(softmax(.)): non-selected gradient is -p_i * dot, NOT zero.
        # must recompute full softmax for the complete Jacobian.
        _topk_over_softmax_bwd_kernel[T,](
            router_logits,
            dlogits_full,
            dtopk_score,
            topk_router_indices,
            topk_router_score,
            router_logits.stride(0),
            router_logits.stride(1),
            dlogits_full.stride(0),
            dlogits_full.stride(1),
            dtopk_score.stride(0),
            dtopk_score.stride(1),
            topk_router_indices.stride(0),
            topk_router_indices.stride(1),
            topk_router_score.stride(0),
            topk_router_score.stride(1),
            E,
            K,
            triton.next_power_of_2(E),
            triton.next_power_of_2(K),
            norm_topk_probs,
        )


@triton.jit
def _topk_bwd_scatter_small_kernel(
    dlogits_full_ptr,
    dscore_ptr,
    idx_ptr,
    stride_dm: tl.constexpr,
    stride_dn: tl.constexpr,
    stride_gm: tl.constexpr,
    stride_gk: tl.constexpr,
    stride_im: tl.constexpr,
    stride_ik: tl.constexpr,
    K: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    row = tl.program_id(axis=0)

    # tl.assume(K <= BLOCK_K)
    k_offs = tl.arange(0, BLOCK_K)
    k_mask = k_offs < K

    idx = tl.load(idx_ptr + row * stride_im + k_offs * stride_ik, mask=k_mask, other=0).to(tl.int32)
    g_sel = tl.load(dscore_ptr + row * stride_gm + k_offs * stride_gk, mask=k_mask, other=0).to(tl.float32)

    # scatter-only: dx[idx] += y_sel * (g_sel - dot)
    add_vals = g_sel

    indices = row * stride_dm + idx * stride_dn
    tl.store(dlogits_full_ptr + indices, add_vals, mask=k_mask)


@torch.library.custom_op(add_op_namespace_prefix("_topk_bwd"), mutates_args={"dlogits_full"})
def _topk_bwd(
    dlogits_full: torch.Tensor,
    dtopk_values: torch.Tensor,
    topk_indices: torch.Tensor,
    K: int,
) -> None:
    T = dtopk_values.shape[0]

    _topk_bwd_scatter_small_kernel[T,](
        dlogits_full,
        dtopk_values,
        topk_indices,
        dlogits_full.stride(0),
        dlogits_full.stride(1),
        dtopk_values.stride(0),
        dtopk_values.stride(1),
        topk_indices.stride(0),
        topk_indices.stride(1),
        K,
        triton.next_power_of_2(K),
    )