/**
 * MPS Linear Attention - Metal Kernels for Apple Silicon
 *
 * Implements:
 *   - causal_conv1d_fwd: depthwise causal 1D convolution + SiLU
 *   - delta_rule_recurrent_step: single-token gated delta rule update
 */

#include <torch/extension.h>
#include <ATen/mps/MPSStream.h>
#include <ATen/native/mps/OperationUtils.h>

#import <Metal/Metal.h>
#import <Foundation/Foundation.h>

#include <string>
#include <mutex>
#include <unordered_map>
#include <dlfcn.h>

// ============================================================================
// Metal library setup
// ============================================================================

static id<MTLLibrary>  g_library  = nil;
static std::mutex      g_lib_mutex;
static std::unordered_map<std::string, id<MTLComputePipelineState>> g_pipelines;

static NSString* KERNEL_SOURCE = nil;

static void load_kernel_source() {
    // Embed the Metal source at compile time via xxd trick / string literal
    // We load from file path relative to this .mm at build time
    NSString* path = [[NSBundle mainBundle] pathForResource:@"kernels" ofType:@"metal"];
    if (!path) {
        // Try relative to module dir
        Dl_info info;
        if (dladdr((void*)load_kernel_source, &info) && info.dli_fname) {
            NSString* dir = [[[NSString stringWithUTF8String:info.dli_fname]
                              stringByDeletingLastPathComponent]
                             stringByDeletingLastPathComponent];
            path = [dir stringByAppendingPathComponent:@"csrc/kernels.metal"];
        }
    }
    if (path && [[NSFileManager defaultManager] fileExistsAtPath:path]) {
        KERNEL_SOURCE = [NSString stringWithContentsOfFile:path
                                                  encoding:NSUTF8StringEncoding
                                                     error:nil];
    }
}

// Inline fallback source (embedded at compile time)
static const char* INLINE_METAL_SOURCE = R"METAL(
#include <metal_stdlib>
using namespace metal;

kernel void causal_conv1d_fwd(
    device const half* x       [[buffer(0)]],
    device const half* weight  [[buffer(1)]],
    device const half* bias    [[buffer(2)]],
    device       half* out     [[buffer(3)]],
    constant int& B            [[buffer(4)]],
    constant int& D            [[buffer(5)]],
    constant int& T            [[buffer(6)]],
    constant int& W            [[buffer(7)]],
    constant int& has_bias     [[buffer(8)]],
    constant int& apply_silu   [[buffer(9)]],
    uint3 gid [[thread_position_in_grid]])
{
    int b = gid.z;
    int d = gid.y;
    int t = gid.x;
    if (b >= B || d >= D || t >= T) return;
    float acc = 0.0f;
    for (int w = 0; w < W; w++) {
        int src_t = t - (W - 1 - w);
        if (src_t >= 0)
            acc += float(x[b * D * T + d * T + src_t]) * float(weight[d * W + w]);
    }
    if (has_bias) acc += float(bias[d]);
    if (apply_silu) {
        float sig = 1.0f / (1.0f + exp(-acc));
        acc = acc * sig;
    }
    out[b * D * T + d * T + t] = half(acc);
}

kernel void causal_conv1d_fwd_f32(
    device const float* x       [[buffer(0)]],
    device const float* weight  [[buffer(1)]],
    device const float* bias    [[buffer(2)]],
    device       float* out     [[buffer(3)]],
    constant int& B             [[buffer(4)]],
    constant int& D             [[buffer(5)]],
    constant int& T             [[buffer(6)]],
    constant int& W             [[buffer(7)]],
    constant int& has_bias      [[buffer(8)]],
    constant int& apply_silu    [[buffer(9)]],
    uint3 gid [[thread_position_in_grid]])
{
    int b = gid.z;
    int d = gid.y;
    int t = gid.x;
    if (b >= B || d >= D || t >= T) return;
    float acc = 0.0f;
    for (int w = 0; w < W; w++) {
        int src_t = t - (W - 1 - w);
        if (src_t >= 0)
            acc += x[b * D * T + d * T + src_t] * weight[d * W + w];
    }
    if (has_bias) acc += bias[d];
    if (apply_silu) {
        float sig = 1.0f / (1.0f + exp(-acc));
        acc = acc * sig;
    }
    out[b * D * T + d * T + t] = acc;
}

