// Fused AllReduce CUDA Kernels for Distributed Training
//
// These kernels provide optimized implementations of collective operations
// that fuse multiple operations to reduce memory bandwidth and kernel launch overhead.
//
// Features:
// - Fused gradient scaling + AllReduce
// - Fused AllReduce + compression
// - Ring AllReduce with pipelining
// - Tree AllReduce for hierarchical topologies
// - Mixed precision support (FP32/FP16/BF16)

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>

namespace cg = cooperative_groups;

// =============================================================================
// Constants and Configuration
// =============================================================================

constexpr int WARP_SIZE = 32;
constexpr int MAX_BLOCK_SIZE = 1024;
constexpr int REDUCE_THREADS = 256;

// Reduction operation types
enum class ReduceOp : int {
    Sum = 0,
    Product = 1,
    Min = 2,
    Max = 3
};

// =============================================================================
// Utility Functions
// =============================================================================

template<typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

template<typename T>
__device__ __forceinline__ T warp_reduce_max(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val = max(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

template<typename T>
__device__ __forceinline__ T warp_reduce_min(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val = min(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

template<typename T>
__device__ __forceinline__ T block_reduce_sum(T val) {
    __shared__ T shared[WARP_SIZE];

    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    val = warp_reduce_sum(val);

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : T(0);
    if (wid == 0) val = warp_reduce_sum(val);

    return val;
}

// =============================================================================
// Fused Gradient Scaling + Reduction Kernel
// =============================================================================

// Fuses gradient scaling (divide by world_size) with local reduction
template<typename T>
__global__ void fused_scale_reduce_kernel(
    T* __restrict__ output,
    const T* __restrict__ input,
    const float scale,
    const int n,
    const ReduceOp op
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        T val = input[idx] * T(scale);

        // Write scaled value
        output[idx] = val;
    }
}

// Vectorized version for better memory bandwidth
__global__ void fused_scale_reduce_f32x4_kernel(
    float4* __restrict__ output,
    const float4* __restrict__ input,
    const float scale,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        float4 val = input[idx];
        val.x *= scale;
        val.y *= scale;
        val.z *= scale;
        val.w *= scale;
        output[idx] = val;
    }
}

// =============================================================================
// Fused AllReduce + FP16 Compression Kernel
// =============================================================================

// Converts FP32 gradients to FP16 during reduction
__global__ void fused_allreduce_fp16_compress_kernel(
    __half* __restrict__ output,
    const float* __restrict__ input,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        output[idx] = __float2half(input[idx]);
    }
}

// Vectorized FP32 to FP16 conversion (2x throughput)
__global__ void fused_fp32_to_fp16_x2_kernel(
    __half2* __restrict__ output,
    const float2* __restrict__ input,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        float2 val = input[idx];
        output[idx] = __floats2half2_rn(val.x, val.y);
    }
}

// FP16 to FP32 decompression
__global__ void fused_fp16_to_fp32_kernel(
    float* __restrict__ output,
    const __half* __restrict__ input,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        output[idx] = __half2float(input[idx]);
    }
}

// =============================================================================
// Fused AllReduce + INT8 Quantization Kernel
// =============================================================================

// Fused quantize + pack for INT8 compression
__global__ void fused_quantize_int8_kernel(
    int8_t* __restrict__ output,
    const float* __restrict__ input,
    float* __restrict__ scale_out,
    float* __restrict__ zero_point_out,
    const int n
) {
    __shared__ float s_min, s_max;

    // Phase 1: Find min/max using reduction
    float local_min = FLT_MAX;
    float local_max = -FLT_MAX;

    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float val = input[i];
        local_min = min(local_min, val);
        local_max = max(local_max, val);
    }

    // Reduce within block
    local_min = -warp_reduce_max(-local_min);  // min via negated max
    local_max = warp_reduce_max(local_max);

    if (threadIdx.x == 0) {
        s_min = local_min;
        s_max = local_max;
    }
    __syncthreads();

    float scale = (s_max - s_min) / 255.0f;
    float zero_point = s_min;

    if (threadIdx.x == 0) {
        *scale_out = scale;
        *zero_point_out = zero_point;
    }

    // Phase 2: Quantize
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
        float val = input[i];
        int8_t quantized = (scale > 0) ?
            (int8_t)__float2int_rn((val - zero_point) / scale) : 0;
        output[i] = max((int8_t)-128, min((int8_t)127, quantized));
    }
}

// Dequantize INT8 back to FP32
__global__ void fused_dequantize_int8_kernel(
    float* __restrict__ output,
    const int8_t* __restrict__ input,
    const float scale,
    const float zero_point,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        output[idx] = (float)input[idx] * scale + zero_point;
    }
}

