Kernel Dev 训练营：以 FlashInfer 为主线

// C1 · Vector add + fused bias-GELU — reference solution.
//
// Two kernels, two indexing styles:
//   vector_add_kernel : one thread per element, grid sized to cover n exactly (+ bounds check).
//   bias_gelu_kernel  : grid-stride loop, grid capped at ~32 blocks per SM, so one launch
//                       works for any size and each thread does several elements.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

#include <algorithm>

#define CHECK_F32_CUDA(t)                                                   \
  TORCH_CHECK((t).is_cuda(), #t " must be a CUDA tensor");                  \
  TORCH_CHECK((t).is_contiguous(), #t " must be contiguous");               \
  TORCH_CHECK((t).scalar_type() == at::kFloat, #t " must be float32")

// ---------------------------------------------------------------- vector add
// __restrict__ promises the compiler that a, b, c never alias, so it may keep values in
// registers and use the read-only data path for a and b.
__global__ void vector_add_kernel(const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ c,
                                  int64_t n) {
  // Global index of this thread. Cast before multiplying: blockIdx.x * blockDim.x is
  // computed in 32-bit unsigned arithmetic and overflows once n > 2^32.
  int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {            // the last block is usually only partly full
    c[i] = a[i] + b[i];   // consecutive threads touch consecutive floats -> coalesced
  }
}

torch::Tensor vector_add(torch::Tensor a, torch::Tensor b) {
  CHECK_F32_CUDA(a);
  CHECK_F32_CUDA(b);
  TORCH_CHECK(a.numel() == b.numel(), "a and b must have the same number of elements");
  auto c = torch::empty_like(a);
  const int64_t n = a.numel();
  if (n == 0) return c;  // a launch with 0 blocks is an error

  const int threads = 256;                               // a multiple of 32 (warp size)
  const int64_t blocks = (n + threads - 1) / threads;    // ceil(n / threads)
  // Launch on PyTorch's current stream so the kernel is ordered with other torch ops.
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  vector_add_kernel<<<blocks, threads, 0, stream>>>(
      a.data_ptr<float>(), b.data_ptr<float>(), c.data_ptr<float>(), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();  // surfaces bad launch configs immediately
  return c;
}

// ---------------------------------------------------------------- fused bias + GELU
// GELU, tanh approximation (what torch.nn.functional.gelu(approximate="tanh") computes):
//   0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 x^3)))
__device__ __forceinline__ float gelu_tanh(float x) {
  const float k0 = 0.7978845608028654f;  // sqrt(2 / pi)
  const float k1 = 0.044715f;
  return 0.5f * x * (1.0f + tanhf(k0 * (x + k1 * x * x * x)));
}

// x: [rows, cols] row-major, bias: [cols], y: [rows, cols].
// Unfused PyTorch does `t = x + bias` (read x, write t) then `gelu(t)` (read t, write y):
// 4 passes over memory. Fused: read x once, write y once -> about half the traffic.
__global__ void bias_gelu_kernel(const float* __restrict__ x,
                                 const float* __restrict__ bias,
                                 float* __restrict__ y,
                                 int64_t total,
                                 int64_t cols) {
  const int64_t stride = (int64_t)blockDim.x * gridDim.x;  // threads in the whole grid
  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) {
    // Row-major: element i sits in column i % cols. bias is tiny and stays in L1/L2.
    y[i] = gelu_tanh(x[i] + bias[i % cols]);
  }
}

torch::Tensor bias_gelu(torch::Tensor x, torch::Tensor bias) {
  CHECK_F32_CUDA(x);
  CHECK_F32_CUDA(bias);
  TORCH_CHECK(x.dim() == 2, "x must be 2-D [rows, cols]");
  TORCH_CHECK(bias.dim() == 1 && bias.size(0) == x.size(1), "bias must be [cols]");
  auto y = torch::empty_like(x);
  const int64_t total = x.numel();
  if (total == 0) return y;

  const int threads = 256;
  // Enough blocks to fill every SM several times over, but no more: the grid-stride loop
  // covers the rest, and fewer blocks means less scheduling overhead for huge tensors.
  const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
  const int64_t blocks = std::min<int64_t>((total + threads - 1) / threads, (int64_t)sms * 32);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  bias_gelu_kernel<<<blocks, threads, 0, stream>>>(
      x.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total, x.size(1));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return y;
}

// C2 · Reduction sum — reference solution.
//
// Three levels, each using the cheapest mechanism available at that level:
//   1. thread : a grid-stride loop keeps a private running sum in a register;
//   2. warp   : 5 rounds of __shfl_down_sync combine 32 lanes without touching memory;
//   3. block  : each warp's lane 0 writes one partial to shared memory, warp 0 reduces those;
//   4. grid   : one atomicAdd per block onto the single output float.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

#include <algorithm>

// Sum of v across the 32 lanes of a warp; the full result ends up in lane 0.
// Round 1: lane i adds lane i+16; round 2: lane i adds lane i+8; ... 16, 8, 4, 2, 1.
// 0xffffffff = all 32 lanes take part (they must all reach this call).
__device__ __forceinline__ float warp_reduce_sum(float v) {
#pragma unroll
  for (int offset = 16; offset > 0; offset >>= 1) {
    v += __shfl_down_sync(0xffffffff, v, offset);
  }
  return v;
}

template <int BLOCK>
__global__ void reduce_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int64_t n) {
  static_assert(BLOCK % 32 == 0 && BLOCK <= 1024, "BLOCK must be a multiple of 32, at most 1024");
  constexpr int WARPS = BLOCK / 32;
  __shared__ float warp_sums[WARPS];  // one slot per warp in this block

  // 1. Thread level. Consecutive threads read consecutive elements on every iteration,
  //    so each warp-wide load is coalesced. Out-of-range threads simply add nothing.
  float v = 0.0f;
  const int64_t stride = (int64_t)BLOCK * gridDim.x;
  for (int64_t i = (int64_t)blockIdx.x * BLOCK + threadIdx.x; i < n; i += stride) {
    v += x[i];
  }

  // 2. Warp level.
  v = warp_reduce_sum(v);

  // 3. Block level: lane 0 of every warp publishes its warp's sum ...
  const int lane = threadIdx.x & 31;  // threadIdx.x % 32
  const int warp = threadIdx.x >> 5;  // threadIdx.x / 32
  if (lane == 0) warp_sums[warp] = v;
  __syncthreads();  // ... and every warp must have written before warp 0 reads

  // ... then the first warp reduces the WARPS partials (lanes >= WARPS contribute 0).
  if (warp == 0) {
    v = (lane < WARPS) ? warp_sums[lane] : 0.0f;
    v = warp_reduce_sum(v);
    // 4. Grid level: one atomic per block. Float addition is not associative, so the
    //    order in which blocks arrive can change the last bits of the result.
    if (lane == 0) atomicAdd(out, v);
  }
}

torch::Tensor reduce_sum(torch::Tensor x) {
  TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kFloat,
              "x must be a contiguous float32 CUDA tensor");
  auto out = torch::zeros({}, x.options());  // atomicAdd accumulates into this, so start at 0
  const int64_t n = x.numel();
  if (n == 0) return out;

  constexpr int BLOCK = 256;
  // A few blocks per SM is enough to saturate bandwidth; more blocks only add atomics.
  const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
  const int64_t blocks = std::min<int64_t>((n + BLOCK - 1) / BLOCK, (int64_t)sms * 4);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  reduce_sum_kernel<BLOCK><<<blocks, BLOCK, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return out;
}

// C3 · Matrix transpose — reference solution.
//
// y[c][r] = x[r][c] for x: [M, N] -> y: [N, M], both row-major float32.
//
// naive : each thread reads x[r][c] (coalesced: consecutive threads, consecutive c) and
//         writes y[c][r] (strided by M floats: every lane of the warp hits a different
//         memory segment). Writes become ~32 separate transactions per warp.
// tiled : a block stages a 32x32 tile in shared memory, then reads it back "sideways"
//         so that the global write is also coalesced. Shared memory is padded to 33
//         columns so the sideways read does not hit the same bank 32 times.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

constexpr int TILE = 32;  // tile edge = warp size: one warp covers one tile row
constexpr int ROWS = 8;   // threads per block in y; each thread copies TILE / ROWS = 4 rows

__global__ void transpose_naive_kernel(const float* __restrict__ x, float* __restrict__ y, int M, int N) {
  const int c = blockIdx.x * blockDim.x + threadIdx.x;  // column in x (fast index)
  const int r = blockIdx.y * blockDim.y + threadIdx.y;  // row in x
  if (r < M && c < N) {
    y[(int64_t)c * M + r] = x[(int64_t)r * N + c];
  }
}

__global__ void transpose_tiled_kernel(const float* __restrict__ x, float* __restrict__ y, int M, int N) {
  // +1 column of padding: element tile[i][j] lives in bank (i * 33 + j) % 32 = (i + j) % 32,
  // so a warp reading a column tile[0..31][j] touches 32 different banks.
  // Without padding, (i * 32 + j) % 32 = j: all 32 lanes hit one bank -> 32-way conflict.
  __shared__ float tile[TILE][TILE + 1];

  // Phase 1: coalesced read of x's tile (blockIdx.y, blockIdx.x) into shared memory.
  const int c = blockIdx.x * TILE + threadIdx.x;   // column in x
  const int r0 = blockIdx.y * TILE + threadIdx.y;  // first row in x handled by this thread
#pragma unroll
  for (int j = 0; j < TILE; j += ROWS) {
    const int r = r0 + j;
    if (r < M && c < N) tile[threadIdx.y + j][threadIdx.x] = x[(int64_t)r * N + c];
  }

  __syncthreads();  // all of the tile must be in shared memory before anyone reads it back

  // Phase 2: x's tile (by, bx) becomes y's tile (bx, by). Swap the block coordinates and
  // let threadIdx.x run along y's rows again, so the global write is coalesced too.
  const int yc = blockIdx.y * TILE + threadIdx.x;   // column in y  (= row in x)
  const int yr0 = blockIdx.x * TILE + threadIdx.y;  // row in y     (= column in x)
#pragma unroll
  for (int j = 0; j < TILE; j += ROWS) {
    const int yr = yr0 + j;
    // y[yr][yc] = x[yc][yr] = tile[yc - by*TILE][yr - bx*TILE] = tile[threadIdx.x][threadIdx.y + j]
    if (yr < N && yc < M) y[(int64_t)yr * M + yc] = tile[threadIdx.x][threadIdx.y + j];
  }
}

static void check_input(const torch::Tensor& x) {
  TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kFloat,
              "x must be a contiguous float32 CUDA tensor");
  TORCH_CHECK(x.dim() == 2, "x must be 2-D");
}

torch::Tensor transpose_naive(torch::Tensor x) {
  check_input(x);
  const int M = x.size(0), N = x.size(1);
  auto y = torch::empty({N, M}, x.options());
  if (M == 0 || N == 0) return y;
  dim3 block(32, 8);
  dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
  transpose_naive_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
      x.data_ptr<float>(), y.data_ptr<float>(), M, N);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return y;
}

torch::Tensor transpose_tiled(torch::Tensor x) {
  check_input(x);
  const int M = x.size(0), N = x.size(1);
  auto y = torch::empty({N, M}, x.options());
  if (M == 0 || N == 0) return y;
  dim3 block(TILE, ROWS);                                      // 256 threads
  dim3 grid((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);     // one block per 32x32 tile of x
  transpose_tiled_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
      x.data_ptr<float>(), y.data_ptr<float>(), M, N);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return y;
}

// C4 · GEMM: naive -> shared-memory tiling -> register tiling — reference solution.
//
// C[M, N] = A[M, K] @ B[K, N], all row-major float32, fp32 accumulation.
//
// Arithmetic intensity is what changes between versions:
//   naive   : every multiply-add loads one A and one B value from global memory (2 loads / FMA).
//   smem    : a 32x32 block loads a tile once into shared memory and reuses each value 32 times.
//   regtile : each thread computes a 4x4 patch of C from registers; every value read from
//             shared memory is used 4 times, so shared-memory traffic drops by 4x again.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

// ---------------------------------------------------------------- v0: naive
__global__ void matmul_naive_kernel(const float* __restrict__ A, const float* __restrict__ B,
                                    float* __restrict__ C, int M, int N, int K) {
  const int row = blockIdx.y * blockDim.y + threadIdx.y;
  const int col = blockIdx.x * blockDim.x + threadIdx.x;  // threadIdx.x -> column: B and C reads coalesce
  if (row >= M || col >= N) return;
  float acc = 0.0f;
  for (int k = 0; k < K; ++k) acc += A[(int64_t)row * K + k] * B[(int64_t)k * N + col];
  C[(int64_t)row * N + col] = acc;
}

// ---------------------------------------------------------------- v1: shared-memory tiling
constexpr int T = 32;  // block = T x T threads = one T x T tile of C