// Fused T-step recurrent kernel: loops over T inside the shader.
// One threadgroup per (b, h). No barriers needed — sequential T within TG.
// q, k: [B, H, T, K], v: [B, H, T, V], g: [B, H, T], beta: [B, H, T]
// state: [B, H, K, V] float32 (in-place), out: [B, H, T, V] half
kernel void delta_rule_recurrent_fused(
    device const half* q      [[buffer(0)]],
    device const half* k      [[buffer(1)]],
    device const half* v      [[buffer(2)]],
    device const float* g     [[buffer(3)]],
    device const half* beta   [[buffer(4)]],
    device       float* state [[buffer(5)]],
    device       half* out    [[buffer(6)]],
    constant int& B           [[buffer(7)]],
    constant int& H           [[buffer(8)]],
    constant int& T_val       [[buffer(9)]],
    constant int& K           [[buffer(10)]],
    constant int& V           [[buffer(11)]],
    constant float& scale     [[buffer(12)]],
    uint2 tgid    [[threadgroup_position_in_grid]],
    uint2 tid2    [[thread_position_in_threadgroup]],
    uint2 tg_dim  [[threads_per_threadgroup]])
{
    uint tid = tid2.x;
    uint tg_size = tg_dim.x;
    int b = tgid.y;
    int h = tgid.x;
    if (b >= B || h >= H) return;

    int bh = b * H + h;
    device float* s_bh = state + bh * K * V;
    float sc = scale;

    for (int t = 0; t < T_val; t++) {
        int bht = bh * T_val + t;
        device const half* q_t    = q    + bht * K;
        device const half* k_t    = k    + bht * K;
        device const half* v_t    = v    + bht * V;
        device       half* out_t  = out  + bht * V;
        float g_val    = exp(g[bht]);
        float beta_val = float(beta[bht]);

        for (int vi = tid; vi < V; vi += tg_size) {
            float kv_mem = 0.0f;
            for (int ki = 0; ki < K; ki++) {
                float s = s_bh[ki * V + vi] * g_val;
                s_bh[ki * V + vi] = s;
                kv_mem += s * float(k_t[ki]);
            }
            float delta = (float(v_t[vi]) - kv_mem) * beta_val;
            for (int ki = 0; ki < K; ki++) {
                s_bh[ki * V + vi] += float(k_t[ki]) * delta;
            }
            float o = 0.0f;
            for (int ki = 0; ki < K; ki++) {
                o += s_bh[ki * V + vi] * float(q_t[ki]);
            }
            out_t[vi] = half(o * sc);
        }
    }
}

// causal_conv1d_update: fused state shift + conv + optional SiLU
// One thread per (b, d). Processes L new tokens.
// conv_state: (B, D, W) — mutated in-place (sliding window)
// hidden_states: (B, D, L), weight: (D, W), bias: (D,), out: (B, D, L)
kernel void causal_conv1d_update_kernel(
    device const half* hidden   [[buffer(0)]],   // (B, D, L)
    device       half* conv_st  [[buffer(1)]],   // (B, D, W) in-place
    device const half* weight   [[buffer(2)]],   // (D, W)
    device const half* bias     [[buffer(3)]],   // (D,)
    device       half* out      [[buffer(4)]],   // (B, D, L)
    constant int& B             [[buffer(5)]],
    constant int& D             [[buffer(6)]],
    constant int& L             [[buffer(7)]],
    constant int& W             [[buffer(8)]],
    constant int& has_bias      [[buffer(9)]],
    constant int& apply_silu    [[buffer(10)]],
    uint3 gid [[thread_position_in_grid]])
{
    int b = gid.y;
    int d = gid.x;
    if (b >= B || d >= D) return;

    device half* st = conv_st + (b * D + d) * W;
    device const half* h = hidden + (b * D + d) * L;
    device half* o = out + (b * D + d) * L;
    device const half* w = weight + d * W;

    // Process each new token
    for (int l = 0; l < L; l++) {
        // Shift state left by 1, push new token
        for (int i = 0; i < W - 1; i++) {
            st[i] = st[i + 1];
        }
        st[W - 1] = h[l];

        // Depthwise conv: dot(state, weight)
        float acc = 0.0f;
        for (int i = 0; i < W; i++) {
            acc += float(st[i]) * float(w[i]);
        }
        if (has_bias) acc += float(bias[d]);
        if (apply_silu) {
            float sig = 1.0f / (1.0f + exp(-acc));
            acc = acc * sig;
        }
        o[l] = half(acc);
    }
}

