# NEGATIVE (non-default stream) run with the static gate OFF: all work (reference ops + a launched no-op kernel)
# is enqueued on a side stream, so cuda events on the default stream see ~nothing; the guarded timer must not.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
__global__ void noop_kernel(float* y) { if (threadIdx.x == 1024) y[0] = 0.f; }
torch::Tensor noop(torch::Tensor y) { noop_kernel<<<1, 32, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>()); return y; }
"""
cpp_src = "torch::Tensor noop(torch::Tensor y);"
ext = load_inline(name="eval_neg_stream_rt", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["noop"])


class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, multiplier, negative_slope):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.multiplier = multiplier
        self.negative_slope = negative_slope
        self.s = torch.cuda.Stream()

    def forward(self, x):
        with torch.cuda.stream(self.s):
            y = ext.noop(F.leaky_relu(F.linear(x, self.gemm.weight, self.gemm.bias) * self.multiplier, self.negative_slope))
        return y