from typing import Any import deep_gemm import tile_kernels import torch from deepinfer import PLATFORM 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() 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()=}" x = tile_kernels.quant.per_token_cast(x, **get_fp8_cast_params()) # 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)