/*
 * Isolated aligned-shape VEA-B numerical prototype.
 *
 * The MMA, online-softmax, FP8 conversion, PV accumulation, and output fragment
 * mappings are adapted from SageAttention 2.2.0 at commit
 * d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 (Apache-2.0), whose helpers record
 * prior adaptation from FlashInfer 0.1.5. This prototype supports only
 * B=H=1, Q=128, KV=192, D=128 and is not production code.
 */

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>

#include <cuda/barrier>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <array>
#include <cstdint>
#include <vector>

#define CUDA_CHECK(call) do { \
  cudaError_t error_ = (call); \
  TORCH_CHECK(error_ == cudaSuccess, cudaGetErrorString(error_)); \
} while (0)

namespace {

constexpr int kQ = 128;
constexpr int kKv = 192;
constexpr int kD = 128;
constexpr int kEpochs = 3;
constexpr int kOwners = 4;
constexpr int kThreads = 320;
constexpr int kReferenceThreads = 128;
constexpr int kDynamicShared = 50 * 1024;
constexpr int kReferenceShared = 32 * 1024;
constexpr float kSoftmaxScale = 0.08838834764831845f;
constexpr float kLog2e = 1.44269504088896340736f;
constexpr float kFp8Offset = 8.807f;

struct Checkpoints {
  float* qk;
  float* m;
  float* d;
  uint32_t* p;
  float* rescale;
  uint32_t* pv;
  float* ro;
  float* reciprocal;
  int64_t* clocks;
};

__device__ __forceinline__ float ptx_exp2(float value) {
  float result;
  asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value));
  return result;
}

__device__ __forceinline__ float ptx_rcp(float value) {
  float result;
  asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value));
  return result;
}

__device__ __forceinline__ void unpack_half2(float* output, uint32_t value) {
  const uint16_t low = value & 0xffff;
  const uint16_t high = value >> 16;
  asm("cvt.f32.f16 %0, %1;" : "=f"(output[0]) : "h"(low));
  asm("cvt.f32.f16 %0, %1;" : "=f"(output[1]) : "h"(high));
}

__device__ __forceinline__ uint32_t fp8x4(float* first, float* second) {
  uint32_t output;
  asm volatile(
      "{ .reg .b16 lo; .reg .b16 hi;\n"
      "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
      "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
      "mov.b32 %0, {lo, hi}; }"
      : "=r"(output)
      : "f"(first[0]), "f"(first[1]), "f"(second[0]), "f"(second[1]));
  return output;
}

__device__ __forceinline__ void ldmatrix_x4(uint32_t* output, const void* pointer) {
  const uint32_t address = static_cast<uint32_t>(__cvta_generic_to_shared(pointer));
  asm volatile(
      "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
      : "=r"(output[0]), "=r"(output[1]), "=r"(output[2]), "=r"(output[3])
      : "r"(address));
}

__device__ __forceinline__ uint2 ldmatrix_x2(const void* pointer) {
  const uint32_t address = static_cast<uint32_t>(__cvta_generic_to_shared(pointer));
  uint2 output;
  asm volatile(
      "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
      : "=r"(output.x), "=r"(output.y)
      : "r"(address));
  return output;
}

__device__ __forceinline__ void int8_mma_init(int32_t* c, uint32_t* a, uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
      "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};"
      : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(0), "r"(0), "r"(0), "r"(0));
}

__device__ __forceinline__ void int8_mma(int32_t* c, uint32_t* a, uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
      "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};"
      : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
}

__device__ __forceinline__ void fp8_mma_init(uint32_t* c, uint32_t* a, uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
      "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};"
      : "=r"(c[0]), "=r"(c[1])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(0), "r"(0));
}

__device__ __forceinline__ void fp8_mma(uint32_t* c, uint32_t* a, uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
      "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};"
      : "=r"(c[0]), "=r"(c[1])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(c[0]), "r"(c[1]));
}

__device__ __forceinline__ void fp8_mma_init_scalar(
    uint32_t& c0, uint32_t& c1, uint4 a, uint2 b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
      "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};"
      : "=r"(c0), "=r"(c1)
      : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y),
        "r"(0), "r"(0));
}

__device__ __forceinline__ void fp8_mma_scalar(
    uint32_t& c0, uint32_t& c1, uint4 a, uint2 b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
      "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};"
      : "=r"(c0), "=r"(c1)
      : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y),
        "r"(c0), "r"(c1));
}

struct Smem128 {
  uint4* base;
  __device__ __forceinline__ uint32_t offset(int row, int column) const {
    return row * 8 + (column ^ (row & 7));
  }
};

struct Smem64 {
  uint4* base;
  __device__ __forceinline__ uint32_t offset(int row, int column) const {
    return row * 4 + (column ^ ((row / 2) & 3));
  }
};