__global__ void matmul_smem_kernel(const float* __restrict__ A, const float* __restrict__ B,
                                   float* __restrict__ C, int M, int N, int K) {
  __shared__ float As[T][T];
  __shared__ float Bs[T][T];
  const int tx = threadIdx.x, ty = threadIdx.y;
  const int row = blockIdx.y * T + ty;
  const int col = blockIdx.x * T + tx;

  float acc = 0.0f;
  for (int k0 = 0; k0 < K; k0 += T) {
    // Every thread loads one element of each tile. Out-of-range elements become 0, which
    // contributes nothing to the dot product, so the inner loop needs no bounds checks.
    As[ty][tx] = (row < M && k0 + tx < K) ? A[(int64_t)row * K + k0 + tx] : 0.0f;
    Bs[ty][tx] = (k0 + ty < K && col < N) ? B[(int64_t)(k0 + ty) * N + col] : 0.0f;
    __syncthreads();  // tiles fully loaded before anyone uses them

#pragma unroll
    for (int k = 0; k < T; ++k) {
      // As[ty][k]: same address for the whole warp (same ty) -> broadcast, no conflict.
      // Bs[k][tx]: 32 consecutive floats -> 32 different banks, no conflict.
      acc += As[ty][k] * Bs[k][tx];
    }
    __syncthreads();  // everyone done with this tile before the next iteration overwrites it
  }
  if (row < M && col < N) C[(int64_t)row * N + col] = acc;
}

// ---------------------------------------------------------------- v2: register tiling
// Block tile BM x BN of C, K in steps of BK, each thread owns a TM x TN patch.
constexpr int BM = 64, BN = 64, BK = 16, TM = 4, TN = 4;
constexpr int THREADS = (BM / TM) * (BN / TN);  // 16 * 16 = 256

__global__ void __launch_bounds__(THREADS)
matmul_regtile_kernel(const float* __restrict__ A, const float* __restrict__ B,
                      float* __restrict__ C, int M, int N, int K) {
  // A's tile is stored transposed (As[k][m]) so the inner loop reads TM consecutive m values
  // for a fixed k; B's tile keeps its natural layout Bs[k][n].
  __shared__ float As[BK][BM];
  __shared__ float Bs[BK][BN];

  const int tid = threadIdx.x;
  const int tx = tid % (BN / TN);  // which TN-wide column strip of the block tile
  const int ty = tid / (BN / TN);  // which TM-tall row strip
  const int block_row = blockIdx.y * BM;
  const int block_col = blockIdx.x * BN;

  float acc[TM][TN] = {};  // the thread's TM x TN results live in registers
  float a_frag[TM], b_frag[TN];

  for (int k0 = 0; k0 < K; k0 += BK) {
    // Cooperative load of A[block_row : +BM, k0 : +BK] (BM*BK = 1024 floats, 4 per thread).
    // e -> (m, k) with k fastest, so consecutive threads read consecutive addresses of A.
    for (int e = tid; e < BM * BK; e += THREADS) {
      const int m = e / BK, k = e % BK;
      const int gm = block_row + m, gk = k0 + k;
      As[k][m] = (gm < M && gk < K) ? A[(int64_t)gm * K + gk] : 0.0f;
    }
    // Cooperative load of B[k0 : +BK, block_col : +BN], n fastest -> coalesced.
    for (int e = tid; e < BK * BN; e += THREADS) {
      const int k = e / BN, n = e % BN;
      const int gk = k0 + k, gn = block_col + n;
      Bs[k][n] = (gk < K && gn < N) ? B[(int64_t)gk * N + gn] : 0.0f;
    }
    __syncthreads();

#pragma unroll
    for (int k = 0; k < BK; ++k) {
      // Load a column of A's tile and a row of B's tile into registers ...
#pragma unroll
      for (int i = 0; i < TM; ++i) a_frag[i] = As[k][ty * TM + i];
#pragma unroll
      for (int j = 0; j < TN; ++j) b_frag[j] = Bs[k][tx * TN + j];
      // ... and take their outer product: TM*TN FMAs for TM+TN shared-memory reads.
#pragma unroll
      for (int i = 0; i < TM; ++i)
#pragma unroll
        for (int j = 0; j < TN; ++j) acc[i][j] += a_frag[i] * b_frag[j];
    }
    __syncthreads();
  }

#pragma unroll
  for (int i = 0; i < TM; ++i) {
    const int gm = block_row + ty * TM + i;
    if (gm >= M) continue;
#pragma unroll
    for (int j = 0; j < TN; ++j) {
      const int gn = block_col + tx * TN + j;
      if (gn < N) C[(int64_t)gm * N + gn] = acc[i][j];
    }
  }
}

// ---------------------------------------------------------------- host wrappers
static void check_inputs(const torch::Tensor& A, const torch::Tensor& B) {
  for (const auto* t : {&A, &B}) {
    TORCH_CHECK(t->is_cuda() && t->is_contiguous() && t->scalar_type() == at::kFloat,
                "inputs must be contiguous float32 CUDA tensors");
    TORCH_CHECK(t->dim() == 2, "inputs must be 2-D");
  }
  TORCH_CHECK(A.size(1) == B.size(0), "inner dimensions must match");
}

torch::Tensor matmul_naive(torch::Tensor A, torch::Tensor B) {
  check_inputs(A, B);
  const int M = A.size(0), K = A.size(1), N = B.size(1);
  auto C = torch::empty({M, N}, A.options());
  if (M == 0 || N == 0) return C;
  dim3 block(32, 8);
  dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
  matmul_naive_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
      A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return C;
}

torch::Tensor matmul_smem(torch::Tensor A, torch::Tensor B) {
  check_inputs(A, B);
  const int M = A.size(0), K = A.size(1), N = B.size(1);
  auto C = torch::empty({M, N}, A.options());
  if (M == 0 || N == 0) return C;
  dim3 block(T, T);  // 1024 threads: the per-block maximum
  dim3 grid((N + T - 1) / T, (M + T - 1) / T);
  matmul_smem_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
      A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return C;
}

torch::Tensor matmul_regtile(torch::Tensor A, torch::Tensor B) {
  check_inputs(A, B);
  const int M = A.size(0), K = A.size(1), N = B.size(1);
  auto C = torch::empty({M, N}, A.options());
  if (M == 0 || N == 0) return C;
  dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
  matmul_regtile_kernel<<<grid, THREADS, 0, at::cuda::getCurrentCUDAStream()>>>(
      A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return C;
}

"""T1 · Fused add + GELU in Triton — reference solution.

out = gelu_tanh(x + y), any shape, fp32 / fp16 / bf16. Math in fp32, result in x.dtype.

Triton's programming model in one kernel:
  * a "program" (≈ a CUDA thread block) handles BLOCK_SIZE consecutive elements;
  * tl.arange gives the program a *vector* of offsets — you write block-level code,
    the compiler spreads the vector across the block's threads;
  * every load/store past the end is disabled with a mask.
"""
import torch
import triton
import triton.language as tl


@triton.jit
def add_gelu_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)                       # which block of elements this program owns
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)  # [BLOCK_SIZE] element offsets
    mask = offs < n_elements                          # the last program is usually partial

    # Load, upcast to fp32 (fp16 x^3 overflows/loses precision quickly).
    x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32)
    y = tl.load(y_ptr + offs, mask=mask, other=0.0).to(tl.float32)
    z = x + y

    # GELU (tanh approximation): 0.5 z (1 + tanh(sqrt(2/pi) (z + 0.044715 z^3))).
    # tanh(u) = 2 sigmoid(2u) - 1; tl.sigmoid is available on every backend.
    inner = 0.7978845608028654 * (z + 0.044715 * z * z * z)
    t = 2.0 * tl.sigmoid(2.0 * inner) - 1.0
    out = 0.5 * z * (1.0 + t)

    # Cast back to the output's element type before storing.
    tl.store(out_ptr + offs, out.to(out_ptr.dtype.element_ty), mask=mask)


def add_gelu(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024) -> torch.Tensor:
    assert x.shape == y.shape and x.dtype == y.dtype and x.device == y.device
    # The kernel indexes memory as a flat array, so it needs contiguous inputs.
    x, y = x.contiguous(), y.contiguous()
    out = torch.empty_like(x)
    n = x.numel()
    if n == 0:
        return out
    # One program per BLOCK_SIZE elements. grid can also be a lambda over the meta-params:
    #   grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
    grid = (triton.cdiv(n, block_size),)
    add_gelu_kernel[grid](x, y, out, n, BLOCK_SIZE=block_size)
    return out

"""T2 · Fused row softmax in Triton — reference solution.

softmax(x) over the last dim of a 2-D tensor [M, N]. One program per row; the whole row
fits in one block (BLOCK_SIZE = next_power_of_2(N)), so x is read once and y written once.
PyTorch's unfused version reads x for max, again for exp/sum, again to normalise.
"""
import torch
import triton
import triton.language as tl


@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_row_stride, out_row_stride, n_cols, BLOCK_SIZE: tl.constexpr):
    # program_id is int32; promote before multiplying by a stride so row * stride cannot
    # overflow for tensors with more than 2^31 elements.
    row = tl.program_id(0).to(tl.int64)
    cols = tl.arange(0, BLOCK_SIZE)        # BLOCK_SIZE is a power of two >= n_cols
    mask = cols < n_cols

    # Padding lanes load -inf: they lose the max, and exp(-inf) = 0 adds nothing to the sum.
    x = tl.load(in_ptr + row * in_row_stride + cols, mask=mask, other=-float("inf")).to(tl.float32)

    # Numerically stable softmax: softmax(x) = softmax(x - max(x)); the largest exponent is 0,
    # so exp never overflows even for fp16 inputs around 1e4.
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    den = tl.sum(num, axis=0)
    y = num / den

    tl.store(out_ptr + row * out_row_stride + cols, y.to(out_ptr.dtype.element_ty), mask=mask)


def softmax(x: torch.Tensor) -> torch.Tensor:
    assert x.dim() == 2, "expects [M, N]"
    M, N = x.shape
    # The kernel walks a row with unit stride; the rows themselves may be strided.
    if x.stride(1) != 1:
        x = x.contiguous()
    y = torch.empty((M, N), dtype=x.dtype, device=x.device)
    if M == 0 or N == 0:
        return y
    BLOCK_SIZE = triton.next_power_of_2(N)
    # More warps for wider rows, so each thread holds a reasonable number of elements.
    num_warps = 4 if BLOCK_SIZE <= 2048 else (8 if BLOCK_SIZE <= 8192 else 16)
    softmax_kernel[(M,)](y, x, x.stride(0), y.stride(0), N, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps)
    return y

"""T3 · LayerNorm / RMSNorm forward in Triton — reference solution.

layer_norm(x[M, N], w, b, eps) -> (y, mean[M], rstd[M])
    y = (x - mean) * rstd * w + b,  rstd = 1 / sqrt(var + eps),  var is the biased variance.
rms_norm(x[M, N], w, eps) -> y
    y = x * rsqrt(mean(x^2) + eps) * w

One program per row, the full row in one block (N <= 16384). mean and rstd are returned
because the backward pass needs exactly these two numbers per row.
"""
import torch
import triton
import triton.language as tl


@triton.jit
def layer_norm_kernel(x_ptr, w_ptr, b_ptr, y_ptr, mean_ptr, rstd_ptr,
                      x_row_stride, y_row_stride, N, eps, BLOCK_N: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < N

    x = tl.load(x_ptr + row * x_row_stride + cols, mask=mask, other=0.0).to(tl.float32)

    # Statistics in fp32. Padding lanes loaded 0, so they do not disturb the sum ...
    mean = tl.sum(x, axis=0) / N
    # ... but x - mean is NOT zero there: zero it again before squaring.
    xc = tl.where(mask, x - mean, 0.0)
    var = tl.sum(xc * xc, axis=0) / N          # two-pass variance: stabler than E[x^2] - E[x]^2
    rstd = 1.0 / tl.sqrt(var + eps)

    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    b = tl.load(b_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    y = xc * rstd * w + b

    tl.store(y_ptr + row * y_row_stride + cols, y.to(y_ptr.dtype.element_ty), mask=mask)
    tl.store(mean_ptr + row, mean)             # scalar stores: one value per program
    tl.store(rstd_ptr + row, rstd)


@triton.jit
def rms_norm_kernel(x_ptr, w_ptr, y_ptr, x_row_stride, y_row_stride, N, eps, BLOCK_N: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < N
    x = tl.load(x_ptr + row * x_row_stride + cols, mask=mask, other=0.0).to(tl.float32)
    rstd = 1.0 / tl.sqrt(tl.sum(x * x, axis=0) / N + eps)   # no mean subtraction, no bias
    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    tl.store(y_ptr + row * y_row_stride + cols, (x * rstd * w).to(y_ptr.dtype.element_ty), mask=mask)


def _launch_params(N: int) -> tuple[int, int]:
    BLOCK_N = triton.next_power_of_2(N)
    if BLOCK_N > 16384:
        raise ValueError("this single-block kernel supports N <= 16384; loop over the row for more")
    num_warps = min(max(BLOCK_N // 256, 1), 8)   # ~256 elements per warp, 1..8 warps
    return BLOCK_N, num_warps


def layer_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor, eps: float = 1e-5):
    shape = x.shape
    x2 = x.reshape(-1, shape[-1])                # [..., N] -> [M, N]
    if x2.stride(-1) != 1:
        x2 = x2.contiguous()
    M, N = x2.shape
    y = torch.empty_like(x2)
    mean = torch.empty(M, dtype=torch.float32, device=x.device)
    rstd = torch.empty(M, dtype=torch.float32, device=x.device)
    if M > 0:
        BLOCK_N, num_warps = _launch_params(N)
        layer_norm_kernel[(M,)](x2, w, b, y, mean, rstd, x2.stride(0), y.stride(0), N, eps,
                                BLOCK_N=BLOCK_N, num_warps=num_warps)
    return y.reshape(shape), mean, rstd


def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    shape = x.shape
    x2 = x.reshape(-1, shape[-1])
    if x2.stride(-1) != 1:
        x2 = x2.contiguous()
    M, N = x2.shape
    y = torch.empty_like(x2)
    if M > 0:
        BLOCK_N, num_warps = _launch_params(N)
        rms_norm_kernel[(M,)](x2, w, y, x2.stride(0), y.stride(0), N, eps, BLOCK_N=BLOCK_N, num_warps=num_warps)
    return y.reshape(shape)

"""T4 · Tiled GEMM with tl.dot, grouped ordering and autotuning — reference solution.

C[M, N] = A[M, K] @ B[K, N]; A, B fp16 or bf16, fp32 accumulation, C in the input dtype.
Each program computes one BLOCK_M x BLOCK_N tile of C, walking K in BLOCK_K steps.
tl.dot lowers to Tensor Core instructions (mma / wgmma) on NVIDIA GPUs.
"""
import os

import torch
import triton
import triton.language as tl


@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
                  GROUP_M: tl.constexpr):
    # ---- 1. which C tile does this program own? (grouped / "super-grouping" order)
    # Plain row-major order (pid_m = pid // num_pid_n) makes consecutive programs sweep a
    # whole row of tiles: they share one A row-panel but each needs a different B column-panel.
    # Grouping GROUP_M rows of tiles together and going down the columns first means programs
    # running at the same time reuse both A and B panels while they are still in L2.
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)   # the last group may be shorter
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ---- 2. pointers to the first A and B tiles (2-D blocks of pointers via broadcasting)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak   # [BLOCK_M, BLOCK_K]
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn   # [BLOCK_K, BLOCK_N]

    # ---- 3. main loop over K: accumulate in fp32 registers
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        k_remaining = K - k * BLOCK_K
        # Masked loads return 0 outside the matrix; zeros add nothing to the dot product.
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)          # acc += a @ b on Tensor Cores
        a_ptrs += BLOCK_K * stride_ak    # slide both tiles along K
        b_ptrs += BLOCK_K * stride_bk

    # ---- 4. epilogue: cast and store the C tile
    c = acc.to(c_ptr.dtype.element_ty)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, c, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


# Candidate configs for the autotuner. num_stages = software-pipelining depth (how many K tiles
# are prefetched into shared memory ahead of use); num_warps = threads per program / 32.
AUTOTUNE_CONFIGS = [
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8}, num_stages=3, num_warps=8),
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_stages=4, num_warps=4),
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_stages=4, num_warps=4),
    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_stages=4, num_warps=4),
    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_stages=5, num_warps=2),
]

