linear.py: ReplicatedLinear and the FP8 GEMM path

class ReplicatedLinear(nn.Linear):
    """Complete linear weights, optionally prepared as FP8 on meta."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        *,
        bias: bool = False,
        use_fp8_gemm: bool = False,
    ) -> None:
        super().__init__(in_features, out_features, bias=bias)
        self.use_fp8_gemm = use_fp8_gemm


class ParallelRoutedLinear(nn.Linear):
    """Expert-major grouped GEMM metadata: (n_local_experts * h_in, h_out).

    HF materialization allocates packed FP4 directly. Dummy initialization uses
    this 2D representation before converting to the runtime layout.
    """

    def __init__(self, h_in: int, h_out: int, n_local_experts: int, *, use_fp8_gemm: bool = False) -> None:
        super().__init__(h_out, h_in * n_local_experts, bias=False)
        self.use_fp8_gemm = use_fp8_gemm
        self.n_local_experts = n_local_experts
        self.weight.n_experts = n_local_experts

    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 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_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 _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