__device__ __forceinline__ void load_q_fragments(
    Smem128 q_smem, int owner, int lane, uint32_t rq[2][4][4]) {
  #pragma unroll
  for (int inner = 0; inner < 4; ++inner) {
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq) {
      const int row = owner * 32 + fq * 16 + lane % 16;
      const int column = lane / 16 + inner * 2;
      ldmatrix_x4(rq[fq][inner], q_smem.base + q_smem.offset(row, column));
    }
  }
}

__device__ __forceinline__ void compute_qk(
    Smem128 k_smem, int lane, uint32_t rq[2][4][4], int32_t rs[2][4][8]) {
  #pragma unroll
  for (int inner = 0; inner < 4; ++inner) {
    #pragma unroll
    for (int fk = 0; fk < 4; ++fk) {
      uint32_t rk[4];
      const int row = lane % 8 + (lane / 16) * 8 + fk * 16;
      const int column = (lane / 8) % 2 + inner * 2;
      ldmatrix_x4(rk, k_smem.base + k_smem.offset(row, column));
      #pragma unroll
      for (int fq = 0; fq < 2; ++fq) {
        if (inner == 0) {
          int8_mma_init(rs[fq][fk], rq[fq][inner], rk);
          int8_mma_init(rs[fq][fk] + 4, rq[fq][inner], rk + 2);
        } else {
          int8_mma(rs[fq][fk], rq[fq][inner], rk);
          int8_mma(rs[fq][fk] + 4, rq[fq][inner], rk + 2);
        }
      }
    }
  }
}

__device__ __forceinline__ int qk_checkpoint_index(
    int epoch, int owner, int lane, int fq, int fk, int element) {
  return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 4 + fk) * 8 + element);
}

__device__ __forceinline__ int state_checkpoint_index(
    int epoch, int owner, int lane, int fq, int state) {
  return ((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 2 + state);
}

__device__ __forceinline__ int p_checkpoint_index(
    int epoch, int owner, int lane, int fq, int fk, int element) {
  return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 2 + fk) * 4 + element);
}

__device__ __forceinline__ int pv_checkpoint_index(
    int epoch, int owner, int lane, int fq, int fv, int element) {
  return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 8 + fv) * 4 + element);
}

__device__ __forceinline__ int ro_checkpoint_index(
    int epoch, int owner, int lane, int fq, int fv, int element) {
  return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 8 + fv) * 8 + element);
}

template <bool Capture>
__device__ __forceinline__ void softmax_update(
    float rs[2][4][8], float m[2][2], float d[2][2], float old_scale[2][2],
    uint32_t rs_f8[2][2][4], float scale, int epoch, int owner, int lane,
    Checkpoints checkpoints) {
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq) {
    #pragma unroll
    for (int state = 0; state < 2; ++state) {
      const float previous = m[fq][state];
      float tile_max = -5000000.0f;
      #pragma unroll
      for (int fk = 0; fk < 4; ++fk) {
        const float local = fmaxf(
            fmaxf(rs[fq][fk][state * 2], rs[fq][fk][state * 2 + 1]),
            fmaxf(rs[fq][fk][state * 2 + 4], rs[fq][fk][state * 2 + 5]));
        tile_max = fmaxf(tile_max, local);
      }
      tile_max = fmaf(tile_max, scale, -kFp8Offset);
      tile_max = fmaxf(tile_max, __shfl_xor_sync(0xffffffff, tile_max, 1));
      tile_max = fmaxf(tile_max, __shfl_xor_sync(0xffffffff, tile_max, 2));
      m[fq][state] = fmaxf(m[fq][state], tile_max);
      old_scale[fq][state] = ptx_exp2(previous - m[fq][state]);
      d[fq][state] *= old_scale[fq][state];
      const float negative_m = -m[fq][state];
      #pragma unroll
      for (int fk = 0; fk < 4; ++fk) {
        rs[fq][fk][state * 2] = ptx_exp2(fmaf(rs[fq][fk][state * 2], scale, negative_m));
        rs[fq][fk][state * 2 + 1] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 1], scale, negative_m));
        rs[fq][fk][state * 2 + 4] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 4], scale, negative_m));
        rs[fq][fk][state * 2 + 5] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 5], scale, negative_m));
      }
    }
  }
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq) {
    #pragma unroll
    for (int fk = 0; fk < 4; ++fk) {
      d[fq][0] += rs[fq][fk][0] + rs[fq][fk][1] + rs[fq][fk][4] + rs[fq][fk][5];
      d[fq][1] += rs[fq][fk][2] + rs[fq][fk][3] + rs[fq][fk][6] + rs[fq][fk][7];
    }
    #pragma unroll
    for (int fk = 0; fk < 2; ++fk) {
      rs_f8[fq][fk][0] = fp8x4(rs[fq][fk * 2], rs[fq][fk * 2] + 4);
      rs_f8[fq][fk][1] = fp8x4(rs[fq][fk * 2] + 2, rs[fq][fk * 2] + 6);
      rs_f8[fq][fk][2] = fp8x4(rs[fq][fk * 2 + 1], rs[fq][fk * 2 + 1] + 4);
      rs_f8[fq][fk][3] = fp8x4(rs[fq][fk * 2 + 1] + 2, rs[fq][fk * 2 + 1] + 6);
    }
  }
  if constexpr (Capture) {
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq) {
      #pragma unroll
      for (int state = 0; state < 2; ++state) {
        const int index = state_checkpoint_index(epoch, owner, lane, fq, state);
        checkpoints.m[index] = m[fq][state];
        checkpoints.d[index] = d[fq][state];
        checkpoints.rescale[index] = old_scale[fq][state];
      }
      #pragma unroll
      for (int fk = 0; fk < 2; ++fk) {
        #pragma unroll
        for (int element = 0; element < 4; ++element) {
          checkpoints.p[p_checkpoint_index(epoch, owner, lane, fq, fk, element)] = rs_f8[fq][fk][element];
        }
      }
    }
  }
}

