"""NVFP4 linear helpers built on vLLM's own kernels (same code path on B200 sm_100 and Spark sm_121).

Two GEMM flavours for a bf16 weight W [N, K]:
  * "marlin": weight-only NVFP4 (W4A16, Marlin kernel) - bf16 activations, best suited to decode GEMV;
  * "cutlass": W4A4 (activations quantized to NVFP4 per call, CUTLASS block-scaled FP4 MMA).
"""
import torch
import torch.nn as nn

from vllm import _custom_ops as ops

FP4_MAX, E4M3_MAX = 6.0, 448.0


class NVFP4Linear(nn.Module):
    def __init__(self, W: torch.Tensor, kind: str = "marlin"):
        super().__init__()
        assert W.dim() == 2 and W.shape[1] % 16 == 0
        self.kind = kind
        N, K = W.shape
        self.output_size_per_partition, self.input_size_per_partition = N, K
        self.params_dtype = W.dtype
        gs = (E4M3_MAX * FP4_MAX / W.abs().amax().float().clamp(min=1e-12)).reshape(1)
        if kind == "cutlass":
            w4, wsf = ops.scaled_fp4_quant(W.contiguous(), gs)
            self.register_buffer("w4", w4, persistent=False)
            self.register_buffer("wsf", wsf, persistent=False)
            self.register_buffer("w_gs", gs, persistent=False)
        elif kind == "marlin":
            from vllm.model_executor.layers.quantization.utils.marlin_utils import marlin_make_workspace_new
            from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import prepare_fp4_layer_for_marlin
            w4, wsf = ops.scaled_fp4_quant(W.contiguous(), gs, is_sf_swizzled_layout=False)
            self.weight = nn.Parameter(w4.view(torch.uint8).reshape(N, K // 2), requires_grad=False)
            self.weight_scale = nn.Parameter(wsf.reshape(N, K // 16), requires_grad=False)
            self.weight_global_scale = nn.Parameter((1.0 / gs).float(), requires_grad=False)
            prepare_fp4_layer_for_marlin(self)
            self.workspace = marlin_make_workspace_new(W.device)
        else:
            raise ValueError(kind)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shp = x.shape
        x2 = x.reshape(-1, shp[-1])
        N = self.output_size_per_partition
        if self.kind == "cutlass":
            a_gs = (E4M3_MAX * FP4_MAX / x2.abs().amax().float().clamp(min=1e-12)).reshape(1)
            a4, asf = ops.scaled_fp4_quant(x2.contiguous(), a_gs)
            alpha = 1.0 / (a_gs * self.w_gs)
            y = ops.cutlass_scaled_fp4_mm(a4, self.w4, asf, self.wsf, alpha, x.dtype)
        else:
            from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import apply_fp4_marlin_linear
            y = apply_fp4_marlin_linear(x2, self.weight, self.weight_scale, self.weight_global_scale,
                                        self.workspace, N, self.input_size_per_partition)
        return y.reshape(*shp[:-1], N)