kernel void causal_conv1d_update_kernel_f32(
    device const float* hidden  [[buffer(0)]],
    device       float* conv_st [[buffer(1)]],
    device const float* weight  [[buffer(2)]],
    device const float* bias    [[buffer(3)]],
    device       float* out     [[buffer(4)]],
    constant int& B             [[buffer(5)]],
    constant int& D             [[buffer(6)]],
    constant int& L             [[buffer(7)]],
    constant int& W             [[buffer(8)]],
    constant int& has_bias      [[buffer(9)]],
    constant int& apply_silu    [[buffer(10)]],
    uint3 gid [[thread_position_in_grid]])
{
    int b = gid.y;
    int d = gid.x;
    if (b >= B || d >= D) return;

    device float* st = conv_st + (b * D + d) * W;
    device const float* h = hidden + (b * D + d) * L;
    device float* o = out + (b * D + d) * L;
    device const float* w = weight + d * W;

    for (int l = 0; l < L; l++) {
        for (int i = 0; i < W - 1; i++) {
            st[i] = st[i + 1];
        }
        st[W - 1] = h[l];

        float acc = 0.0f;
        for (int i = 0; i < W; i++) {
            acc += st[i] * w[i];
        }
        if (has_bias) acc += bias[d];
        if (apply_silu) {
            float sig = 1.0f / (1.0f + exp(-acc));
            acc = acc * sig;
        }
        o[l] = acc;
    }
}

// State is float32 for precision. q/k/v/beta are half, output is half.
// Uses threadgroup memory to cache k and q vectors (read once, shared by all threads).
kernel void delta_rule_recurrent_step(
    device const half* q      [[buffer(0)]],   // [B, H, K]
    device const half* k      [[buffer(1)]],   // [B, H, K]
    device const half* v      [[buffer(2)]],   // [B, H, V]
    device const float* g     [[buffer(3)]],   // [B, H] log-gate
    device const half* beta   [[buffer(4)]],   // [B, H]
    device       float* state [[buffer(5)]],   // [B, H, K, V] float32!
    device       half* out    [[buffer(6)]],   // [B, H, V]
    constant int& B           [[buffer(7)]],
    constant int& H           [[buffer(8)]],
    constant int& K           [[buffer(9)]],
    constant int& V           [[buffer(10)]],
    constant float& scale     [[buffer(11)]],
    uint2 tgid    [[threadgroup_position_in_grid]],
    uint2 tid2    [[thread_position_in_threadgroup]],
    uint2 tg_dim  [[threads_per_threadgroup]])
{
    uint tid = tid2.x;
    uint tg_size = tg_dim.x;
    int b = tgid.y;
    int h = tgid.x;
    if (b >= B || h >= H) return;

    int bh = b * H + h;
    device const half* q_bh    = q    + bh * K;
    device const half* k_bh    = k    + bh * K;
    device const half* v_bh    = v    + bh * V;
    device       float* s_bh   = state + bh * K * V;
    device       half* out_bh  = out  + bh * V;

    float g_val    = exp(g[bh]);
    float beta_val = float(beta[bh]);
    float sc = scale;

    for (int vi = tid; vi < V; vi += tg_size) {
        float kv_mem = 0.0f;
        for (int ki = 0; ki < K; ki++) {
            float s = s_bh[ki * V + vi] * g_val;
            s_bh[ki * V + vi] = s;
            kv_mem += s * float(k_bh[ki]);
        }
        float delta = (float(v_bh[vi]) - kv_mem) * beta_val;
        for (int ki = 0; ki < K; ki++) {
            s_bh[ki * V + vi] += float(k_bh[ki]) * delta;
        }
        float o = 0.0f;
        for (int ki = 0; ki < K; ki++) {
            o += s_bh[ki * V + vi] * float(q_bh[ki]);
        }
        out_bh[vi] = half(o * sc);
    }
}
)METAL";

static id<MTLDevice> get_device() {
    return at::mps::MPSDevice::getInstance()->device();
}

static bool ensure_library() {
    std::lock_guard<std::mutex> lock(g_lib_mutex);
    if (g_library) return true;

    id<MTLDevice> device = get_device();
    NSError* error = nil;

    // Try file-based source first
    if (!KERNEL_SOURCE) load_kernel_source();
    NSString* src = KERNEL_SOURCE ?: [NSString stringWithUTF8String:INLINE_METAL_SOURCE];

    MTLCompileOptions* opts = [[MTLCompileOptions alloc] init];
    opts.mathMode = MTLMathModeFast;

    g_library = [device newLibraryWithSource:src options:opts error:&error];
    if (!g_library) {
        NSLog(@"[mla] Metal compile error: %@", error);
        return false;
    }
    return true;
}