template <bool Capture>
__device__ __forceinline__ void compute_pv(
    Smem64 v_smem, uint32_t rs_f8[2][2][4], float ro[2][8][8],
    int epoch, int owner, int lane, Checkpoints checkpoints) {
  const int row_base = lane % 8 + (lane / 16) * 8;
  const int column_base = (lane / 8) % 2;
  #pragma unroll
  for (int fv = 0; fv < 8; ++fv) {
    uint32_t instant[2][4];
    #pragma unroll
    for (int fk = 0; fk < 2; ++fk) {
      uint32_t rv[4];
      const uint32_t offset = v_smem.offset(row_base + fv * 16, column_base + fk * 2);
      ldmatrix_x4(rv, v_smem.base + offset);
      #pragma unroll
      for (int fq = 0; fq < 2; ++fq) {
        if (fk == 0) {
          fp8_mma_init(instant[fq], rs_f8[fq][fk], rv);
          fp8_mma_init(instant[fq] + 2, rs_f8[fq][fk], rv + 2);
        } else {
          fp8_mma(instant[fq], rs_f8[fq][fk], rv);
          fp8_mma(instant[fq] + 2, rs_f8[fq][fk], rv + 2);
        }
      }
    }
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq) {
      #pragma unroll
      for (int element = 0; element < 4; ++element) {
        float pair[2];
        unpack_half2(pair, instant[fq][element]);
        ro[fq][fv][element * 2] += pair[0];
        ro[fq][fv][element * 2 + 1] += pair[1];
        if constexpr (Capture) {
          checkpoints.pv[pv_checkpoint_index(epoch, owner, lane, fq, fv, element)] = instant[fq][element];
        }
      }
      if constexpr (Capture) {
        #pragma unroll
        for (int element = 0; element < 8; ++element) {
          checkpoints.ro[ro_checkpoint_index(epoch, owner, lane, fq, fv, element)] = ro[fq][fv][element];
        }
      }
    }
  }
}

template <bool Capture>
__device__ __forceinline__ void compute_pv_from_shared(
    Smem64 v_smem, const uint4* score_slot, float ro[2][8][8],
    volatile float* retire_slot, int epoch, int owner, int lane,
    Checkpoints checkpoints) {
  const int row_base = lane % 8 + (lane / 16) * 8;
  const int column_base = (lane / 8) % 2;
  #pragma unroll
  for (int fv = 0; fv < 8; ++fv) {
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq) {
      #pragma unroll
      for (int output_pair = 0; output_pair < 2; ++output_pair) {
        uint32_t instant[2];
        #pragma unroll
        for (int fk = 0; fk < 2; ++fk) {
          const uint4 packed = score_slot[((owner * 4 + fq * 2 + fk) * 32 + lane)];
          const uint32_t offset = v_smem.offset(
              row_base + fv * 16 + output_pair * 8, column_base + fk * 2);
          const uint2 rv = ldmatrix_x2(v_smem.base + offset);
          if (fk == 0)
            fp8_mma_init_scalar(instant[0], instant[1], packed, rv);
          else
            fp8_mma_scalar(instant[0], instant[1], packed, rv);
        }
        #pragma unroll
        for (int element = 0; element < 2; ++element) {
          float pair[2];
          unpack_half2(pair, instant[element]);
          const int ro_element = output_pair * 4 + element * 2;
          ro[fq][fv][ro_element] += pair[0];
          ro[fq][fv][ro_element + 1] += pair[1];
          if constexpr (Capture)
            checkpoints.pv[pv_checkpoint_index(
                epoch, owner, lane, fq, fv, output_pair * 2 + element)] = instant[element];
        }
      }
      if constexpr (!Capture) {
        #pragma unroll
        for (int element = 0; element < 8; ++element)
          *retire_slot = ro[fq][fv][element];
      }
      __syncwarp();
      if constexpr (Capture) {
        #pragma unroll
        for (int element = 0; element < 8; ++element)
          checkpoints.ro[ro_checkpoint_index(epoch, owner, lane, fq, fv, element)] = ro[fq][fv][element];
      }
    }
  }
}

