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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void mix_quant_kernel(const float* x, const float* logits, float* mixed, float scale, int B, int T, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * T * D;
    if (idx >= total) return;
    int bt = idx / D;
    float alpha = 1.0f / (1.0f + expf(-logits[bt]));
    float val = x[idx];
    float q = roundf(val / scale);
    q = fmaxf(fminf(q, 127.0f), -127.0f);
    float xq = q * scale;
    mixed[idx] = alpha * xq + (1.0f - alpha) * val;
}

torch::Tensor selective_mix(torch::Tensor x, torch::Tensor layout_logits, float scale) {
    auto B = x.size(0);
    auto T = x.size(1);
    auto D = x.size(2);
    auto mixed = torch::empty_like(x);
    int total = B * T * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    mix_quant_kernel<<<blocks, threads>>>(x.data_ptr<float>(), layout_logits.data_ptr<float>(), mixed.data_ptr<float>(), scale, B, T, D);
    return mixed;
}
'''

cpp_source = r'''
torch::Tensor selective_mix(torch::Tensor x, torch::Tensor layout_logits, float scale);
'''

module = load_inline(
    name="selective_quant",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["selective_mix"],
    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:
        B, T, D = x.shape
        amax = x.detach().abs().max().clamp(min=1e-8)
        scale = amax.item() / 127.0
        mixed = module.selective_mix(x, layout_logits, scale)
        return self.out_proj(mixed)