submission 806844

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

import os
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t


# Active routes (E55):
#   n in {1024, 2048}        -> qr_blocked16_gemm_smem_panel  (SMEM-resident panel + BF16 cuBLAS trailing)
#   n == 512                 -> qr_blocked16_gemm_higham_concat (SMEM panel + K=48 BF16 Higham concat GEMM)
#   n in {176, 352}          -> qr_blocked16_gemm             (panel + 2x BF16 cuBLAS trailing)
#   n == 32 (and other small)-> qr_cublas_geqrf via cublasSgeqrfBatched
#   else (n=4096)            -> torch.geqrf fallback
#
# Dead experiments (qr512_split, qr_blocked16 SIMT, cluster panel/DSM, Higham 3-GEMM,
# qr_block16_panel_pack bmm fallback, qr512_panel) were trimmed in the E58 cleanup;
# see NOTES.md for rationale. Git history (E9, E16-E18, E30, E46-E48, E55) preserves
# the kernels themselves.
CPP_SRC = """
void qr_blocked16_gemm(torch::Tensor h, torch::Tensor tau, torch::Tensor work, torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp);
void qr_blocked16_gemm_smem_panel(torch::Tensor h, torch::Tensor tau, torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp, int warps_per_matrix);
void qr_blocked16_gemm_cluster_smem(torch::Tensor h, torch::Tensor tau, torch::Tensor work,
                                    torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp,
                                    torch::Tensor scratch);
void qr_blocked16_gemm_higham_concat(torch::Tensor h, torch::Tensor tau, torch::Tensor work,
                                     torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp,
                                     torch::Tensor y_concat48, torch::Tensor tmp_concat48);
void qr_cublas_geqrf(torch::Tensor h, torch::Tensor tau);
"""