__device__ __forceinline__ void save_output(
    nv_bfloat16* output, float ro[2][8][8], float reciprocal[2][2],
    const float* v_scale, int owner, int lane) {
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq) {
    #pragma unroll
    for (int fv = 0; fv < 8; ++fv) {
      float scales[4];
      const int base = (lane % 4) * 2 + fv * 16;
      scales[0] = v_scale[base];
      scales[1] = v_scale[base + 1];
      scales[2] = v_scale[base + 8];
      scales[3] = v_scale[base + 9];
      #pragma unroll
      for (int pair = 0; pair < 4; ++pair) {
        const int element = pair * 2;
        const int state = (element % 4) / 2;
        ro[fq][fv][element] *= reciprocal[fq][state];
        ro[fq][fv][element + 1] *= reciprocal[fq][state];
        const int scale_pair = pair >= 2 ? 2 : 0;
        ro[fq][fv][element] *= scales[scale_pair];
        ro[fq][fv][element + 1] *= scales[scale_pair + 1];
        const int row = owner * 32 + fq * 16 + lane / 4 + 8 * state;
        const int column = fv * 16 + (lane % 4) * 2 + (pair >= 2 ? 8 : 0);
        const nv_bfloat162 packed = __float22bfloat162_rn(
            make_float2(ro[fq][fv][element], ro[fq][fv][element + 1]));
        *reinterpret_cast<nv_bfloat162*>(output + row * 128 + column) = packed;
      }
    }
  }
}

__device__ __forceinline__ void save_output_streamed(
    nv_bfloat16* output, float ro[2][8][8], const float* reciprocal,
    const float* v_scale, int owner, int lane) {
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq) {
    #pragma unroll
    for (int fv = 0; fv < 8; ++fv) {
      #pragma unroll
      for (int pair = 0; pair < 4; ++pair) {
        const int element = pair * 2;
        const int component = (element % 4) / 2;
        const float normalizer = reciprocal[owner * 32 + fq * 16 + component * 8 + lane / 4];
        ro[fq][fv][element] *= normalizer;
        ro[fq][fv][element + 1] *= normalizer;
        const int column = fv * 16 + (lane % 4) * 2 + (pair >= 2 ? 8 : 0);
        ro[fq][fv][element] *= v_scale[column];
        ro[fq][fv][element + 1] *= v_scale[column + 1];
        const int row = owner * 32 + fq * 16 + lane / 4 + 8 * component;
        const nv_bfloat162 packed = __float22bfloat162_rn(
            make_float2(ro[fq][fv][element], ro[fq][fv][element + 1]));
        *reinterpret_cast<nv_bfloat162*>(output + row * 128 + column) = packed;
      }
    }
  }
}