static id<MTLComputePipelineState> get_pipeline(const std::string& name) {
    {
        std::lock_guard<std::mutex> lock(g_lib_mutex);
        auto it = g_pipelines.find(name);
        if (it != g_pipelines.end()) return it->second;
    }

    if (!ensure_library()) return nil;

    id<MTLFunction> fn = [g_library newFunctionWithName:[NSString stringWithUTF8String:name.c_str()]];
    if (!fn) {
        NSLog(@"[mla] Function not found: %s", name.c_str());
        return nil;
    }

    NSError* error = nil;
    id<MTLDevice> device = get_device();
    id<MTLComputePipelineState> pso = [device newComputePipelineStateWithFunction:fn error:&error];
    if (!pso) {
        NSLog(@"[mla] Pipeline error: %@", error);
        return nil;
    }

    std::lock_guard<std::mutex> lock(g_lib_mutex);
    g_pipelines[name] = pso;
    return pso;
}

// ============================================================================
// Buffer helpers
// ============================================================================

static id<MTLBuffer> get_buffer(const torch::Tensor& t) {
    return at::native::mps::getMTLBufferStorage(t);
}

static int64_t byte_offset(const torch::Tensor& t) {
    return t.storage_offset() * t.element_size();
}

// ============================================================================
// causal_conv1d forward
// ============================================================================

torch::Tensor metal_causal_conv1d(
    torch::Tensor x,       // (B, D, T)
    torch::Tensor weight,  // (D, W)
    torch::Tensor bias,    // (D,) or empty
    bool apply_silu)
{
    TORCH_CHECK(x.is_mps(), "x must be MPS");
    TORCH_CHECK(x.dim() == 3, "x must be (B, D, T)");

    x = x.contiguous();
    weight = weight.contiguous();

    int B = x.size(0), D = x.size(1), T = x.size(2);
    int W = weight.size(1);
    int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;

    auto out = torch::empty_like(x);

    bool is_f32 = (x.scalar_type() == torch::kFloat);
    std::string kernel_name = is_f32 ? "causal_conv1d_fwd_f32" : "causal_conv1d_fwd";

    auto pso = get_pipeline(kernel_name);
    TORCH_CHECK(pso, "Failed to get pipeline: ", kernel_name);

    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
        [enc setComputePipelineState:pso];

        [enc setBuffer:get_buffer(x)      offset:byte_offset(x)      atIndex:0];
        [enc setBuffer:get_buffer(weight) offset:byte_offset(weight)  atIndex:1];

        auto bias_c = has_bias ? bias.contiguous() : out;
        [enc setBuffer:get_buffer(bias_c) offset:(has_bias ? byte_offset(bias_c) : 0) atIndex:2];

        [enc setBuffer:get_buffer(out) offset:byte_offset(out) atIndex:3];
        [enc setBytes:&B          length:sizeof(int) atIndex:4];
        [enc setBytes:&D          length:sizeof(int) atIndex:5];
        [enc setBytes:&T          length:sizeof(int) atIndex:6];
        [enc setBytes:&W          length:sizeof(int) atIndex:7];
        [enc setBytes:&has_bias   length:sizeof(int) atIndex:8];
        int silu_flag = apply_silu ? 1 : 0;
        [enc setBytes:&silu_flag  length:sizeof(int) atIndex:9];

        MTLSize grid = MTLSizeMake(T, D, B);
        NSUInteger tg_total = pso.maxTotalThreadsPerThreadgroup;
        MTLSize tg = MTLSizeMake(MIN((NSUInteger)T, tg_total), 1, 1);
        [enc dispatchThreads:grid threadsPerThreadgroup:tg];
    }

    return out;
}

// ============================================================================
// causal_conv1d_update (fused state shift + conv + SiLU)
// ============================================================================