CUDA_SRC = r"""
#include <ATen/cuda/CUDABlas.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cstdio>
#include <cstdlib>
#include <torch/extension.h>

static void check_status(cublasStatus_t status, const char *where);
static cublasHandle_t raw_blas_handle();

__device__ __forceinline__ float qr_y_at(const float *a, int n, int panel, int local_col, int row) {
  int col = panel + local_col;
  if (row < col) {
    return 0.0f;
  }
  if (row == col) {
    return 1.0f;
  }
  return a[row + col * n];
}

// E16-derived single-CTA-per-matrix panel kernel. Factors a width-16 panel
// in-place in h[], builds the compact W in `work`, then directly emits
// column-major Y / W into ygemm/wgemm for the cuBLAS trailing GEMMs.
// (E52: "DirectPack" is now unconditional — the bmm fallback that needed
// the non-packing variant was dropped in the E58 cleanup.)
__global__ void qr_block16_panel_kernel(
    float *h,
    float *tau,
    float *work,
    float *ygemm,
    float *wgemm,
    int64_t h_step,
    int64_t tau_step,
    int n,
    int panel) {
  constexpr int nb = 16;
  constexpr int threads = 256;
  __shared__ float reduce[threads];
  __shared__ float shared_tau;
  __shared__ float shared_scale;
  __shared__ float ypy[nb * nb];

  int width = min(nb, n - panel);
  float *a = h + static_cast<int64_t>(blockIdx.x) * h_step;
  float *t = tau + static_cast<int64_t>(blockIdx.x) * tau_step;
  float *wmat = work + (static_cast<int64_t>(blockIdx.x) * n * nb);

  for (int jj = 0; jj < width; ++jj) {
    int k = panel + jj;
    float local = 0.0f;
    for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
      float x = a[i + k * n];
      local += x * x;
    }
    reduce[threadIdx.x] = local;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
      if (threadIdx.x < stride) {
        reduce[threadIdx.x] += reduce[threadIdx.x + stride];
      }
      __syncthreads();
    }

    if (threadIdx.x == 0) {
      float alpha = a[k + k * n];
      float tail_ss = reduce[0];
      if (tail_ss == 0.0f) {
        t[k] = 0.0f;
        shared_tau = 0.0f;
        shared_scale = 0.0f;
      } else {
        float norm = sqrtf(alpha * alpha + tail_ss);
        float beta = (alpha >= 0.0f) ? -norm : norm;
        float tau_k = (beta - alpha) / beta;
        t[k] = tau_k;
        shared_tau = tau_k;
        shared_scale = 1.0f / (alpha - beta);
        a[k + k * n] = beta;
      }
    }
    __syncthreads();

    float tau_k = shared_tau;
    if (tau_k != 0.0f) {
      for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
        a[i + k * n] *= shared_scale;
      }
    }
    __syncthreads();

    for (int j = k + 1; j < panel + width; ++j) {
      float dot = 0.0f;
      if (threadIdx.x == 0) {
        dot = a[k + j * n];
      }
      if (tau_k != 0.0f) {
        for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
          dot += a[i + k * n] * a[i + j * n];
        }
      }
      reduce[threadIdx.x] = dot;
      __syncthreads();

      for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
          reduce[threadIdx.x] += reduce[threadIdx.x + stride];
        }
        __syncthreads();
      }

      float update = tau_k * reduce[0];
      if (threadIdx.x == 0) {
        a[k + j * n] -= update;
      }
      if (tau_k != 0.0f) {
        for (int i = k + 1 + threadIdx.x; i < n; i += blockDim.x) {
          a[i + j * n] -= a[i + k * n] * update;
        }
      }
      __syncthreads();
    }
  }

  for (int idx = threadIdx.x; idx < nb * nb; idx += blockDim.x) {
    ypy[idx] = 0.0f;
  }
  __syncthreads();

  for (int l = 0; l < width; ++l) {
    for (int j = l + 1; j < width; ++j) {
      float local = 0.0f;
      for (int row = panel + threadIdx.x; row < n; row += blockDim.x) {
        local += qr_y_at(a, n, panel, l, row) * qr_y_at(a, n, panel, j, row);
      }
      reduce[threadIdx.x] = local;
      __syncthreads();

      for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
          reduce[threadIdx.x] += reduce[threadIdx.x + stride];
        }
        __syncthreads();
      }
      if (threadIdx.x == 0) {
        ypy[l * nb + j] = reduce[0];
      }
      __syncthreads();
    }
  }

  for (int j = 0; j < width; ++j) {
    float tau_j = t[panel + j];
    for (int row = panel + threadIdx.x; row < n; row += blockDim.x) {
      float y = qr_y_at(a, n, panel, j, row);
      float accum = y;
      for (int l = 0; l < j; ++l) {
        accum += wmat[row * nb + l] * ypy[l * nb + j];
      }
      wmat[row * nb + j] = -tau_j * accum;
    }
    __syncthreads();
  }

  // E52: direct pack into cuBLAS-ready column-major Y/W panels.
  float *yout = ygemm + (static_cast<int64_t>(blockIdx.x) * n * nb);
  float *wout = wgemm + (static_cast<int64_t>(blockIdx.x) * n * nb);
  for (int row = panel + threadIdx.x; row < n; row += blockDim.x) {
    for (int j = 0; j < width; ++j) {
      yout[row + j * n] = qr_y_at(a, n, panel, j, row);
      wout[row + j * n] = wmat[row * nb + j];
    }
  }
}

__device__ __forceinline__ float warp_sum_smem_panel(float value) {
  #pragma unroll
  for (int shift = 16; shift > 0; shift >>= 1) {
    value += __shfl_xor_sync(0xffffffff, value, shift);
  }
  return value;
}

// E54: shared-memory-resident panel factorization. The serial Householder
// dependency chain still sets the lower bound, but for large panels the
// single-CTA SIMT panel repeatedly reloads the just-factored panel from HBM
// to pack Y/W. This kernel keeps the active n-by-16 panel in SMEM, factors
// there, writes H once, and emits column-major Y/W directly. Dynamic SMEM:
// (m*nb + 2*(warps+1) + nb*nb) * sizeof(float). The two reduction
// scratch banks are ping-ponged so a following reduction cannot alias a slow
// reader from the previous one; this is the race E62's full-CTA tree avoided
// more conservatively.
__global__ void qr_block16_panel_smem_kernel(
    float *h,
    float *tau,
    float *ygemm,
    float *wgemm,
    int64_t h_step,
    int64_t tau_step,
    int64_t y_step,
    int n,
    int panel) {
  constexpr int nb = 16;
  const int lane = threadIdx.x;
  const int warp_id = threadIdx.y;
  const int warps = blockDim.y;
  const int tid = warp_id * 32 + lane;
  const int total_threads = warps * 32;
  const int batch_id = static_cast<int>(blockIdx.x);

  float *a = h + static_cast<int64_t>(batch_id) * h_step;
  float *t = tau + static_cast<int64_t>(batch_id) * tau_step;
  float *yptr = ygemm + static_cast<int64_t>(batch_id) * y_step;
  float *wptr = wgemm + static_cast<int64_t>(batch_id) * y_step;

  const int width = min(nb, n - panel);
  const int m = n - panel;

  extern __shared__ float smem[];
  float *panel_s = smem;
  float *partials = panel_s + m * nb;
  float *ypy = partials + 2 * (warps + 1);
  int reduce_phase = 0;

  #define CTA_SUM(input_value, output_value)                         \
    do {                                                             \
      int __base = (reduce_phase++ & 1) * (warps + 1);               \
      float *__partials = partials + __base;                         \
      float __v = warp_sum_smem_panel((input_value));                \
      if (lane == 0) __partials[warp_id] = __v;                      \
      __syncthreads();                                               \
      if (warp_id == 0) {                                            \
        float __r = (lane < warps) ? __partials[lane] : 0.0f;        \
        __r = warp_sum_smem_panel(__r);                              \
        if (lane == 0) __partials[warps] = __r;                      \
      }                                                              \
      __syncthreads();                                               \
      (output_value) = __partials[warps];                            \
      __syncthreads();                                               \
    } while (0)

  for (int idx = tid; idx < m * width; idx += total_threads) {
    int j = idx / m;
    int i = idx - j * m;
    panel_s[i + j * m] = a[(panel + i) + static_cast<int64_t>(panel + j) * n];
  }
  __syncthreads();

  for (int k = 0; k < width; ++k) {
    float local_ss = 0.0f;
    for (int i = k + 1 + tid; i < m; i += total_threads) {
      float x = panel_s[i + k * m];
      local_ss += x * x;
    }
    float tail_ss;
    CTA_SUM(local_ss, tail_ss);

    float alpha = panel_s[k + k * m];
    float beta;
    float tau_k;
    float scale;
    if (tail_ss == 0.0f) {
      beta = alpha;
      tau_k = 0.0f;
      scale = 0.0f;
    } else {
      float norm = sqrtf(alpha * alpha + tail_ss);
      beta = (alpha >= 0.0f) ? -norm : norm;
      tau_k = (beta - alpha) / beta;
      scale = 1.0f / (alpha - beta);
    }

    if (tid == 0) {
      panel_s[k + k * m] = beta;
      t[panel + k] = tau_k;
    }

    if (tau_k != 0.0f) {
      for (int i = k + 1 + tid; i < m; i += total_threads) {
        panel_s[i + k * m] *= scale;
      }
      __syncthreads();

      for (int j = k + 1; j < width; ++j) {
        float dot = (tid == 0) ? panel_s[k + j * m] : 0.0f;
        for (int i = k + 1 + tid; i < m; i += total_threads) {
          dot += panel_s[i + k * m] * panel_s[i + j * m];
        }
        float dot_all;
        CTA_SUM(dot, dot_all);

        float update = tau_k * dot_all;
        if (tid == 0) {
          panel_s[k + j * m] -= update;
        }
        for (int i = k + 1 + tid; i < m; i += total_threads) {
          panel_s[i + j * m] -= panel_s[i + k * m] * update;
        }
        __syncthreads();
      }
    } else {
      __syncthreads();
    }
  }

  for (int idx = tid; idx < m * width; idx += total_threads) {
    int j = idx / m;
    int i = idx - j * m;
    a[(panel + i) + static_cast<int64_t>(panel + j) * n] = panel_s[i + j * m];
  }

  for (int idx = tid; idx < nb * nb; idx += total_threads) {
    ypy[idx] = 0.0f;
  }
  __syncthreads();

  for (int l = 0; l < width; ++l) {
    for (int j = l + 1; j < width; ++j) {
      float dot = 0.0f;
      for (int i = j + tid; i < m; i += total_threads) {
        float yl = panel_s[i + l * m];
        float yj = (i == j) ? 1.0f : panel_s[i + j * m];
        dot += yl * yj;
      }
      float dot_all;
      CTA_SUM(dot, dot_all);
      if (tid == 0) {
        ypy[l * nb + j] = dot_all;
      }
      __syncthreads();
    }
  }

  for (int row = panel + tid; row < n; row += total_threads) {
    int i = row - panel;
    #pragma unroll
    for (int j = 0; j < nb; ++j) {
      float y = 0.0f;
      if (j < width) {
        if (i < j) {
          y = 0.0f;
        } else if (i == j) {
          y = 1.0f;
        } else {
          y = panel_s[i + j * m];
        }
      }
      yptr[row + static_cast<int64_t>(j) * n] = y;
    }

    #pragma unroll
    for (int j = 0; j < nb; ++j) {
      if (j >= width) {
        wptr[row + static_cast<int64_t>(j) * n] = 0.0f;
        continue;
      }
      float y;
      if (i < j) {
        y = 0.0f;
      } else if (i == j) {
        y = 1.0f;
      } else {
        y = panel_s[i + j * m];
      }
      float accum = y;
      for (int l = 0; l < j; ++l) {
        accum += wptr[row + static_cast<int64_t>(l) * n] * ypy[l * nb + j];
      }
      wptr[row + static_cast<int64_t>(j) * n] = -t[panel + j] * accum;
    }
  }

  #undef CTA_SUM
}

static bool qr_deep_timing_enabled() {
  const char *flag = std::getenv("QR_DEEP_TIMING");
  return flag != nullptr && flag[0] == '1';
}

static void qr_deep_begin(cudaEvent_t event) {
  C10_CUDA_CHECK(cudaEventRecord(event));
}

static void qr_deep_end(cudaEvent_t begin, cudaEvent_t end, float *accum_ms) {
  C10_CUDA_CHECK(cudaEventRecord(end));
  C10_CUDA_CHECK(cudaEventSynchronize(end));
  float ms = 0.0f;
  C10_CUDA_CHECK(cudaEventElapsedTime(&ms, begin, end));
  *accum_ms += ms;
}

// E72: cluster-barrier version of the multi-CTA panel kernel. Same algorithm
// as E71 but the per-phase grid sync uses Blackwell's hardware cluster
// barrier (PTX `barrier.cluster.arrive/wait`) instead of an atomic spin on
// a per-matrix counter. The cluster barrier is ~hundreds of ns vs the
// atomic-spin ~2 us per sync. Cluster size is limited to 8 CTAs per matrix
// (portable max). Each cluster covers one matrix.
//
// Workspace per matrix (laid out contiguously in `scratch`):
//   [0 .. CTAS * nb)              partials  -- one slot per CTA per dot
//   [CTAS*nb .. +32)              scalars   -- tau_k, scale, broadcasted updates
//   [+32 .. +32 + nb*nb)          ypy       -- 16x16 Y^T Y triangular
__device__ __forceinline__ float qr_warp_sum_mc(float v) {
  for (int s = 16; s > 0; s >>= 1) {
    v += __shfl_xor_sync(0xffffffff, v, s);
  }
  return v;
}

__device__ __forceinline__ float qr_block_sum_mc(float v, float *smem_buf, int tid, int nwarps) {
  v = qr_warp_sum_mc(v);
  int lane = tid & 31;
  int wid = tid >> 5;
  if (lane == 0) smem_buf[wid] = v;
  __syncthreads();
  if (wid == 0) {
    float t = (lane < nwarps) ? smem_buf[lane] : 0.0f;
    t = qr_warp_sum_mc(t);
    if (lane == 0) smem_buf[0] = t;
  }
  __syncthreads();
  return smem_buf[0];
}


// E73: SMEM-cached cluster panel. Same algorithm + cluster sync as E72 but
// each CTA holds its row chunk of the panel AND of the W matrix in SMEM,
// so the column factor loop, T construction, and W construction are all
// SMEM-resident. Only inter-CTA partial sums and YPY go through global
// scratch. HBM is touched twice: load panel chunk at kernel start, write
// panel back + pack Y/W at kernel end.
//
// Per CTA SMEM (for CHUNK_SIZE=512, CTAS=8):
//   panel_s[CHUNK_SIZE * nb]  32 KB  -- column-major panel chunk
//   wmat_s [CHUNK_SIZE * nb]  32 KB  -- column-major W chunk
//   smem_buf[16]              tiny   -- block-reduce scratch
//   sm_scalars[20]            tiny   -- tau, scale, sm_updates
// Total ~64 KB per CTA: panel_s is static, wmat_s is dynamic SMEM
// (cudaFuncSetAttribute opt-in to allow >48 KB total per CTA).
//
// Global scratch per matrix:
//   partials[CTAS * nb]   -- partial-sum slots
//   scalars [32]          -- master broadcasts (tau, scale, updates)
//   ypy     [nb * nb]     -- T construction output (read by every CTA in W loop)
template <int CTAS, int CHUNK_SIZE>
__global__
__cluster_dims__(1, CTAS, 1)
void qr_block16_panel_cluster_smem_kernel(
    float *h,
    float *tau,
    float *work,             // (batch, n, 16) row-major W (fallback only)
    float *ygemm,            // (batch, n, 16) col-major packed Y
    float *wgemm,            // (batch, n, 16) col-major packed W
    float *scratch,
    int64_t h_step,
    int64_t tau_step,
    int64_t y_step,
    int n,
    int panel) {
  constexpr int nb = 16;
  const int bid = blockIdx.x;
  const int cta = blockIdx.y;
  const int tid = threadIdx.x;
  const int nthread = blockDim.x;
  const int nwarps = nthread >> 5;

  __shared__ float panel_s[CHUNK_SIZE * nb];
  __shared__ float smem_buf[16];
  extern __shared__ float dyn_smem[];
  float *wmat_s = dyn_smem;  // CHUNK_SIZE * nb floats
  __shared__ float sm_tau;
  __shared__ float sm_scale;
  __shared__ float sm_updates[16];

#define CLUSTER_SYNC()                                                \
  do {                                                                \
    __syncthreads();                                                  \
    asm volatile("barrier.cluster.arrive.release.aligned;");          \
    asm volatile("barrier.cluster.wait.acquire.aligned;");            \
    __syncthreads();                                                  \
  } while (0)

  float *a = h + (int64_t)bid * h_step;
  float *t = tau + (int64_t)bid * tau_step;
  float *wmat = work + (int64_t)bid * n * nb;
  float *yout = ygemm + (int64_t)bid * y_step;
  float *wout = wgemm + (int64_t)bid * y_step;

  // Scratch layout per matrix:
  //   partials[CTAS * nb]
  //   scalars [32]          (per-col: tau, scale, 15 effective_updates)
  //   scales  [nb]          (persistent per-col scale array)
  //   ypy     [nb * nb]
  const int part_stride = CTAS * nb + 32 + nb + nb * nb;
  float *part = scratch + (int64_t)bid * part_stride;
  float *scalars = part + CTAS * nb;
  float *ypy = scalars + 32 + nb;

  const int row_start = cta * CHUNK_SIZE;
  const int row_end = min(n, row_start + CHUNK_SIZE);
  const int local_size = row_end - row_start;
  const int width = min(nb, n - panel);

  // ----- Load panel chunk from HBM to SMEM -----
  // rows [row_start, row_end), cols [panel, panel + width)
  for (int idx = tid; idx < CHUNK_SIZE * nb; idx += nthread) {
    int local_col = idx / CHUNK_SIZE;
    int local_row = idx - local_col * CHUNK_SIZE;
    int global_row = row_start + local_row;
    int global_col = panel + local_col;
    float v = 0.0f;
    if (local_col < width && local_row < local_size && global_col < n) {
      v = a[global_row + (int64_t)global_col * n];
    }
    panel_s[local_row + local_col * CHUNK_SIZE] = v;
  }
  __syncthreads();

  // Master CTA owns the panel rows [panel, panel+width). For aligned panels
  // (n is a multiple of CHUNK_SIZE * CTAS won't matter since panel steps by nb=16),
  // all 16 panel rows are in master_cta's chunk.
  const int master_cta = panel / CHUNK_SIZE;
  const int master_local_panel = panel - master_cta * CHUNK_SIZE;

  // ----- Column factorization loop (combined-phase variant) -----
  // E76: 2 cluster syncs per column instead of 4. Strategy: keep v_raw
  // (unscaled) in panel_s throughout the loop, fuse norm+dot partials into
  // a single reduction pass, and have the master fold the scale into the
  // effective_update broadcast. After the column loop a single scale-back
  // pass multiplies each v_raw by scale[col]. T construction also reads
  // unscaled v's and the master applies scale[l]*scale[j] to ypy[l,j].
  //
  // Persistent scales: scratch[scales_off + col] holds scale[col] for the
  // duration of this kernel invocation. Layout extension is documented in
  // the driver.
  const int scales_off = CTAS * nb + 32;  // 16 floats after the per-col scalars
  float *scales_arr = part + scales_off;
  for (int jj = 0; jj < width; ++jj) {
    int k = panel + jj;
    int n_trailing = panel + width - 1 - k;

    // Phase A: combined partial sum-of-squares (col k) + partial unscaled dot
    // products against the trailing panel cols. Single pass over rows.
    int rs_g = max(row_start, k + 1);
    int rs_l = rs_g - row_start;
    float pnorm = 0.0f;
    float pdots[15];
    #pragma unroll
    for (int j = 0; j < 15; ++j) pdots[j] = 0.0f;

    if (rs_l < local_size) {
      for (int li = rs_l + tid; li < local_size; li += nthread) {
        float v = panel_s[li + jj * CHUNK_SIZE];
        pnorm += v * v;
        #pragma unroll
        for (int j = 0; j < 15; ++j) {
          if (j < n_trailing) {
            pdots[j] += v * panel_s[li + (jj + 1 + j) * CHUNK_SIZE];
          }
        }
      }
    }

    // Reduce norm into slot 0, dots into slots 1..n_trailing.
    pnorm = qr_block_sum_mc(pnorm, smem_buf, tid, nwarps);
    if (tid == 0) part[cta * nb + 0] = pnorm;
    for (int j = 0; j < n_trailing; ++j) {
      float r = qr_block_sum_mc(pdots[j], smem_buf, tid, nwarps);
      if (tid == 0) part[cta * nb + 1 + j] = r;
    }

    CLUSTER_SYNC();

    // Phase B: master computes tau, scale, effective_updates in one shot.
    if (cta == master_cta && tid == 0) {
      float tail_ss = 0.0f;
      for (int c = 0; c < CTAS; ++c) tail_ss += part[c * nb + 0];
      int diag_local = master_local_panel + jj;
      float alpha = panel_s[diag_local + jj * CHUNK_SIZE];
      float tau_k = 0.0f, scale = 0.0f;
      if (tail_ss != 0.0f) {
        float norm = sqrtf(alpha * alpha + tail_ss);
        float beta = (alpha >= 0.0f) ? -norm : norm;
        tau_k = (beta - alpha) / beta;
        scale = 1.0f / (alpha - beta);
        panel_s[diag_local + jj * CHUNK_SIZE] = beta;  // R[k,k]
      }
      t[k] = tau_k;
      scalars[0] = tau_k;
      scalars[1] = scale;
      scales_arr[jj] = scale;  // persistent across the col loop

      for (int j = 0; j < n_trailing; ++j) {
        // unscaled partial-dot sum; the v[k]=1 contribution comes from the
        // on-diagonal a[k, k+1+j] value still in panel_s.
        float pd_sum = 0.0f;
        for (int c = 0; c < CTAS; ++c) pd_sum += part[c * nb + 1 + j];
        float dot = panel_s[diag_local + (jj + 1 + j) * CHUNK_SIZE] + scale * pd_sum;
        float update = tau_k * dot;
        float effective = scale * update;
        scalars[2 + j] = effective;
        panel_s[diag_local + (jj + 1 + j) * CHUNK_SIZE] -= update;  // on-diagonal R update
      }
    }

    CLUSTER_SYNC();

    if (tid < 16 && tid < n_trailing) sm_updates[tid] = scalars[2 + tid];
    __syncthreads();

    // Phase C: apply trailing update using v_raw * effective_update.
    // (No cluster sync needed afterwards; each CTA only modifies its own
    //  SMEM, and the next column's reads are intra-CTA.)
    if (n_trailing > 0) {
      int rs_l4 = max(0, k + 1 - row_start);
      if (rs_l4 < local_size) {
        for (int li = rs_l4 + tid; li < local_size; li += nthread) {
          float v = panel_s[li + jj * CHUNK_SIZE];
          for (int j = 0; j < n_trailing; ++j) {
            panel_s[li + (jj + 1 + j) * CHUNK_SIZE] -= v * sm_updates[j];
          }
        }
      }
    }
    __syncthreads();
  }

  // After the column loop, panel_s holds UNSCALED v_raw values in the
  // sub-diagonal of each column. T / W / Y consumers need scaled v's, so
  // multiply each below-diagonal v_raw by scales_arr[col]. Each CTA scales
  // its own row chunk in SMEM; the last column's Phase B CLUSTER_SYNC
  // already made the master's scales_arr writes globally visible.
  for (int col = 0; col < width; ++col) {
    float s = scales_arr[col];
    if (s != 0.0f) {
      int rs_l = max(0, (panel + col + 1) - row_start);
      if (rs_l < local_size) {
        for (int li = rs_l + tid; li < local_size; li += nthread) {
          panel_s[li + col * CHUNK_SIZE] *= s;
        }
      }
    }
  }
  __syncthreads();

  // ----- T construction (ypy upper triangle), partial dots from SMEM -----
  // E77: one cluster sync per l iteration. Master's ypy writes don't need to
  // be visible to other CTAs until W construction reads them, so we drop the
  // post-master-write sync and add a single CLUSTER_SYNC before the W loop.
  for (int l = 0; l < width - 1; ++l) {
    int nj = width - 1 - l;
    float pdots[15];
    #pragma unroll
    for (int j = 0; j < 15; ++j) pdots[j] = 0.0f;

    int col_l = panel + l;
    int rs_g = max(row_start, panel);
    int rs_l = rs_g - row_start;
    if (rs_l < local_size) {
      for (int li = rs_l + tid; li < local_size; li += nthread) {
        int global_row = row_start + li;
        float yl;
        if (global_row < col_l) yl = 0.0f;
        else if (global_row == col_l) yl = 1.0f;
        else yl = panel_s[li + l * CHUNK_SIZE];
        if (yl == 0.0f) continue;

        #pragma unroll
        for (int j = 0; j < 15; ++j) {
          if (j < nj) {
            int col_j = panel + l + 1 + j;
            float yj;
            if (global_row < col_j) yj = 0.0f;
            else if (global_row == col_j) yj = 1.0f;
            else yj = panel_s[li + (l + 1 + j) * CHUNK_SIZE];
            pdots[j] += yl * yj;
          }
        }
      }
    }

    for (int j = 0; j < nj; ++j) {
      float r = qr_block_sum_mc(pdots[j], smem_buf, tid, nwarps);
      if (tid == 0) part[cta * nb + j] = r;
    }

    CLUSTER_SYNC();

    if (cta == master_cta && tid == 0) {
      for (int j = 0; j < nj; ++j) {
        float s = 0.0f;
        for (int c = 0; c < CTAS; ++c) s += part[c * nb + j];
        ypy[l * nb + (l + 1 + j)] = s;
      }
    }
    // No post-write sync; W reads happen after a single CLUSTER_SYNC below.
  }

  CLUSTER_SYNC();

  // ----- W construction (per-CTA row chunk, SMEM-resident wmat_s) -----
  // For each j, wmat_s[row, j] = -tau_j * (Y[row, j] + Σ_{l<j} wmat_s[row, l] * ypy[l, j]).
  // Each CTA computes its row chunk independently (no cross-CTA dependency).
  // wmat_s layout matches panel_s: column-major within (CHUNK_SIZE, nb).
  for (int j = 0; j < width; ++j) {
    float tau_j = t[panel + j];
    int col_j = panel + j;
    int rs_g = max(row_start, panel);
    int rs_l = rs_g - row_start;
    if (rs_l < local_size) {
      for (int li = rs_l + tid; li < local_size; li += nthread) {
        int global_row = row_start + li;
        float y;
        if (global_row < col_j) y = 0.0f;
        else if (global_row == col_j) y = 1.0f;
        else y = panel_s[li + j * CHUNK_SIZE];
        float accum = y;
        for (int l = 0; l < j; ++l) {
          accum += wmat_s[li + l * CHUNK_SIZE] * ypy[l * nb + j];
        }
        wmat_s[li + j * CHUNK_SIZE] = -tau_j * accum;
      }
    }
    __syncthreads();
  }

  // ----- Write panel back to HBM and pack Y / W to ygemm / wgemm. -----
  // Panel write: each CTA writes its chunk's panel columns to h[row, col].
  for (int idx = tid; idx < local_size * nb; idx += nthread) {
    int local_col = idx / local_size;
    int local_row = idx - local_col * local_size;
    if (local_col >= width) break;
    int global_row = row_start + local_row;
    int global_col = panel + local_col;
    a[global_row + (int64_t)global_col * n] = panel_s[local_row + local_col * CHUNK_SIZE];
  }

  // Y / W pack: each CTA writes its rows in [max(row_start, panel), row_end).
  {
    int rs_g = max(row_start, panel);
    int rs_l = rs_g - row_start;
    if (rs_l < local_size) {
      for (int li = rs_l + tid; li < local_size; li += nthread) {
        int global_row = row_start + li;
        for (int j = 0; j < width; ++j) {
          int col = panel + j;
          float yval;
          if (global_row < col) yval = 0.0f;
          else if (global_row == col) yval = 1.0f;
          else yval = panel_s[li + j * CHUNK_SIZE];
          yout[global_row + j * n] = yval;
          wout[global_row + j * n] = wmat_s[li + j * CHUNK_SIZE];
        }
      }
    }
  }

#undef CLUSTER_SYNC
}

void qr_blocked16_gemm_cluster_smem(
    torch::Tensor h,
    torch::Tensor tau,
    torch::Tensor work,
    torch::Tensor ygemm,
    torch::Tensor wgemm,
    torch::Tensor tmp,
    torch::Tensor scratch) {
  TORCH_CHECK(h.is_cuda() && tau.is_cuda() && work.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(ygemm.is_cuda() && wgemm.is_cuda() && tmp.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(scratch.is_cuda(), "CUDA workspace required");
  TORCH_CHECK(h.scalar_type() == torch::kFloat32, "h must be float32");
  TORCH_CHECK(scratch.scalar_type() == torch::kFloat32, "scratch must be float32");
  TORCH_CHECK(h.stride(1) == 1 && h.stride(2) == h.size(1), "h column-major");
  TORCH_CHECK(ygemm.stride(1) == 1 && ygemm.stride(2) == h.size(1), "ygemm column-major");
  TORCH_CHECK(wgemm.stride(1) == 1 && wgemm.stride(2) == h.size(1), "wgemm column-major");
  const int batch = static_cast<int>(h.size(0));
  const int n = static_cast<int>(h.size(1));
  TORCH_CHECK(h.size(2) == n, "h square");
  TORCH_CHECK(n == 4096, "qr_blocked16_gemm_cluster_smem only supports n=4096 currently");

  constexpr int CTAS = 16;
  constexpr int CHUNK_SIZE = 256;
  static_assert(CTAS * CHUNK_SIZE == 4096, "CTAS*CHUNK_SIZE must equal n");
  constexpr int nb = 16;
  const int expected = CTAS * nb + 32 + nb + nb * nb;
  TORCH_CHECK(scratch.numel() >= (int64_t)batch * expected, "scratch too small");

  constexpr float one = 1.0f;
  constexpr float zero = 0.0f;
  cublasHandle_t handle = raw_blas_handle();
  cublasComputeType_t compute = CUBLAS_COMPUTE_32F_FAST_16BF;
  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;

  // Dynamic SMEM is just wmat_s (panel_s is static).
  constexpr int dyn_smem_bytes = CHUNK_SIZE * nb * (int)sizeof(float);
  static bool smem_optin_done = false;
  if (!smem_optin_done) {
    int dev = 0;
    cudaGetDevice(&dev);
    int max_optin = 0;
    cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
    TORCH_CHECK(max_optin >= dyn_smem_bytes,
                "cluster_smem dyn SMEM (", dyn_smem_bytes, ") > device opt-in cap (", max_optin, ")");
    auto kernel_ptr = &qr_block16_panel_cluster_smem_kernel<CTAS, CHUNK_SIZE>;
    cudaError_t err = cudaFuncSetAttribute(
        reinterpret_cast<const void*>(kernel_ptr),
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        dyn_smem_bytes);
    TORCH_CHECK(err == cudaSuccess,
                "cudaFuncSetAttribute(MaxDynamicSharedMemorySize) failed: ",
                cudaGetErrorString(err), " (requested=", dyn_smem_bytes,
                ", cap=", max_optin, ")");
    // CTAS > 8 requires non-portable cluster size attribute.
    if (CTAS > 8) {
      err = cudaFuncSetAttribute(
          reinterpret_cast<const void*>(kernel_ptr),
          cudaFuncAttributeNonPortableClusterSizeAllowed, 1);
      TORCH_CHECK(err == cudaSuccess,
                  "cudaFuncSetAttribute(NonPortableClusterSizeAllowed) failed: ",
                  cudaGetErrorString(err));
    }
    smem_optin_done = true;
  }

  const int threads = 256;
  for (int panel = 0; panel < n; panel += nb) {
    int width = std::min(nb, n - panel);
    int m = n - panel;

    dim3 panel_grid(batch, CTAS, 1);
    qr_block16_panel_cluster_smem_kernel<CTAS, CHUNK_SIZE>
        <<<panel_grid, threads, dyn_smem_bytes>>>(
        h.data_ptr<float>(), tau.data_ptr<float>(), work.data_ptr<float>(),
        ygemm.data_ptr<float>(), wgemm.data_ptr<float>(),
        scratch.data_ptr<float>(),
        h.stride(0), tau.stride(0), ygemm.stride(0),
        n, panel);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int trailing = n - panel - width;
    if (trailing > 0) {
      float *cptr = h.data_ptr<float>() + panel + static_cast<int64_t>(panel + width) * n;
      float *yptr = ygemm.data_ptr<float>() + panel;
      float *wptr = wgemm.data_ptr<float>() + panel;
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_T, CUBLAS_OP_N, width, trailing, m,
              &one, wptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              &zero, tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              batch, compute, algo),
          "qr_blocked16_gemm_cluster_smem GEMM1");
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_N, CUBLAS_OP_N, m, trailing, width,
              &one, yptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              &one, cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              batch, compute, algo),
          "qr_blocked16_gemm_cluster_smem GEMM2");
    }
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}


// n=176, n=352: panel + 2x BF16 cuBLAS trailing GEMM. Both shapes have only
// dense / well-conditioned public residual tests, so BF16 (FP32 accumulate)
// for both GEMMs is within tolerance. The n=512 band test that requires
// FP32 GEMM2 routes through qr_blocked16_gemm_higham_concat instead.
void qr_blocked16_gemm(torch::Tensor h, torch::Tensor tau, torch::Tensor work, torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp) {
  TORCH_CHECK(h.is_cuda() && tau.is_cuda() && work.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(ygemm.is_cuda() && wgemm.is_cuda() && tmp.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(h.scalar_type() == torch::kFloat32, "h must be float32");
  TORCH_CHECK(tau.scalar_type() == torch::kFloat32, "tau must be float32");
  TORCH_CHECK(work.scalar_type() == torch::kFloat32, "work must be float32");
  TORCH_CHECK(ygemm.scalar_type() == torch::kFloat32 && wgemm.scalar_type() == torch::kFloat32, "gemm work must be float32");
  TORCH_CHECK(tmp.scalar_type() == torch::kFloat32, "tmp must be float32");
  TORCH_CHECK(h.dim() == 3 && tau.dim() == 2 && work.dim() == 3, "bad ranks");
  TORCH_CHECK(ygemm.dim() == 3 && wgemm.dim() == 3 && tmp.dim() == 3, "bad ranks");
  const int batch = static_cast<int>(h.size(0));
  const int n = static_cast<int>(h.size(1));
  TORCH_CHECK(h.size(2) == n, "h must be square");
  TORCH_CHECK(work.size(1) == n && work.size(2) == 16, "bad work shape");
  TORCH_CHECK(ygemm.size(0) == batch && ygemm.size(1) == n && ygemm.size(2) == 16, "bad ygemm shape");
  TORCH_CHECK(wgemm.size(0) == batch && wgemm.size(1) == n && wgemm.size(2) == 16, "bad wgemm shape");
  TORCH_CHECK(tmp.size(0) == batch && tmp.size(1) == 16 && tmp.size(2) == n, "bad tmp shape");
  TORCH_CHECK(h.stride(1) == 1 && h.stride(2) == n, "h must be compact column-major");
  TORCH_CHECK(work.stride(2) == 1 && work.stride(1) == 16, "work must be compact");
  TORCH_CHECK(ygemm.stride(1) == 1 && ygemm.stride(2) == n, "ygemm must be column-major");
  TORCH_CHECK(wgemm.stride(1) == 1 && wgemm.stride(2) == n, "wgemm must be column-major");
  TORCH_CHECK(tmp.stride(1) == 1 && tmp.stride(2) == 16, "tmp must be column-major");

  constexpr int nb = 16;
  constexpr float one = 1.0f;
  constexpr float zero = 0.0f;
  cublasHandle_t handle = raw_blas_handle();
  cublasComputeType_t compute = CUBLAS_COMPUTE_32F_FAST_16BF;
  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;

  for (int panel = 0; panel < n; panel += nb) {
    int width = min(nb, n - panel);
    int m = n - panel;
    dim3 panel_grid(batch, 1, 1);
    qr_block16_panel_kernel<<<panel_grid, 256>>>(
        h.data_ptr<float>(), tau.data_ptr<float>(), work.data_ptr<float>(),
        ygemm.data_ptr<float>(), wgemm.data_ptr<float>(),
        h.stride(0), tau.stride(0), n, panel);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int trailing = n - panel - width;
    if (trailing > 0) {
      float *cptr = h.data_ptr<float>() + panel + (panel + width) * n;
      float *yptr = ygemm.data_ptr<float>() + panel;
      float *wptr = wgemm.data_ptr<float>() + panel;
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_T, CUBLAS_OP_N, width, trailing, m,
              &one, wptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              &zero, tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              batch, compute, algo),
          "qr_blocked16_gemm GEMM1");
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_N, CUBLAS_OP_N, m, trailing, width,
              &one, yptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              &one, cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              batch, compute, algo),
          "qr_blocked16_gemm GEMM2");
    }
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void qr_blocked16_gemm_smem_panel(
    torch::Tensor h,
    torch::Tensor tau,
    torch::Tensor ygemm,
    torch::Tensor wgemm,
    torch::Tensor tmp,
    int warps_per_matrix) {
  TORCH_CHECK(h.is_cuda() && tau.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(ygemm.is_cuda() && wgemm.is_cuda() && tmp.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(h.scalar_type() == torch::kFloat32, "h must be float32");
  TORCH_CHECK(tau.scalar_type() == torch::kFloat32, "tau must be float32");
  TORCH_CHECK(ygemm.scalar_type() == torch::kFloat32 && wgemm.scalar_type() == torch::kFloat32, "gemm work must be float32");
  TORCH_CHECK(tmp.scalar_type() == torch::kFloat32, "tmp must be float32");
  TORCH_CHECK(h.dim() == 3 && tau.dim() == 2, "bad ranks");
  TORCH_CHECK(ygemm.dim() == 3 && wgemm.dim() == 3 && tmp.dim() == 3, "bad ranks");
  TORCH_CHECK(warps_per_matrix >= 1 && warps_per_matrix <= 8, "bad warp count");
  const int batch = static_cast<int>(h.size(0));
  const int n = static_cast<int>(h.size(1));
  TORCH_CHECK(h.size(2) == n, "h must be square");
  TORCH_CHECK(tau.size(0) == batch && tau.size(1) == n, "bad tau shape");
  TORCH_CHECK(ygemm.size(0) == batch && ygemm.size(1) == n && ygemm.size(2) == 16, "bad ygemm shape");
  TORCH_CHECK(wgemm.size(0) == batch && wgemm.size(1) == n && wgemm.size(2) == 16, "bad wgemm shape");
  TORCH_CHECK(tmp.size(0) == batch && tmp.size(1) == 16 && tmp.size(2) == n, "bad tmp shape");
  TORCH_CHECK(h.stride(1) == 1 && h.stride(2) == n, "h must be compact column-major");
  TORCH_CHECK(ygemm.stride(1) == 1 && ygemm.stride(2) == n, "ygemm must be column-major");
  TORCH_CHECK(wgemm.stride(1) == 1 && wgemm.stride(2) == n, "wgemm must be column-major");
  TORCH_CHECK(tmp.stride(1) == 1 && tmp.stride(2) == 16, "tmp must be column-major");

  constexpr int nb = 16;
  constexpr float one = 1.0f;
  constexpr float zero = 0.0f;
  cublasHandle_t handle = raw_blas_handle();

  static int max_smem_optin = 0;
  if (max_smem_optin == 0) {
    int dev = 0;
    cudaGetDevice(&dev);
    cudaDeviceGetAttribute(&max_smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
    cudaFuncSetAttribute(
        qr_block16_panel_smem_kernel,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        max_smem_optin);
  }

  // n in {1024, 2048} have no FP32-IEEE-only stress case in the public
  // benchmark inputs (only n=512 has the `band` constraint), so both
  // trailing GEMMs run on BF16 tensor cores.
  cublasComputeType_t compute = CUBLAS_COMPUTE_32F_FAST_16BF;
  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;

  for (int panel = 0; panel < n; panel += nb) {
    int width = min(nb, n - panel);
    int m = n - panel;
    int smem_bytes = static_cast<int>((static_cast<int64_t>(m) * nb + 2 * (warps_per_matrix + 1) + nb * nb) * sizeof(float));
    TORCH_CHECK(smem_bytes <= max_smem_optin, "panel smem request too large");

    dim3 panel_grid(batch, 1, 1);
    dim3 panel_block(32, warps_per_matrix, 1);
    qr_block16_panel_smem_kernel<<<panel_grid, panel_block, smem_bytes>>>(
        h.data_ptr<float>(),
        tau.data_ptr<float>(),
        ygemm.data_ptr<float>(),
        wgemm.data_ptr<float>(),
        h.stride(0),
        tau.stride(0),
        ygemm.stride(0),
        n,
        panel);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int trailing = n - panel - width;
    if (trailing > 0) {
      float *cptr = h.data_ptr<float>() + panel + static_cast<int64_t>(panel + width) * n;
      float *yptr = ygemm.data_ptr<float>() + panel;
      float *wptr = wgemm.data_ptr<float>() + panel;
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_T, CUBLAS_OP_N, width, trailing, m,
              &one, wptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              &zero, tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              batch, compute, algo),
          "qr_blocked16_gemm_smem_panel GEMM1");
      check_status(
          cublasGemmStridedBatchedEx(
              handle, CUBLAS_OP_N, CUBLAS_OP_N, m, trailing, width,
              &one, yptr, CUDA_R_32F, n, static_cast<long long>(n * nb),
                    tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
              &one, cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
              batch, compute, algo),
          "qr_blocked16_gemm_smem_panel GEMM2");
    }
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

// E49: K=48 concat-correction Higham GEMM2 for n=512.
//
// The n=512 `band` stress case requires GEMM2 (C += Y @ tmp) to be ~FP32-IEEE
// accurate; a plain BF16 GEMM2 produces a scaled residual ~23.6× the budget.
// The honest precision/throughput unlock is Higham one-pass refinement:
//
//   Y = bf16(Y) + Y_res        (Y_res < ulp(Y))
//   tmp = bf16(tmp) + tmp_res
//   Y * tmp = bf16(Y)*bf16(tmp) + Y_res*bf16(tmp) + bf16(Y)*tmp_res + O(eps^2)
//
// E48 ran these as three separate BF16 GEMMs but tripled the trailing-C
// HBM traffic. E49 collapses them into a single K=48 BF16 GEMM:
//
//   Y_concat48 layout (batch, n, 48), column-major BF16:
//     columns  0..15 : bf16(Y)
//     columns 16..31 : Y_res
//     columns 32..47 : bf16(Y) (duplicate)
//   tmp_concat48 layout (batch, 48, n), column-major BF16:
//     rows  0..15 : bf16(tmp)
//     rows 16..31 : bf16(tmp) (duplicate)
//     rows 32..47 : tmp_res
//
// The duplication of bf16(Y) and bf16(tmp) wastes BF16 storage but reads/
// writes the trailing C block ONCE per panel instead of three times.

__global__ void cast_y_to_concat48_kernel(
    const float *y_src,                  // ygemm (batch, n, 16) FP32
    __nv_bfloat16 *y_concat_dst,         // y_concat48 (batch, n, 48) BF16
    int batch,
    int n,
    int panel) {
  constexpr int nb = 16;
  int b = blockIdx.y;
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int m = n - panel;
  if (b >= batch || idx >= m * nb) return;
  int r_local = idx % m;
  int r = panel + r_local;
  int c = idx / m;  // c in [0, 16)

  int64_t src_off = static_cast<int64_t>(b) * n * nb + r + c * n;
  float v = y_src[src_off];
  __nv_bfloat16 v_bf = __float2bfloat16(v);
  float v_fp = __bfloat162float(v_bf);
  __nv_bfloat16 v_res = __float2bfloat16(v - v_fp);

  int64_t dst_base = static_cast<int64_t>(b) * n * 48 + r;
  y_concat_dst[dst_base + c * n] = v_bf;
  y_concat_dst[dst_base + (c + 16) * n] = v_res;
  y_concat_dst[dst_base + (c + 32) * n] = v_bf;
}

__global__ void cast_tmp_to_concat48_kernel(
    const float *tmp_src,                  // tmp (batch, 16, n) FP32
    __nv_bfloat16 *tmp_concat_dst,         // tmp_concat48 (batch, 48, n) BF16
    int batch,
    int n,
    int trailing) {
  constexpr int nb = 16;
  int b = blockIdx.y;
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (b >= batch || idx >= nb * trailing) return;
  int r = idx % nb;  // r in [0, 16)
  int c = idx / nb;  // c in [0, trailing)

  int64_t src_off = static_cast<int64_t>(b) * nb * n + r + c * nb;
  float v = tmp_src[src_off];
  __nv_bfloat16 v_bf = __float2bfloat16(v);
  float v_fp = __bfloat162float(v_bf);
  __nv_bfloat16 v_res = __float2bfloat16(v - v_fp);

  int64_t dst_base = static_cast<int64_t>(b) * 48 * n + c * 48;
  tmp_concat_dst[dst_base + r] = v_bf;
  tmp_concat_dst[dst_base + (r + 16)] = v_bf;
  tmp_concat_dst[dst_base + (r + 32)] = v_res;
}

void qr_blocked16_gemm_higham_concat(torch::Tensor h, torch::Tensor tau, torch::Tensor work,
                                     torch::Tensor ygemm, torch::Tensor wgemm, torch::Tensor tmp,
                                     torch::Tensor y_concat48, torch::Tensor tmp_concat48) {
  TORCH_CHECK(h.is_cuda() && tau.is_cuda() && work.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(y_concat48.is_cuda() && tmp_concat48.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(h.scalar_type() == torch::kFloat32, "h must be float32");
  TORCH_CHECK(ygemm.scalar_type() == torch::kFloat32 && wgemm.scalar_type() == torch::kFloat32, "ygemm/wgemm must be float32");
  TORCH_CHECK(tmp.scalar_type() == torch::kFloat32, "tmp must be float32");
  TORCH_CHECK(y_concat48.scalar_type() == torch::kBFloat16, "y_concat48 must be bfloat16");
  TORCH_CHECK(tmp_concat48.scalar_type() == torch::kBFloat16, "tmp_concat48 must be bfloat16");
  const int batch = static_cast<int>(h.size(0));
  const int n = static_cast<int>(h.size(1));
  TORCH_CHECK(h.size(2) == n, "h must be square");
  TORCH_CHECK(y_concat48.size(0) == batch && y_concat48.size(1) == n && y_concat48.size(2) == 48, "bad y_concat48 shape");
  TORCH_CHECK(tmp_concat48.size(0) == batch && tmp_concat48.size(1) == 48 && tmp_concat48.size(2) == n, "bad tmp_concat48 shape");
  TORCH_CHECK(h.stride(1) == 1 && h.stride(2) == n, "h must be compact column-major");
  TORCH_CHECK(y_concat48.stride(1) == 1 && y_concat48.stride(2) == n, "y_concat48 must be column-major (ld=n)");
  TORCH_CHECK(tmp_concat48.stride(1) == 1 && tmp_concat48.stride(2) == 48, "tmp_concat48 must be column-major (ld=48)");

  constexpr int nb = 16;
  constexpr float one = 1.0f;
  constexpr float zero = 0.0f;
  cublasHandle_t handle = raw_blas_handle();
  cublasComputeType_t compute = CUBLAS_COMPUTE_32F_FAST_16BF;
  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;
  bool deep_timing = qr_deep_timing_enabled();
  float panel_ms = 0.0f;
  float gemm1_ms = 0.0f;
  float cast_y_ms = 0.0f;
  float cast_tmp_ms = 0.0f;
  float gemm2_ms = 0.0f;
  cudaEvent_t deep_begin, deep_end;
  if (deep_timing) {
    C10_CUDA_CHECK(cudaEventCreate(&deep_begin));
    C10_CUDA_CHECK(cudaEventCreate(&deep_end));
  }

  for (int panel = 0; panel < n; panel += nb) {
    int width = std::min(nb, n - panel);
    int m = n - panel;

    // E62: use the SMEM-resident panel kernel here too. The deep E61/E62
    // timing split showed the old global-memory panel consuming ~11.7 ms of
    // the ~18 ms n=512 QR path, much more than either trailing GEMM.
    constexpr int warps_per_matrix = 4;
    static int max_smem_optin = 0;
    if (max_smem_optin == 0) {
      int dev = 0;
      cudaGetDevice(&dev);
      cudaDeviceGetAttribute(&max_smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
      cudaFuncSetAttribute(
          qr_block16_panel_smem_kernel,
          cudaFuncAttributeMaxDynamicSharedMemorySize,
          max_smem_optin);
    }
    int smem_bytes = static_cast<int>((static_cast<int64_t>(m) * nb + 2 * (warps_per_matrix + 1) + nb * nb) * sizeof(float));
    TORCH_CHECK(smem_bytes <= max_smem_optin, "panel smem request too large");
    dim3 panel_grid(batch, 1, 1);
    dim3 panel_block(32, warps_per_matrix, 1);
    if (deep_timing) qr_deep_begin(deep_begin);
    qr_block16_panel_smem_kernel<<<panel_grid, panel_block, smem_bytes>>>(
        h.data_ptr<float>(),
        tau.data_ptr<float>(),
        ygemm.data_ptr<float>(),
        wgemm.data_ptr<float>(),
        h.stride(0),
        tau.stride(0),
        ygemm.stride(0),
        n,
        panel);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    if (deep_timing) qr_deep_end(deep_begin, deep_end, &panel_ms);

    int trailing = n - panel - width;
    if (trailing <= 0) continue;

    float *cptr = h.data_ptr<float>() + panel + (panel + width) * n;
    float *wptr_fp = wgemm.data_ptr<float>() + panel;

    // GEMM1: tmp = W^T @ C (BF16, FP32 accumulate).
    if (deep_timing) qr_deep_begin(deep_begin);
    check_status(
        cublasGemmStridedBatchedEx(
            handle, CUBLAS_OP_T, CUBLAS_OP_N, width, trailing, m,
            &one, wptr_fp, CUDA_R_32F, n, static_cast<long long>(n * nb),
                  cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
            &zero, tmp.data_ptr<float>(), CUDA_R_32F, nb, static_cast<long long>(nb * n),
            batch, compute, algo),
        "qr_blocked16_gemm_higham_concat GEMM1");
    if (deep_timing) qr_deep_end(deep_begin, deep_end, &gemm1_ms);

    // Cast ygemm into Y_concat48 (full 48 cols populated).
    {
      int threads_per_block = 256;
      int blocks_x = (m * nb + threads_per_block - 1) / threads_per_block;
      dim3 grid(blocks_x, batch);
      if (deep_timing) qr_deep_begin(deep_begin);
      cast_y_to_concat48_kernel<<<grid, threads_per_block>>>(
          ygemm.data_ptr<float>(),
          reinterpret_cast<__nv_bfloat16 *>(y_concat48.data_ptr<at::BFloat16>()),
          batch, n, panel);
      C10_CUDA_KERNEL_LAUNCH_CHECK();
      if (deep_timing) qr_deep_end(deep_begin, deep_end, &cast_y_ms);
    }

    // Cast tmp into tmp_concat48.
    {
      int threads_per_block = 256;
      int blocks_x = (nb * trailing + threads_per_block - 1) / threads_per_block;
      dim3 grid(blocks_x, batch);
      if (deep_timing) qr_deep_begin(deep_begin);
      cast_tmp_to_concat48_kernel<<<grid, threads_per_block>>>(
          tmp.data_ptr<float>(),
          reinterpret_cast<__nv_bfloat16 *>(tmp_concat48.data_ptr<at::BFloat16>()),
          batch, n, trailing);
      C10_CUDA_KERNEL_LAUNCH_CHECK();
      if (deep_timing) qr_deep_end(deep_begin, deep_end, &cast_tmp_ms);
    }

    // GEMM2 concat: C += Y_concat48 @ tmp_concat48 (K=48, BF16 tensor core,
    // FP32 accumulator). One C update for all three Higham terms.
    __nv_bfloat16 *yptr_concat = reinterpret_cast<__nv_bfloat16 *>(y_concat48.data_ptr<at::BFloat16>()) + panel;
    __nv_bfloat16 *tptr_concat = reinterpret_cast<__nv_bfloat16 *>(tmp_concat48.data_ptr<at::BFloat16>());
    if (deep_timing) qr_deep_begin(deep_begin);
    check_status(
        cublasGemmStridedBatchedEx(
            handle, CUBLAS_OP_N, CUBLAS_OP_N, m, trailing, 48,
            &one, yptr_concat, CUDA_R_16BF, n, static_cast<long long>(n * 48),
                  tptr_concat, CUDA_R_16BF, 48, static_cast<long long>(48 * n),
            &one, cptr, CUDA_R_32F, n, static_cast<long long>(h.stride(0)),
            batch, compute, algo),
        "qr_blocked16_gemm_higham_concat GEMM2_concat");
    if (deep_timing) qr_deep_end(deep_begin, deep_end, &gemm2_ms);
  }
  if (deep_timing) {
    float total_ms = panel_ms + gemm1_ms + cast_y_ms + cast_tmp_ms + gemm2_ms;
    std::printf(
        "[deep] route=blocked16_higham_concat batch=%d n=%d "
        "panel=%.3f gemm1=%.3f cast_y=%.3f cast_tmp=%.3f gemm2=%.3f sum=%.3f\n",
        batch, n, panel_ms, gemm1_ms, cast_y_ms, cast_tmp_ms, gemm2_ms, total_ms);
    std::fflush(stdout);
    C10_CUDA_CHECK(cudaEventDestroy(deep_begin));
    C10_CUDA_CHECK(cudaEventDestroy(deep_end));
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

__global__ void fill_ptrs(
    float **h_ptrs,
    float **tau_ptrs,
    float *h,
    float *tau,
    int64_t h_step,
    int64_t tau_step,
    int batch) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < batch) {
    h_ptrs[i] = h + i * h_step;
    tau_ptrs[i] = tau + i * tau_step;
  }
}

static void check_status(cublasStatus_t status, const char *where) {
  if (status != CUBLAS_STATUS_SUCCESS) {
    TORCH_CHECK(false, where, " failed with status ", static_cast<int>(status));
  }
}

static cublasHandle_t raw_blas_handle() {
  static cublasHandle_t handle = nullptr;
  if (handle == nullptr) {
    check_status(cublasCreate(&handle), "cublasCreate");
  }
  return handle;
}

void qr_cublas_geqrf(torch::Tensor h, torch::Tensor tau) {
  TORCH_CHECK(h.is_cuda() && tau.is_cuda(), "CUDA tensors required");
  TORCH_CHECK(h.scalar_type() == torch::kFloat32, "h must be float32");
  TORCH_CHECK(tau.scalar_type() == torch::kFloat32, "tau must be float32");
  TORCH_CHECK(h.dim() == 3 && tau.dim() == 2, "bad ranks");
  const int batch = static_cast<int>(h.size(0));
  const int n = static_cast<int>(h.size(1));
  TORCH_CHECK(h.size(2) == n, "h must be square");
  TORCH_CHECK(tau.size(0) == batch && tau.size(1) == n, "bad tau shape");
  TORCH_CHECK(h.stride(1) == 1, "h must be column-major per matrix");

  auto ptr_opts = torch::TensorOptions().device(h.device()).dtype(torch::kUInt64);
  auto h_ptrs = torch::empty({batch}, ptr_opts);
  auto tau_ptrs = torch::empty({batch}, ptr_opts);
  const int threads = 256;
  const int blocks = (batch + threads - 1) / threads;
  fill_ptrs<<<blocks, threads>>>(
      reinterpret_cast<float **>(h_ptrs.data_ptr<uint64_t>()),
      reinterpret_cast<float **>(tau_ptrs.data_ptr<uint64_t>()),
      h.data_ptr<float>(),
      tau.data_ptr<float>(),
      h.stride(0),
      tau.stride(0),
      batch);
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  int info = 0;
  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
  check_status(
      cublasSgeqrfBatched(
          handle,
          n,
          n,
          reinterpret_cast<float **>(h_ptrs.data_ptr<uint64_t>()),
          static_cast<int>(h.stride(2)),
          reinterpret_cast<float **>(tau_ptrs.data_ptr<uint64_t>()),
          &info,
          batch),
      "cublasSgeqrfBatched");
  TORCH_CHECK(info == 0, "bad geqrf argument ", info);
}
"""