template <bool Capture>
__global__ __launch_bounds__(kThreads, 1) void vea_b_numeric_kernel(
    const int8_t* q, const int8_t* k, const int8_t* v,
    const float* q_scale, const float* k_scale, const float* v_scale,
    nv_bfloat16* output, Checkpoints checkpoints) {
  extern __shared__ uint8_t shared[];
  Smem128 q_smem{reinterpret_cast<uint4*>(shared)};
  Smem128 k_slots[2] = {
      {reinterpret_cast<uint4*>(shared + 16 * 1024)},
      {reinterpret_cast<uint4*>(shared + 24 * 1024)}};
  Smem64 v_slots[2] = {
      {reinterpret_cast<uint4*>(shared + 32 * 1024)},
      {reinterpret_cast<uint4*>(shared + 40 * 1024)}};
  using BlockBarrier = cuda::barrier<cuda::thread_scope_block>;
  BlockBarrier* barriers = reinterpret_cast<BlockBarrier*>(shared + 48 * 1024);
  float* scale_slots = reinterpret_cast<float*>(shared + 48 * 1024 + 128);
  float* final_reciprocal = scale_slots + 2 * 128;
  const int tid = threadIdx.x;
  const int warp = tid / 32;
  const int lane = tid % 32;
  const bool qk_owner = warp >= 2 && warp <= 5;
  const bool pv_owner = warp >= 6;
  const int owner = qk_owner ? warp - 2 : (pv_owner ? warp - 6 : -1);

  if (tid < 6) {
    const int count = tid < 4 ? 320 : 160;
    init(barriers + tid, count);
  }
  for (int chunk = tid; chunk < 1024; chunk += kThreads) {
    const int row = chunk / 8;
    const int column = chunk % 8;
    q_smem.base[q_smem.offset(row, column)] = reinterpret_cast<const uint4*>(q + row * 128)[column];
  }
  __syncthreads();

  uint32_t rq[2][4][4];
  if (qk_owner) load_q_fragments(q_smem, owner, lane, rq);
  __syncthreads();

  if (qk_owner) {
    float m[2][2] = {{-5000000.0f, -5000000.0f}, {-5000000.0f, -5000000.0f}};
    float d[2][2] = {{1.0f, 1.0f}, {1.0f, 1.0f}};
    #pragma unroll
    for (int epoch = 0; epoch < kEpochs; ++epoch) {
      const int slot = epoch & 1;
      if (epoch >= 2) barriers[2 + slot].arrive_and_wait();
      barriers[4 + slot].arrive_and_wait();
      const uint64_t started = clock64();
      union ScoreStorage { int32_t integer[2][4][8]; float fp32[2][4][8]; } scores;
      compute_qk(k_slots[slot], lane, rq, scores.integer);
      const float dequant = q_scale[owner] * k_scale[epoch];
      const bool final_pair = epoch + 2 >= kEpochs;
      #pragma unroll
      for (int fq = 0; fq < 2; ++fq) {
        #pragma unroll
        for (int fk = 0; fk < 4; ++fk) {
          #pragma unroll
          for (int element = 0; element < 8; ++element) {
            const float converted = __int2float_rz(scores.integer[fq][fk][element]);
            scores.fp32[fq][fk][element] = final_pair ? converted * dequant : converted;
            if constexpr (Capture) {
              checkpoints.qk[qk_checkpoint_index(epoch, owner, lane, fq, fk, element)] = scores.fp32[fq][fk][element];
            }
          }
        }
      }
      float old_scale[2][2];
      uint32_t probability[2][2][4];
      softmax_update<Capture>(scores.fp32, m, d, old_scale, probability,
                              final_pair ? kSoftmaxScale * kLog2e : kSoftmaxScale * kLog2e * dequant,
                              epoch, owner, lane, checkpoints);
      if ((lane & 3) == 0) {
        #pragma unroll
        for (int fq = 0; fq < 2; ++fq)
          #pragma unroll
          for (int state = 0; state < 2; ++state)
            scale_slots[slot * 128 + owner * 32 + fq * 16 + state * 8 + lane / 4] = old_scale[fq][state];
      }
      uint4* score_slot = reinterpret_cast<uint4*>(shared + slot * 8 * 1024);
      #pragma unroll
      for (int fq = 0; fq < 2; ++fq)
        #pragma unroll
        for (int fk = 0; fk < 2; ++fk)
          score_slot[((owner * 4 + fq * 2 + fk) * 32 + lane)] = *reinterpret_cast<uint4*>(probability[fq][fk]);
      if constexpr (Capture) {
        if (lane == 0) {
          checkpoints.clocks[((epoch * 8 + owner) * 2)] = started;
          checkpoints.clocks[((epoch * 8 + owner) * 2) + 1] = clock64();
        }
      }
      if (epoch == kEpochs - 1) {
        #pragma unroll
        for (int fq = 0; fq < 2; ++fq) {
          #pragma unroll
          for (int component = 0; component < 2; ++component) {
            d[fq][component] += __shfl_xor_sync(0xffffffff, d[fq][component], 1);
            d[fq][component] += __shfl_xor_sync(0xffffffff, d[fq][component], 2);
            const float reciprocal = ptx_rcp(d[fq][component]);
            if constexpr (Capture)
              checkpoints.reciprocal[state_checkpoint_index(0, owner, lane, fq, component)] = reciprocal;
            if ((lane & 3) == 0)
              final_reciprocal[owner * 32 + fq * 16 + component * 8 + lane / 4] = reciprocal;
          }
        }
      }
      (void)barriers[slot].arrive();
    }
    barriers[2].arrive_and_wait();
    barriers[3].arrive_and_wait();
  } else if (pv_owner) {
    float ro[2][8][8];
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq)
      #pragma unroll
      for (int fv = 0; fv < 8; ++fv)
        #pragma unroll
        for (int element = 0; element < 8; ++element) ro[fq][fv][element] = 0.0f;
    #pragma unroll
    for (int epoch = 0; epoch < kEpochs; ++epoch) {
      const int slot = epoch & 1;
      barriers[slot].arrive_and_wait();
      const uint64_t started = clock64();
      #pragma unroll
      for (int fq = 0; fq < 2; ++fq) {
        #pragma unroll
        for (int component = 0; component < 2; ++component) {
          const float old_scale = scale_slots[slot * 128 + owner * 32 + fq * 16 + component * 8 + lane / 4];
          #pragma unroll
          for (int fv = 0; fv < 8; ++fv) {
            ro[fq][fv][component * 2] *= old_scale;
            ro[fq][fv][component * 2 + 1] *= old_scale;
            ro[fq][fv][component * 2 + 4] *= old_scale;
            ro[fq][fv][component * 2 + 5] *= old_scale;
          }
        }
      }
      __syncwarp();
      compute_pv_from_shared<Capture>(
          v_slots[slot], reinterpret_cast<uint4*>(shared + slot * 8 * 1024),
          ro, scale_slots + slot * 128 + owner * 32 + lane,
          epoch, owner, lane, checkpoints);
      if constexpr (Capture) {
        if (lane == 0) {
          checkpoints.clocks[((epoch * 8 + 4 + owner) * 2)] = started;
          checkpoints.clocks[((epoch * 8 + 4 + owner) * 2) + 1] = clock64();
        }
      }
      (void)barriers[2 + slot].arrive();
    }
    save_output_streamed(output, ro, final_reciprocal, v_scale, owner, lane);
  } else if (warp == 0) {
    #pragma unroll
    for (int epoch = 0; epoch < kEpochs; ++epoch) {
      const int slot = epoch & 1;
      if (epoch >= 2) barriers[2 + slot].arrive_and_wait();
      for (int chunk = lane; chunk < 512; chunk += 32) {
        const int row = chunk / 8;
        const int column = chunk % 8;
        k_slots[slot].base[k_slots[slot].offset(row, column)] =
            reinterpret_cast<const uint4*>(k + (epoch * 64 + row) * 128)[column];
      }
      (void)barriers[4 + slot].arrive();
      (void)barriers[slot].arrive();
    }
    barriers[2].arrive_and_wait();
    barriers[3].arrive_and_wait();
  } else {
    #pragma unroll
    for (int epoch = 0; epoch < kEpochs; ++epoch) {
      const int slot = epoch & 1;
      if (epoch >= 2) barriers[2 + slot].arrive_and_wait();
      for (int chunk = lane; chunk < 512; chunk += 32) {
        const int row = chunk / 4;
        const int column = chunk % 4;
        v_slots[slot].base[v_slots[slot].offset(row, column)] =
            reinterpret_cast<const uint4*>(v + row * 192 + epoch * 64)[column];
      }
      (void)barriers[slot].arrive();
    }
    barriers[2].arrive_and_wait();
    barriers[3].arrive_and_wait();
  }
}