// =============================================================================
// Ring AllReduce Kernels
// =============================================================================

// Ring reduce-scatter phase: partial reduction of chunks
template<typename T>
__global__ void ring_reduce_scatter_kernel(
    T* __restrict__ data,
    const T* __restrict__ recv_buf,
    const int chunk_size,
    const int chunk_offset,
    const ReduceOp op
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < chunk_size) {
        int global_idx = chunk_offset + idx;
        T local_val = data[global_idx];
        T recv_val = recv_buf[idx];

        T result;
        switch (op) {
            case ReduceOp::Sum:
                result = local_val + recv_val;
                break;
            case ReduceOp::Product:
                result = local_val * recv_val;
                break;
            case ReduceOp::Max:
                result = max(local_val, recv_val);
                break;
            case ReduceOp::Min:
                result = min(local_val, recv_val);
                break;
            default:
                result = local_val + recv_val;
        }

        data[global_idx] = result;
    }
}

// Ring all-gather phase: copy received chunks
template<typename T>
__global__ void ring_allgather_kernel(
    T* __restrict__ data,
    const T* __restrict__ recv_buf,
    const int chunk_size,
    const int chunk_offset
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < chunk_size) {
        data[chunk_offset + idx] = recv_buf[idx];
    }
}

// =============================================================================
// Tree AllReduce Kernels (for hierarchical topologies)
// =============================================================================

// Hierarchical reduce: first reduce within node, then across nodes
template<typename T>
__global__ void tree_reduce_local_kernel(
    T* __restrict__ output,
    const T* __restrict__ inputs,  // Array of pointers from local GPUs
    const int n,
    const int num_gpus,
    const ReduceOp op
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        T result = inputs[idx];  // First GPU's value

        for (int gpu = 1; gpu < num_gpus; gpu++) {
            T val = inputs[gpu * n + idx];
            switch (op) {
                case ReduceOp::Sum:
                    result += val;
                    break;
                case ReduceOp::Product:
                    result *= val;
                    break;
                case ReduceOp::Max:
                    result = max(result, val);
                    break;
                case ReduceOp::Min:
                    result = min(result, val);
                    break;
            }
        }

        output[idx] = result;
    }
}

// =============================================================================
// Fused TopK Sparsification Kernel
// =============================================================================

// Find top-k values by absolute magnitude
__global__ void topk_threshold_kernel(
    float* __restrict__ output,
    uint32_t* __restrict__ indices,
    int* __restrict__ count,
    const float* __restrict__ input,
    const float threshold,
    const int n,
    const int max_k
) {
    __shared__ int s_count;
    if (threadIdx.x == 0) s_count = 0;
    __syncthreads();

    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        float val = input[idx];
        if (fabsf(val) >= threshold) {
            int pos = atomicAdd(&s_count, 1);
            if (pos < max_k) {
                int global_pos = atomicAdd(count, 1);
                if (global_pos < max_k) {
                    output[global_pos] = val;
                    indices[global_pos] = idx;
                }
            }
        }
    }
}

// Sparse decompress: scatter values back to dense tensor
__global__ void sparse_decompress_kernel(
    float* __restrict__ output,
    const float* __restrict__ values,
    const uint32_t* __restrict__ indices,
    const int nnz,
    const int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // Zero out the output first (separate kernel or memset)
    if (idx < nnz) {
        uint32_t out_idx = indices[idx];
        if (out_idx < n) {
            output[out_idx] = values[idx];
        }
    }
}

// =============================================================================
// Fused 1-bit SGD Kernels
// =============================================================================