# autotune wraps the same jitted kernel: the first call for each new (M, N, K) benchmarks all
# configs and caches the fastest one; later calls with that key reuse it.
matmul_kernel_autotuned = triton.autotune(configs=AUTOTUNE_CONFIGS, key=["M", "N", "K"])(matmul_kernel)


def matmul(a: torch.Tensor, b: torch.Tensor, config: dict | None = None) -> torch.Tensor:
    """config=None -> autotuned on GPU; pass e.g. {"BLOCK_M": 64, ...} to force one config."""
    assert a.dim() == 2 and b.dim() == 2 and a.shape[1] == b.shape[0], "shape mismatch"
    assert a.dtype == b.dtype and a.dtype in (torch.float16, torch.bfloat16)
    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), dtype=a.dtype, device=a.device)
    if M == 0 or N == 0:
        return c
    args = (a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1))

    # The CPU interpreter cannot time configs, so it always uses one fixed config.
    if config is None and os.environ.get("TRITON_INTERPRET") == "1":
        config = {"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}
    if config is not None:
        grid = (triton.cdiv(M, config["BLOCK_M"]) * triton.cdiv(N, config["BLOCK_N"]),)
        matmul_kernel[grid](*args, **config)
    else:
        # A 1-D grid sized from whichever config the autotuner is trying.
        grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
        matmul_kernel_autotuned[grid](*args)
    return c

"""T5 · FlashAttention forward (tiled + online softmax) in Triton — reference solution.

o = softmax(q k^T * sm_scale [+ causal mask]) v, q/k/v/o: [B, H, S, D] fp16 or bf16.
One program = BLOCK_M queries of one (batch, head). It streams K/V in BLOCK_N chunks and
never materialises the S x S score matrix: only a BLOCK_M x BLOCK_N tile exists at a time.

Online softmax state per query row i (all fp32):
    m_i  running max of the scores seen so far
    l_i  running sum of exp(score - m_i)
    acc  running sum of exp(score - m_i) * v
When a new chunk raises the max to m_new, old terms are rescaled by alpha = exp(m_i - m_new).
At the end o_i = acc_i / l_i.
"""
import math

import torch
import triton
import triton.language as tl


@triton.jit
def flash_attn_fwd_kernel(Q, K, V, O, sm_scale,
                          stride_b, stride_h, stride_s,       # same strides for q, k, v, o; stride_d == 1
                          H, S,
                          CAUSAL: tl.constexpr, HEAD_DIM: tl.constexpr,
                          BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    start_m = tl.program_id(0)                 # which block of queries
    bh = tl.program_id(1)                      # which (batch, head)
    off_b = bh // H
    off_h = bh % H
    base = off_b.to(tl.int64) * stride_b + off_h.to(tl.int64) * stride_h

    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)    # query positions
    offs_n = tl.arange(0, BLOCK_N)                        # key positions within a chunk
    offs_d = tl.arange(0, HEAD_DIM)

    # The query tile is loaded once and stays in registers for the whole loop.
    q = tl.load(Q + base + offs_m[:, None] * stride_s + offs_d[None, :],
                mask=offs_m[:, None] < S, other=0.0)

    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)

    # Causal: queries in this block never look past their own last position, so key chunks
    # beyond (start_m + 1) * BLOCK_M are skipped entirely (about half the work).
    if CAUSAL:
        hi = tl.minimum((start_m + 1) * BLOCK_M, S)
    else:
        hi = S

    for start_n in range(0, hi, BLOCK_N):
        cols = start_n + offs_n
        # Load K already transposed: a [HEAD_DIM, BLOCK_N] tile, so q @ k is a plain tl.dot.
        k = tl.load(K + base + cols[None, :] * stride_s + offs_d[:, None],
                    mask=cols[None, :] < S, other=0.0)
        s = tl.dot(q, k) * sm_scale                        # [BLOCK_M, BLOCK_N] fp32 scores

        valid = cols[None, :] < S                          # keys past the end of the sequence
        if CAUSAL:
            valid = valid & (offs_m[:, None] >= cols[None, :])
        s = tl.where(valid, s, float("-inf"))

        # Online softmax update. Every real query row has at least one valid key (key 0) in
        # the first chunk, so m_new is finite from the first iteration on and no
        # (-inf) - (-inf) = NaN can appear.
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        p = tl.exp(s - m_new[:, None])                     # [BLOCK_M, BLOCK_N]
        alpha = tl.exp(m_i - m_new)                        # rescales the old state; 0 on the first chunk
        l_i = l_i * alpha + tl.sum(p, axis=1)

        v = tl.load(V + base + cols[:, None] * stride_s + offs_d[None, :],
                    mask=cols[:, None] < S, other=0.0)
        acc = acc * alpha[:, None]
        acc = tl.dot(p.to(v.dtype), v, acc)                # acc += p @ v (p in fp16/bf16 for Tensor Cores)
        m_i = m_new

    o = acc / l_i[:, None]
    tl.store(O + base + offs_m[:, None] * stride_s + offs_d[None, :],
             o.to(O.dtype.element_ty), mask=offs_m[:, None] < S)


def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False,
                    sm_scale: float | None = None, block_m: int = 64, block_n: int = 64) -> torch.Tensor:
    assert q.shape == k.shape == v.shape and q.dim() == 4, "q, k, v: [B, H, S, D]"
    assert q.dtype == k.dtype == v.dtype and q.dtype in (torch.float16, torch.bfloat16)
    B, H, S, D = q.shape
    assert D in (16, 32, 64, 128), "HEAD_DIM must be a power of two in [16, 128]"
    q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
    o = torch.empty_like(q)
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(D)
    grid = (triton.cdiv(S, block_m), B * H)
    flash_attn_fwd_kernel[grid](q, k, v, o, sm_scale, q.stride(0), q.stride(1), q.stride(2), H, S,
                                CAUSAL=causal, HEAD_DIM=D, BLOCK_M=block_m, BLOCK_N=block_n,
                                num_warps=4, num_stages=2)
    return o

"""L1 · Paged KV cache: slot mapping + cache write (vLLM's reshape_and_cache) — reference solution.

Cache layout (one tensor each for K and V):
    k_cache, v_cache : [num_blocks, block_size, num_kv_heads, head_dim]
A sequence's token at logical position p lives in physical block block_table[seq][p // block_size]
at offset p % block_size. Flattening (block, offset) gives a "slot":
    slot = physical_block * block_size + offset
and because the cache is contiguous, slot s starts at element s * num_kv_heads * head_dim.
"""
import torch
import triton
import triton.language as tl


def slot_mapping(block_table: torch.Tensor, seq_idx: torch.Tensor, positions: torch.Tensor,
                 block_size: int) -> torch.Tensor:
    """block_table [B, max_blocks] int32, seq_idx [T], positions [T] -> slots [T] int64.

    Pure PyTorch, vectorised: no Python loop over tokens.
    """
    seq_idx = seq_idx.long()
    positions = positions.long()
    logical_block = positions // block_size                                  # which page of the sequence
    physical_block = block_table[seq_idx, logical_block].long()              # gather: [T]
    return physical_block * block_size + positions % block_size


@triton.jit
def reshape_and_cache_kernel(key_ptr, value_ptr, k_cache_ptr, v_cache_ptr, slot_ptr,
                             key_stride, value_stride, HD, BLOCK: tl.constexpr):
    t = tl.program_id(0).to(tl.int64)          # one program per new token
    slot = tl.load(slot_ptr + t)               # int64; -1 marks padding tokens to skip
    offs = tl.arange(0, BLOCK)                 # the token's H*D values, flattened
    # Padding token -> 0 valid elements, every lane masked off. (Written as one comparison
    # rather than `(offs < HD) & (slot >= 0)`: same meaning, and the CPU interpreter rejects
    # a mask built by and-ing a scalar with a block.)
    n_valid = tl.where(slot >= 0, HD, 0)
    mask = offs < n_valid

    # key/value rows may be strided (e.g. a view into a fused qkv tensor), but each row's
    # [H, D] part is contiguous.
    k = tl.load(key_ptr + t * key_stride + offs, mask=mask)
    v = tl.load(value_ptr + t * value_stride + offs, mask=mask)
    # The cache is contiguous: slot s occupies elements [s * HD, (s + 1) * HD).
    tl.store(k_cache_ptr + slot * HD + offs, k, mask=mask)
    tl.store(v_cache_ptr + slot * HD + offs, v, mask=mask)


def reshape_and_cache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor,
                      v_cache: torch.Tensor, slots: torch.Tensor) -> None:
    """key/value [T, H, D] -> written in place into k_cache/v_cache [NB, BS, H, D] at `slots`."""
    T, H, D = key.shape
    assert value.shape == key.shape
    assert k_cache.is_contiguous() and v_cache.is_contiguous()
    assert k_cache.shape[2:] == (H, D) and k_cache.dtype == key.dtype
    for t in (key, value):
        assert t.stride(2) == 1 and t.stride(1) == D, "each token's [H, D] must be contiguous"
    if T == 0:
        return
    HD = H * D
    reshape_and_cache_kernel[(T,)](key, value, k_cache, v_cache, slots.long(),
                                   key.stride(0), value.stride(0), HD, BLOCK=triton.next_power_of_2(HD))

"""L2 · Paged decode attention with GQA — reference solution.

Decode step: each sequence has ONE new query token per head, and attends over its whole
history stored in the paged KV cache.
    q            [B, Hq, D]
    k/v_cache    [num_blocks, PAGE, Hkv, D]       (Hq = GROUP * Hkv)
    block_tables [B, max_blocks] int32, context_lens [B] int32 (>= 1)
    out          [B, Hq, D]

One program per (sequence, kv head). It handles the GROUP query heads that share this kv head
together, so every K/V page is read from memory once for the whole group — decode is
memory-bound, and this is the main saving GQA offers. The group is padded to >= 16 rows so
q @ k^T can use tl.dot (Tensor Cores need M >= 16).
"""
import math

import torch
import triton
import triton.language as tl