template <bool Capture>
__global__ __launch_bounds__(kReferenceThreads, 1) void sage_checkpoint_kernel(
    const int8_t* q, const int8_t* k, const int8_t* v,
    const float* q_scale, const float* k_scale, const float* v_scale,
    nv_bfloat16* output, Checkpoints checkpoints) {
  extern __shared__ uint8_t shared[];
  Smem128 q_smem{reinterpret_cast<uint4*>(shared)};
  Smem128 k_smem{reinterpret_cast<uint4*>(shared + 16 * 1024)};
  Smem64 v_smem{reinterpret_cast<uint4*>(shared + 24 * 1024)};
  const int tid = threadIdx.x;
  const int owner = tid / 32;
  const int lane = tid % 32;
  for (int chunk = tid; chunk < 1024; chunk += kReferenceThreads) {
    const int row = chunk / 8;
    const int column = chunk % 8;
    q_smem.base[q_smem.offset(row, column)] = reinterpret_cast<const uint4*>(q + row * 128)[column];
  }
  __syncthreads();
  uint32_t rq[2][4][4];
  load_q_fragments(q_smem, owner, lane, rq);
  float m[2][2] = {{-5000000.0f, -5000000.0f}, {-5000000.0f, -5000000.0f}};
  float d[2][2] = {{1.0f, 1.0f}, {1.0f, 1.0f}};
  float ro[2][8][8];
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq)
    #pragma unroll
    for (int fv = 0; fv < 8; ++fv)
      #pragma unroll
      for (int element = 0; element < 8; ++element) ro[fq][fv][element] = 0.0f;

  #pragma unroll
  for (int epoch = 0; epoch < kEpochs; ++epoch) {
    for (int chunk = tid; chunk < 512; chunk += kReferenceThreads) {
      const int row = chunk / 8;
      const int column = chunk % 8;
      k_smem.base[k_smem.offset(row, column)] = reinterpret_cast<const uint4*>(k + (epoch * 64 + row) * 128)[column];
      const int vrow = chunk / 4;
      const int vcolumn = chunk % 4;
      v_smem.base[v_smem.offset(vrow, vcolumn)] = reinterpret_cast<const uint4*>(v + vrow * 192 + epoch * 64)[vcolumn];
    }
    __syncthreads();
    union ScoreStorage { int32_t integer[2][4][8]; float fp32[2][4][8]; } scores;
    compute_qk(k_smem, lane, rq, scores.integer);
    const float dequant = q_scale[owner] * k_scale[epoch];
    const bool final_pair = epoch + 2 >= kEpochs;
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq)
      #pragma unroll
      for (int fk = 0; fk < 4; ++fk)
        #pragma unroll
        for (int element = 0; element < 8; ++element) {
          const float converted = __int2float_rz(scores.integer[fq][fk][element]);
          scores.fp32[fq][fk][element] = final_pair ? converted * dequant : converted;
          if constexpr (Capture)
            checkpoints.qk[qk_checkpoint_index(epoch, owner, lane, fq, fk, element)] = scores.fp32[fq][fk][element];
        }
    float old_scale[2][2];
    uint32_t probability[2][2][4];
    softmax_update<Capture>(scores.fp32, m, d, old_scale, probability,
                            final_pair ? kSoftmaxScale * kLog2e : kSoftmaxScale * kLog2e * dequant,
                            epoch, owner, lane, checkpoints);
    #pragma unroll
    for (int fq = 0; fq < 2; ++fq)
      #pragma unroll
      for (int fv = 0; fv < 8; ++fv)
        #pragma unroll
        for (int element = 0; element < 8; ++element)
          ro[fq][fv][element] *= old_scale[fq][(element % 4) / 2];
    compute_pv<Capture>(v_smem, probability, ro, epoch, owner, lane, checkpoints);
    __syncthreads();
  }
  float reciprocal[2][2];
  #pragma unroll
  for (int fq = 0; fq < 2; ++fq)
    #pragma unroll
    for (int state = 0; state < 2; ++state) {
      d[fq][state] += __shfl_xor_sync(0xffffffff, d[fq][state], 1);
      d[fq][state] += __shfl_xor_sync(0xffffffff, d[fq][state], 2);
      reciprocal[fq][state] = ptx_rcp(d[fq][state]);
      if constexpr (Capture) checkpoints.reciprocal[state_checkpoint_index(0, owner, lane, fq, state)] = reciprocal[fq][state];
    }
  save_output(output, ro, reciprocal, v_scale, owner, lane);
}