// Compress gradients to 1-bit with error feedback
__global__ void onebit_compress_kernel(
    uint32_t* __restrict__ output,       // Packed bits (32 values per uint32)
    float* __restrict__ error_feedback,  // Error accumulated for next iteration
    float* __restrict__ pos_mean_out,
    float* __restrict__ neg_mean_out,
    const float* __restrict__ input,
    const int n
) {
    __shared__ float s_pos_sum, s_neg_sum;
    __shared__ int s_pos_count, s_neg_count;
    __shared__ float s_mean;

    if (threadIdx.x == 0) {
        s_pos_sum = 0; s_neg_sum = 0;
        s_pos_count = 0; s_neg_count = 0;
    }
    __syncthreads();

    // Phase 1: Compute mean
    float local_sum = 0;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        local_sum += input[i];
    }
    local_sum = block_reduce_sum(local_sum);

    if (threadIdx.x == 0) {
        s_mean = local_sum / n;
    }
    __syncthreads();

    float mean = s_mean;

    // Phase 2: Compute positive/negative means and pack bits
    float local_pos_sum = 0, local_neg_sum = 0;
    int local_pos_count = 0, local_neg_count = 0;

    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float val = input[i];
        if (val >= mean) {
            local_pos_sum += val;
            local_pos_count++;
        } else {
            local_neg_sum += val;
            local_neg_count++;
        }
    }

    atomicAdd(&s_pos_sum, local_pos_sum);
    atomicAdd(&s_neg_sum, local_neg_sum);
    atomicAdd(&s_pos_count, local_pos_count);
    atomicAdd(&s_neg_count, local_neg_count);
    __syncthreads();

    float pos_mean = (s_pos_count > 0) ? s_pos_sum / s_pos_count : mean;
    float neg_mean = (s_neg_count > 0) ? s_neg_sum / s_neg_count : mean;

    if (threadIdx.x == 0) {
        *pos_mean_out = pos_mean;
        *neg_mean_out = neg_mean;
    }

    // Phase 3: Pack bits and compute error
    int word_idx = blockIdx.x * blockDim.x + threadIdx.x;
    int start_idx = word_idx * 32;

    if (start_idx < n) {
        uint32_t packed = 0;
        for (int bit = 0; bit < 32 && start_idx + bit < n; bit++) {
            int idx = start_idx + bit;
            float val = input[idx];
            float reconstructed;

            if (val >= mean) {
                packed |= (1u << bit);
                reconstructed = pos_mean;
            } else {
                reconstructed = neg_mean;
            }

            // Update error feedback
            if (error_feedback != nullptr) {
                error_feedback[idx] += val - reconstructed;
            }
        }
        output[word_idx] = packed;
    }
}

// Decompress 1-bit gradients
__global__ void onebit_decompress_kernel(
    float* __restrict__ output,
    const uint32_t* __restrict__ packed,
    const float pos_mean,
    const float neg_mean,
    const int n
) {
    int word_idx = blockIdx.x * blockDim.x + threadIdx.x;
    int start_idx = word_idx * 32;

    if (start_idx < n) {
        uint32_t bits = packed[word_idx];
        for (int bit = 0; bit < 32 && start_idx + bit < n; bit++) {
            int idx = start_idx + bit;
            output[idx] = (bits & (1u << bit)) ? pos_mean : neg_mean;
        }
    }
}

// =============================================================================
// Multi-GPU Gradient Accumulation Kernel
// =============================================================================

// Accumulate gradients from multiple GPUs (for data parallel)
template<typename T, int NUM_GPUS>
__global__ void multi_gpu_accumulate_kernel(
    T* __restrict__ output,
    const T* const* __restrict__ inputs,  // Array of input pointers
    const int n,
    const ReduceOp op
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        T result = inputs[0][idx];

        #pragma unroll
        for (int gpu = 1; gpu < NUM_GPUS; gpu++) {
            T val = inputs[gpu][idx];
            switch (op) {
                case ReduceOp::Sum:
                    result += val;
                    break;
                case ReduceOp::Product:
                    result *= val;
                    break;
                case ReduceOp::Max:
                    result = max(result, val);
                    break;
                case ReduceOp::Min:
                    result = min(result, val);
                    break;
            }
        }

        output[idx] = result;
    }
}

// =============================================================================
// Stream-Overlapped Copy Kernel (for compute/comm overlap)
// =============================================================================

// Async copy with notification flag
__global__ void async_copy_with_flag_kernel(
    float* __restrict__ dst,
    const float* __restrict__ src,
    volatile int* __restrict__ ready_flag,
    const int n,
    const int flag_value
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // Coalesced copy
    if (idx < n) {
        dst[idx] = src[idx];
    }

    // Last thread sets flag
    __syncthreads();
    if (idx == 0) {
        __threadfence_system();  // Ensure all writes visible
        *ready_flag = flag_value;
    }
}

// =============================================================================
// Bucket Fusion Kernel
// =============================================================================

// Fuse multiple small tensors into one large buffer for AllReduce
__global__ void bucket_pack_kernel(
    float* __restrict__ bucket,
    const float* const* __restrict__ tensors,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    const float* src = tensors[tensor_id];

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        bucket[offset + i] = src[i];
    }
}

// Unpack from bucket back to individual tensors
__global__ void bucket_unpack_kernel(
    float* const* __restrict__ tensors,
    const float* __restrict__ bucket,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    float* dst = tensors[tensor_id];

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        dst[i] = bucket[offset + i];
    }
}

// =============================================================================
// Host-Callable Wrapper Functions
// =============================================================================

