import torch
import torch.nn as nn

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


def _fake_quant_symmetric_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Symmetric int8 fake-quant: round(x/scale) clamped to [-127, 127], times scale."""
    q = torch.clamp(torch.round(x / scale), -127.0, 127.0)
    return q * scale


class Model(nn.Module):
    """
    **QuantVLA selective layout fusion** (reference target for a fused CUDA/Triton kernel).

    Vision–language–action stacks often use **token-wise mixed precision**: some
    timesteps or modalities tolerate aggressive quantization while others stay in
    full precision. This module fuses in **one forward**:

      1. **Layout routing** — ``sigmoid(layout_logits)`` yields a per-token mix
         coefficient in ``(0, 1)`` (quant vs FP residual).
      2. **Quantitative path** — global symmetric **int8-style fake quantization**
         (scale from the tensor max, ``/ 127``).
      3. **Selective blend** — ``α · x_quant + (1 - α) · x``.
      4. **Head** — linear projection on the mixed activations.

    This matches a “quantitative selective layout” block: the layout logits (from a
    router, modality tags, or calibration) choose how strongly each position uses the
    quantized layout vs the identity layout.
    """

    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, D) VLA token activations (fused modalities).
            layout_logits: (B, T) unbounded logits; larger → more mass on quantized path.
        Returns:
            (B, T, D) updated activations.
        """
        B, T, D = x.shape
        assert D == self.dim
        assert layout_logits.shape == (B, T)

        alpha = torch.sigmoid(layout_logits).unsqueeze(-1)
        with torch.no_grad():
            amax = x.detach().abs().max().clamp(min=1e-8)
        scale = amax / 127.0
        x_q = _fake_quant_symmetric_int8(x, scale)
        mixed = alpha * x_q + (1.0 - alpha) * x
        return self.out_proj(mixed)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self, dim: int):
        super().__init__()
        self._impl = Model(dim)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        return self._impl(x, layout_logits)


batch_size = 8
seq_len = 64
dim = 128


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), dim)
    layout_logits = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len))
    return [x, layout_logits]


def get_init_inputs():
    return [dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor qvla_forward(torch::Tensor x, torch::Tensor layout_logits, torch::Tensor weight, torch::Tensor bias);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__inline__ __device__ float warp_max(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
    return v;
}

__global__ void reduce_amax_kernel(const float* __restrict__ x, float* __restrict__ scale_out, int n) {
    __shared__ float smem[8];
    int tid = threadIdx.x;
    float m = 0.0f;
    int n4 = n >> 2;
    const float4* x4 = reinterpret_cast<const float4*>(x);
    for (int i = tid; i < n4; i += 256) {
        float4 v = x4[i];
        m = fmaxf(m, fabsf(v.x));
        m = fmaxf(m, fabsf(v.y));
        m = fmaxf(m, fabsf(v.z));
        m = fmaxf(m, fabsf(v.w));
    }
    for (int i = (n4 << 2) + tid; i < n; i += 256) m = fmaxf(m, fabsf(x[i]));
    m = warp_max(m);
    if ((tid & 31) == 0) smem[tid >> 5] = m;
    __syncthreads();
    m = (tid < 8) ? smem[tid] : 0.0f;
    if (tid < 32) {
        m = warp_max(m);
        if (tid == 0) scale_out[0] = fmaxf(m, 1.0e-8f) * (1.0f / 127.0f);
    }
}

__global__ void fused_matmul128_kernel(const float* __restrict__ x,
                                       const float* __restrict__ logits,
                                       const float* __restrict__ weight,
                                       const float* __restrict__ bias,
                                       const float* __restrict__ scale_ptr,
                                       float* __restrict__ out) {
    __shared__ float As[16][17];
    __shared__ float Bs[16][33];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row = blockIdx.y * 16 + ty;
    int col0 = blockIdx.x * 32 + tx;
    int col1 = col0 + 16;
    float scale = scale_ptr[0];
    float inv_scale = 1.0f / scale;
    float l = logits[row];
    float alpha = 1.0f / (1.0f + expf(-l));
    float acc0 = 0.0f, acc1 = 0.0f;
    #pragma unroll
    for (int k0 = 0; k0 < 128; k0 += 16) {
        int kA = k0 + tx;
        float xi = x[row * 128 + kA];
        float q = nearbyintf(xi * inv_scale);
        q = fminf(127.0f, fmaxf(-127.0f, q));
        As[ty][tx] = fmaf(alpha, q * scale - xi, xi);
        int kB = k0 + ty;
        Bs[ty][tx] = weight[col0 * 128 + kB];
        Bs[ty][tx + 16] = weight[col1 * 128 + kB];
        __syncthreads();
        #pragma unroll
        for (int k = 0; k < 16; ++k) {
            float a = As[ty][k];
            acc0 = fmaf(a, Bs[k][tx], acc0);
            acc1 = fmaf(a, Bs[k][tx + 16], acc1);
        }
        __syncthreads();
    }
    out[row * 128 + col0] = acc0 + bias[col0];
    out[row * 128 + col1] = acc1 + bias[col1];
}

torch::Tensor qvla_forward(torch::Tensor x, torch::Tensor layout_logits, torch::Tensor weight, torch::Tensor bias) {
    auto out = torch::empty_like(x);
    auto scale = torch::empty({1}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    reduce_amax_kernel<<<1, 256, 0, stream>>>(x.data_ptr<float>(), scale.data_ptr<float>(), (int)x.numel());
    fused_matmul128_kernel<<<dim3(4, 32), dim3(16, 16), 0, stream>>>(x.data_ptr<float>(), layout_logits.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), scale.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
"""
_ext = load_inline(name="qvla_selective_fused_ext_v4", cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=["qvla_forward"], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)
    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        return _ext.qvla_forward(x, layout_logits, self.out_proj.weight, self.out_proj.bias)