std::vector<torch::Tensor> make_outputs(torch::Device device) {
  auto byte = torch::TensorOptions().device(device).dtype(torch::kInt32);
  auto fp32 = torch::TensorOptions().device(device).dtype(torch::kFloat32);
  auto output = torch::empty({1, 128, 1, 128}, torch::TensorOptions().device(device).dtype(torch::kBFloat16));
  auto qk = torch::empty({3, 4, 32, 2, 4, 8}, fp32);
  auto m = torch::empty({3, 4, 32, 2, 2}, fp32);
  auto d = torch::empty({3, 4, 32, 2, 2}, fp32);
  auto p = torch::empty({3, 4, 32, 2, 2, 4}, byte);
  auto rescale = torch::empty({3, 4, 32, 2, 2}, fp32);
  auto pv = torch::empty({3, 4, 32, 2, 8, 4}, byte);
  auto ro = torch::empty({3, 4, 32, 2, 8, 8}, fp32);
  auto reciprocal = torch::empty({4, 32, 2, 2}, fp32);
  auto clocks = torch::zeros({3, 8, 2}, torch::TensorOptions().device(device).dtype(torch::kInt64));
  return {output, qk, m, d, p, rescale, pv, ro, reciprocal, clocks};
}

Checkpoints pointers(std::vector<torch::Tensor>& outputs) {
  return {outputs[1].data_ptr<float>(), outputs[2].data_ptr<float>(), outputs[3].data_ptr<float>(),
          reinterpret_cast<uint32_t*>(outputs[4].data_ptr<int32_t>()), outputs[5].data_ptr<float>(),
          reinterpret_cast<uint32_t*>(outputs[6].data_ptr<int32_t>()), outputs[7].data_ptr<float>(),
          outputs[8].data_ptr<float>(), outputs[9].data_ptr<int64_t>()};
}