torch::Tensor metal_causal_conv1d_update(
    torch::Tensor hidden_states,  // (B, D, L)
    torch::Tensor conv_state,     // (B, D, W) in-place
    torch::Tensor weight,         // (D, W)
    torch::Tensor bias,           // (D,) or empty
    bool apply_silu)
{
    TORCH_CHECK(hidden_states.is_mps(), "hidden_states must be MPS");
    TORCH_CHECK(conv_state.is_contiguous(), "conv_state must be contiguous for in-place update");

    hidden_states = hidden_states.contiguous();
    weight = weight.contiguous();

    int B = hidden_states.size(0), D = hidden_states.size(1), L = hidden_states.size(2);
    int W = weight.size(1);
    int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;

    auto out = torch::empty_like(hidden_states);

    bool is_f32 = (hidden_states.scalar_type() == torch::kFloat);
    std::string kernel_name = is_f32 ? "causal_conv1d_update_kernel_f32" : "causal_conv1d_update_kernel";

    auto pso = get_pipeline(kernel_name);
    TORCH_CHECK(pso, "Failed to get pipeline: ", kernel_name);

    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
        [enc setComputePipelineState:pso];

        [enc setBuffer:get_buffer(hidden_states) offset:byte_offset(hidden_states) atIndex:0];
        [enc setBuffer:get_buffer(conv_state)    offset:byte_offset(conv_state)    atIndex:1];
        [enc setBuffer:get_buffer(weight)        offset:byte_offset(weight)        atIndex:2];

        auto bias_c = has_bias ? bias.contiguous() : out;
        [enc setBuffer:get_buffer(bias_c) offset:(has_bias ? byte_offset(bias_c) : 0) atIndex:3];

        [enc setBuffer:get_buffer(out) offset:byte_offset(out) atIndex:4];
        [enc setBytes:&B        length:sizeof(int) atIndex:5];
        [enc setBytes:&D        length:sizeof(int) atIndex:6];
        [enc setBytes:&L        length:sizeof(int) atIndex:7];
        [enc setBytes:&W        length:sizeof(int) atIndex:8];
        [enc setBytes:&has_bias length:sizeof(int) atIndex:9];
        int silu_flag = apply_silu ? 1 : 0;
        [enc setBytes:&silu_flag length:sizeof(int) atIndex:10];

        MTLSize grid = MTLSizeMake(D, B, 1);
        NSUInteger tg_total = pso.maxTotalThreadsPerThreadgroup;
        MTLSize tg = MTLSizeMake(MIN((NSUInteger)D, tg_total), 1, 1);
        [enc dispatchThreads:grid threadsPerThreadgroup:tg];
    }

    return out;
}

// ============================================================================
// delta_rule_recurrent_step
// ============================================================================

torch::Tensor metal_delta_rule_step(
    torch::Tensor q,      // (B, H, K)
    torch::Tensor k,      // (B, H, K)
    torch::Tensor v,      // (B, H, V)
    torch::Tensor g,      // (B, H) float32 log-gate
    torch::Tensor beta,   // (B, H)
    torch::Tensor state,  // (B, H, K, V)  in-place update
    float scale)
{
    TORCH_CHECK(q.is_mps(), "Tensors must be on MPS");
    TORCH_CHECK(state.is_contiguous(), "state must be contiguous for in-place update");
    TORCH_CHECK(state.scalar_type() == torch::kFloat, "state must be float32");
    q = q.contiguous(); k = k.contiguous(); v = v.contiguous();
    g = g.contiguous().to(torch::kFloat);
    beta = beta.contiguous();

    int B = q.size(0), H = q.size(1), K = q.size(2), V = v.size(2);
    auto out = torch::empty({B, H, V}, q.options());

    auto pso = get_pipeline("delta_rule_recurrent_step");
    TORCH_CHECK(pso, "Failed to get pipeline: delta_rule_recurrent_step");

    auto stream = at::mps::getCurrentMPSStream();
    @autoreleasepool {
        id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
        [enc setComputePipelineState:pso];

        [enc setBuffer:get_buffer(q)     offset:byte_offset(q)     atIndex:0];
        [enc setBuffer:get_buffer(k)     offset:byte_offset(k)     atIndex:1];
        [enc setBuffer:get_buffer(v)     offset:byte_offset(v)     atIndex:2];
        [enc setBuffer:get_buffer(g)     offset:byte_offset(g)     atIndex:3];
        [enc setBuffer:get_buffer(beta)  offset:byte_offset(beta)  atIndex:4];
        [enc setBuffer:get_buffer(state) offset:byte_offset(state) atIndex:5];
        [enc setBuffer:get_buffer(out)   offset:byte_offset(out)   atIndex:6];
        [enc setBytes:&B     length:sizeof(int)   atIndex:7];
        [enc setBytes:&H     length:sizeof(int)   atIndex:8];
        [enc setBytes:&K     length:sizeof(int)   atIndex:9];
        [enc setBytes:&V     length:sizeof(int)   atIndex:10];
        [enc setBytes:&scale length:sizeof(float) atIndex:11];

        NSUInteger tg_size = MIN((NSUInteger)V, pso.maxTotalThreadsPerThreadgroup);
        MTLSize tg_count = MTLSizeMake(H, B, 1);
        MTLSize tg       = MTLSizeMake(tg_size, 1, 1);
        [enc dispatchThreadgroups:tg_count threadsPerThreadgroup:tg];
    }

    return out;
}