@triton.jit
def paged_decode_kernel(Q, K_cache, V_cache, block_tables, context_lens, O, scale,
                        stride_qb, stride_qh,
                        stride_cb, stride_ct, stride_ch,          # cache strides: block, token, head (dim stride 1)
                        stride_btb,                               # block_tables row stride
                        stride_ob, stride_oh,
                        GROUP: tl.constexpr, GROUP_PAD: tl.constexpr,
                        PAGE: tl.constexpr, HEAD_DIM: tl.constexpr):
    b = tl.program_id(0).to(tl.int64)
    kvh = tl.program_id(1).to(tl.int64)
    ctx = tl.load(context_lens + b)

    offs_g = tl.arange(0, GROUP_PAD)                 # query heads within the group (padded)
    offs_t = tl.arange(0, PAGE)                      # token offsets within a page
    offs_d = tl.arange(0, HEAD_DIM)
    g_mask = offs_g < GROUP
    q_heads = kvh * GROUP + offs_g                   # query heads served by this kv head

    q = tl.load(Q + b * stride_qb + q_heads[:, None] * stride_qh + offs_d[None, :],
                mask=g_mask[:, None], other=0.0)     # [GROUP_PAD, D]; padded rows are 0

    m_i = tl.full([GROUP_PAD], float("-inf"), dtype=tl.float32)
    l_i = tl.zeros([GROUP_PAD], dtype=tl.float32)
    acc = tl.zeros([GROUP_PAD, HEAD_DIM], dtype=tl.float32)

    num_pages = tl.cdiv(ctx, PAGE)
    for p in range(0, num_pages):
        # Indirection: logical page p of this sequence -> physical block in the cache.
        phys = tl.load(block_tables + b * stride_btb + p).to(tl.int64)
        tok = p * PAGE + offs_t
        t_mask = tok < ctx                           # the last page is usually partial

        kv_off = phys * stride_cb + offs_t[:, None] * stride_ct + kvh * stride_ch + offs_d[None, :]
        k = tl.load(K_cache + kv_off, mask=t_mask[:, None], other=0.0)   # [PAGE, D]
        v = tl.load(V_cache + kv_off, mask=t_mask[:, None], other=0.0)   # [PAGE, D]

        s = tl.dot(q, tl.trans(k)) * scale           # [GROUP_PAD, PAGE] fp32
        s = tl.where(t_mask[None, :], s, float("-inf"))

        # Same online-softmax update as FlashAttention (T5). Page 0 always has token 0, so
        # m_new is finite from the first iteration.
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        p_ = tl.exp(s - m_new[:, None])
        alpha = tl.exp(m_i - m_new)
        l_i = l_i * alpha + tl.sum(p_, axis=1)
        acc = acc * alpha[:, None]
        acc = tl.dot(p_.to(v.dtype), v, acc)
        m_i = m_new

    out = acc / l_i[:, None]
    tl.store(O + b * stride_ob + q_heads[:, None] * stride_oh + offs_d[None, :],
             out.to(O.dtype.element_ty), mask=g_mask[:, None])


def paged_decode_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                           block_tables: torch.Tensor, context_lens: torch.Tensor,
                           scale: float | None = None) -> torch.Tensor:
    B, Hq, D = q.shape
    NB, PAGE, Hkv, D2 = k_cache.shape
    assert D == D2 and v_cache.shape == k_cache.shape and Hq % Hkv == 0
    assert k_cache.is_contiguous() and v_cache.is_contiguous() and q.stride(-1) == 1
    assert PAGE >= 16 and PAGE & (PAGE - 1) == 0, "page size: a power of two >= 16 (tl.dot)"
    assert D >= 16 and D & (D - 1) == 0, "head_dim: a power of two >= 16"
    GROUP = Hq // Hkv
    GROUP_PAD = max(16, triton.next_power_of_2(GROUP))
    if scale is None:
        scale = 1.0 / math.sqrt(D)
    out = torch.empty_like(q)
    paged_decode_kernel[(B, Hkv)](
        q, k_cache, v_cache, block_tables, context_lens, out, scale,
        q.stride(0), q.stride(1),
        k_cache.stride(0), k_cache.stride(1), k_cache.stride(2),
        block_tables.stride(0),
        out.stride(0), out.stride(1),
        GROUP=GROUP, GROUP_PAD=GROUP_PAD, PAGE=PAGE, HEAD_DIM=D, num_warps=4)
    return out

"""L3 · Split-KV decode attention (Flash-Decoding) + log-sum-exp merge — reference solution.

L2 launches B * Hkv programs. With a small batch and a long context that is far fewer
programs than the GPU has SMs, and each one loops over thousands of tokens: the GPU sits
mostly idle. Split-KV also splits the context of each sequence into NUM_SPLITS ranges:

  pass 1 (grid B x Hkv x NUM_SPLITS): attention over one range of pages -> a partial output
          o_s (already normalised within the range) and its log-sum-exp lse_s = m_s + log(l_s)
  pass 2 (grid B x Hq): merge the partials exactly:
          w_s = exp(lse_s - max_s lse_s),  out = sum_s w_s * o_s / sum_s w_s
An empty range (short sequences) gets lse = -inf and weight 0.
"""
import math

import torch
import triton
import triton.language as tl


@triton.jit
def split_kv_partial_kernel(Q, K_cache, V_cache, block_tables, context_lens, O_part, LSE, scale,
                            stride_qb, stride_qh,
                            stride_cb, stride_ct, stride_ch,
                            stride_btb,
                            stride_pb, stride_ph, stride_ps,      # O_part [B, Hq, NUM_SPLITS, D] fp32
                            stride_lb, stride_lh,                 # LSE    [B, Hq, NUM_SPLITS]    fp32
                            NUM_SPLITS: tl.constexpr,
                            GROUP: tl.constexpr, GROUP_PAD: tl.constexpr,
                            PAGE: tl.constexpr, HEAD_DIM: tl.constexpr):
    b = tl.program_id(0).to(tl.int64)
    kvh = tl.program_id(1).to(tl.int64)
    split = tl.program_id(2)
    ctx = tl.load(context_lens + b)

    # This program's contiguous range of pages: [p_start, p_end).
    num_pages = tl.cdiv(ctx, PAGE)
    pages_per_split = tl.cdiv(num_pages, NUM_SPLITS)
    p_start = split * pages_per_split
    p_end = tl.minimum(p_start + pages_per_split, num_pages)   # may be <= p_start: empty split

    offs_g = tl.arange(0, GROUP_PAD)
    offs_t = tl.arange(0, PAGE)
    offs_d = tl.arange(0, HEAD_DIM)
    g_mask = offs_g < GROUP
    q_heads = kvh * GROUP + offs_g
    q = tl.load(Q + b * stride_qb + q_heads[:, None] * stride_qh + offs_d[None, :],
                mask=g_mask[:, None], other=0.0)

    m_i = tl.full([GROUP_PAD], float("-inf"), dtype=tl.float32)
    l_i = tl.zeros([GROUP_PAD], dtype=tl.float32)
    acc = tl.zeros([GROUP_PAD, HEAD_DIM], dtype=tl.float32)

    for p in range(p_start, p_end):                  # identical body to L2
        phys = tl.load(block_tables + b * stride_btb + p).to(tl.int64)
        tok = p * PAGE + offs_t
        t_mask = tok < ctx
        kv_off = phys * stride_cb + offs_t[:, None] * stride_ct + kvh * stride_ch + offs_d[None, :]
        k = tl.load(K_cache + kv_off, mask=t_mask[:, None], other=0.0)
        v = tl.load(V_cache + kv_off, mask=t_mask[:, None], other=0.0)
        s = tl.dot(q, tl.trans(k)) * scale
        s = tl.where(t_mask[None, :], s, float("-inf"))
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        p_ = tl.exp(s - m_new[:, None])
        alpha = tl.exp(m_i - m_new)
        l_i = l_i * alpha + tl.sum(p_, axis=1)
        acc = acc * alpha[:, None]
        acc = tl.dot(p_.to(v.dtype), v, acc)
        m_i = m_new

    # Finalise this range. An empty range has l_i == 0: write o = 0 and lse = -inf so the
    # merge gives it weight exp(-inf) = 0 (and never divide by zero).
    has = l_i > 0
    o = acc / tl.where(has, l_i, 1.0)[:, None]
    lse = tl.where(has, m_i + tl.log(tl.where(has, l_i, 1.0)), float("-inf"))

    tl.store(O_part + b * stride_pb + q_heads[:, None] * stride_ph + split * stride_ps + offs_d[None, :],
             o, mask=g_mask[:, None])
    tl.store(LSE + b * stride_lb + q_heads * stride_lh + split, lse, mask=g_mask)


@triton.jit
def merge_kernel(O_part, LSE, O,
                 stride_pb, stride_ph, stride_ps, stride_lb, stride_lh, stride_ob, stride_oh,
                 NUM_SPLITS: tl.constexpr, SPLIT_PAD: tl.constexpr, HEAD_DIM: tl.constexpr):
    b = tl.program_id(0).to(tl.int64)
    h = tl.program_id(1).to(tl.int64)
    offs_s = tl.arange(0, SPLIT_PAD)
    offs_d = tl.arange(0, HEAD_DIM)
    s_mask = offs_s < NUM_SPLITS

    lse = tl.load(LSE + b * stride_lb + h * stride_lh + offs_s, mask=s_mask, other=float("-inf"))
    m = tl.max(lse, axis=0)                  # finite: context_len >= 1, so some split is non-empty
    w = tl.exp(lse - m)                      # [SPLIT_PAD]; empty or padded splits -> 0
    o = tl.load(O_part + b * stride_pb + h * stride_ph + offs_s[:, None] * stride_ps + offs_d[None, :],
                mask=s_mask[:, None], other=0.0)                       # [SPLIT_PAD, D]
    out = tl.sum(w[:, None] * o, axis=0) / tl.sum(w, axis=0)
    tl.store(O + b * stride_ob + h * stride_oh + offs_d, out.to(O.dtype.element_ty))


def split_kv_decode_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                              block_tables: torch.Tensor, context_lens: torch.Tensor,
                              num_splits: int = 4, scale: float | None = None) -> torch.Tensor:
    B, Hq, D = q.shape
    NB, PAGE, Hkv, D2 = k_cache.shape
    assert D == D2 and v_cache.shape == k_cache.shape and Hq % Hkv == 0
    assert k_cache.is_contiguous() and v_cache.is_contiguous() and q.stride(-1) == 1
    assert PAGE >= 16 and PAGE & (PAGE - 1) == 0 and D >= 16 and D & (D - 1) == 0
    assert num_splits >= 1
    GROUP = Hq // Hkv
    GROUP_PAD = max(16, triton.next_power_of_2(GROUP))
    if scale is None:
        scale = 1.0 / math.sqrt(D)

    # Partial results stay in fp32: they are combined again, so keep the extra precision.
    o_part = torch.empty((B, Hq, num_splits, D), dtype=torch.float32, device=q.device)
    lse = torch.empty((B, Hq, num_splits), dtype=torch.float32, device=q.device)
    split_kv_partial_kernel[(B, Hkv, num_splits)](
        q, k_cache, v_cache, block_tables, context_lens, o_part, lse, scale,
        q.stride(0), q.stride(1),
        k_cache.stride(0), k_cache.stride(1), k_cache.stride(2),
        block_tables.stride(0),
        o_part.stride(0), o_part.stride(1), o_part.stride(2),
        lse.stride(0), lse.stride(1),
        NUM_SPLITS=num_splits, GROUP=GROUP, GROUP_PAD=GROUP_PAD, PAGE=PAGE, HEAD_DIM=D, num_warps=4)

    out = torch.empty_like(q)
    merge_kernel[(B, Hq)](
        o_part, lse, out,
        o_part.stride(0), o_part.stride(1), o_part.stride(2), lse.stride(0), lse.stride(1),
        out.stride(0), out.stride(1),
        NUM_SPLITS=num_splits, SPLIT_PAD=max(2, triton.next_power_of_2(num_splits)), HEAD_DIM=D, num_warps=1)
    return out

"""F1 · FlashInfer's paged-KV format (CSR) + plan/run decode — reference solution.

vLLM-style block table (what L1-L3 used):
    block_tables [B, max_blocks] int32 (padded), context_lens [B] int32
FlashInfer's CSR format (no padding; each sequence's pages are contiguous in `indices`):
    indptr        [B + 1] int32   pages of sequence b are indices[indptr[b] : indptr[b+1]]
    indices       [nnz]   int32   physical page ids, concatenated
    last_page_len [B]     int32   tokens used in each sequence's last page, in [1, page_size]
so  context_len[b] = (indptr[b+1] - indptr[b] - 1) * page_size + last_page_len[b].
"""
import torch


def block_table_to_csr(block_tables: torch.Tensor, context_lens: torch.Tensor, page_size: int):
    """Vectorised (no Python loop over sequences). Requires context_lens >= 1."""
    B, max_blocks = block_tables.shape
    ctx = context_lens.long()
    num_pages = (ctx + page_size - 1) // page_size                       # [B]
    indptr = torch.zeros(B + 1, dtype=torch.int32, device=block_tables.device)
    indptr[1:] = torch.cumsum(num_pages, 0).to(torch.int32)
    # Keep the first num_pages[b] entries of row b. Boolean indexing walks the tensor in
    # row-major order, so each sequence's pages stay together and in order.
    keep = torch.arange(max_blocks, device=block_tables.device)[None, :] < num_pages[:, None]
    indices = block_tables[keep].to(torch.int32)
    last_page_len = (ctx - (num_pages - 1) * page_size).to(torch.int32)   # == (ctx - 1) % P + 1
    return indptr, indices, last_page_len


def csr_to_block_table(indptr: torch.Tensor, indices: torch.Tensor, last_page_len: torch.Tensor,
                       page_size: int):
    """Inverse of block_table_to_csr -> (block_tables padded with 0, context_lens)."""
    num_pages = (indptr[1:] - indptr[:-1]).long()                         # [B]
    B = num_pages.numel()
    max_blocks = int(num_pages.max().item()) if B else 0
    col = torch.arange(max_blocks, device=indices.device)[None, :]
    keep = col < num_pages[:, None]
    block_tables = torch.zeros(B, max_blocks, dtype=torch.int32, device=indices.device)
    block_tables[keep] = indices                                          # same row-major order
    context_lens = ((num_pages - 1) * page_size + last_page_len.long()).to(torch.int32)
    return block_tables, context_lens