# E59 step 1: tcgen05 smoke kernel is its own load_inline extension because:
#   1. It needs -gencode=arch=compute_100a,code=sm_100a (the architecture-specific
#      `a` variant) to expose tcgen05 instructions, but compiling the existing
#      cuBLAS-heavy submission with that flag triggered a 10+ minute ptxas hang
#      (probably a known compute_100a / cuBLAS-header interaction).
#   2. Isolating the tcgen05 work keeps the main extension's cold compile fast
#      (~30-60 s, same as E58) so iteration on the active route stays cheap.
#   3. Long-term this also makes the dev/prod split cleaner: the smoke is a
#      debug scaffold, not part of the active dispatch.
TCGEN05_CPP_SRC = """
void qr_tcgen05_smoke(torch::Tensor out);
void qr_tcgen05_bf16_gemm(torch::Tensor A, torch::Tensor B, torch::Tensor D);
void qr_tcgen05_bf16_gemm_k32(torch::Tensor A, torch::Tensor B, torch::Tensor D);
void qr_tcgen05_bf16_gemm_batched(torch::Tensor A, torch::Tensor B, torch::Tensor D);
void qr_tcgen05_bf16_gemm_persistent(torch::Tensor A, torch::Tensor B, torch::Tensor D, torch::Tensor counter);
"""