extern "C" {

void launch_fused_scale_reduce(
    float* output,
    const float* input,
    float scale,
    int n,
    cudaStream_t stream
) {
    int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;
    fused_scale_reduce_kernel<float><<<grid_size, block_size, 0, stream>>>(
        output, input, scale, n, ReduceOp::Sum
    );
}

void launch_fused_fp16_compress(
    __half* output,
    const float* input,
    int n,
    cudaStream_t stream
) {
    int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;
    fused_allreduce_fp16_compress_kernel<<<grid_size, block_size, 0, stream>>>(
        output, input, n
    );
}

void launch_fused_fp16_decompress(
    float* output,
    const __half* input,
    int n,
    cudaStream_t stream
) {
    int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;
    fused_fp16_to_fp32_kernel<<<grid_size, block_size, 0, stream>>>(
        output, input, n
    );
}

void launch_ring_reduce_scatter(
    float* data,
    const float* recv_buf,
    int chunk_size,
    int chunk_offset,
    int op,
    cudaStream_t stream
) {
    int block_size = 256;
    int grid_size = (chunk_size + block_size - 1) / block_size;
    ring_reduce_scatter_kernel<float><<<grid_size, block_size, 0, stream>>>(
        data, recv_buf, chunk_size, chunk_offset, static_cast<ReduceOp>(op)
    );
}

void launch_ring_allgather(
    float* data,
    const float* recv_buf,
    int chunk_size,
    int chunk_offset,
    cudaStream_t stream
) {
    int block_size = 256;
    int grid_size = (chunk_size + block_size - 1) / block_size;
    ring_allgather_kernel<float><<<grid_size, block_size, 0, stream>>>(
        data, recv_buf, chunk_size, chunk_offset
    );
}

void launch_onebit_compress(
    uint32_t* output,
    float* error_feedback,
    float* pos_mean,
    float* neg_mean,
    const float* input,
    int n,
    cudaStream_t stream
) {
    int num_words = (n + 31) / 32;
    int block_size = 256;
    int grid_size = (num_words + block_size - 1) / block_size;
    onebit_compress_kernel<<<grid_size, block_size, 0, stream>>>(
        output, error_feedback, pos_mean, neg_mean, input, n
    );
}

void launch_onebit_decompress(
    float* output,
    const uint32_t* packed,
    float pos_mean,
    float neg_mean,
    int n,
    cudaStream_t stream
) {
    int num_words = (n + 31) / 32;
    int block_size = 256;
    int grid_size = (num_words + block_size - 1) / block_size;
    onebit_decompress_kernel<<<grid_size, block_size, 0, stream>>>(
        output, packed, pos_mean, neg_mean, n
    );
}

void launch_sparse_decompress(
    float* output,
    const float* values,
    const uint32_t* indices,
    int nnz,
    int n,
    cudaStream_t stream
) {
    // First zero the output
    cudaMemsetAsync(output, 0, n * sizeof(float), stream);

    int block_size = 256;
    int grid_size = (nnz + block_size - 1) / block_size;
    sparse_decompress_kernel<<<grid_size, block_size, 0, stream>>>(
        output, values, indices, nnz, n
    );
}

void launch_bucket_pack(
    float* bucket,
    const float* const* tensors,
    const int* offsets,
    const int* sizes,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_pack_kernel<<<grid, block, 0, stream>>>(
        bucket, tensors, offsets, sizes, num_tensors
    );
}

void launch_bucket_unpack(
    float* const* tensors,
    const float* bucket,
    const int* offsets,
    const int* sizes,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_unpack_kernel<<<grid, block, 0, stream>>>(
        tensors, bucket, offsets, sizes, num_tensors
    );
}

} // extern "C"

// =============================================================================
// Advanced Gradient Packing Kernels
// =============================================================================

// Fused pack + scale kernel - copies gradients and applies scaling in one pass
__global__ void bucket_pack_scale_kernel(
    float* __restrict__ bucket,
    const float* const* __restrict__ tensors,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const float scale,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    const float* src = tensors[tensor_id];

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        bucket[offset + i] = src[i] * scale;
    }
}

