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 <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void reduce_absmax_kernel(const float* __restrict__ x, float* __restrict__ out, long long total) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    float m = 0.0f;
    for (long long i = tid; i < total; i += blockDim.x) {
        float v = fabsf(x[i]);
        m = fmaxf(m, v);
    }
    sdata[tid] = m;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);
        __syncthreads();
    }
    if (tid == 0) out[0] = fmaxf(sdata[0], 1.0e-8f);
}

#define TILE 16

__global__ void qvla_linear_kernel(
    const float* __restrict__ x,
    const float* __restrict__ logits,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ amax_ptr,
    float* __restrict__ y,
    int N, int D) {

    __shared__ float As[TILE][TILE + 1];
    __shared__ float Bs[TILE][TILE + 1];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row = blockIdx.y * TILE + ty;
    int col = blockIdx.x * TILE + tx;

    float alpha = 0.0f;
    if (row < N) {
        float l = logits[row];
        alpha = 1.0f / (1.0f + expf(-l));
    }

    float scale = (*amax_ptr) * (1.0f / 127.0f);
    float inv_scale = 127.0f / (*amax_ptr);
    float acc = 0.0f;

    for (int tile = 0; tile < D; tile += TILE) {
        int ka = tile + tx;
        if (row < N && ka < D) {
            float v = x[row * D + ka];
            float q = nearbyintf(v * inv_scale);
            q = fminf(127.0f, fmaxf(-127.0f, q));
            float qv = q * scale;
            As[ty][tx] = v + alpha * (qv - v);
        } else {
            As[ty][tx] = 0.0f;
        }

        int kb = tile + ty;
        if (col < D && kb < D) {
            Bs[ty][tx] = weight[col * D + kb];
        } else {
            Bs[ty][tx] = 0.0f;
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE; ++k) {
            acc += As[ty][k] * Bs[k][tx];
        }
        __syncthreads();
    }

    if (row < N && col < D) {
        y[row * D + col] = acc + bias[col];
    }
}

torch::Tensor qvla_forward(torch::Tensor x, torch::Tensor layout_logits, torch::Tensor weight, torch::Tensor bias) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    auto y = torch::empty_like(x);
    auto tmp = torch::empty({1}, x.options());

    const int B = (int)x.size(0);
    const int T = (int)x.size(1);
    const int D = (int)x.size(2);
    const int N = B * T;
    const long long total = (long long)N * (long long)D;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int rthreads = 256;
    reduce_absmax_kernel<<<1, rthreads, rthreads * sizeof(float), stream>>>(
        x.data_ptr<float>(), tmp.data_ptr<float>(), total);

    dim3 block(TILE, TILE);
    dim3 grid((D + TILE - 1) / TILE, (N + TILE - 1) / TILE);
    qvla_linear_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        layout_logits.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        tmp.data_ptr<float>(),
        y.data_ptr<float>(),
        N, D);
    return y;
}
"""

_qvla_ext = load_inline(
    name="qvla_selective_layout_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["qvla_forward"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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 _qvla_ext.qvla_forward(x, layout_logits, self.out_proj.weight, self.out_proj.bias)