DeepInfer/deepinfer/model_executor/layers/linear.py

from typing import Any

import deep_gemm
import tile_kernels
import torch

from deepinfer import PLATFORM
from deepinfer.neuron_ops import nki_gemm
from deepinfer.model_executor.utils import reconstruct_param

_OUTPUT_ALIGNMENT_BYTES = 32


def get_fp8_casting_granularity() -> int:
    return 32


def get_fp8_cast_params() -> dict[str, Any]:
    return {
        "fmt": "e4m3",
        "num_per_channels": get_fp8_casting_granularity(),
        "round_sf": True,
        "use_packed_ue8m0": True,
        "use_tma_aligned_col_major_sf": True,
    }


def bf16_gemm_nt(
    a: torch.Tensor,
    b: torch.Tensor,
    d: torch.Tensor,
    c: torch.Tensor | None = None,
) -> None:
    deep_gemm.bf16_gemm_nt(a, b, d, c)


class ReplicatedLinear(torch.nn.Module):
    """Replicated linear layer.

    Args:
        input_size: input dimension of the linear layer.
        output_size: output dimension of the linear layer.
        use_fp32_output: If true, output fp32 instead of input dtype (bf16 weight only).
        weight: weight parameter of shape (output_size, input_size).
        bias: optional bias parameter.
    """

    def __init__(
        self,
        input_size: int,
        output_size: int,
        use_fp32_output: bool = False,
        *,
        weight: torch.nn.Parameter,
        bias: torch.nn.Parameter | None = None,
    ) -> None:
        super().__init__()

        # Keep input parameters
        self.input_size = input_size
        self.output_size = output_size
        if weight.dtype != torch.bfloat16:
            assert not use_fp32_output
        self.use_fp32_output = use_fp32_output

        assert weight.shape == (output_size, input_size)
        self.weight = reconstruct_param(weight)
        if self.weight.element_size() == 1:
            self.weight.sf = weight.sf
            if self.weight.sf.dtype != PLATFORM.fp8_sf_type:
                self.weight.sf = deep_gemm.transform_sf_into_required_layout(
                    self.weight.sf,
                    self.weight.size(0),
                    self.weight.size(1),
                    recipe=(1, get_fp8_casting_granularity(), get_fp8_casting_granularity()),
                    is_sfa=False,
                )

        if bias is None:
            self.register_parameter("bias", None)
        else:
            self.bias = bias

    def _allocate_output(
        self,
        shape: tuple[int, ...],
        dtype: torch.dtype,
        device: torch.device,
    ) -> torch.Tensor:
        assert len(shape) >= 2 and shape[-1] == self.output_size, f"{shape=}, {self.output_size=}"
        itemsize = dtype.itemsize
        assert _OUTPUT_ALIGNMENT_BYTES % itemsize == 0, f"{dtype=}"
        outputs_per_alignment = _OUTPUT_ALIGNMENT_BYTES // itemsize
        physical_output_size = (self.output_size + outputs_per_alignment - 1) // outputs_per_alignment * outputs_per_alignment
        storage = torch.empty((*shape[:-1], physical_output_size), dtype=dtype, device=device)
        return storage[..., : self.output_size]

    def _has_aligned_output_stride(self, output: torch.Tensor) -> bool:
        assert output.ndim >= 2 and output.shape[-1] == self.output_size and output.stride(-1) == 1
        return output.stride(-2) * output.element_size() % _OUTPUT_ALIGNMENT_BYTES == 0

    def _flatten_output(self, output: torch.Tensor) -> torch.Tensor:
        assert output.ndim >= 2 and output.shape[-1] == self.output_size and output.stride(-1) == 1
        assert all(output.stride(i) == output.stride(i + 1) * output.shape[i + 1] for i in range(output.ndim - 2)), (
            f"{output.shape=}, {output.stride()=}"
        )
        m = output.shape[:-1].numel()
        if PLATFORM.is_neuron:
            return output.view(m, self.output_size)  # a view, not as_strided: the GEMM writes through it under torch.compile
        return output.as_strided((m, self.output_size), (output.stride(-2), 1))

    def _align_output(self, output: torch.Tensor) -> torch.Tensor:
        if self._has_aligned_output_stride(output):
            return output
        aligned = self._allocate_output(tuple(output.shape), output.dtype, output.device)
        aligned.copy_(output)
        return aligned

    def apply_weight_bf16(
        self,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
        out_acc: torch.Tensor | None = None,
    ) -> torch.Tensor:
        weight = self.weight
        m_shape = x.shape[:-1]
        m = m_shape.numel()
        n, k = weight.shape
        if out_acc is None:
            out = self._allocate_output((*m_shape, n), torch.float if self.use_fp32_output else x.dtype, x.device)
        else:
            out = out_acc
            assert out.shape == (*m_shape, n)
            assert out.dtype == (torch.float if self.use_fp32_output else x.dtype)
            assert out.device == x.device
            out_acc = self._flatten_output(out_acc)
        assert self._has_aligned_output_stride(out), f"unaligned output: {out.shape=}, {out.stride()=}, {out.dtype=}"
        bf16_gemm_nt(x.view(m, k), weight, self._flatten_output(out), out_acc)
        if bias is not None:
            out += bias
        return out

    def apply_weight_fp8(
        self,
        x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        bias: torch.Tensor | None = None,
        out_acc: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert bias is None

        if isinstance(x, torch.Tensor):
            assert x.dim() == 2, f"{x.shape=}, {x.stride()=}"
            if out_acc is None and PLATFORM.is_neuron and nki_gemm.has_kernel(self.weight):
                # The NKI kernel quantizes the bf16 input itself (same 1x32 fp8 round trip as per_token_cast).
                return nki_gemm.linear_qdq(x, self.weight)
            x = tile_kernels.quant.per_token_cast(x, **get_fp8_cast_params())
        elif out_acc is None and PLATFORM.is_neuron and nki_gemm.has_kernel(self.weight) and x[0].shape[-1] % 128 == 0:
            # An fp8 activation (e.g. Engram's rows fetched from the host): decode it exactly to bf16 and take the NKI
            # kernel, whose 1x32 re-quantization is idempotent on it. Through the deep_gemm shim the whole fp8
            # weight was dequantized in the graph on every call (9 ms of a prefill layer graph, issue #29).
            k = x[0].shape[-1]
            return nki_gemm.linear_qdq(nki_gemm._activation_bf16(x[0], x[1], k).reshape(-1, k), self.weight)

        # FP8 GEMM.
        weight = (self.weight, self.weight.sf)
        out_shape = x[0].shape[:-1] + (self.weight.shape[0],)
        if out_acc is None:
            out = self._allocate_output(out_shape, torch.bfloat16, x[0].device)
        else:
            out = out_acc
            assert out.shape == out_shape
            assert out.dtype == torch.bfloat16
            assert out.device == x[0].device

        assert self._has_aligned_output_stride(out), f"unaligned output: {out.shape=}, {out.stride()=}, {out.dtype=}"
        assert x[1].dtype == weight[1].dtype == PLATFORM.fp8_sf_type, f"{x[1].dtype=}, {weight[1].dtype=}"
        deep_gemm.fp8_fp4_gemm_nt(x, weight, out, out_acc, recipe=(1, 1, get_fp8_casting_granularity()))
        return out

    def apply_weight(
        self,
        x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        bias: torch.Tensor | None,
        out_acc: torch.Tensor | None,
    ) -> torch.Tensor:
        if self.weight.dtype == torch.bfloat16:
            return self.apply_weight_bf16(x, bias, out_acc)
        if self.weight.element_size() == 1:
            return self.apply_weight_fp8(x, bias, out_acc)
        raise TypeError(f"unsupported weight dtype: {self.weight.dtype}")

    def forward(
        self,
        x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        out_acc: torch.Tensor | None = None,
    ) -> torch.Tensor:
        return self.apply_weight(x, self.bias, out_acc)