void validate_inputs(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                     torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale) {
  TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "inputs must be CUDA tensors");
  TORCH_CHECK(q.scalar_type() == torch::kInt8 && k.scalar_type() == torch::kInt8, "Q/K must be INT8");
  TORCH_CHECK(q.sizes() == torch::IntArrayRef({1, 128, 1, 128}), "Q shape must be [1,128,1,128]");
  TORCH_CHECK(k.sizes() == torch::IntArrayRef({1, 192, 1, 128}), "K shape must be [1,192,1,128]");
  TORCH_CHECK(v.numel() == 128 * 192 && v.element_size() == 1, "V must contain [1,128,1,192] FP8 bytes");
  TORCH_CHECK(q_scale.numel() == 4 && k_scale.numel() == 3 && v_scale.numel() == 128, "invalid scale shapes");
  TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "inputs must be contiguous");
}

template <typename Kernel>
std::array<int64_t, 8> attributes(Kernel kernel, int shared_bytes, int threads) {
  cudaFuncAttributes value{};
  CUDA_CHECK(cudaFuncGetAttributes(&value, kernel));
  int active = 0;
  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, kernel, threads, shared_bytes));
  return {value.numRegs, static_cast<int64_t>(value.sharedSizeBytes), static_cast<int64_t>(value.localSizeBytes),
          value.maxThreadsPerBlock, value.binaryVersion, value.ptxVersion,
          static_cast<int64_t>(value.maxDynamicSharedSizeBytes), active};
}

}  // namespace

std::vector<torch::Tensor> launch_vea_b_numeric(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale,
    bool capture) {
  validate_inputs(q, k, v, q_scale, k_scale, v_scale);
  c10::cuda::CUDAGuard guard(q.device());
  auto outputs = make_outputs(q.device());
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
  const auto stream = at::cuda::getDefaultCUDAStream();
  if (capture) {
    vea_b_numeric_kernel<true><<<1, kThreads, kDynamicShared, stream>>>(
        q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
        q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
        reinterpret_cast<nv_bfloat16*>(outputs[0].data_ptr()), pointers(outputs));
  } else {
    vea_b_numeric_kernel<false><<<1, kThreads, kDynamicShared, stream>>>(
        q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
        q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
        reinterpret_cast<nv_bfloat16*>(outputs[0].data_ptr()), pointers(outputs));
  }
  CUDA_CHECK(cudaGetLastError());
  return outputs;
}

std::vector<torch::Tensor> launch_sage_checkpoint(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale) {
  validate_inputs(q, k, v, q_scale, k_scale, v_scale);
  c10::cuda::CUDAGuard guard(q.device());
  auto outputs = make_outputs(q.device());
  CUDA_CHECK(cudaFuncSetAttribute(sage_checkpoint_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared));
  sage_checkpoint_kernel<true><<<1, kReferenceThreads, kReferenceShared, at::cuda::getDefaultCUDAStream()>>>(
      q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
      q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
      reinterpret_cast<nv_bfloat16*>(outputs[0].data_ptr()), pointers(outputs));
  CUDA_CHECK(cudaGetLastError());
  return outputs;
}

torch::Tensor launch_vea_b_numeric_into(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale,
    torch::Tensor output) {
  validate_inputs(q, k, v, q_scale, k_scale, v_scale);
  TORCH_CHECK(output.is_cuda() && output.scalar_type() == torch::kBFloat16 &&
              output.sizes() == torch::IntArrayRef({1, 128, 1, 128}) && output.is_contiguous(),
              "output must be contiguous CUDA BF16 [1,128,1,128]");
  c10::cuda::CUDAGuard guard(q.device());
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
  Checkpoints empty{};
  vea_b_numeric_kernel<false><<<1, kThreads, kDynamicShared, at::cuda::getDefaultCUDAStream()>>>(
      q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
      q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
      reinterpret_cast<nv_bfloat16*>(output.data_ptr()), empty);
  CUDA_CHECK(cudaGetLastError());
  return output;
}

torch::Tensor numeric_kernel_attributes() {
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
  CUDA_CHECK(cudaFuncSetAttribute(sage_checkpoint_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared));
  std::array<std::array<int64_t, 8>, 3> rows = {
      attributes(vea_b_numeric_kernel<false>, kDynamicShared, kThreads),
      attributes(vea_b_numeric_kernel<true>, kDynamicShared, kThreads),
      attributes(sage_checkpoint_kernel<true>, kReferenceShared, kReferenceThreads)};
  auto output = torch::empty({3, 8}, torch::TensorOptions().dtype(torch::kInt64));
  auto access = output.accessor<int64_t, 2>();
  for (int row = 0; row < 3; ++row)
    for (int column = 0; column < 8; ++column) access[row][column] = rows[row][column];
  return output;
}