TCGEN05_CUDA_SRC = r"""
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <torch/extension.h>

// E59 step 1: tcgen05 BF16 smoke kernel — phase B.
//
// Phase A (alloc + dealloc) verified that compute_100a compiles, that the
// kernel launches without trap, and that whole-warp sync.aligned semantics
// are respected. Phase B adds the actual MMA: a single tcgen05.mma kind::f16
// of A (M=64 × K=16 BF16 ones) * B (K=16 × N=64 BF16 ones), producing a
// FP32 64×64 accumulator in TMEM. Expected: every entry = 16.0.
//
// Verifies:
//   1. tcgen05.mma.cta_group::1.kind::f16 emits UTCHMMA in SASS (binary
//      go/no-go for the entire E59+ direction).
//   2. tcgen05.commit on an mbarrier + mbarrier.try_wait.parity.acquire
//      forms a working MMA-completion signal.
//   3. tcgen05.ld.32x32b.x8 reads back the FP32 accumulator correctly.
//   4. The SMEM-descriptor layout (LBO/SBO/no-swizzle bit-46) matches what
//      tcgen05.mma expects for contiguous (M, 8) BF16 blocks.
//
// CRITICAL warp-discipline rules (learned the hard way in phase A):
//   - tcgen05.alloc / dealloc / ld         : WHOLE WARP (sync.aligned)
//   - tcgen05.mma / commit                 : single thread (elect.sync)
//   - mbarrier.init / arrive               : single thread (elect.sync)
//   - mbarrier.try_wait.parity.acquire     : all threads (per-thread spin)
//   - tcgen05.fence::after_thread_sync     : all threads (it's a fence)

__device__ __forceinline__ uint32_t qr_e59_elect_sync() {
  uint32_t pred = 0;
  asm volatile(
      "{\n\t"
      ".reg .pred %%px;\n\t"
      "elect.sync _|%%px, %1;\n\t"
      "@%%px mov.s32 %0, 1;\n\t"
      "}"
      : "+r"(pred) : "r"(0xFFFFFFFFu));
  return pred;
}

// Encode a byte offset / address into the low 14 bits of (x >> 4) for
// tcgen05 SMEM descriptors. PTX 8.7 §9.7.16.1.
__device__ __forceinline__ uint64_t qr_e59_desc_encode(uint64_t x) {
  return (x & 0x3FFFFULL) >> 4ULL;
}

// 128 threads (4 warps): M=128 is matmul_v1's canonical "Layout D" — rows 0..127
// laid out one-per-lane across all 128 TMEM lanes, so each warp's 32-lane
// tcgen05.ld reads a clean 32-row slab. (M=64 has a doubled layout that packs
// 2 rows per lane in lanes 0..31; that's harder to get right on the first try
// and is not what we'll use for QR.)
__global__ __launch_bounds__(128) void qr_tcgen05_smoke_kernel(__nv_bfloat16 *out) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 16;  // == MMA_K for BF16

  __shared__ __align__(1024) __nv_bfloat16 A_smem_arr[BLOCK_M * BLOCK_K];
  __shared__ __align__(1024) __nv_bfloat16 B_smem_arr[BLOCK_K * BLOCK_N];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ uint64_t mbar_mem[1];
  __shared__ uint32_t tmem_addr_mem[1];

  const int tid = threadIdx.x;
  const int warp_id = tid >> 5;
  const int lane_id = tid & 31;

  // Fill A, B with 1.0. With BF16 ones and K=16, A*B -> D[i,j] = 16.0.
  for (int idx = tid; idx < BLOCK_M * BLOCK_K; idx += blockDim.x) {
    A_smem_arr[idx] = __float2bfloat16(1.0f);
  }
  for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += blockDim.x) {
    B_smem_arr[idx] = __float2bfloat16(1.0f);
  }

  const uint32_t A_smem = static_cast<uint32_t>(__cvta_generic_to_shared(A_smem_arr));
  const uint32_t B_smem = static_cast<uint32_t>(__cvta_generic_to_shared(B_smem_arr));
  const uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_mem));
  const uint32_t tmem_addr_smem = static_cast<uint32_t>(__cvta_generic_to_shared(tmem_addr_mem));

  // Whole-warp alloc on warp 0, single-thread mbarrier init on warp 1.
  if (warp_id == 0) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                 :: "r"(tmem_addr_smem), "n"(BLOCK_N));
  } else if (warp_id == 1 && qr_e59_elect_sync()) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(mbar_addr));
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  __syncthreads();
  const uint32_t taddr = tmem_addr_mem[0];

  asm volatile("tcgen05.fence::after_thread_sync;");

  // Single-thread MMA issue. Descriptors: contiguous (M, 8) BF16 blocks,
  // LBO = stride between K-blocks (= M * 16 bytes), SBO = 128 (= 8 cols *
  // 16 bytes), bit-46 set per matmul_v1.cu convention.
  if (warp_id == 0 && qr_e59_elect_sync()) {
    const uint64_t LBO_A = static_cast<uint64_t>(BLOCK_M) * 16;  // 1024
    const uint64_t LBO_B = static_cast<uint64_t>(BLOCK_N) * 16;  // 1024
    const uint64_t SBO   = 8ULL * 16ULL;                          // 128
    const uint64_t a_desc = qr_e59_desc_encode(A_smem)
                          | (qr_e59_desc_encode(LBO_A) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    const uint64_t b_desc = qr_e59_desc_encode(B_smem)
                          | (qr_e59_desc_encode(LBO_B) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    // i_desc: dtype=FP32 acc, atype=BF16, btype=BF16, MMA_N/8, MMA_M/16.
    constexpr uint32_t i_desc = (1U << 4U)
                              | (1U << 7U)
                              | (1U << 10U)
                              | (((uint32_t)BLOCK_N >> 3U) << 17U)
                              | (((uint32_t)BLOCK_M >> 4U) << 24U);
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, 0, 0;\n\t"  // p = (0 != 0) = false -> overwrite D (do NOT accumulate)
        "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
        "}"
        :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc));
    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mbar_addr) : "memory");
  }

  // All threads wait for MMA completion via the mbarrier (per-thread spin
  // with .try_wait.parity.acquire). matmul_v1's "ticks" 0x989680 is just the
  // try_wait suspend duration, not a loop count — the loop bails on @P1.
  {
    const uint32_t ticks = 0x989680u;
    const int phase = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT_E59:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE_E59;\n\t"
        "bra.uni LAB_WAIT_E59;\n\t"
        "DONE_E59:\n\t"
        "}"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks));
  }

  asm volatile("tcgen05.fence::after_thread_sync;");

  // Whole-warp tcgen05.ld. Each warp covers 32 TMEM lanes; with 4 warps and
  // M=128 we cover all 128 lanes. The lane offset goes in the upper 16 bits
  // of the TMEM address.
  for (int n = 0; n < BLOCK_N / 8; n++) {
    float tmp[8];
    const uint32_t addr = taddr + ((uint32_t)(warp_id * 32) << 16) + (uint32_t)(n * 8);
    asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                 : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                   "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                 : "r"(addr));
    asm volatile("tcgen05.wait::ld.sync.aligned;");

    const int row = warp_id * 32 + lane_id;
    const int col_base = n * 8;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
      out[row * BLOCK_N + col_base + i] = __float2bfloat16(tmp[i]);
    }
  }

  __syncthreads();

  // Whole-warp dealloc (only warp 0 needs to issue, but every lane in that
  // warp must participate — sync.aligned).
  if (warp_id == 0) {
    asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
                 :: "r"(taddr), "n"(BLOCK_N));
  }
}

void qr_tcgen05_smoke(torch::Tensor out) {
  TORCH_CHECK(out.is_cuda(), "out must be CUDA");
  TORCH_CHECK(out.scalar_type() == torch::kBFloat16, "out must be bfloat16");
  TORCH_CHECK(out.numel() == 128 * 64, "out must have 128*64 elements");
  TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
  qr_tcgen05_smoke_kernel<<<1, 128>>>(
      reinterpret_cast<__nv_bfloat16 *>(out.data_ptr<at::BFloat16>()));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

// E59 step 2: QR-shaped tcgen05 BF16 GEMM tile.
//
// Same MMA primitives as the step-1 smoke, but now with real input matrices:
// A (M=128, K=16) BF16 + B (K=16, N=64) BF16, both row-major in global memory,
// loaded into SMEM in the (M, 8) / (N, 8) contiguous-block layouts the MMA
// descriptors expect. Output D (M=128, N=64) BF16. Verifies that we get the
// right answer on non-trivial values (i.e., that the SMEM layout reshape
// is correct), which is the precondition for slotting this into QR.
//
// Shape choice: M=128 N=64 K=16 matches one trailing-apply tile of a QR
// panel (panel width = 16 = K, the QR convention; trailing rows = M, trailing
// cols = N). For (640, 512) panel 0, the trailing-apply C += Y @ tmp has
// 496 rows and 496 cols; we'd tile it as M=128 x N=64 chunks.
__global__ __launch_bounds__(128) void qr_tcgen05_bf16_gemm_kernel(
    const __nv_bfloat16 *A,  // (M, K) row-major BF16 in global memory
    const __nv_bfloat16 *B,  // (K, N) row-major BF16 in global memory
    __nv_bfloat16 *D         // (M, N) row-major BF16 in global memory (output)
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 16;  // == MMA_K for BF16

  __shared__ __align__(1024) __nv_bfloat16 A_smem_arr[BLOCK_M * BLOCK_K];
  __shared__ __align__(1024) __nv_bfloat16 B_smem_arr[BLOCK_K * BLOCK_N];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ uint64_t mbar_mem[1];
  __shared__ uint32_t tmem_addr_mem[1];

  const int tid = threadIdx.x;
  const int warp_id = tid >> 5;
  const int lane_id = tid & 31;

  // Load A from global into SMEM in (M, 8)-block layout.
  //
  // A_global row-major (M, K): A[r,c] at A_global[r*K + c].
  // A_smem (M, 8) blocks: block_b = c/8, c_in_block = c%8.
  //   dst index = block_b * (M*8) + r*8 + c_in_block.
  // (matmul_v1.cu loads this layout from TMA descriptors; we do it from
  // ordinary global loads since the smoke doesn't pull in TMA setup.)
  for (int idx = tid; idx < BLOCK_M * BLOCK_K; idx += blockDim.x) {
    int r = idx / BLOCK_K;
    int c = idx - r * BLOCK_K;
    int block_b = c / 8;
    int c_in_block = c & 7;
    int dst = block_b * BLOCK_M * 8 + r * 8 + c_in_block;
    A_smem_arr[dst] = A[r * BLOCK_K + c];
  }

  // Load B from global into SMEM in (N, 8)-block layout.
  //
  // B_global row-major (K, N): B[k,n] at B_global[k*N + n].
  // B_smem (N, 8) blocks of 8 K-rows each: block_b = k/8, k_in_block = k%8.
  //   dst index = block_b * (N*8) + n*8 + k_in_block.
  for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += blockDim.x) {
    int k = idx / BLOCK_N;
    int n = idx - k * BLOCK_N;
    int block_b = k / 8;
    int k_in_block = k & 7;
    int dst = block_b * BLOCK_N * 8 + n * 8 + k_in_block;
    B_smem_arr[dst] = B[k * BLOCK_N + n];
  }

  const uint32_t A_smem = static_cast<uint32_t>(__cvta_generic_to_shared(A_smem_arr));
  const uint32_t B_smem = static_cast<uint32_t>(__cvta_generic_to_shared(B_smem_arr));
  const uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_mem));
  const uint32_t tmem_addr_smem = static_cast<uint32_t>(__cvta_generic_to_shared(tmem_addr_mem));

  // Whole-warp alloc on warp 0, single-thread mbarrier init on warp 1.
  if (warp_id == 0) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                 :: "r"(tmem_addr_smem), "n"(BLOCK_N));
  } else if (warp_id == 1 && qr_e59_elect_sync()) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(mbar_addr));
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  __syncthreads();
  const uint32_t taddr = tmem_addr_mem[0];

  asm volatile("tcgen05.fence::after_thread_sync;");

  if (warp_id == 0 && qr_e59_elect_sync()) {
    const uint64_t LBO_A = static_cast<uint64_t>(BLOCK_M) * 16;
    const uint64_t LBO_B = static_cast<uint64_t>(BLOCK_N) * 16;
    const uint64_t SBO   = 8ULL * 16ULL;
    const uint64_t a_desc = qr_e59_desc_encode(A_smem)
                          | (qr_e59_desc_encode(LBO_A) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    const uint64_t b_desc = qr_e59_desc_encode(B_smem)
                          | (qr_e59_desc_encode(LBO_B) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    constexpr uint32_t i_desc = (1U << 4U)
                              | (1U << 7U)
                              | (1U << 10U)
                              | (((uint32_t)BLOCK_N >> 3U) << 17U)
                              | (((uint32_t)BLOCK_M >> 4U) << 24U);
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, 0, 0;\n\t"  // p = (0 != 0) = false -> overwrite D (fresh D = A@B)
        "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
        "}"
        :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc));
    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mbar_addr) : "memory");
  }

  // All threads wait for MMA completion.
  {
    const uint32_t ticks = 0x989680u;
    const int phase = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT_E59_GEMM:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE_E59_GEMM;\n\t"
        "bra.uni LAB_WAIT_E59_GEMM;\n\t"
        "DONE_E59_GEMM:\n\t"
        "}"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks));
  }

  asm volatile("tcgen05.fence::after_thread_sync;");

  // Whole-warp tcgen05.ld: 4 warps cover 128 lanes; each warp gets 32 rows
  // of the M=128 output. Each ld.32x32b.x8 reads 8 cols, distributed 8 cols
  // per thread (so 1 lane = 1 row of D, each thread holds 8 col values).
  for (int n_iter = 0; n_iter < BLOCK_N / 8; n_iter++) {
    float tmp[8];
    const uint32_t addr = taddr + ((uint32_t)(warp_id * 32) << 16) + (uint32_t)(n_iter * 8);
    asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                 : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                   "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                 : "r"(addr));
    asm volatile("tcgen05.wait::ld.sync.aligned;");

    const int row = warp_id * 32 + lane_id;
    const int col_base = n_iter * 8;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
      D[row * BLOCK_N + col_base + i] = __float2bfloat16(tmp[i]);
    }
  }

  __syncthreads();

  if (warp_id == 0) {
    asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
                 :: "r"(taddr), "n"(BLOCK_N));
  }
}

void qr_tcgen05_bf16_gemm(torch::Tensor A, torch::Tensor B, torch::Tensor D) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda() && D.is_cuda(), "A, B, D must be CUDA");
  TORCH_CHECK(A.scalar_type() == torch::kBFloat16, "A must be bfloat16");
  TORCH_CHECK(B.scalar_type() == torch::kBFloat16, "B must be bfloat16");
  TORCH_CHECK(D.scalar_type() == torch::kBFloat16, "D must be bfloat16");
  TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && D.is_contiguous(), "must be contiguous");
  TORCH_CHECK(A.numel() == 128 * 16, "A must be (128, 16)");
  TORCH_CHECK(B.numel() == 16 * 64, "B must be (16, 64)");
  TORCH_CHECK(D.numel() == 128 * 64, "D must be (128, 64)");
  qr_tcgen05_bf16_gemm_kernel<<<1, 128>>>(
      reinterpret_cast<const __nv_bfloat16 *>(A.data_ptr<at::BFloat16>()),
      reinterpret_cast<const __nv_bfloat16 *>(B.data_ptr<at::BFloat16>()),
      reinterpret_cast<__nv_bfloat16 *>(D.data_ptr<at::BFloat16>()));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

// E59 step 3: 2-MMA K-loop. K=32 = 2 * MMA_K = the same inner-loop structure
// matmul_v1.cu uses for its mainloop. Validates that we can correctly:
//   1. Issue MMA #0 with enable_input_d=0 (fresh D = A0*B0)
//   2. Issue MMA #1 with enable_input_d=1 (accumulate D += A1*B1)
//   3. Both MMAs hit the same TMEM region with no fence / cleanup needed
//      between them — tcgen05's MMA pipeline handles ordering.
//   4. The single tcgen05.commit at the end of the loop arms the mbarrier
//      only after BOTH MMAs have committed to TMEM.
//
// A is (M=128, K=32) BF16 row-major; B is (K=32, N=64) BF16 row-major.
// SMEM layout: 4 (M, 8) blocks for A (covering K-cols 0..7, 8..15, 16..23, 24..31)
//              4 (N, 8) blocks for B (covering K-rows 0..7, 8..15, 16..23, 24..31)
//
// Descriptor offsets between the two MMAs:
//   - MMA 0 a_desc.base = A_smem            (covers blocks 0-1 = K=0..15)
//   - MMA 1 a_desc.base = A_smem + M*16 BF16  (covers blocks 2-3 = K=16..31)
//     i.e. + MMA_K * sizeof(BF16) = +32 bytes per K-row of A, times M rows = +M*32 bytes
//   - Same offset for B descriptors (+ N*32 bytes between MMAs).
__global__ __launch_bounds__(128) void qr_tcgen05_bf16_gemm_k32_kernel(
    const __nv_bfloat16 *A,  // (M=128, K=32) row-major BF16
    const __nv_bfloat16 *B,  // (K=32, N=64) row-major BF16
    __nv_bfloat16 *D         // (M=128, N=64) row-major BF16 (output)
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 32;       // total K
  constexpr int MMA_K = 16;         // tcgen05.mma kind::f16 K per call
  constexpr int NUM_MMA = BLOCK_K / MMA_K;  // 2

  __shared__ __align__(1024) __nv_bfloat16 A_smem_arr[BLOCK_M * BLOCK_K];
  __shared__ __align__(1024) __nv_bfloat16 B_smem_arr[BLOCK_K * BLOCK_N];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ uint64_t mbar_mem[1];
  __shared__ uint32_t tmem_addr_mem[1];

  const int tid = threadIdx.x;
  const int warp_id = tid >> 5;
  const int lane_id = tid & 31;

  // A_global (M, K) row-major -> A_smem (M, 8)-blocks (4 of them, K=32=4*8).
  for (int idx = tid; idx < BLOCK_M * BLOCK_K; idx += blockDim.x) {
    int r = idx / BLOCK_K;
    int c = idx - r * BLOCK_K;
    int block_b = c / 8;
    int c_in_block = c & 7;
    int dst = block_b * BLOCK_M * 8 + r * 8 + c_in_block;
    A_smem_arr[dst] = A[r * BLOCK_K + c];
  }
  // B_global (K, N) row-major -> B_smem (N, 8)-blocks of 8 K-rows each (4 of them).
  for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += blockDim.x) {
    int k = idx / BLOCK_N;
    int n = idx - k * BLOCK_N;
    int block_b = k / 8;
    int k_in_block = k & 7;
    int dst = block_b * BLOCK_N * 8 + n * 8 + k_in_block;
    B_smem_arr[dst] = B[k * BLOCK_N + n];
  }

  const uint32_t A_smem = static_cast<uint32_t>(__cvta_generic_to_shared(A_smem_arr));
  const uint32_t B_smem = static_cast<uint32_t>(__cvta_generic_to_shared(B_smem_arr));
  const uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_mem));
  const uint32_t tmem_addr_smem = static_cast<uint32_t>(__cvta_generic_to_shared(tmem_addr_mem));

  if (warp_id == 0) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                 :: "r"(tmem_addr_smem), "n"(BLOCK_N));
  } else if (warp_id == 1 && qr_e59_elect_sync()) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(mbar_addr));
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  __syncthreads();
  const uint32_t taddr = tmem_addr_mem[0];

  asm volatile("tcgen05.fence::after_thread_sync;");

  if (warp_id == 0 && qr_e59_elect_sync()) {
    const uint64_t LBO_A = static_cast<uint64_t>(BLOCK_M) * 16;
    const uint64_t LBO_B = static_cast<uint64_t>(BLOCK_N) * 16;
    const uint64_t SBO   = 8ULL * 16ULL;

    constexpr uint32_t i_desc = (1U << 4U)
                              | (1U << 7U)
                              | (1U << 10U)
                              | (((uint32_t)BLOCK_N >> 3U) << 17U)
                              | (((uint32_t)BLOCK_M >> 4U) << 24U);

    // MMA 0: a_desc base = A_smem,                b_desc base = B_smem.
    //        enable_input_d = 0 (overwrite D).
    const uint64_t a_desc0 = qr_e59_desc_encode(A_smem)
                           | (qr_e59_desc_encode(LBO_A) << 16ULL)
                           | (qr_e59_desc_encode(SBO)   << 32ULL)
                           | (1ULL << 46ULL);
    const uint64_t b_desc0 = qr_e59_desc_encode(B_smem)
                           | (qr_e59_desc_encode(LBO_B) << 16ULL)
                           | (qr_e59_desc_encode(SBO)   << 32ULL)
                           | (1ULL << 46ULL);
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, 0, 0;\n\t"  // p = false -> overwrite D
        "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
        "}"
        :: "r"(taddr), "l"(a_desc0), "l"(b_desc0), "r"(i_desc));

    // MMA 1: descriptor bases advance by MMA_K * sizeof(BF16) per K-row,
    //        times BLOCK_M (or BLOCK_N) rows = BLOCK_M * 32 byte address offset.
    //        We add (MMA_K * 2 / 16) = 4 to the desc.base 14-bit encoded field;
    //        equivalently, recompute the descriptor with the new SMEM byte addr.
    const uint32_t A_smem_1 = A_smem + (uint32_t)(BLOCK_M * MMA_K * 2);  // +M*32 bytes
    const uint32_t B_smem_1 = B_smem + (uint32_t)(BLOCK_N * MMA_K * 2);  // +N*32 bytes
    const uint64_t a_desc1 = qr_e59_desc_encode(A_smem_1)
                           | (qr_e59_desc_encode(LBO_A) << 16ULL)
                           | (qr_e59_desc_encode(SBO)   << 32ULL)
                           | (1ULL << 46ULL);
    const uint64_t b_desc1 = qr_e59_desc_encode(B_smem_1)
                           | (qr_e59_desc_encode(LBO_B) << 16ULL)
                           | (qr_e59_desc_encode(SBO)   << 32ULL)
                           | (1ULL << 46ULL);
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, 1, 0;\n\t"  // p = true -> accumulate (D += A*B)
        "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
        "}"
        :: "r"(taddr), "l"(a_desc1), "l"(b_desc1), "r"(i_desc));

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mbar_addr) : "memory");
  }

  {
    const uint32_t ticks = 0x989680u;
    const int phase = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT_E59_K32:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE_E59_K32;\n\t"
        "bra.uni LAB_WAIT_E59_K32;\n\t"
        "DONE_E59_K32:\n\t"
        "}"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks));
  }

  asm volatile("tcgen05.fence::after_thread_sync;");

  for (int n_iter = 0; n_iter < BLOCK_N / 8; n_iter++) {
    float tmp[8];
    const uint32_t addr = taddr + ((uint32_t)(warp_id * 32) << 16) + (uint32_t)(n_iter * 8);
    asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                 : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                   "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                 : "r"(addr));
    asm volatile("tcgen05.wait::ld.sync.aligned;");

    const int row = warp_id * 32 + lane_id;
    const int col_base = n_iter * 8;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
      D[row * BLOCK_N + col_base + i] = __float2bfloat16(tmp[i]);
    }
  }

  __syncthreads();

  if (warp_id == 0) {
    asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
                 :: "r"(taddr), "n"(BLOCK_N));
  }
}

void qr_tcgen05_bf16_gemm_k32(torch::Tensor A, torch::Tensor B, torch::Tensor D) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda() && D.is_cuda(), "A, B, D must be CUDA");
  TORCH_CHECK(A.scalar_type() == torch::kBFloat16, "A must be bfloat16");
  TORCH_CHECK(B.scalar_type() == torch::kBFloat16, "B must be bfloat16");
  TORCH_CHECK(D.scalar_type() == torch::kBFloat16, "D must be bfloat16");
  TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && D.is_contiguous(), "must be contiguous");
  TORCH_CHECK(A.numel() == 128 * 32, "A must be (128, 32)");
  TORCH_CHECK(B.numel() == 32 * 64, "B must be (32, 64)");
  TORCH_CHECK(D.numel() == 128 * 64, "D must be (128, 64)");
  qr_tcgen05_bf16_gemm_k32_kernel<<<1, 128>>>(
      reinterpret_cast<const __nv_bfloat16 *>(A.data_ptr<at::BFloat16>()),
      reinterpret_cast<const __nv_bfloat16 *>(B.data_ptr<at::BFloat16>()),
      reinterpret_cast<__nv_bfloat16 *>(D.data_ptr<at::BFloat16>()));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

// E60a: batched grid-launched version of the step-2 single-tile kernel.
//
// Same MMA + read-back as qr_tcgen05_bf16_gemm_kernel, but the CTA picks
// up its (tile_m, tile_n, batch) from blockIdx.{x,y,z} instead of doing
// one fixed tile from a single launch. This is the simplest "scale out"
// version — no persistent CTAs, no work stealing, just one CTA per tile.
// Good enough to test the throughput hypothesis: if at QR-realistic scale
// the per-tile launch overhead dilutes well, we should land within ~2x of
// cuBLAS strided batched. If we don't, we know we need Mufeez v2+ (persistent
// + warp specialization) before attempting integration.
//
// Inputs (all BF16, row-major contiguous):
//   A : (batch, M_total, K=16)
//   B : (batch, K=16, N_total)
//   D : (batch, M_total, N_total)  (output; fresh, not accumulated)
//
// Requires M_total % 128 == 0 and N_total % 64 == 0. Caller pads if needed.
__global__ __launch_bounds__(128) void qr_tcgen05_bf16_gemm_batched_kernel(
    const __nv_bfloat16 *A_base,
    const __nv_bfloat16 *B_base,
    __nv_bfloat16 *D_base,
    int M_total,
    int N_total
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 16;

  const int tile_m = blockIdx.x;
  const int tile_n = blockIdx.y;
  const int b      = blockIdx.z;
  const int row_base = tile_m * BLOCK_M;
  const int col_base = tile_n * BLOCK_N;

  // Per-batch base pointers.
  const __nv_bfloat16 *A_b = A_base + static_cast<int64_t>(b) * M_total * BLOCK_K;
  const __nv_bfloat16 *B_b = B_base + static_cast<int64_t>(b) * BLOCK_K * N_total;
  __nv_bfloat16 *D_b       = D_base + static_cast<int64_t>(b) * M_total * N_total;

  __shared__ __align__(1024) __nv_bfloat16 A_smem_arr[BLOCK_M * BLOCK_K];
  __shared__ __align__(1024) __nv_bfloat16 B_smem_arr[BLOCK_K * BLOCK_N];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ uint64_t mbar_mem[1];
  __shared__ uint32_t tmem_addr_mem[1];

  const int tid = threadIdx.x;
  const int warp_id = tid >> 5;
  const int lane_id = tid & 31;

  // Load A_tile (128, 16) from global into SMEM (M, 8)-blocks.
  // A's leading dim is K=16 (row stride within one batch), so A_b row R is at A_b + R*K.
  for (int idx = tid; idx < BLOCK_M * BLOCK_K; idx += blockDim.x) {
    int r = idx / BLOCK_K;
    int c = idx - r * BLOCK_K;
    int block_b = c / 8;
    int c_in_block = c & 7;
    int dst = block_b * BLOCK_M * 8 + r * 8 + c_in_block;
    A_smem_arr[dst] = A_b[(row_base + r) * BLOCK_K + c];
  }
  // Load B_tile (16, 64) from global into SMEM (N, 8)-blocks.
  // B's leading dim is N_total (row stride), so B_b[k, col] is at B_b + k*N_total + col.
  for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += blockDim.x) {
    int k = idx / BLOCK_N;
    int n = idx - k * BLOCK_N;
    int block_b = k / 8;
    int k_in_block = k & 7;
    int dst = block_b * BLOCK_N * 8 + n * 8 + k_in_block;
    B_smem_arr[dst] = B_b[k * N_total + (col_base + n)];
  }

  const uint32_t A_smem = static_cast<uint32_t>(__cvta_generic_to_shared(A_smem_arr));
  const uint32_t B_smem = static_cast<uint32_t>(__cvta_generic_to_shared(B_smem_arr));
  const uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_mem));
  const uint32_t tmem_addr_smem = static_cast<uint32_t>(__cvta_generic_to_shared(tmem_addr_mem));

  if (warp_id == 0) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                 :: "r"(tmem_addr_smem), "n"(BLOCK_N));
  } else if (warp_id == 1 && qr_e59_elect_sync()) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(mbar_addr));
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  __syncthreads();
  const uint32_t taddr = tmem_addr_mem[0];

  asm volatile("tcgen05.fence::after_thread_sync;");

  if (warp_id == 0 && qr_e59_elect_sync()) {
    const uint64_t LBO_A = static_cast<uint64_t>(BLOCK_M) * 16;
    const uint64_t LBO_B = static_cast<uint64_t>(BLOCK_N) * 16;
    const uint64_t SBO   = 8ULL * 16ULL;
    const uint64_t a_desc = qr_e59_desc_encode(A_smem)
                          | (qr_e59_desc_encode(LBO_A) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    const uint64_t b_desc = qr_e59_desc_encode(B_smem)
                          | (qr_e59_desc_encode(LBO_B) << 16ULL)
                          | (qr_e59_desc_encode(SBO)   << 32ULL)
                          | (1ULL << 46ULL);
    constexpr uint32_t i_desc = (1U << 4U)
                              | (1U << 7U)
                              | (1U << 10U)
                              | (((uint32_t)BLOCK_N >> 3U) << 17U)
                              | (((uint32_t)BLOCK_M >> 4U) << 24U);
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, 0, 0;\n\t"  // p = false -> overwrite D
        "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
        "}"
        :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc));
    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mbar_addr) : "memory");
  }

  {
    const uint32_t ticks = 0x989680u;
    const int phase = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT_E60_BATCH:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE_E60_BATCH;\n\t"
        "bra.uni LAB_WAIT_E60_BATCH;\n\t"
        "DONE_E60_BATCH:\n\t"
        "}"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks));
  }

  asm volatile("tcgen05.fence::after_thread_sync;");

  for (int n_iter = 0; n_iter < BLOCK_N / 8; n_iter++) {
    float tmp[8];
    const uint32_t addr = taddr + ((uint32_t)(warp_id * 32) << 16) + (uint32_t)(n_iter * 8);
    asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                 : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                   "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                 : "r"(addr));
    asm volatile("tcgen05.wait::ld.sync.aligned;");

    const int row = warp_id * 32 + lane_id;
    const int col_local = n_iter * 8;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
      D_b[(row_base + row) * N_total + (col_base + col_local + i)] = __float2bfloat16(tmp[i]);
    }
  }

  __syncthreads();

  if (warp_id == 0) {
    asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
                 :: "r"(taddr), "n"(BLOCK_N));
  }
}

void qr_tcgen05_bf16_gemm_batched(torch::Tensor A, torch::Tensor B, torch::Tensor D) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda() && D.is_cuda(), "A, B, D must be CUDA");
  TORCH_CHECK(A.scalar_type() == torch::kBFloat16 &&
              B.scalar_type() == torch::kBFloat16 &&
              D.scalar_type() == torch::kBFloat16, "all tensors must be bfloat16");
  TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && D.is_contiguous(), "all must be contiguous");
  TORCH_CHECK(A.dim() == 3 && B.dim() == 3 && D.dim() == 3, "all must be 3D");
  const int batch = A.size(0);
  const int M_total = A.size(1);
  const int N_total = D.size(2);
  TORCH_CHECK(A.size(2) == 16, "A's K must be 16");
  TORCH_CHECK(B.size(0) == batch && B.size(1) == 16 && B.size(2) == N_total, "B shape mismatch");
  TORCH_CHECK(D.size(0) == batch && D.size(1) == M_total, "D shape mismatch");
  TORCH_CHECK(M_total % 128 == 0, "M_total must be a multiple of 128");
  TORCH_CHECK(N_total % 64 == 0, "N_total must be a multiple of 64");
  dim3 grid(M_total / 128, N_total / 64, batch);
  qr_tcgen05_bf16_gemm_batched_kernel<<<grid, 128>>>(
      reinterpret_cast<const __nv_bfloat16 *>(A.data_ptr<at::BFloat16>()),
      reinterpret_cast<const __nv_bfloat16 *>(B.data_ptr<at::BFloat16>()),
      reinterpret_cast<__nv_bfloat16 *>(D.data_ptr<at::BFloat16>()),
      M_total, N_total);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

// E61: persistent version of the E60a batched tcgen05 GEMM.
//
// E60a launched one CTA per output tile. That exposed the correct tcgen05
// instruction path but paid CTA setup and TMEM allocation on ~20K tiny tiles.
// This variant launches roughly one CTA per SM; each CTA repeatedly claims a
// tile from a global counter. The math schedule is intentionally unchanged so
// the measurement isolates the persistent-scheduling question.
__global__ __launch_bounds__(128) void qr_tcgen05_bf16_gemm_persistent_kernel(
    const __nv_bfloat16 *A_base,
    const __nv_bfloat16 *B_base,
    __nv_bfloat16 *D_base,
    int *work_counter,
    int M_total,
    int N_total,
    int tiles_m,
    int tiles_n,
    int total_tiles
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 16;

  __shared__ __align__(1024) __nv_bfloat16 A_smem_arr[BLOCK_M * BLOCK_K];
  __shared__ __align__(1024) __nv_bfloat16 B_smem_arr[BLOCK_K * BLOCK_N];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ uint64_t mbar_mem[1];
  __shared__ uint32_t tmem_addr_mem[1];
  __shared__ int shared_work_idx;

  const int tid = threadIdx.x;
  const int warp_id = tid >> 5;
  const int lane_id = tid & 31;

  while (true) {
    if (tid == 0) {
      shared_work_idx = atomicAdd(work_counter, 1);
    }
    __syncthreads();

    const int work_idx = shared_work_idx;
    if (work_idx >= total_tiles) {
      break;
    }

    const int tiles_per_batch = tiles_m * tiles_n;
    const int b = work_idx / tiles_per_batch;
    const int rem = work_idx - b * tiles_per_batch;
    const int tile_m = rem / tiles_n;
    const int tile_n = rem - tile_m * tiles_n;
    const int row_base = tile_m * BLOCK_M;
    const int col_base = tile_n * BLOCK_N;

    const __nv_bfloat16 *A_b = A_base + static_cast<int64_t>(b) * M_total * BLOCK_K;
    const __nv_bfloat16 *B_b = B_base + static_cast<int64_t>(b) * BLOCK_K * N_total;
    __nv_bfloat16 *D_b       = D_base + static_cast<int64_t>(b) * M_total * N_total;

    for (int idx = tid; idx < BLOCK_M * BLOCK_K; idx += blockDim.x) {
      int r = idx / BLOCK_K;
      int c = idx - r * BLOCK_K;
      int block_b = c / 8;
      int c_in_block = c & 7;
      int dst = block_b * BLOCK_M * 8 + r * 8 + c_in_block;
      A_smem_arr[dst] = A_b[(row_base + r) * BLOCK_K + c];
    }
    for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += blockDim.x) {
      int k = idx / BLOCK_N;
      int n = idx - k * BLOCK_N;
      int block_b = k / 8;
      int k_in_block = k & 7;
      int dst = block_b * BLOCK_N * 8 + n * 8 + k_in_block;
      B_smem_arr[dst] = B_b[k * N_total + (col_base + n)];
    }
    __syncthreads();

    const uint32_t A_smem = static_cast<uint32_t>(__cvta_generic_to_shared(A_smem_arr));
    const uint32_t B_smem = static_cast<uint32_t>(__cvta_generic_to_shared(B_smem_arr));
    const uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(mbar_mem));
    const uint32_t tmem_addr_smem = static_cast<uint32_t>(__cvta_generic_to_shared(tmem_addr_mem));

    if (warp_id == 0) {
      asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                   :: "r"(tmem_addr_smem), "n"(BLOCK_N));
    } else if (warp_id == 1 && qr_e59_elect_sync()) {
      asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(mbar_addr));
      asm volatile("fence.mbarrier_init.release.cluster;");
    }
    __syncthreads();
    const uint32_t taddr = tmem_addr_mem[0];

    asm volatile("tcgen05.fence::after_thread_sync;");

    if (warp_id == 0 && qr_e59_elect_sync()) {
      const uint64_t LBO_A = static_cast<uint64_t>(BLOCK_M) * 16;
      const uint64_t LBO_B = static_cast<uint64_t>(BLOCK_N) * 16;
      const uint64_t SBO   = 8ULL * 16ULL;
      const uint64_t a_desc = qr_e59_desc_encode(A_smem)
                            | (qr_e59_desc_encode(LBO_A) << 16ULL)
                            | (qr_e59_desc_encode(SBO)   << 32ULL)
                            | (1ULL << 46ULL);
      const uint64_t b_desc = qr_e59_desc_encode(B_smem)
                            | (qr_e59_desc_encode(LBO_B) << 16ULL)
                            | (qr_e59_desc_encode(SBO)   << 32ULL)
                            | (1ULL << 46ULL);
      constexpr uint32_t i_desc = (1U << 4U)
                                | (1U << 7U)
                                | (1U << 10U)
                                | (((uint32_t)BLOCK_N >> 3U) << 17U)
                                | (((uint32_t)BLOCK_M >> 4U) << 24U);
      asm volatile(
          "{\n\t"
          ".reg .pred p;\n\t"
          "setp.ne.b32 p, 0, 0;\n\t"
          "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t"
          "}"
          :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc));
      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                   :: "r"(mbar_addr) : "memory");
    }

    {
      const uint32_t ticks = 0x989680u;
      const int phase = 0;
      asm volatile(
          "{\n\t"
          ".reg .pred P1;\n\t"
          "LAB_WAIT_E61_PERSIST:\n\t"
          "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
          "@P1 bra.uni DONE_E61_PERSIST;\n\t"
          "bra.uni LAB_WAIT_E61_PERSIST;\n\t"
          "DONE_E61_PERSIST:\n\t"
          "}"
          :: "r"(mbar_addr), "r"(phase), "r"(ticks));
    }

    asm volatile("tcgen05.fence::after_thread_sync;");

    for (int n_iter = 0; n_iter < BLOCK_N / 8; n_iter++) {
      float tmp[8];
      const uint32_t addr = taddr + ((uint32_t)(warp_id * 32) << 16) + (uint32_t)(n_iter * 8);
      asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                   : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                     "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                   : "r"(addr));
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      const int row = warp_id * 32 + lane_id;
      const int col_local = n_iter * 8;
      #pragma unroll
      for (int i = 0; i < 8; i++) {
        D_b[(row_base + row) * N_total + (col_base + col_local + i)] = __float2bfloat16(tmp[i]);
      }
    }

    __syncthreads();

    if (warp_id == 0) {
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
                   :: "r"(taddr), "n"(BLOCK_N));
    }
    __syncthreads();
  }
}

void qr_tcgen05_bf16_gemm_persistent(torch::Tensor A, torch::Tensor B, torch::Tensor D, torch::Tensor counter) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda() && D.is_cuda() && counter.is_cuda(), "A, B, D, counter must be CUDA");
  TORCH_CHECK(A.scalar_type() == torch::kBFloat16 &&
              B.scalar_type() == torch::kBFloat16 &&
              D.scalar_type() == torch::kBFloat16, "all tensors must be bfloat16");
  TORCH_CHECK(counter.scalar_type() == torch::kInt32 && counter.numel() == 1, "counter must be int32[1]");
  TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && D.is_contiguous() && counter.is_contiguous(), "all must be contiguous");
  TORCH_CHECK(A.dim() == 3 && B.dim() == 3 && D.dim() == 3, "all must be 3D");
  const int batch = A.size(0);
  const int M_total = A.size(1);
  const int N_total = D.size(2);
  TORCH_CHECK(A.size(2) == 16, "A's K must be 16");
  TORCH_CHECK(B.size(0) == batch && B.size(1) == 16 && B.size(2) == N_total, "B shape mismatch");
  TORCH_CHECK(D.size(0) == batch && D.size(1) == M_total, "D shape mismatch");
  TORCH_CHECK(M_total % 128 == 0, "M_total must be a multiple of 128");
  TORCH_CHECK(N_total % 64 == 0, "N_total must be a multiple of 64");
  const int tiles_m = M_total / 128;
  const int tiles_n = N_total / 64;
  const int total_tiles = batch * tiles_m * tiles_n;
  int sm_count = 0;
  C10_CUDA_CHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, A.get_device()));
  const int blocks = min(total_tiles, sm_count);
  qr_tcgen05_bf16_gemm_persistent_kernel<<<blocks, 128>>>(
      reinterpret_cast<const __nv_bfloat16 *>(A.data_ptr<at::BFloat16>()),
      reinterpret_cast<const __nv_bfloat16 *>(B.data_ptr<at::BFloat16>()),
      reinterpret_cast<__nv_bfloat16 *>(D.data_ptr<at::BFloat16>()),
      counter.data_ptr<int>(),
      M_total, N_total, tiles_m, tiles_n, total_tiles);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}
"""