// Fused pack + INT8 quantize kernel
__global__ void bucket_pack_quantize_int8_kernel(
    int8_t* __restrict__ bucket,
    float* __restrict__ scales,  // Per-tensor scales
    const float* const* __restrict__ tensors,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    const float* src = tensors[tensor_id];

    // First pass: find max abs value using warp reduction
    __shared__ float s_max[32];
    float local_max = 0.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        local_max = fmaxf(local_max, fabsf(src[i]));
    }

    // Warp reduce
    for (int offset = 16; offset > 0; offset /= 2) {
        local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
    }

    // Block reduce
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;
    if (lane == 0) s_max[wid] = local_max;
    __syncthreads();

    if (threadIdx.x < 32) {
        local_max = (threadIdx.x < blockDim.x / 32) ? s_max[threadIdx.x] : 0.0f;
        for (int offset = 16; offset > 0; offset /= 2) {
            local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
        }
        if (threadIdx.x == 0) {
            atomicMax((int*)&scales[tensor_id], __float_as_int(local_max));
        }
    }
    __syncthreads();

    // Second pass: quantize
    float scale = scales[tensor_id] / 127.0f;
    if (scale == 0.0f) scale = 1.0f;  // Avoid division by zero

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        float val = src[i] / scale;
        bucket[offset + i] = (int8_t)fmaxf(-127.0f, fminf(127.0f, roundf(val)));
    }
}

// Fused unpack + dequantize INT8 kernel
__global__ void bucket_unpack_dequantize_int8_kernel(
    float* const* __restrict__ tensors,
    const int8_t* __restrict__ bucket,
    const float* __restrict__ scales,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    float* dst = tensors[tensor_id];
    float scale = scales[tensor_id] / 127.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        dst[i] = (float)bucket[offset + i] * scale;
    }
}

// Fused pack + gradient clipping kernel
__global__ void bucket_pack_clip_kernel(
    float* __restrict__ bucket,
    const float* const* __restrict__ tensors,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const float max_norm,
    const float* __restrict__ grad_norms,  // Pre-computed per-tensor norms
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    const float* src = tensors[tensor_id];
    float clip_coef = (grad_norms[tensor_id] > max_norm) ? (max_norm / grad_norms[tensor_id]) : 1.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        bucket[offset + i] = src[i] * clip_coef;
    }
}

// Fused unpack + weight update kernel (SGD)
__global__ void bucket_unpack_sgd_kernel(
    float* const* __restrict__ params,
    float* const* __restrict__ grads,
    const float* __restrict__ bucket,
    const int* __restrict__ offsets,
    const int* __restrict__ sizes,
    const float lr,
    const float weight_decay,
    const int num_tensors
) {
    int tensor_id = blockIdx.y;
    if (tensor_id >= num_tensors) return;

    int offset = offsets[tensor_id];
    int size = sizes[tensor_id];
    float* param = params[tensor_id];
    float* grad = grads[tensor_id];

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size; i += gridDim.x * blockDim.x) {
        float g = bucket[offset + i];
        grad[i] = g;  // Store the averaged gradient
        param[i] -= lr * (g + weight_decay * param[i]);
    }
}

// =============================================================================
// Host Wrappers for Advanced Kernels
// =============================================================================

extern "C" {

void launch_bucket_pack_scale(
    float* bucket,
    const float* const* tensors,
    const int* offsets,
    const int* sizes,
    float scale,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_pack_scale_kernel<<<grid, block, 0, stream>>>(
        bucket, tensors, offsets, sizes, scale, num_tensors
    );
}

void launch_bucket_pack_quantize_int8(
    int8_t* bucket,
    float* scales,
    const float* const* tensors,
    const int* offsets,
    const int* sizes,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    // Initialize scales to 0
    cudaMemsetAsync(scales, 0, num_tensors * sizeof(float), stream);

    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_pack_quantize_int8_kernel<<<grid, block, 0, stream>>>(
        bucket, scales, tensors, offsets, sizes, num_tensors
    );
}

void launch_bucket_unpack_dequantize_int8(
    float* const* tensors,
    const int8_t* bucket,
    const float* scales,
    const int* offsets,
    const int* sizes,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_unpack_dequantize_int8_kernel<<<grid, block, 0, stream>>>(
        tensors, bucket, scales, offsets, sizes, num_tensors
    );
}

void launch_bucket_pack_clip(
    float* bucket,
    const float* const* tensors,
    const int* offsets,
    const int* sizes,
    float max_norm,
    const float* grad_norms,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_pack_clip_kernel<<<grid, block, 0, stream>>>(
        bucket, tensors, offsets, sizes, max_norm, grad_norms, num_tensors
    );
}

void launch_bucket_unpack_sgd(
    float* const* params,
    float* const* grads,
    const float* bucket,
    const int* offsets,
    const int* sizes,
    float lr,
    float weight_decay,
    int num_tensors,
    int max_size,
    cudaStream_t stream
) {
    dim3 block(256);
    dim3 grid((max_size + 255) / 256, num_tensors);
    bucket_unpack_sgd_kernel<<<grid, block, 0, stream>>>(
        params, grads, bucket, offsets, sizes, lr, weight_decay, num_tensors
    );
}

} // extern "C"