vllm.models.inkling.nvidia.ops.silu_and_mul ¶

def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor:
    """SiLU-and-mul for the interleaved fused gate/up layout.

    Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP).
    """
    assert gateup_output.is_contiguous(), (
        f"{gateup_output.shape=} {gateup_output.stride()=}"
    )
    assert gateup_output.ndim == 2, f"{gateup_output.shape=}"

    M = gateup_output.shape[0]
    hidden_size = gateup_output.shape[1]
    assert hidden_size % 2 == 0, f"{hidden_size=}"
    N = hidden_size // 2

    down_input = torch.empty(
        (M, N), device=gateup_output.device, dtype=gateup_output.dtype
    )
    if M == 0:
        return down_input

    BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N)))
    if M <= 1:
        BLOCK_SIZE_M = 4
    elif M <= 256:
        BLOCK_SIZE_M = 2
    elif M < 4096:
        BLOCK_SIZE_M = 4
    else:
        BLOCK_SIZE_M = 16
        BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N)))
    max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
    num_sms = torch.cuda.get_device_properties(
        gateup_output.device
    ).multi_processor_count
    grid_size = min(num_sms * 4, max_grid_size)

    _silu_and_mul_triton_kernel[(grid_size,)](
        gateup_out_ptr=gateup_output,
        down_inp_ptr=down_input,
        M=M,
        N=N,
        GRID_SIZE=grid_size,
        NUM_STAGES=1,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        EVEN_N=N % BLOCK_SIZE_N == 0,
        INT64_INDEX=gateup_output.nbytes >= 2**31,
        num_warps=8,
    )

    return down_input

def sink_silu_mul_epilogue(
    raw: torch.Tensor,  # [T, S * 2F] gemm1 output (interleaved gate/up rows)
    alphas: torch.Tensor,  # [S] fp32
    gammas: torch.Tensor,  # [T, S] fp32
    ratios: torch.Tensor,  # [S] fp32
    n_experts: int,
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e).

    One kernel replaces the per-expert dequant column scale, the SwiGLU, and
    the per-token gamma multiply between the two sink GEMMs.
    """
    tokens = raw.shape[0]
    f = raw.shape[1] // (2 * n_experts)
    out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype)
    if tokens == 0:
        return out
    # raw may be a column-slice of a padded GEMM output (rows strided).
    assert raw.stride(1) == 1 and gammas.stride(1) == 1
    # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512.
    block_f = min(512, f & (-f))
    _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))](
        raw,
        alphas,
        gammas,
        ratios,
        out,
        tokens,
        raw.stride(0),
        gammas.stride(0),
        F=f,
        S=n_experts,
        BLOCK_F=block_f,
    )
    return out