_qr_native = load_inline(
    name="qr_batched_householder",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=[
        "qr_blocked16_gemm",
        "qr_blocked16_gemm_smem_panel",
        "qr_blocked16_gemm_cluster_smem",
        "qr_blocked16_gemm_higham_concat",
        "qr_cublas_geqrf",
    ],
    extra_cuda_cflags=["-arch=sm_100a", "-std=c++20"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)

# E59 step 1: tcgen05 smoke kernel as its own extension (see TCGEN05_CUDA_SRC
# comment for why this is split out). Use -gencode=arch=compute_100a,code=sm_100a
# so the PTX virtual target picks up the architecture-specific `a` variant —
# plain -arch=sm_100a leaves the virtual target at compute_100 and ptxas rejects
# all tcgen05 instructions. Pattern from learn-cuda/02e_matmul_sm100/main.py.
_qr_e59_native = None


def _qr_e59():
    """Lazy-load dev-only tcgen05 kernels.

    The active QR route does not call these kernels. Keeping the extension lazy
    prevents leaderboard imports from paying the E59/E61 compile cost.
    """
    global _qr_e59_native
    if _qr_e59_native is None:
        _qr_e59_native = load_inline(
            name="qr_e59_tcgen05",
            cpp_sources=[TCGEN05_CPP_SRC],
            cuda_sources=[TCGEN05_CUDA_SRC],
            functions=[
                "qr_tcgen05_smoke",
                "qr_tcgen05_bf16_gemm",
                "qr_tcgen05_bf16_gemm_k32",
                "qr_tcgen05_bf16_gemm_batched",
                "qr_tcgen05_bf16_gemm_persistent",
            ],
            extra_cuda_cflags=["-gencode=arch=compute_100a,code=sm_100a", "-std=c++20", "-Xptxas=-v"],
            extra_ldflags=[],
            verbose=False,
        )
    return _qr_e59_native

_CUBLAS_MAX_TRY_N = 512
_bad_cublas_n: set[int] = set()
_PHASE_TIMING = os.environ.get("QR_PHASE_TIMING") == "1"


def _qr_e60_batched_correctness() -> None:
    """E60a step 1: verify the batched grid-launched tcgen05 GEMM matches
    torch.matmul on a small case before scaling up. (batch=4, M=256, N=128,
    K=16) — exercises multiple tiles per matrix AND multiple matrices.
    """
    torch.manual_seed(4)
    batch, M, N, K = 4, 256, 128, 16
    A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16).contiguous()
    B = torch.randn(batch, K, N, device="cuda", dtype=torch.bfloat16).contiguous()
    D = torch.zeros(batch, M, N, device="cuda", dtype=torch.bfloat16)

    print(f"[E60a correctness] batch={batch} M={M} N={N} K={K}", flush=True)
    _qr_e59().qr_tcgen05_bf16_gemm_batched(A, B, D)
    torch.cuda.synchronize()

    D_test = D.to(torch.float32)
    D_ref = (A.to(torch.float32) @ B.to(torch.float32)).to(torch.bfloat16).to(torch.float32)
    abs_diff = (D_test - D_ref).abs()
    n_exact = int((D_test == D_ref).sum().item())
    n = batch * M * N
    print(
        f"[E60a correctness] max_abs_diff={abs_diff.max().item():.6f} "
        f"n_exact={n_exact}/{n}",
        flush=True,
    )