def flashinfer_decode(q, k_cache, v_cache, block_tables, context_lens, workspace=None, return_lse=False):
    """Run FlashInfer's batch decode on the L2-style inputs.

    q [B, Hq, D]; k_cache/v_cache [num_pages, page_size, Hkv, D] ("NHD" layout, passed as a tuple).
    plan() does the CPU-side scheduling once per batch shape; run() launches the kernel and can
    be called again (e.g. once per layer) with the same plan.
    """
    import flashinfer

    B, Hq, D = q.shape
    _, page_size, Hkv, _ = k_cache.shape
    indptr, indices, last = block_table_to_csr(block_tables, context_lens, page_size)
    if workspace is None:
        workspace = torch.empty(128 << 20, dtype=torch.uint8, device=q.device)   # scratch for split-KV partials
    wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(workspace, kv_layout="NHD")
    wrapper.plan(indptr, indices, last, Hq, Hkv, D, page_size,
                 q_data_type=q.dtype, kv_data_type=k_cache.dtype)
    return wrapper.run(q, (k_cache, v_cache), return_lse=return_lse)

"""F2 · merge_state / merge_states with FlashInfer's base-2 LSE — reference solution.

An attention "state" over a subset of keys is (v, s): v is the softmax-normalised output over
that subset, s the log-sum-exp of its scores. FlashInfer stores s in BASE 2:
    s = log2( sum_j exp(score_j) )          (natural-log LSE = s * ln 2)
Merging two states (same queries, disjoint key sets):
    m   = max(s_a, s_b)
    w_a = 2^(s_a - m),  w_b = 2^(s_b - m)
    v   = (w_a v_a + w_b v_b) / (w_a + w_b),   s = m + log2(w_a + w_b)
Shapes follow flashinfer.merge_state / merge_states:
    merge_state : v_a, v_b [S, H, D]; s_a, s_b [S, H] fp32  -> (v [S, H, D], s [S, H])
    merge_states: v [S, N, H, D];     s [S, N, H]   fp32    -> (v [S, H, D], s [S, H])
"""
import torch
import triton
import triton.language as tl


@triton.jit
def merge_state_kernel(va_ptr, sa_ptr, vb_ptr, sb_ptr, vo_ptr, so_ptr, H, D,
                       BLOCK_D: tl.constexpr):
    pid = tl.program_id(0).to(tl.int64)       # one program per (position, head) = pid
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D
    sa = tl.load(sa_ptr + pid)
    sb = tl.load(sb_ptr + pid)
    m = tl.maximum(sa, sb)
    wa = tl.exp2(sa - m)                      # base 2 throughout: exp2 / log2
    wb = tl.exp2(sb - m)
    va = tl.load(va_ptr + pid * D + offs, mask=mask, other=0.0).to(tl.float32)
    vb = tl.load(vb_ptr + pid * D + offs, mask=mask, other=0.0).to(tl.float32)
    v = (wa * va + wb * vb) / (wa + wb)
    tl.store(vo_ptr + pid * D + offs, v.to(vo_ptr.dtype.element_ty), mask=mask)
    tl.store(so_ptr + pid, m + tl.log2(wa + wb))


@triton.jit
def merge_states_kernel(v_ptr, s_ptr, vo_ptr, so_ptr, N, H, D,
                        BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr):
    pid = tl.program_id(0).to(tl.int64)       # (position, head)
    pos = pid // H
    h = pid % H
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_D)
    n_mask = offs_n < N
    # s[pos, n, h] lives at (pos * N + n) * H + h
    s = tl.load(s_ptr + (pos * N + offs_n) * H + h, mask=n_mask, other=float("-inf"))
    m = tl.max(s, axis=0)
    w = tl.exp2(s - m)                        # padded states -> 2^-inf = 0
    v = tl.load(v_ptr + ((pos * N + offs_n[:, None]) * H + h) * D + offs_d[None, :],
                mask=n_mask[:, None] & (offs_d[None, :] < D), other=0.0).to(tl.float32)
    w_sum = tl.sum(w, axis=0)
    out = tl.sum(w[:, None] * v, axis=0) / w_sum
    tl.store(vo_ptr + pid * D + offs_d, out.to(vo_ptr.dtype.element_ty), mask=offs_d < D)
    tl.store(so_ptr + pid, m + tl.log2(w_sum))


def merge_state(v_a, s_a, v_b, s_b):
    S, H, D = v_a.shape
    v_a, v_b = v_a.contiguous(), v_b.contiguous()
    s_a, s_b = s_a.float().contiguous(), s_b.float().contiguous()
    v_out = torch.empty_like(v_a)
    s_out = torch.empty((S, H), dtype=torch.float32, device=v_a.device)
    if S * H:
        merge_state_kernel[(S * H,)](v_a, s_a, v_b, s_b, v_out, s_out, H, D,
                                     BLOCK_D=triton.next_power_of_2(D))
    return v_out, s_out


def merge_states(v, s):
    S, N, H, D = v.shape
    v, s = v.contiguous(), s.float().contiguous()
    v_out = torch.empty((S, H, D), dtype=v.dtype, device=v.device)
    s_out = torch.empty((S, H), dtype=torch.float32, device=v.device)
    if S * H:
        merge_states_kernel[(S * H,)](v, s, v_out, s_out, N, H, D,
                                      BLOCK_N=max(2, triton.next_power_of_2(N)),
                                      BLOCK_D=triton.next_power_of_2(D))
    return v_out, s_out

"""F3 · fused_add_rmsnorm (in place) + silu_and_mul, with FlashInfer's semantics — reference solution.

fused_add_rmsnorm(x, residual, w, eps)   -> None, both tensors updated IN PLACE:
    h = x + residual          (computed in fp32)
    residual <- h             (stored in residual's dtype: the next layer's residual stream)
    x        <- h * rsqrt(mean(h^2) + eps) * w
This is the pre-norm Transformer step "add the sublayer output to the residual, then normalise
for the next sublayer" in one pass: read x, residual once; write both once.

silu_and_mul(x[..., 2d]) -> [..., d]:  silu(x[..., :d]) * x[..., d:]   (the SwiGLU MLP gate)
"""
import torch
import triton
import triton.language as tl