// ============================================================================
// delta_rule_recurrent_fused (T-step, single dispatch)
// ============================================================================

torch::Tensor metal_delta_rule_fused(
    torch::Tensor q,      // (B, H, T, K)
    torch::Tensor k,      // (B, H, T, K)
    torch::Tensor v,      // (B, H, T, V)
    torch::Tensor g,      // (B, H, T) float32
    torch::Tensor beta,   // (B, H, T) half
    torch::Tensor state,  // (B, H, K, V) float32 in-place
    float scale)
{
    TORCH_CHECK(q.is_mps(), "Tensors must be on MPS");
    TORCH_CHECK(state.is_contiguous(), "state must be contiguous");
    TORCH_CHECK(state.scalar_type() == torch::kFloat, "state must be float32");
    q = q.contiguous(); k = k.contiguous(); v = v.contiguous();
    g = g.contiguous().to(torch::kFloat);
    beta = beta.contiguous();

    int B = q.size(0), H = q.size(1), T = q.size(2), K = q.size(3), V = v.size(3);
    auto out = torch::empty({B, H, T, V}, q.options());

    auto pso = get_pipeline("delta_rule_recurrent_fused");
    TORCH_CHECK(pso, "Failed to get pipeline: delta_rule_recurrent_fused");

    auto stream = at::mps::getCurrentMPSStream();
    @autoreleasepool {
        id<MTLComputeCommandEncoder> enc = stream->commandEncoder();
        [enc setComputePipelineState:pso];

        [enc setBuffer:get_buffer(q)     offset:byte_offset(q)     atIndex:0];
        [enc setBuffer:get_buffer(k)     offset:byte_offset(k)     atIndex:1];
        [enc setBuffer:get_buffer(v)     offset:byte_offset(v)     atIndex:2];
        [enc setBuffer:get_buffer(g)     offset:byte_offset(g)     atIndex:3];
        [enc setBuffer:get_buffer(beta)  offset:byte_offset(beta)  atIndex:4];
        [enc setBuffer:get_buffer(state) offset:byte_offset(state) atIndex:5];
        [enc setBuffer:get_buffer(out)   offset:byte_offset(out)   atIndex:6];
        [enc setBytes:&B     length:sizeof(int)   atIndex:7];
        [enc setBytes:&H     length:sizeof(int)   atIndex:8];
        [enc setBytes:&T     length:sizeof(int)   atIndex:9];
        [enc setBytes:&K     length:sizeof(int)   atIndex:10];
        [enc setBytes:&V     length:sizeof(int)   atIndex:11];
        [enc setBytes:&scale length:sizeof(float) atIndex:12];

        NSUInteger tg_size = MIN((NSUInteger)V, pso.maxTotalThreadsPerThreadgroup);
        MTLSize tg_count = MTLSizeMake(H, B, 1);
        MTLSize tg       = MTLSizeMake(tg_size, 1, 1);
        [enc dispatchThreadgroups:tg_count threadsPerThreadgroup:tg];
    }

    return out;
}

// ============================================================================
// pybind11
// ============================================================================

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("causal_conv1d", &metal_causal_conv1d,
          "Causal depthwise conv1d + SiLU (MPS)",
          py::arg("x"), py::arg("weight"), py::arg("bias"), py::arg("apply_silu")=true);

    m.def("causal_conv1d_update", &metal_causal_conv1d_update,
          "Causal conv1d update with state (MPS)",
          py::arg("hidden_states"), py::arg("conv_state"), py::arg("weight"),
          py::arg("bias"), py::arg("apply_silu")=true);

    m.def("delta_rule_step", &metal_delta_rule_step,
          "Single-token gated delta rule recurrent step (MPS)",
          py::arg("q"), py::arg("k"), py::arg("v"), py::arg("g"),
          py::arg("beta"), py::arg("state"), py::arg("scale"));

    m.def("delta_rule_fused", &metal_delta_rule_fused,
          "Fused T-step gated delta rule recurrent (MPS)",
          py::arg("q"), py::arg("k"), py::arg("v"), py::arg("g"),
          py::arg("beta"), py::arg("state"), py::arg("scale"));
}