def _qr_e61_persistent_correctness() -> None:
    """E61 step 1: verify the persistent tcgen05 GEMM produces the same BF16
    output as torch.matmul on the small multi-tile case used for E60a.
    """
    torch.manual_seed(6)
    batch, M, N, K = 4, 256, 128, 16
    A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16).contiguous()
    B = torch.randn(batch, K, N, device="cuda", dtype=torch.bfloat16).contiguous()
    D = torch.zeros(batch, M, N, device="cuda", dtype=torch.bfloat16)
    counter = torch.zeros(1, device="cuda", dtype=torch.int32)

    print(f"[E61 correctness] persistent batch={batch} M={M} N={N} K={K}", flush=True)
    _qr_e59().qr_tcgen05_bf16_gemm_persistent(A, B, D, counter)
    torch.cuda.synchronize()

    D_test = D.to(torch.float32)
    D_ref = (A.to(torch.float32) @ B.to(torch.float32)).to(torch.bfloat16).to(torch.float32)
    abs_diff = (D_test - D_ref).abs()
    n_exact = int((D_test == D_ref).sum().item())
    n = batch * M * N
    print(
        f"[E61 correctness] max_abs_diff={abs_diff.max().item():.6f} "
        f"n_exact={n_exact}/{n}",
        flush=True,
    )