@triton.jit
def fused_add_rmsnorm_kernel(x_ptr, r_ptr, w_ptr, x_stride, r_stride, N, eps, BLOCK_N: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < N
    x = tl.load(x_ptr + row * x_stride + cols, mask=mask, other=0.0).to(tl.float32)
    r = tl.load(r_ptr + row * r_stride + cols, mask=mask, other=0.0).to(tl.float32)
    h = x + r
    tl.store(r_ptr + row * r_stride + cols, h.to(r_ptr.dtype.element_ty), mask=mask)   # new residual
    rstd = 1.0 / tl.sqrt(tl.sum(h * h, axis=0) / N + eps)
    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    tl.store(x_ptr + row * x_stride + cols, (h * rstd * w).to(x_ptr.dtype.element_ty), mask=mask)


@triton.jit
def silu_and_mul_kernel(x_ptr, out_ptr, x_stride, out_stride, d, BLOCK: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)
    offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)   # 2-D grid: rows x column blocks
    mask = offs < d
    gate = tl.load(x_ptr + row * x_stride + offs, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(x_ptr + row * x_stride + d + offs, mask=mask, other=0.0).to(tl.float32)
    out = gate * tl.sigmoid(gate) * up                        # silu(g) = g * sigmoid(g)
    tl.store(out_ptr + row * out_stride + offs, out.to(out_ptr.dtype.element_ty), mask=mask)


def fused_add_rmsnorm(x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> None:
    assert x.shape == residual.shape and x.shape[-1] == weight.shape[0]
    assert x.stride(-1) == 1 and residual.stride(-1) == 1
    N = x.shape[-1]
    x2, r2 = x.view(-1, N), residual.view(-1, N)       # view, not reshape: we must write in place
    M = x2.shape[0]
    if M == 0:
        return
    BLOCK_N = triton.next_power_of_2(N)
    fused_add_rmsnorm_kernel[(M,)](x2, r2, weight, x2.stride(0), r2.stride(0), N, eps,
                                   BLOCK_N=BLOCK_N, num_warps=min(max(BLOCK_N // 256, 1), 8))


def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
    assert x.shape[-1] % 2 == 0 and x.stride(-1) == 1
    d = x.shape[-1] // 2
    x2 = x.reshape(-1, 2 * d)
    out = torch.empty((x2.shape[0], d), dtype=x.dtype, device=x.device)
    if x2.shape[0]:
        BLOCK = min(1024, triton.next_power_of_2(d))
        silu_and_mul_kernel[(x2.shape[0], triton.cdiv(d, BLOCK))](x2, out, x2.stride(0), out.stride(0), d, BLOCK=BLOCK)
    return out.reshape(*x.shape[:-1], d)

"""F4 · Sorting-free top-k / top-p renormalisation — reference solution.

top_k_renorm_probs(probs [B, V], k)  : keep the k largest probabilities per row, renormalise.
top_p_renorm_probs(probs [B, V], p)  : keep the smallest set of largest probabilities whose
                                       mass reaches p (nucleus), renormalise.
Same semantics as flashinfer.top_k_renorm_probs / top_p_renorm_probs (for rows without ties).

Sorting a 150k-entry row per request per step is the slow part of naive samplers. FlashInfer
never sorts: it repeatedly reduces (mass, count) of the entries above a pivot and moves the
pivot. Here the pivot search is an exact bisection over the float's bit pattern: for x >= 0,
the int32 view of an fp32 value increases with the value, so 31 halvings of [0, bits(max)+1)
pin down the exact threshold. Each round is one pass over the row.
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _renorm_kernel(probs_ptr, out_ptr, k_ptr, p_ptr, V, IS_TOP_K: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)
    base = probs_ptr + row * V
    offs = tl.arange(0, BLOCK)

    # hi = bits(max) + 1 so that "count(x >= hi) == 0"; lo = 0 so that "count(x >= lo) == V".
    mx = tl.zeros([BLOCK], dtype=tl.float32)
    for start in range(0, V, BLOCK):
        x = tl.load(base + start + offs, mask=start + offs < V, other=0.0)
        mx = tl.maximum(mx, x)
    lo = tl.zeros([], dtype=tl.int32)
    hi = tl.max(mx, axis=0).to(tl.int32, bitcast=True) + 1

    if IS_TOP_K:
        k = tl.load(k_ptr + row)
    else:
        p = tl.load(p_ptr + row)

    # Invariant: f(lo) satisfies the condition, f(hi) does not. 31 rounds shrink hi - lo to 1.
    for _ in range(31):
        mid = lo + (hi - lo) // 2
        acc = tl.zeros([BLOCK], dtype=tl.float32)
        for start in range(0, V, BLOCK):
            x = tl.load(base + start + offs, mask=start + offs < V, other=0.0)
            above = x.to(tl.int32, bitcast=True) >= mid      # x >= value(mid), compared as bits
            if IS_TOP_K:
                acc += above.to(tl.float32)                  # count above the pivot
            else:
                acc += tl.where(above, x, 0.0)               # mass above the pivot
        total = tl.sum(acc, axis=0)
        if IS_TOP_K:
            ok = total >= k
        else:
            ok = total >= p
        lo = tl.where(ok, mid, lo)
        hi = tl.where(ok, hi, mid)

    # Keep x >= value(lo): exactly k entries (top-k) / the nucleus (top-p). Renormalise.
    kept = tl.zeros([BLOCK], dtype=tl.float32)
    for start in range(0, V, BLOCK):
        x = tl.load(base + start + offs, mask=start + offs < V, other=0.0)
        kept += tl.where(x.to(tl.int32, bitcast=True) >= lo, x, 0.0)
    denom = tl.sum(kept, axis=0)
    for start in range(0, V, BLOCK):
        m = start + offs < V
        x = tl.load(base + start + offs, mask=m, other=0.0)
        y = tl.where(x.to(tl.int32, bitcast=True) >= lo, x / denom, 0.0)
        tl.store(out_ptr + row * V + start + offs, y, mask=m)


def _per_row(val, B, dtype, device):
    if isinstance(val, torch.Tensor):
        return val.to(device=device, dtype=dtype).expand(B).contiguous()
    return torch.full((B,), val, dtype=dtype, device=device)


def top_k_renorm_probs(probs: torch.Tensor, top_k) -> torch.Tensor:
    probs = probs.float().contiguous()
    B, V = probs.shape
    out = torch.empty_like(probs)
    k = _per_row(top_k, B, torch.float32, probs.device)       # compared with a float count
    _renorm_kernel[(B,)](probs, out, k, k, V, IS_TOP_K=True, BLOCK=4096, num_warps=16)
    return out


def top_p_renorm_probs(probs: torch.Tensor, top_p) -> torch.Tensor:
    probs = probs.float().contiguous()
    B, V = probs.shape
    out = torch.empty_like(probs)
    p = _per_row(top_p, B, torch.float32, probs.device)
    _renorm_kernel[(B,)](probs, out, p, p, V, IS_TOP_K=False, BLOCK=4096, num_warps=16)
    return out

"""G1 · Layout calculator for Triton / Gluon BlockedLayout and SliceLayout — reference solution.

Pure Python, runs anywhere. A BlockedLayout says how the elements of a block-level tensor are
spread over the registers, lanes (threads of a warp) and warps of one program:

    size_per_thread  [d] : contiguous elements a thread owns along dim d (one "chunk")
    threads_per_warp [d] : lanes laid out along dim d          (product must be 32)
    warps_per_cta    [d] : warps laid out along dim d          (product = num_warps)
    order                : dims from fastest- to slowest-varying when numbering lanes,
                           warps and registers (order[0] is the fastest)

One "CTA tile" covers tile[d] = spt[d] * tpw[d] * wpc[d] elements along each dim:
    coord[d] = rep[d] * tile[d] + warp_c[d] * (spt[d] * tpw[d]) + lane_c[d] * spt[d] + s[d]
If the tensor is larger than the tile, the pattern repeats (each thread gets more registers);
if it is smaller, coordinates wrap around (several threads hold the same element: broadcast).

Register numbering used in this drill (it matches how Triton numbers registers for blocked
layouts): the elements of one size_per_thread chunk come first (linearised by `order`),
then the repetitions (also linearised by `order`).
"""
from dataclasses import dataclass
from itertools import product
from math import log2, prod


@dataclass(frozen=True)
class BlockedLayout:
    size_per_thread: tuple[int, ...]
    threads_per_warp: tuple[int, ...]
    warps_per_cta: tuple[int, ...]
    order: tuple[int, ...]

    def __post_init__(self):
        r = len(self.size_per_thread)
        assert len(self.threads_per_warp) == len(self.warps_per_cta) == len(self.order) == r
        assert sorted(self.order) == list(range(r)), "order must be a permutation of the dims"
        assert prod(self.threads_per_warp) == 32, "a warp has 32 lanes"

    @property
    def rank(self) -> int:
        return len(self.size_per_thread)

    @property
    def num_warps(self) -> int:
        return prod(self.warps_per_cta)

    def tile(self) -> tuple[int, ...]:
        """Elements covered by one CTA tile along each dim."""
        return tuple(s * t * w for s, t, w in zip(self.size_per_thread, self.threads_per_warp, self.warps_per_cta))


def _delinearize(index: int, sizes: tuple[int, ...], order: tuple[int, ...]) -> list[int]:
    """Inverse of numbering a multi-index with order[0] as the fastest-varying dim."""
    coords = [0] * len(sizes)
    for d in order:
        coords[d] = index % sizes[d]
        index //= sizes[d]
    return coords


def _linearize(coords, sizes, order) -> int:
    index, mult = 0, 1
    for d in order:
        index += coords[d] * mult
        mult *= sizes[d]
    return index


def reps(layout: BlockedLayout, shape: tuple[int, ...]) -> tuple[int, ...]:
    """How many times the CTA tile repeats along each dim (1 if the tensor is smaller)."""
    out = []
    for n, t in zip(shape, layout.tile()):
        assert n % t == 0 or t % n == 0, "shape and tile must divide each other (powers of two)"
        out.append(max(1, n // t))
    return tuple(out)


def thread_elements(layout: BlockedLayout, shape: tuple[int, ...], warp: int, lane: int) -> list[tuple[int, ...]]:
    """Coordinates held by (warp, lane), indexed by register number."""
    spt, tpw, wpc, order = layout.size_per_thread, layout.threads_per_warp, layout.warps_per_cta, layout.order
    lane_c = _delinearize(lane, tpw, order)    # this lane's position inside the warp
    warp_c = _delinearize(warp, wpc, order)    # this warp's position inside the CTA
    tile = layout.tile()
    r = reps(layout, shape)
    out = []
    for rep_i in range(prod(r)):                       # repetitions: outer
        rep_c = _delinearize(rep_i, r, order)
        for s_i in range(prod(spt)):                   # elements of one chunk: inner
            s_c = _delinearize(s_i, spt, order)
            coord = tuple(
                (rep_c[d] * tile[d] + warp_c[d] * spt[d] * tpw[d] + lane_c[d] * spt[d] + s_c[d]) % shape[d]
                for d in range(layout.rank)
            )
            out.append(coord)
    return out


def owners(layout: BlockedLayout, shape: tuple[int, ...], coord: tuple[int, ...]) -> list[tuple[int, int, int]]:
    """Every (warp, lane, register) holding element `coord` (more than one when broadcast)."""
    found = []
    for warp in range(layout.num_warps):
        for lane in range(32):
            for reg, c in enumerate(thread_elements(layout, shape, warp, lane)):
                if c == tuple(coord):
                    found.append((warp, lane, reg))
    return found


def slice_elements(parent: BlockedLayout, parent_shape: tuple[int, ...], dim: int,
                   warp: int, lane: int) -> list[tuple[int, ...]]:
    """SliceLayout(dim, parent): the layout of a tensor obtained by removing `dim`, e.g. the
    result of tl.sum(x, axis=dim). A thread holds the remaining coordinates of its parent
    elements; duplicates collapse, first occurrence keeps its register order."""
    seen, out = set(), []
    for c in thread_elements(parent, parent_shape, warp, lane):
        rest = c[:dim] + c[dim + 1:]
        if rest not in seen:
            seen.add(rest)
            out.append(rest)
    return out


def reduce_plan(layout: BlockedLayout, shape: tuple[int, ...], axis: int) -> dict:
    """What tl.sum / tl.max along `axis` costs with this layout.

    in_thread      values each thread combines in its own registers for one output
    shuffle_rounds butterfly __shfl_xor rounds across the lanes laid out along `axis`
    cross_warp     True if warps are laid out along `axis`: partials go through shared
                   memory plus a barrier, the expensive part
    """
    r = reps(layout, shape)
    return {
        "in_thread": layout.size_per_thread[axis] * r[axis],
        "shuffle_rounds": int(log2(layout.threads_per_warp[axis])),
        "cross_warp": layout.warps_per_cta[axis] > 1,
    }


def warp_load_sectors(layout: BlockedLayout, shape: tuple[int, ...], dtype_bytes: int, warp: int = 0) -> tuple[int, int]:
    """Memory efficiency of the first vectorised load a warp issues on a row-major tensor.

    Each lane loads the first `vec` registers of its chunk in one instruction, where
    vec = min(spt[order[0]], 16 // dtype_bytes) (at most 16 bytes per lane per instruction).
    Returns (32-byte sectors touched, bytes actually requested). Ideal: sectors * 32 == bytes.
    """
    vec = max(1, min(layout.size_per_thread[layout.order[0]], 16 // dtype_bytes))
    strides = [prod(shape[d + 1:]) for d in range(len(shape))]     # row-major element strides
    sectors = set()
    for lane in range(32):
        for c in thread_elements(layout, shape, warp, lane)[:vec]:
            addr = sum(ci * si for ci, si in zip(c, strides)) * dtype_bytes
            sectors.add(addr // 32)
    return len(sectors), 32 * vec * dtype_bytes


def all_coords(shape):
    return list(product(*(range(n) for n in shape)))

"""G2 · CuTe layout algebra in pure Python — reference solution.

A CuTe Layout is a pair (Shape, Stride) of (nested) integer tuples. It is a FUNCTION from
coordinates to offsets:   offset = sum_i coord_i * stride_i.
A plain integer index is first turned into a coordinate colexicographically (first mode fastest).
Everything in CUTLASS 3 / CuTe DSL — thread-value layouts, tiled MMA/copy, smem swizzles,
FlashAttention-4 and FlashInfer's CuTe-DSL kernels — is built from these few operations.

    Layout((4, 2), (1, 4))   "4x2 column-major"   written (4,2):(1,4)
    Layout((4, 2), (2, 1))   "4x2 row-major"      written (4,2):(2,1)
"""
from dataclasses import dataclass
from math import prod


def flatten(t):
    if isinstance(t, tuple):
        return tuple(x for e in t for x in flatten(e))
    return (t,)


def _product(shape):
    return prod(flatten(shape))


def _idx2crd(idx, shape):
    """Colexicographic: the first mode varies fastest. Follows the nesting of `shape`."""
    if isinstance(shape, tuple):
        crd = []
        for s in shape:
            n = _product(s)
            crd.append(_idx2crd(idx % n, s))
            idx //= n
        return tuple(crd)
    return idx


def _inner(crd, shape, stride):
    if isinstance(shape, tuple):
        if not isinstance(crd, tuple):            # an integer for a nested mode: unflatten it first
            crd = _idx2crd(crd, shape)
        return sum(_inner(c, s, d) for c, s, d in zip(crd, shape, stride))
    return crd * stride


@dataclass(frozen=True)
class Layout:
    shape: object
    stride: object

    def __call__(self, crd):
        return _inner(crd, self.shape, self.stride)

    def size(self):                     # number of coordinates (the domain)
        return _product(self.shape)

    def cosize(self):                   # 1 + largest offset (the codomain the layout touches)
        return self(self.size() - 1) + 1 if self.size() else 0

    def __repr__(self):
        return f"{self.shape}:{self.stride}"


def make_layout_right(shape):
    """Compact row-major (LayoutRight): the LAST mode has stride 1."""
    fl = flatten(shape)
    strides, acc = [0] * len(fl), 1
    for i in reversed(range(len(fl))):
        strides[i] = acc
        acc *= fl[i]
    return Layout(shape, _unflatten(shape, iter(strides)))


def _unflatten(shape, it):
    if isinstance(shape, tuple):
        return tuple(_unflatten(s, it) for s in shape)
    return next(it)


def make_layout_left(shape):
    """Compact column-major (LayoutLeft): stride = exclusive prefix product of the flattened shape."""
    strides, acc = [], 1
    for s in flatten(shape):
        strides.append(acc)
        acc *= s
    return Layout(shape, _unflatten(shape, iter(strides)))


def coalesce(layout: Layout) -> Layout:
    """Same function, fewest modes: drop size-1 modes, merge (s0:d0, s1:d1) when d1 == s0 * d0."""
    shapes, strides = [], []
    for s, d in zip(flatten(layout.shape), flatten(layout.stride)):
        if s == 1:
            continue
        if shapes and strides[-1] * shapes[-1] == d:
            shapes[-1] *= s                       # the new mode continues the previous one
        else:
            shapes.append(s)
            strides.append(d)
    if not shapes:
        return Layout(1, 0)
    if len(shapes) == 1:
        return Layout(shapes[0], strides[0])
    return Layout(tuple(shapes), tuple(strides))


def composition(a: Layout, b: Layout) -> Layout:
    """R = A o B, i.e. R(i) = A(B(i)) for i < size(B). B is applied mode by mode.

    For one mode s:d of B, walk A's (coalesced) modes: skip the first d elements ("divide"),
    then keep the next s ("mod"). Requires the usual CuTe divisibility conditions."""
    if isinstance(b.shape, tuple):
        modes = [composition(a, Layout(s, d)) for s, d in zip(b.shape, b.stride)]
        return Layout(tuple(m.shape for m in modes), tuple(m.stride for m in modes))
    if b.stride == 0:
        return Layout(b.shape, 0)
    fa = coalesce(a)
    a_shapes, a_strides = flatten(fa.shape), flatten(fa.stride)
    rest_shape, rest_stride = b.shape, b.stride
    shapes, strides = [], []
    for s, d in zip(a_shapes[:-1], a_strides[:-1]):
        assert s % rest_stride == 0 or rest_stride % s == 0, "stride divisibility"
        new_shape = min(max(1, s // rest_stride), rest_shape)
        if new_shape != 1:
            shapes.append(new_shape)
            strides.append(rest_stride * d)
        assert rest_shape % new_shape == 0, "shape divisibility"
        rest_shape //= new_shape
        rest_stride = -(-rest_stride // s)        # ceil division
    if rest_shape != 1 or not shapes:             # whatever is left runs along A's last mode
        shapes.append(rest_shape)
        strides.append(rest_stride * a_strides[-1])
    if len(shapes) == 1:
        return Layout(shapes[0], strides[0])
    return Layout(tuple(shapes), tuple(strides))


def complement(a: Layout, cotarget: int) -> Layout:
    """The layout of the "holes" of A in [0, cotarget): (A, complement(A)) together cover it.

    Sort A's modes by stride; each gap between where one mode ends and the next begins becomes
    a mode of the complement, plus a final mode up to cotarget."""
    modes = sorted((d, s) for s, d in zip(flatten(a.shape), flatten(a.stride)) if s != 1 and d != 0)
    shapes, strides, current = [], [], 1
    for d, s in modes:
        assert d % current == 0, "A must be injective with divisible strides"
        shapes.append(d // current)
        strides.append(current)
        current = s * d
    shapes.append(-(-cotarget // current))
    strides.append(current)
    return coalesce(Layout(tuple(shapes), tuple(strides)))


def logical_divide(a: Layout, tiler: Layout) -> Layout:
    """Split A into (tile, rest): mode 0 = the elements the tiler picks, mode 1 = repetitions.
    logical_divide(A, B) = A o (B, complement(B, size(A)))."""
    c = complement(tiler, a.size())
    return composition(a, Layout((tiler.shape, c.shape), (tiler.stride, c.stride)))

"""M1 · MoE routing + moe_align_block_size — reference solution.

route(logits [T, E], top_k)  -> (topk_weights [T, K] fp32, topk_ids [T, K] int32)
    softmax over experts, keep the top_k, renormalise the kept weights to sum to 1
    (Mixtral / Qwen-MoE style; DeepSeek-V3 uses sigmoid scores + group-limited routing).

moe_align_block_size(topk_ids [T, K], num_experts, block_m)
    -> sorted_token_ids [P] int32, expert_ids [P // block_m] int32, num_tokens_post_padded int
The (token, slot) pairs are numbered flat: f = t * K + k. Grouping them by expert and padding each
expert's group to a multiple of block_m lets ONE grouped GEMM launch (M2) process every expert:
each block of block_m rows belongs to exactly one expert. Padding rows hold the value T * K
(an out-of-range id the GEMM masks off). Same contract as vLLM's moe_align_block_size, except
that P here is the exact padded length instead of a preallocated upper bound.
"""
import torch


def route(router_logits: torch.Tensor, top_k: int):
    probs = torch.softmax(router_logits.float(), dim=-1)
    w, ids = torch.topk(probs, top_k, dim=-1)
    w = w / w.sum(-1, keepdim=True)
    return w, ids.to(torch.int32)


def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int, block_m: int):
    flat = topk_ids.reshape(-1).long()                       # [T*K] expert of each (token, slot)
    n = flat.numel()
    dev = flat.device
    counts = torch.bincount(flat, minlength=num_experts)     # tokens per expert
    padded = (counts + block_m - 1) // block_m * block_m     # each group rounded up to block_m
    pad_start = torch.cumsum(padded, 0) - padded             # where each expert's group begins
    total = int(padded.sum().item())

    # Stable sort by expert keeps the original token order inside each expert's group.
    order = torch.argsort(flat, stable=True)                 # flat ids, grouped by expert
    sorted_experts = flat[order]
    group_start = torch.cumsum(counts, 0) - counts           # start of each group in `order`
    rank = torch.arange(n, device=dev) - group_start[sorted_experts]   # position within the group
    dest = pad_start[sorted_experts] + rank                  # position in the padded layout

    sorted_token_ids = torch.full((total,), n, dtype=torch.int32, device=dev)   # padding = T*K
    sorted_token_ids[dest] = order.to(torch.int32)
    expert_ids = torch.repeat_interleave(torch.arange(num_experts, device=dev), padded // block_m).to(torch.int32)
    return sorted_token_ids, expert_ids, total

"""M2 · Fused MoE grouped GEMM in Triton (the vLLM fused_moe_kernel structure) — reference solution.

One kernel launch computes the expert matmul for every (token, slot) pair:
    out[f, :] = x[f // K, :] @ W[expert(f)]^T  (* topk_weight[f] if MUL_WEIGHT)
using the padded, expert-grouped order from M1. Grid = (padded blocks of BLOCK_M rows) x (N tiles);
each M-block belongs to exactly one expert, so its B operand is a single expert's weight tile.

moe_layer() is a full SwiGLU MoE MLP on top of it:
    route -> align -> GEMM1 (x @ W13^T, [.., 2I]) -> silu_and_mul -> GEMM2 (h @ W2^T, * weight) -> sum over K
"""
import torch
import torch.nn.functional as F
import triton
import triton.language as tl

from drill import load

_m1 = load("m1_moe_align")   # your M1 (or the reference with DRILL_IMPL=solutions)
route, moe_align_block_size = _m1.route, _m1.moe_align_block_size


@triton.jit
def fused_moe_kernel(a_ptr, w_ptr, out_ptr, topk_w_ptr, sorted_ids_ptr, expert_ids_ptr,
                     N, Kdim, num_valid, top_k,
                     stride_am, stride_we, stride_wn, stride_om,
                     A_ROWS_ARE_TOKENS: tl.constexpr, MUL_WEIGHT: tl.constexpr,
                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)                       # which padded block of BLOCK_M (token, slot) rows
    pid_n = tl.program_id(1)                       # which N tile

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    flat_ids = tl.load(sorted_ids_ptr + offs_m)    # (token, slot) ids; padding rows = num_valid
    row_mask = flat_ids < num_valid
    expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)

    # GEMM1 reads x[token]: several slots of one token share the same input row (gather).
    # GEMM2 reads h[flat_id]: every (token, slot) has its own intermediate row.
    if A_ROWS_ARE_TOKENS:
        a_rows = flat_ids // top_k
    else:
        a_rows = flat_ids
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + a_rows[:, None].to(tl.int64) * stride_am + offs_k[None, :]
    # W[e] is [N, Kdim] row-major; load a [BLOCK_K, BLOCK_N] tile of W[e]^T.
    w_ptrs = w_ptr + expert * stride_we + offs_n[None, :] * stride_wn + offs_k[:, None]

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k0 in range(0, tl.cdiv(Kdim, BLOCK_K)):
        k_left = Kdim - k0 * BLOCK_K
        a = tl.load(a_ptrs, mask=row_mask[:, None] & (offs_k[None, :] < k_left), other=0.0)
        w = tl.load(w_ptrs, mask=(offs_k[:, None] < k_left) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, w, acc)
        a_ptrs += BLOCK_K
        w_ptrs += BLOCK_K

    if MUL_WEIGHT:                                 # fold the router weight into the epilogue
        tw = tl.load(topk_w_ptr + flat_ids, mask=row_mask, other=0.0)
        acc = acc * tw[:, None]
    out_ptrs = out_ptr + flat_ids[:, None].to(tl.int64) * stride_om + offs_n[None, :]
    tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=row_mask[:, None] & (offs_n[None, :] < N))


def fused_moe_gemm(a, w, sorted_ids, expert_ids, num_valid, top_k, topk_weights=None,
                   a_rows_are_tokens=True, block_m=64, block_n=64, block_k=32):
    """a [T, Kdim] (tokens) or [T*K, Kdim]; w [E, N, Kdim] -> out [T*K, N] in a.dtype."""
    E, N, Kdim = w.shape
    assert a.stride(-1) == 1 and w.is_contiguous()
    out = torch.zeros((num_valid, N), dtype=a.dtype, device=a.device)
    tw = topk_weights.reshape(-1).float().contiguous() if topk_weights is not None else out
    grid = (sorted_ids.numel() // block_m, triton.cdiv(N, block_n))
    fused_moe_kernel[grid](a, w, out, tw, sorted_ids, expert_ids, N, Kdim, num_valid, top_k,
                           a.stride(0), w.stride(0), w.stride(1), out.stride(0),
                           A_ROWS_ARE_TOKENS=a_rows_are_tokens, MUL_WEIGHT=topk_weights is not None,
                           BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k)
    return out


def moe_layer(x, router_logits, w13, w2, top_k, block_m=64):
    """x [T, H]; w13 [E, 2I, H] (gate rows then up rows); w2 [E, H, I] -> [T, H]."""
    T, H = x.shape
    E = w13.shape[0]
    topk_w, topk_ids = route(router_logits, top_k)
    sorted_ids, expert_ids, _ = moe_align_block_size(topk_ids, E, block_m)
    n = T * top_k
    h = fused_moe_gemm(x, w13, sorted_ids, expert_ids, n, top_k, block_m=block_m)          # [T*K, 2I]
    I = h.shape[1] // 2
    h = (F.silu(h[:, :I].float()) * h[:, I:].float()).to(x.dtype)                         # F3's silu_and_mul
    y = fused_moe_gemm(h, w2, sorted_ids, expert_ids, n, top_k, topk_weights=topk_w,
                       a_rows_are_tokens=False, block_m=block_m)                          # [T*K, H]
    return y.view(T, top_k, H).float().sum(1).to(x.dtype)                                 # combine slots

"""M3 · FP8 fine-grained (block-scaled) quantisation + GEMM — reference solution.

The DeepSeek-V3 / DeepGEMM recipe, written in Triton:
  activations: per-token-group quantisation, one fp32 scale per 1 x 128 group
  weights:     per-block quantisation,       one fp32 scale per 128 x 128 block
  C[m, n] = sum_kb  (A_q[m, kb] @ B_q[n, kb]^T) * sa[m, kb] * sb[n // 128, kb]
Each 128-wide K step is one fp8 tl.dot (Tensor Cores, fp32 accumulation) followed by a scale
multiply, so the scales change every K block without leaving fp32.
FP8 e4m3: largest finite value 448. Needs sm_89+ (Ada, Hopper, Blackwell incl. RTX 50xx).
"""
import torch
import triton
import triton.language as tl

FP8 = torch.float8_e4m3fn
FP8_MAX = 448.0
GROUP = 128


@triton.jit
def per_token_group_quant_kernel(x_ptr, q_ptr, s_ptr, K, eps, GROUP_SIZE: tl.constexpr):
    # One program per (row, group): the scale makes the group's max |x| map to 448.
    row = tl.program_id(0).to(tl.int64)
    g = tl.program_id(1)
    offs = g * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
    mask = offs < K
    x = tl.load(x_ptr + row * K + offs, mask=mask, other=0.0).to(tl.float32)
    amax = tl.maximum(tl.max(tl.abs(x), axis=0), eps)         # eps: an all-zero group gets a tiny scale, not 0
    scale = amax / 448.0
    q = tl.clamp(x / scale, -448.0, 448.0)
    tl.store(q_ptr + row * K + offs, q.to(q_ptr.dtype.element_ty), mask=mask)
    tl.store(s_ptr + row * tl.cdiv(K, GROUP_SIZE) + g, scale)


def per_token_group_quant_fp8(x: torch.Tensor, group_size: int = GROUP, eps: float = 1e-10):
    """x [M, K] -> (x_q [M, K] fp8 e4m3, scales [M, ceil(K/group)] fp32)."""
    x = x.contiguous()
    M, K = x.shape
    q = torch.empty((M, K), dtype=FP8, device=x.device)
    s = torch.empty((M, triton.cdiv(K, group_size)), dtype=torch.float32, device=x.device)
    per_token_group_quant_kernel[(M, triton.cdiv(K, group_size))](x, q, s, K, eps, GROUP_SIZE=group_size)
    return q, s


def block_quant_fp8(w: torch.Tensor, block: int = GROUP):
    """Weights [N, K] -> (w_q fp8 [N, K], scales [ceil(N/128), ceil(K/128)]). Offline, so plain PyTorch."""
    N, K = w.shape
    nb, kb = triton.cdiv(N, block), triton.cdiv(K, block)
    wp = torch.zeros(nb * block, kb * block, dtype=torch.float32, device=w.device)
    wp[:N, :K] = w.float()
    blocks = wp.view(nb, block, kb, block)
    amax = blocks.abs().amax(dim=(1, 3)).clamp(min=1e-10)              # [nb, kb]
    scale = amax / FP8_MAX
    q = (blocks / scale[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).to(FP8)
    return q.view(nb * block, kb * block)[:N, :K].contiguous(), scale


@triton.jit
def fp8_block_gemm_kernel(a_ptr, b_ptr, c_ptr, sa_ptr, sb_ptr, M, N, K,
                          stride_am, stride_bn, stride_cm, stride_sam, stride_sbn,
                          BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_m[:, None].to(tl.int64) * stride_am + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :].to(tl.int64) * stride_bn + offs_k[:, None]      # B^T tile [BK, BN]
    m_mask = offs_m < M
    n_mask = offs_n < N
    # BLOCK_N divides 128, so all columns of this tile share one weight-scale row: n // 128.
    sb_row = (pid_n * BLOCK_N) // 128

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for kb in range(0, tl.cdiv(K, BLOCK_K)):                  # BLOCK_K == 128 == one scale group
        k_left = K - kb * BLOCK_K
        a = tl.load(a_ptrs, mask=m_mask[:, None] & (offs_k[None, :] < k_left), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_left) & n_mask[None, :], other=0.0)
        sa = tl.load(sa_ptr + offs_m * stride_sam + kb, mask=m_mask, other=0.0)   # [BLOCK_M]
        sb = tl.load(sb_ptr + sb_row * stride_sbn + kb)                             # scalar
        acc += tl.dot(a, b) * (sa[:, None] * sb)              # fp8 x fp8 -> fp32, then rescale
        a_ptrs += BLOCK_K
        b_ptrs += BLOCK_K

    c_ptrs = c_ptr + offs_m[:, None].to(tl.int64) * stride_cm + offs_n[None, :]
    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :])


def fp8_block_gemm(a_q, a_s, b_q, b_s, out_dtype=torch.bfloat16, block_m=64, block_n=128):
    """a_q [M, K] fp8 + a_s [M, K/128]; b_q [N, K] fp8 + b_s [N/128, K/128] -> C [M, N] = A @ B^T."""
    M, K = a_q.shape
    N = b_q.shape[0]
    assert 128 % block_n == 0 and a_q.is_contiguous() and b_q.is_contiguous()
    c = torch.empty((M, N), dtype=out_dtype, device=a_q.device)
    grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n))
    fp8_block_gemm_kernel[grid](a_q, b_q, c, a_s, b_s, M, N, K,
                                a_q.stride(0), b_q.stride(0), c.stride(0), a_s.stride(0), b_s.stride(0),
                                BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=128, num_warps=4, num_stages=3)
    return c

"""X1 · First custom Metal kernel in MLX: fused add + GELU — reference solution.

mx.fast.metal_kernel(name, input_names, output_names, source, header) compiles the BODY of a Metal
kernel: MLX writes the signature for you. Inputs arrive as `const device T* x` (or `constant`
scalars), outputs as `device T* out`, and the Metal attributes you mention in the body
(thread_position_in_grid, threadgroup_position_in_grid, ...) are added as parameters.

Grid semantics differ from CUDA: `grid` is the TOTAL number of threads per dimension (Metal's
dispatchThreads), not the number of threadgroups, and it need not be a multiple of `threadgroup`.
"""
import mlx.core as mx

_HEADER = """
inline float gelu_tanh(float z) {
    // 0.5 z (1 + tanh(sqrt(2/pi) (z + 0.044715 z^3)))
    return 0.5f * z * (1.0f + metal::precise::tanh(0.7978845608028654f * (z + 0.044715f * z * z * z)));
}
"""

# One element per thread. T is a template parameter, so one source serves fp32 / fp16 / bf16.
_SRC = """
    uint i = thread_position_in_grid.x;
    if (i < n) {
        out[i] = T(gelu_tanh(float(x[i]) + float(y[i])));   // math in fp32, store in T
    }
"""

# N_PER elements per thread, contiguous: fewer threads, each issuing wider loads.
_SRC_VEC = """
    uint base = thread_position_in_grid.x * N_PER;
    for (uint j = 0; j < N_PER; ++j) {
        uint i = base + j;
        if (i < n) out[i] = T(gelu_tanh(float(x[i]) + float(y[i])));
    }
"""

_kernel = mx.fast.metal_kernel(name="add_gelu", input_names=["x", "y", "n"], output_names=["out"],
                               source=_SRC, header=_HEADER)
_kernel_vec = mx.fast.metal_kernel(name="add_gelu_vec", input_names=["x", "y", "n"], output_names=["out"],
                                   source=_SRC_VEC, header=_HEADER)


def add_gelu(x: mx.array, y: mx.array, threadgroup: int = 256) -> mx.array:
    assert x.shape == y.shape and x.dtype == y.dtype
    n = x.size
    # ensure_row_contiguous (the default) makes MLX hand the kernel contiguous copies when needed.
    return _kernel(inputs=[x, y, n], template=[("T", x.dtype)],
                   grid=(max(n, 1), 1, 1), threadgroup=(threadgroup, 1, 1),
                   output_shapes=[x.shape], output_dtypes=[x.dtype])[0]


def add_gelu_vec(x: mx.array, y: mx.array, n_per: int = 4, threadgroup: int = 256) -> mx.array:
    assert x.shape == y.shape and x.dtype == y.dtype
    n = x.size
    threads = max((n + n_per - 1) // n_per, 1)
    return _kernel_vec(inputs=[x, y, n], template=[("T", x.dtype), ("N_PER", n_per)],
                       grid=(threads, 1, 1), threadgroup=(threadgroup, 1, 1),
                       output_shapes=[x.shape], output_dtypes=[x.dtype])[0]

"""X2 · Row softmax and RMSNorm in Metal: simdgroup reductions + threadgroup memory — reference solution.

Apple GPU vocabulary <-> CUDA:
    threadgroup            <-> block           threadgroup memory     <-> shared memory
    simdgroup (32 threads) <-> warp            simd_sum / simd_max    <-> warp shuffle reductions
    threadgroup_barrier(mem_flags::mem_threadgroup)                   <-> __syncthreads()
One threadgroup per row: each thread strides over the row, simd_* reduces inside a simdgroup,
then lane 0 of each simdgroup posts its partial to threadgroup memory for the final reduction.
"""
import mlx.core as mx

# Block-wide reduction helper: every thread returns the reduced value. `scratch` holds one
# slot per simdgroup; the trailing barrier lets the caller reuse `scratch` right away.
_HEADER = """
template <bool IS_MAX>
inline float tg_reduce(float v, threadgroup float* scratch, uint lane, uint sg, uint nsg) {
    v = IS_MAX ? simd_max(v) : simd_sum(v);
    if (lane == 0) scratch[sg] = v;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    float init = IS_MAX ? -INFINITY : 0.0f;
    v = lane < nsg ? scratch[lane] : init;
    v = IS_MAX ? simd_max(v) : simd_sum(v);       // every simdgroup reduces the partials itself
    threadgroup_barrier(mem_flags::mem_threadgroup);
    return v;
}
"""

_SOFTMAX = """
    uint row = threadgroup_position_in_grid.x;
    uint tid = thread_position_in_threadgroup.x;
    uint nthr = threads_per_threadgroup.x;
    uint lane = thread_index_in_simdgroup;
    uint sg = simdgroup_index_in_threadgroup;
    uint nsg = (nthr + 31) / 32;
    threadgroup float scratch[32];
    uint n = uint(ncols);
    // `auto`, not `const device T*`: MLX may pass small inputs in the `constant` address space.
    auto xr = x + row * n;

    float m = -INFINITY;                                   // pass 1: row max
    for (uint c = tid; c < n; c += nthr) m = max(m, float(xr[c]));
    m = tg_reduce<true>(m, scratch, lane, sg, nsg);

    float s = 0.0f;                                        // pass 2: sum of exp(x - max)
    for (uint c = tid; c < n; c += nthr) s += metal::exp(float(xr[c]) - m);
    s = tg_reduce<false>(s, scratch, lane, sg, nsg);

    float inv = 1.0f / s;                                  // pass 3: normalise and write
    for (uint c = tid; c < n; c += nthr) out[row * n + c] = T(metal::exp(float(xr[c]) - m) * inv);
"""

_RMSNORM = """
    uint row = threadgroup_position_in_grid.x;
    uint tid = thread_position_in_threadgroup.x;
    uint nthr = threads_per_threadgroup.x;
    uint lane = thread_index_in_simdgroup;
    uint sg = simdgroup_index_in_threadgroup;
    uint nsg = (nthr + 31) / 32;
    threadgroup float scratch[32];
    uint n = uint(ncols);
    // `auto`, not `const device T*`: MLX may pass small inputs in the `constant` address space.
    auto xr = x + row * n;

    float ss = 0.0f;
    for (uint c = tid; c < n; c += nthr) { float v = float(xr[c]); ss += v * v; }
    ss = tg_reduce<false>(ss, scratch, lane, sg, nsg);
    float rstd = metal::precise::rsqrt(ss / float(n) + eps);
    for (uint c = tid; c < n; c += nthr) out[row * n + c] = T(float(xr[c]) * rstd * float(w[c]));
"""

_softmax = mx.fast.metal_kernel(name="row_softmax", input_names=["x", "ncols"], output_names=["out"],
                                source=_SOFTMAX, header=_HEADER)
_rmsnorm = mx.fast.metal_kernel(name="row_rmsnorm", input_names=["x", "w", "ncols", "eps"], output_names=["out"],
                                source=_RMSNORM, header=_HEADER)


def _threads(ncols: int) -> int:
    # one simdgroup per 32..1024 columns, at most 1024 threads per threadgroup
    return min(1024, max(32, 1 << (max(ncols, 1) - 1).bit_length())) if ncols < 1024 else 1024


def softmax(x: mx.array) -> mx.array:
    assert x.ndim == 2
    rows, ncols = x.shape
    t = _threads(ncols)
    return _softmax(inputs=[x, ncols], template=[("T", x.dtype)], grid=(rows * t, 1, 1), threadgroup=(t, 1, 1),
                    output_shapes=[x.shape], output_dtypes=[x.dtype])[0]


def rms_norm(x: mx.array, w: mx.array, eps: float = 1e-6) -> mx.array:
    assert x.ndim == 2 and w.shape == (x.shape[1],)
    rows, ncols = x.shape
    t = _threads(ncols)
    return _rmsnorm(inputs=[x, w, ncols, eps], template=[("T", x.dtype)], grid=(rows * t, 1, 1),
                    threadgroup=(t, 1, 1), output_shapes=[x.shape], output_dtypes=[x.dtype])[0]

"""X3 · Matmul in Metal: threadgroup-memory tiling, then simdgroup_matrix (8x8 hardware MMA) — reference solution.

C[M, N] = A[M, K] @ B[K, N], row-major, fp32 accumulation.

matmul_tiled : C4's v1 on Apple GPUs. 16x16 threadgroup computes a 16x16 tile of C; A and B tiles go
               through threadgroup memory, every value is reused 16 times.
matmul_simd  : 128 threads = 4 simdgroups compute a 32x32 tile; each simdgroup owns a 16x16 quarter
               as 2x2 simdgroup_float8x8 accumulators and uses simdgroup_multiply_accumulate,
               the Apple-GPU matrix instruction (what MLX's own "steel" GEMM kernels are built on).
"""
import mlx.core as mx

_TILED = """
    const uint TS = 16;
    uint2 t = thread_position_in_threadgroup.xy;
    uint2 g = threadgroup_position_in_grid.xy;
    uint row = g.y * TS + t.y;
    uint col = g.x * TS + t.x;
    threadgroup float As[16][16];
    threadgroup float Bs[16][16];
    float acc = 0.0f;
    for (uint k0 = 0; k0 < K; k0 += TS) {
        As[t.y][t.x] = (row < M && k0 + t.x < K) ? float(A[row * K + k0 + t.x]) : 0.0f;   // zero-fill edges
        Bs[t.y][t.x] = (k0 + t.y < K && col < N) ? float(B[(k0 + t.y) * N + col]) : 0.0f;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint k = 0; k < TS; ++k) acc += As[t.y][k] * Bs[k][t.x];
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    if (row < M && col < N) C[row * N + col] = T(acc);
"""

_SIMD = """
    // Tile: BM = BN = 32, BK = 16. 128 threads = 4 simdgroups arranged 2 x 2.
    const uint BM = 32, BN = 32, BK = 16;
    uint tid = thread_index_in_threadgroup;
    uint sg = simdgroup_index_in_threadgroup;
    uint sg_r = sg / 2, sg_c = sg % 2;                 // this simdgroup's 16x16 quarter
    uint2 g = threadgroup_position_in_grid.xy;
    uint row0 = g.y * BM, col0 = g.x * BN;

    threadgroup float As[32 * 16];                     // [BM][BK]
    threadgroup float Bs[16 * 32];                     // [BK][BN]
    threadgroup float Cs[32 * 32];                     // staging for the masked store

    simdgroup_float8x8 c00 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
    simdgroup_float8x8 c01 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
    simdgroup_float8x8 c10 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
    simdgroup_float8x8 c11 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
    simdgroup_float8x8 a0, a1, b0, b1;

    for (uint k0 = 0; k0 < K; k0 += BK) {
        for (uint e = tid; e < BM * BK; e += 128) {    // 512 values each, 4 per thread
            uint r = e / BK, k = e % BK;
            As[e] = (row0 + r < M && k0 + k < K) ? float(A[(row0 + r) * K + k0 + k]) : 0.0f;
            uint kb = e / BN, c = e % BN;
            Bs[e] = (k0 + kb < K && col0 + c < N) ? float(B[(k0 + kb) * N + col0 + c]) : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint kk = 0; kk < BK; kk += 8) {
            // simdgroup_load(matrix, pointer to the 8x8 block's first element, row stride in elements)
            simdgroup_load(a0, As + (sg_r * 16 + 0) * BK + kk, BK);
            simdgroup_load(a1, As + (sg_r * 16 + 8) * BK + kk, BK);
            simdgroup_load(b0, Bs + kk * BN + sg_c * 16 + 0, BN);
            simdgroup_load(b1, Bs + kk * BN + sg_c * 16 + 8, BN);
            simdgroup_multiply_accumulate(c00, a0, b0, c00);   // c += a @ b, 8x8x8
            simdgroup_multiply_accumulate(c01, a0, b1, c01);
            simdgroup_multiply_accumulate(c10, a1, b0, c10);
            simdgroup_multiply_accumulate(c11, a1, b1, c11);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    simdgroup_store(c00, Cs + (sg_r * 16 + 0) * BN + sg_c * 16 + 0, BN);
    simdgroup_store(c01, Cs + (sg_r * 16 + 0) * BN + sg_c * 16 + 8, BN);
    simdgroup_store(c10, Cs + (sg_r * 16 + 8) * BN + sg_c * 16 + 0, BN);
    simdgroup_store(c11, Cs + (sg_r * 16 + 8) * BN + sg_c * 16 + 8, BN);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (uint e = tid; e < BM * BN; e += 128) {        // masked write-back for partial edge tiles
        uint r = e / BN, c = e % BN;
        if (row0 + r < M && col0 + c < N) C[(row0 + r) * N + col0 + c] = T(Cs[e]);
    }
"""

_tiled = mx.fast.metal_kernel(name="matmul_tiled", input_names=["A", "B", "M", "N", "K"], output_names=["C"], source=_TILED)
_simd = mx.fast.metal_kernel(name="matmul_simd", input_names=["A", "B", "M", "N", "K"], output_names=["C"], source=_SIMD,
                             header="#include <metal_simdgroup_matrix>\n")


def _check(a, b):
    assert a.ndim == 2 and b.ndim == 2 and a.shape[1] == b.shape[0] and a.dtype == b.dtype
    return a.shape[0], b.shape[1], a.shape[1]


def matmul_tiled(a: mx.array, b: mx.array) -> mx.array:
    M, N, K = _check(a, b)
    grid = (((N + 15) // 16) * 16, ((M + 15) // 16) * 16, 1)       # threads, not threadgroups
    return _tiled(inputs=[a, b, M, N, K], template=[("T", a.dtype)], grid=grid, threadgroup=(16, 16, 1),
                  output_shapes=[(M, N)], output_dtypes=[a.dtype])[0]


def matmul_simd(a: mx.array, b: mx.array) -> mx.array:
    M, N, K = _check(a, b)
    # one 128-thread threadgroup per 32x32 tile: grid x = 128 * tiles_n, grid y = tiles_m
    grid = (((N + 31) // 32) * 128, (M + 31) // 32, 1)
    return _simd(inputs=[a, b, M, N, K], template=[("T", a.dtype)], grid=grid, threadgroup=(128, 1, 1),
                 output_shapes=[(M, N)], output_dtypes=[a.dtype])[0]