"""Blackwell W4A16 with fused int4 conversion and asymmetric metadata.

Decode uses a native CUDA GEMV laid out as 32 output lanes by 32 quant-group
warps.  Prefill uses CUTLASS 3.x's SM100 cute:: mixed-input UMMA collective;
the collective converts signed int4 to bf16 and applies the per-128-K scale and
additive zero bias in its mainloop before the tensor-core MMA.
"""
from __future__ import annotations

import torch
import torch.nn as nn


GROUP_SIZE = 128


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K == 4096 and group_size == GROUP_SIZE
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        self.register_buffer("w_q", torch.empty((K // 2, N), dtype=torch.uint8))
        self.register_buffer(
            "scales", torch.empty((K // group_size, N), dtype=torch.bfloat16)
        )
        self.register_buffer(
            "zeros", torch.empty((K // group_size, N), dtype=torch.bfloat16)
        )
        self._cutlass_w = None
        self._cutlass_bias = None

    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        # Packed-layout caches must never outlive a weight reload.
        self._cutlass_w = None
        self._cutlass_bias = None
        return super().load_state_dict(state_dict, strict=strict, assign=assign)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        from cutlass_backend import ext

        if self.M == 1:
            return ext.w4_cuda_gemv(x, self.w_q, self.scales, self.zeros)

        # The source packs adjacent K nibbles but stores the resulting bytes
        # with N contiguous.  CUTLASS wants K-contiguous bytes, so transpose
        # those bytes once (still int4; no full dequantized matrix is made).
        # XOR 0x88 recenters each nibble from uint4 q to signed int4 q-8.
        # CUTLASS applies q_signed*scale + bias, hence bias=(8-zero)*scale.
        if self._cutlass_w is None:
            self._cutlass_w = self.w_q.t().contiguous() ^ 0x88
            self._cutlass_bias = ((8 - self.zeros) * self.scales).to(torch.bfloat16)

        if self.M >= 128:
            return ext.w4_cutlass_2sm(
                x, self._cutlass_w, self.scales, self._cutlass_bias
            )
        return ext.w4_cutlass_1sm(
            x, self._cutlass_w, self.scales, self._cutlass_bias
        )


M = 1
N = 12288
K = 4096


def get_inputs():
    return [torch.randn(M, K, dtype=torch.bfloat16)]


def get_init_inputs():
    return [M, N, K]


# ==================================================================
# ===== sidecar: cutlass_backend.py (12270 bytes, loaded by solution.py) =====
# ==================================================================

import os
from pathlib import Path

os.environ["TORCH_CUDA_ARCH_LIST"] = "10.0a"
os.environ["CUDA_HOME"] = "/usr/local/cuda-12.8"
os.environ["PATH"] = "/[REDACTED-home]/kb-hard/.venv/bin:" + os.environ["PATH"]

from torch.utils.cpp_extension import load_inline
_ROOT = Path(__file__).resolve().parent

_CPP = r"""
#include <torch/extension.h>
torch::Tensor w4_cutlass_1sm(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z);
torch::Tensor w4_cutlass_1sm128(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z);
torch::Tensor w4_cutlass_2sm(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z);
torch::Tensor w4_cuda_gemv(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z);
"""

_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/detail/collective/mixed_input_utils.hpp"

using namespace cute;
using ElementA = cutlass::bfloat16_t;
using ElementB = cutlass::int4b_t;
using ElementScale = cutlass::bfloat16_t;
using ElementZero = cutlass::bfloat16_t;
using ElementD = cutlass::bfloat16_t;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::RowMajor;
constexpr int AlignmentA = 8;
constexpr int AlignmentB = 32;
constexpr int AlignmentD = 8;
using ArchTag = cutlass::arch::Sm100;
using OperatorClass = cutlass::arch::OpClassTensorOp;
using ElementAccumulator = float;
using MmaTileShape = Shape<_256,_128,_128>;
using ClusterShape = Shape<_2,_1,_1>;
using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmMixedInputSm100;
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
using ScaleConfig = cutlass::detail::Sm100MixedInputBlockwiseScaleConfig<1, 128>;
using LayoutScale = decltype(ScaleConfig::deduce_layout_scale());

using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OperatorClass, MmaTileShape, ClusterShape,
    cutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementAccumulator,
    void, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    ElementD, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    EpilogueSchedule>::CollectiveOp;

using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    cute::tuple<ElementB, ElementScale, ElementZero>,
    cute::tuple<typename cutlass::layout::LayoutTranspose<LayoutB>::type, LayoutScale>, AlignmentB,
    ElementA, typename cutlass::layout::LayoutTranspose<LayoutA>::type, AlignmentA,
    ElementAccumulator, MmaTileShape, ClusterShape,
    cutlass::gemm::collective::StageCountAutoCarveout<
      static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
    MainloopSchedule>::CollectiveOp;

using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
    Shape<int,int,int,int>, CollectiveMainloop, CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

using MmaTileShape1 = Shape<_128,_64,_128>;
using ClusterShape1 = Shape<_1,_1,_1>;
using CollectiveEpilogue1 = typename cutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OperatorClass, MmaTileShape1, ClusterShape1,
    cutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementAccumulator,
    void, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    ElementD, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    cutlass::epilogue::TmaWarpSpecialized1Sm>::CollectiveOp;
using CollectiveMainloop1 = typename cutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    cute::tuple<ElementB, ElementScale, ElementZero>,
    cute::tuple<typename cutlass::layout::LayoutTranspose<LayoutB>::type, LayoutScale>, AlignmentB,
    ElementA, typename cutlass::layout::LayoutTranspose<LayoutA>::type, AlignmentA,
    ElementAccumulator, MmaTileShape1, ClusterShape1,
    cutlass::gemm::collective::StageCountAutoCarveout<
      static_cast<int>(sizeof(typename CollectiveEpilogue1::SharedStorage))>,
    cutlass::gemm::KernelTmaWarpSpecialized1SmMixedInputSm100>::CollectiveOp;
using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal<
    Shape<int,int,int,int>, CollectiveMainloop1, CollectiveEpilogue1>;
using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel1>;

using MmaTileShape128 = Shape<_128,_128,_128>;
using CollectiveEpilogue128 = typename cutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OperatorClass, MmaTileShape128, ClusterShape1,
    cutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementAccumulator,
    void, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    ElementD, typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
    cutlass::epilogue::TmaWarpSpecialized1Sm>::CollectiveOp;
using CollectiveMainloop128 = typename cutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    cute::tuple<ElementB, ElementScale, ElementZero>,
    cute::tuple<typename cutlass::layout::LayoutTranspose<LayoutB>::type, LayoutScale>, AlignmentB,
    ElementA, typename cutlass::layout::LayoutTranspose<LayoutA>::type, AlignmentA,
    ElementAccumulator, MmaTileShape128, ClusterShape1,
    cutlass::gemm::collective::StageCountAutoCarveout<
      static_cast<int>(sizeof(typename CollectiveEpilogue128::SharedStorage))>,
    cutlass::gemm::KernelTmaWarpSpecialized1SmMixedInputSm100>::CollectiveOp;
using GemmKernel128 = cutlass::gemm::kernel::GemmUniversal<
    Shape<int,int,int,int>, CollectiveMainloop128, CollectiveEpilogue128>;
using Gemm128 = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel128>;

using StrideA = cutlass::detail::TagToStrideA_t<LayoutA>;
using StrideB = cutlass::detail::TagToStrideB_t<LayoutB>;

template <class GemmT>
torch::Tensor run_w4(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z) {
  using StrideC = typename GemmT::GemmKernel::StrideC;
  using StrideD = typename GemmT::GemmKernel::StrideD;
  const int M = x.size(0);
  const int K = x.size(1);
  const int N = s.size(1);
  auto out = torch::empty({M, N}, x.options());

  auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, make_shape(M, K, 1));
  auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, make_shape(N, K, 1));
  auto stride_C = cutlass::make_cute_packed_stride(StrideC{}, make_shape(N, M, 1));
  auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, make_shape(N, M, 1));
  auto layout_S = ScaleConfig::tile_atom_to_shape_scale(make_shape(N, K, 1));

  typename GemmT::Arguments args{
    cutlass::gemm::GemmUniversalMode::kGemm,
    {N, M, K, 1},
    {reinterpret_cast<ElementB*>(w.data_ptr()), stride_B,
     reinterpret_cast<ElementA*>(x.data_ptr()), stride_A,
     reinterpret_cast<ElementScale*>(s.data_ptr()), layout_S,
     reinterpret_cast<ElementZero*>(z.data_ptr())},
    {{1.0f, 0.0f}, nullptr, stride_C,
     reinterpret_cast<ElementD*>(out.data_ptr()), stride_D}
  };

  GemmT gemm;
  auto can = gemm.can_implement(args);
  TORCH_CHECK(can == cutlass::Status::kSuccess, "CUTLASS cannot implement: ", int(can));
  size_t workspace_size = GemmT::get_workspace_size(args);
  auto workspace = torch::empty({static_cast<long long>(workspace_size)},
                                x.options().dtype(torch::kUInt8));
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  auto status = gemm(args, workspace.data_ptr(), stream);
  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS launch failed: ", int(status));
  return out;
}

torch::Tensor w4_cutlass_1sm(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z) {
  return run_w4<Gemm1>(x, w, s, z);
}
torch::Tensor w4_cutlass_1sm128(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z) {
  return run_w4<Gemm128>(x, w, s, z);
}
torch::Tensor w4_cutlass_2sm(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z) {
  return run_w4<Gemm>(x, w, s, z);
}

__device__ __forceinline__ float2 dequant_bf162(uint8_t q, __nv_bfloat162 z, __nv_bfloat162 s) {
  __nv_bfloat162 qb = __floats2bfloat162_rn(float(q & 15), float(q >> 4));
  return __bfloat1622float2(__hmul2(__hsub2(qb, z), s));
}

template <int N>
__global__ __launch_bounds__(1024, 1) void w4_gemv_kernel(
    const __nv_bfloat16* __restrict__ x,
    const uint8_t* __restrict__ w,
    const __nv_bfloat16* __restrict__ scales,
    const __nv_bfloat16* __restrict__ zeros,
    __nv_bfloat16* __restrict__ out) {
  const int lane = threadIdx.x;
  const int split = threadIdx.y;
  const int n = int(blockIdx.x) * 32 + lane;
  __shared__ float partial[32][33];
  float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;

  if (n < N) {
      int g = split;
      __nv_bfloat16 s = scales[g * N + n];
      __nv_bfloat16 z = zeros[g * N + n];
      __nv_bfloat162 s2 = __halves2bfloat162(s, s);
      __nv_bfloat162 z2 = __halves2bfloat162(z, z);
      int base = g * 64;
      #pragma unroll
      for (int p = 0; p < 64; p += 4) {
        uint8_t q0 = w[(base + p + 0) * N + n];
        uint8_t q1 = w[(base + p + 1) * N + n];
        uint8_t q2 = w[(base + p + 2) * N + n];
        uint8_t q3 = w[(base + p + 3) * N + n];
        int k = 2 * (base + p);
        float2 w0 = dequant_bf162(q0, z2, s2);
        float2 w1 = dequant_bf162(q1, z2, s2);
        float2 w2 = dequant_bf162(q2, z2, s2);
        float2 w3 = dequant_bf162(q3, z2, s2);
        float2 x0 = __bfloat1622float2(reinterpret_cast<const __nv_bfloat162*>(x + k)[0]);
        float2 x1 = __bfloat1622float2(reinterpret_cast<const __nv_bfloat162*>(x + k)[1]);
        float2 x2 = __bfloat1622float2(reinterpret_cast<const __nv_bfloat162*>(x + k)[2]);
        float2 x3 = __bfloat1622float2(reinterpret_cast<const __nv_bfloat162*>(x + k)[3]);
        a0 = fmaf(x0.x, w0.x, a0); a0 = fmaf(x0.y, w0.y, a0);
        a1 = fmaf(x1.x, w1.x, a1); a1 = fmaf(x1.y, w1.y, a1);
        a2 = fmaf(x2.x, w2.x, a2); a2 = fmaf(x2.y, w2.y, a2);
        a3 = fmaf(x3.x, w3.x, a3); a3 = fmaf(x3.y, w3.y, a3);
      }
  }
  partial[split][lane] = (a0 + a1) + (a2 + a3);
  __syncthreads();
  if (split == 0 && n < N) {
    float sum = 0.f;
    #pragma unroll
    for (int i = 0; i < 32; ++i) sum += partial[i][lane];
    out[n] = __float2bfloat16_rn(sum);
  }
}

torch::Tensor w4_cuda_gemv(torch::Tensor x, torch::Tensor w, torch::Tensor s, torch::Tensor z) {
  const int N = w.size(1);
  auto out = torch::empty({1, N}, x.options());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  dim3 block(32, 32);
  if (N == 4096) {
    w4_gemv_kernel<4096><<<128, block, 0, stream>>>(
      reinterpret_cast<__nv_bfloat16*>(x.data_ptr()), w.data_ptr<uint8_t>(),
      reinterpret_cast<__nv_bfloat16*>(s.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(z.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(out.data_ptr()));
  } else {
    w4_gemv_kernel<12288><<<384, block, 0, stream>>>(
      reinterpret_cast<__nv_bfloat16*>(x.data_ptr()), w.data_ptr<uint8_t>(),
      reinterpret_cast<__nv_bfloat16*>(s.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(z.data_ptr()),
      reinterpret_cast<__nv_bfloat16*>(out.data_ptr()));
  }
  return out;
}
"""

ext = load_inline(
    name="w4a16_cutlass_sm100_v10",
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    functions=["w4_cutlass_1sm", "w4_cutlass_1sm128", "w4_cutlass_2sm", "w4_cuda_gemv"],
    extra_include_paths=[
        str(_ROOT / "scratch-cutlass" / "include"),
        str(_ROOT / "scratch-cutlass" / "tools" / "util" / "include"),
    ],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "--expt-relaxed-constexpr",
                       "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                       "-U__CUDA_NO_BFLOAT162_OPERATORS__"],
    with_cuda=True,
    verbose=True,
)