def _qr_e61_persistent_bench(reps: int = 20) -> None:
    """E61 step 2: compare persistent tcgen05, E60a grid tcgen05, and cuBLAS
    on the QR-scale padded GEMM shape batch=640, M=N=512, K=16.
    """
    torch.manual_seed(7)
    batch = 640
    M, N, K = 512, 512, 16
    A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16).contiguous()
    B = torch.randn(batch, K, N, device="cuda", dtype=torch.bfloat16).contiguous()
    D_test = torch.zeros(batch, M, N, device="cuda", dtype=torch.bfloat16)
    counter = torch.zeros(1, device="cuda", dtype=torch.int32)

    for _ in range(3):
        counter.zero_()
        _qr_e59().qr_tcgen05_bf16_gemm_persistent(A, B, D_test, counter)
        _qr_e59().qr_tcgen05_bf16_gemm_batched(A, B, D_test)
        _ = torch.bmm(A, B)
    torch.cuda.synchronize()

    p_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    p_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        counter.zero_()
        p_starts[i].record()
        _qr_e59().qr_tcgen05_bf16_gemm_persistent(A, B, D_test, counter)
        p_ends[i].record()
    torch.cuda.synchronize()
    persistent_ms = sorted(p_starts[i].elapsed_time(p_ends[i]) for i in range(reps))

    g_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    g_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        g_starts[i].record()
        _qr_e59().qr_tcgen05_bf16_gemm_batched(A, B, D_test)
        g_ends[i].record()
    torch.cuda.synchronize()
    grid_ms = sorted(g_starts[i].elapsed_time(g_ends[i]) for i in range(reps))

    c_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    c_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        c_starts[i].record()
        _ = torch.bmm(A, B)
        c_ends[i].record()
    torch.cuda.synchronize()
    cublas_ms = sorted(c_starts[i].elapsed_time(c_ends[i]) for i in range(reps))

    flops = 2.0 * batch * M * N * K
    p_med = persistent_ms[len(persistent_ms)//2]
    g_med = grid_ms[len(grid_ms)//2]
    c_med = cublas_ms[len(cublas_ms)//2]
    print(
        f"[E61 bench] batch={batch} M={M} N={N} K={K} reps={reps}\n"
        f"  persistent tcgen05 median {p_med:.3f} ms  "
        f"min {persistent_ms[0]:.3f}  max {persistent_ms[-1]:.3f}  "
        f"({flops / (p_med / 1000.0) / 1e12:.1f} TF/s)\n"
        f"  grid tcgen05       median {g_med:.3f} ms  "
        f"min {grid_ms[0]:.3f}  max {grid_ms[-1]:.3f}  "
        f"({flops / (g_med / 1000.0) / 1e12:.1f} TF/s)\n"
        f"  cuBLAS torch.bmm   median {c_med:.3f} ms  "
        f"min {cublas_ms[0]:.3f}  max {cublas_ms[-1]:.3f}  "
        f"({flops / (c_med / 1000.0) / 1e12:.1f} TF/s)\n"
        f"  ratios: persistent/grid={p_med/g_med:.2f}x, "
        f"persistent/cuBLAS={p_med/c_med:.2f}x",
        flush=True,
    )


def _qr_e60_bench(reps: int = 20) -> None:
    """E60a step 2: bench batched tcgen05 vs cuBLAS strided batched at
    QR-realistic scale (batch=640, M=512, N=512, K=16).

    For (640, 512) panel 0 the actual trailing apply is M=N=496 (not a
    multiple of 128/64), so we round up to padded 512 — this slightly
    overestimates the work but matches what an aligned integration would
    do. Reference is torch.bmm which dispatches to cublasGemmStridedBatchedEx.
    """
    torch.manual_seed(5)
    batch = 640
    M, N, K = 512, 512, 16
    A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16).contiguous()
    B = torch.randn(batch, K, N, device="cuda", dtype=torch.bfloat16).contiguous()
    D_test = torch.zeros(batch, M, N, device="cuda", dtype=torch.bfloat16)

    # Warmup
    for _ in range(3):
        _qr_e59().qr_tcgen05_bf16_gemm_batched(A, B, D_test)
        _ = torch.bmm(A, B)
    torch.cuda.synchronize()

    t_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    t_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        t_starts[i].record()
        _qr_e59().qr_tcgen05_bf16_gemm_batched(A, B, D_test)
        t_ends[i].record()
    torch.cuda.synchronize()
    tcgen05_ms = sorted(t_starts[i].elapsed_time(t_ends[i]) for i in range(reps))

    c_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    c_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        c_starts[i].record()
        _ = torch.bmm(A, B)
        c_ends[i].record()
    torch.cuda.synchronize()
    cublas_ms = sorted(c_starts[i].elapsed_time(c_ends[i]) for i in range(reps))

    flops = 2.0 * batch * M * N * K
    tcgen05_tfs = flops / (tcgen05_ms[len(tcgen05_ms)//2] / 1000.0) / 1e12
    cublas_tfs = flops / (cublas_ms[len(cublas_ms)//2] / 1000.0) / 1e12
    print(
        f"[E60a bench] batch={batch} M={M} N={N} K={K} reps={reps}\n"
        f"  tcgen05 batched  median {tcgen05_ms[len(tcgen05_ms)//2]:.3f} ms  "
        f"min {tcgen05_ms[0]:.3f}  max {tcgen05_ms[-1]:.3f}  ({tcgen05_tfs:.1f} TF/s)\n"
        f"  cuBLAS (torch.bmm) median {cublas_ms[len(cublas_ms)//2]:.3f} ms  "
        f"min {cublas_ms[0]:.3f}  max {cublas_ms[-1]:.3f}  ({cublas_tfs:.1f} TF/s)\n"
        f"  ratio (tcgen05/cuBLAS) median = "
        f"{tcgen05_ms[len(tcgen05_ms)//2]/cublas_ms[len(cublas_ms)//2]:.2f}x",
        flush=True,
    )


def _qr_e59_bench(reps: int = 100) -> None:
    """E59 step 4 prerequisite: time the single-CTA tcgen05 GEMM vs torch
    on identical shape, to gauge whether the tcgen05 path is competitive
    before committing to a tiled integration into qr_blocked16_gemm.

    Shape M=128 N=64 K=16: matches one trailing-apply tile of QR.
    Both kernels do the same math (BF16 inputs, FP32 accumulate, BF16
    output); only the *implementation* differs. The cuBLAS reference is
    via torch.matmul which dispatches to cublasGemmEx under the hood.
    """
    torch.manual_seed(2)
    M, N, K = 128, 64, 16
    A = torch.randn(M, K, device="cuda", dtype=torch.bfloat16).contiguous()
    B = torch.randn(K, N, device="cuda", dtype=torch.bfloat16).contiguous()
    D = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16)

    # Warmup
    for _ in range(5):
        _qr_e59().qr_tcgen05_bf16_gemm(A, B, D)
        _ = torch.matmul(A, B)
    torch.cuda.synchronize()

    # Time tcgen05 (single CTA, 128 threads)
    t_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    t_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        t_starts[i].record()
        _qr_e59().qr_tcgen05_bf16_gemm(A, B, D)
        t_ends[i].record()
    torch.cuda.synchronize()
    tcgen05_us = sorted(t_starts[i].elapsed_time(t_ends[i]) * 1000 for i in range(reps))

    # Time torch.matmul (cuBLAS under the hood)
    s_starts = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    s_ends = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        s_starts[i].record()
        _ = torch.matmul(A, B)
        s_ends[i].record()
    torch.cuda.synchronize()
    torch_us = sorted(s_starts[i].elapsed_time(s_ends[i]) * 1000 for i in range(reps))

    print(
        f"[E59 bench] shape=({M},{N},{K}) reps={reps}\n"
        f"  tcgen05  median {tcgen05_us[len(tcgen05_us)//2]:.2f} us  "
        f"min {tcgen05_us[0]:.2f}  max {tcgen05_us[-1]:.2f}\n"
        f"  torch    median {torch_us[len(torch_us)//2]:.2f} us  "
        f"min {torch_us[0]:.2f}  max {torch_us[-1]:.2f}\n"
        f"  ratio (tcgen05/torch) median = "
        f"{tcgen05_us[len(tcgen05_us)//2]/torch_us[len(torch_us)//2]:.2f}x",
        flush=True,
    )


def _qr_e59_gemm_k32_run() -> torch.Tensor:
    """E59 step 3: 2-MMA accumulating K-loop (K=32).

    Same shapes as step 2 except K=32, which exercises matmul_v1.cu's exact
    inner-loop pattern: MMA #0 with enable_input_d=0 (overwrite, fresh D),
    then MMA #1 with enable_input_d=1 (accumulate, D += A1*B1). Confirms
    that two MMAs can correctly share TMEM with the right predicates.
    Bit-exact pass here is the prerequisite for any K > 16 path.
    """
    torch.manual_seed(1)
    M, N, K = 128, 64, 32
    A = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
    B = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
    D = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16)

    print(f"[E59 step 3] launching K=32 2-MMA accumulate (M={M}, N={N}, K={K})", flush=True)
    _qr_e59().qr_tcgen05_bf16_gemm_k32(A.contiguous(), B.contiguous(), D)
    torch.cuda.synchronize()

    D_test = D.to(torch.float32)
    D_ref = (A.to(torch.float32) @ B.to(torch.float32)).to(torch.bfloat16).to(torch.float32)

    abs_diff = (D_test - D_ref).abs()
    rel_diff = abs_diff / (D_ref.abs() + 1e-6)
    n_exact = int((D_test == D_ref).sum().item())
    n = M * N
    print(
        f"[E59 step 3] D_test[0,0]={D_test[0,0].item():.4f} vs D_ref[0,0]={D_ref[0,0].item():.4f} | "
        f"D_test[64,32]={D_test[64,32].item():.4f} vs D_ref[64,32]={D_ref[64,32].item():.4f} | "
        f"max_abs_diff={abs_diff.max().item():.6f} max_rel_diff={rel_diff.max().item():.6f} | "
        f"n_exact={n_exact}/{n}",
        flush=True,
    )
    return D_test


def _qr_e59_gemm_run() -> torch.Tensor:
    """E59 step 2: QR-shaped tcgen05 BF16 GEMM tile.

    Runs D = A @ B for random BF16 A (128, 16) and B (16, 64). Compares
    against torch.matmul reference (which uses cuBLAS tensor cores
    internally and shares the same BF16-input / FP32-accumulate semantics
    as tcgen05.mma.kind::f16). Passing this confirms the SMEM-layout
    reshape from row-major A/B into the (M, 8) / (N, 8) contiguous-block
    layouts is correct, which is the prerequisite for using this kernel
    on real QR panels.
    """
    torch.manual_seed(0)
    M, N, K = 128, 64, 16
    A = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
    B = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
    D = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16)

    print(f"[E59 step 2] launching BF16 GEMM (M={M}, N={N}, K={K}, random A,B)", flush=True)
    _qr_e59().qr_tcgen05_bf16_gemm(A.contiguous(), B.contiguous(), D)
    torch.cuda.synchronize()

    D_test = D.to(torch.float32)
    # Reference: BF16 inputs accumulated in FP32 via torch.matmul.
    D_ref = (A.to(torch.float32) @ B.to(torch.float32)).to(torch.bfloat16).to(torch.float32)

    abs_diff = (D_test - D_ref).abs()
    rel_diff = abs_diff / (D_ref.abs() + 1e-6)
    max_abs = abs_diff.max().item()
    max_rel = rel_diff.max().item()
    # Tight tolerance because both sides cast to BF16 before comparison.
    n_exact = int((D_test == D_ref).sum().item())
    n = M * N
    print(
        f"[E59 step 2] D_test[0,0]={D_test[0,0].item():.4f} vs D_ref[0,0]={D_ref[0,0].item():.4f} | "
        f"D_test[64,32]={D_test[64,32].item():.4f} vs D_ref[64,32]={D_ref[64,32].item():.4f} | "
        f"max_abs_diff={max_abs:.6f} max_rel_diff={max_rel:.6f} | "
        f"n_exact={n_exact}/{n}",
        flush=True,
    )
    return D_test


def _qr_e59_smoke_run() -> torch.Tensor:
    """E59 step 1 phase B: BF16 tcgen05 MMA smoke (all-ones).

    Launches a single (M=128, N=64, K=16) BF16 MMA on all-ones inputs.
    Expected every output entry = 16.0 (sum_k 1.0 * 1.0 over k=0..15).
    A correct value + UTCHMMA appearing in cuobjdump SASS is the binary
    go/no-go for the entire E59+ fused-trailing-apply direction.
    """
    out = torch.zeros(128 * 64, device="cuda", dtype=torch.bfloat16)
    print("[E59 smoke B] launching BF16 MMA smoke (M=128, N=64, K=16, all-ones)", flush=True)
    _qr_e59().qr_tcgen05_smoke(out)
    torch.cuda.synchronize()
    out_f = out.view(128, 64).to(torch.float32)
    expected = 16.0
    n_correct = int((out_f == expected).sum().item())
    n = 128 * 64
    print(
        f"[E59 smoke B] out[0,0]={out_f[0, 0].item():.4f} "
        f"out[63,63]={out_f[63, 63].item():.4f} "
        f"out[127,63]={out_f[127, 63].item():.4f} | "
        f"min={out_f.min().item():.4f} max={out_f.max().item():.4f} "
        f"mean={out_f.mean().item():.4f} | "
        f"n_eq_16.0={n_correct}/{n} (all_match={n_correct == n})",
        flush=True,
    )
    return out_f


def _evt() -> torch.cuda.Event:
    e = torch.cuda.Event(enable_timing=True)
    e.record()
    return e


def _ms(a: torch.cuda.Event, b: torch.cuda.Event) -> float:
    return a.elapsed_time(b)


def _cublas_geqrf(data: torch.Tensor) -> output_t | None:
    batch, n, _ = data.shape
    if n > _CUBLAS_MAX_TRY_N or n in _bad_cublas_n:
        return None

    h = torch.empty_strided((batch, n, n), (n * n, 1, n), device=data.device, dtype=data.dtype)
    tau = torch.empty((batch, n), device=data.device, dtype=data.dtype)
    if _PHASE_TIMING:
        e0 = _evt()
    h.copy_(data)
    if _PHASE_TIMING:
        e1 = _evt()
    try:
        _qr_native.qr_cublas_geqrf(h, tau)
    except RuntimeError:
        _bad_cublas_n.add(n)
        return None
    if _PHASE_TIMING:
        e2 = _evt()
        torch.cuda.synchronize()
        print(
            f"[phase] route=cublas_geqrf batch={batch} n={n} "
            f"copy={_ms(e0, e1):.3f} qr={_ms(e1, e2):.3f} total={_ms(e0, e2):.3f}",
            flush=True,
        )
    return h, tau


# E54/E55: SMEM-resident panel route for the large low-batch shapes whose
# old single-CTA-per-matrix panel + separate pack was HBM-bound.
_SMEM_PANEL_WARPS_BY_N: dict[int, int] = {1024: 4, 2048: 8}


_MULTI_CTA_BY_N: dict[int, int] = {4096: 32}


def _native_blocked16(data: torch.Tensor) -> output_t | None:
    batch, n, _ = data.shape
    if n not in (176, 352, 512, 1024, 2048, 4096):
        return None

    h = torch.empty_strided((batch, n, n), (n * n, 1, n), device=data.device, dtype=data.dtype)
    tau = torch.empty((batch, n), device=data.device, dtype=data.dtype)
    work = torch.empty((batch, n, 16), device=data.device, dtype=data.dtype)
    if _PHASE_TIMING:
        e0 = _evt()
    h.copy_(data)
    if _PHASE_TIMING:
        e1 = _evt()

    ygemm = torch.empty_strided((batch, n, 16), (n * 16, 1, n), device=data.device, dtype=data.dtype)
    wgemm = torch.empty_strided((batch, n, 16), (n * 16, 1, n), device=data.device, dtype=data.dtype)
    tmp = torch.empty_strided((batch, 16, n), (16 * n, 1, 16), device=data.device, dtype=data.dtype)

    if n in _MULTI_CTA_BY_N:
        # E76: combined-phase column factor (2 cluster syncs / col).
        ctas = 16
        scratch = torch.empty(
            (batch, ctas * 16 + 32 + 16 + 16 * 16), device=data.device, dtype=torch.float32
        )
        if _PHASE_TIMING:
            q0 = _evt()
        _qr_native.qr_blocked16_gemm_cluster_smem(h, tau, work, ygemm, wgemm, tmp, scratch)
        if _PHASE_TIMING:
            q1 = _evt()
            torch.cuda.synchronize()
            print(
                f"[phase] route=blocked16_cluster_smem batch={batch} n={n} ctas={ctas} "
                f"copy={_ms(e0, e1):.3f} qr={_ms(q0, q1):.3f} total={_ms(e0, q1):.3f}",
                flush=True,
            )
        return h, tau

    if n in _SMEM_PANEL_WARPS_BY_N:
        warps = _SMEM_PANEL_WARPS_BY_N[n]
        if _PHASE_TIMING:
            q0 = _evt()
        _qr_native.qr_blocked16_gemm_smem_panel(h, tau, ygemm, wgemm, tmp, warps)
        if _PHASE_TIMING:
            q1 = _evt()
            torch.cuda.synchronize()
            print(
                f"[phase] route=blocked16_smem_panel batch={batch} n={n} warps={warps} "
                f"copy={_ms(e0, e1):.3f} qr={_ms(q0, q1):.3f} total={_ms(e0, q1):.3f}",
                flush=True,
            )
        return h, tau

    if n == 512:
        # E49: K=48 concat-correction Higham GEMM2 (BF16 tensor core,
        # FP32 accumulate) replacing the IEEE-FP32 GEMM2.
        y_concat48 = torch.empty_strided((batch, n, 48), (n * 48, 1, n), device=data.device, dtype=torch.bfloat16)
        tmp_concat48 = torch.empty_strided((batch, 48, n), (48 * n, 1, 48), device=data.device, dtype=torch.bfloat16)
        if _PHASE_TIMING:
            q0 = _evt()
        _qr_native.qr_blocked16_gemm_higham_concat(
            h, tau, work, ygemm, wgemm, tmp, y_concat48, tmp_concat48,
        )
        if _PHASE_TIMING:
            q1 = _evt()
            torch.cuda.synchronize()
            print(
                f"[phase] route=blocked16_higham_concat batch={batch} n={n} "
                f"copy={_ms(e0, e1):.3f} qr={_ms(q0, q1):.3f} total={_ms(e0, q1):.3f}",
                flush=True,
            )
        return h, tau

    # n in {176, 352}: standard BF16 cuBLAS trailing GEMMs.
    if _PHASE_TIMING:
        q0 = _evt()
    _qr_native.qr_blocked16_gemm(h, tau, work, ygemm, wgemm, tmp)
    if _PHASE_TIMING:
        q1 = _evt()
        torch.cuda.synchronize()
        print(
            f"[phase] route=blocked16_native_gemm batch={batch} n={n} "
            f"copy={_ms(e0, e1):.3f} qr={_ms(q0, q1):.3f} total={_ms(e0, q1):.3f}",
            flush=True,
        )
    if n == 176:
        h = h.contiguous()
    return h, tau


def custom_kernel(data: input_t) -> output_t:
    result = _native_blocked16(data)
    if result is not None:
        return result
    result = _cublas_geqrf(data)
    if result is not None:
        return result
    if _PHASE_TIMING:
        e0 = _evt()
    result = torch.geqrf(data)
    if _PHASE_TIMING:
        e1 = _evt()
        torch.cuda.synchronize()
        batch, n, _ = data.shape
        print(f"[phase] route=torch_geqrf batch={batch} n={n} qr={_ms(e0, e1):.3f}", flush=True)
    return result