Batched square compact-Householder QR for the GPU MODE challenge.

The kernel accepts a batch × n × n CUDA FP32 tensor and returns (H, tau) in the same compact convention as torch.geqrf:

• torch.triu(H) is the triangular factor R.
• The strict lower triangle of H stores the Householder vectors.
• torch.linalg.householder_product(H, tau) materializes Q.

The current implementation is a correctness baseline backed by torch.geqrf. It satisfies the compact-factor contract but is not presented as a leaderboard-optimized custom kernel.

Requirements:

• Linux with Python 3.10+ and uv
• An NVIDIA GPU and a CUDA-compatible PyTorch installation

Install the development environment:

UV_TORCH_BACKEND=auto uv sync --group dev

Run the tests:

uv run pytest -q

Run a representative benchmark matrix:

uv run python benchmark.py

Run selected cases or override the default shapes:

uv run python benchmark.py \
  --cases dense mixed rank-deficient near-collinear \
  --shapes 16x512 4x1024 1x2048 \
  --warmup 3 --repeats 10

Each shape is written as batch×n. Results are emitted as JSON lines so they can be redirected to a file and compared mechanically.

from solution import custom_kernel

H, tau = custom_kernel(A)
Q = torch.linalg.householder_product(H, tau)
R = torch.triu(H)

A must have shape (batch, n, n), dtype torch.float32, and reside on a CUDA device. H and tau are FP32 tensors on the same device with shapes (batch, n, n) and (batch, n).

The hot path intentionally performs no Python validation. This avoids adding dispatch overhead; invalid inputs are rejected by PyTorch.

The local checker mirrors the challenge's normwise hard gates. It materializes Q, extracts R = triu(H), and computes in FP64:

factor residual     = ||QᵀA - R||₁ / ||A||₁
orthogonality error = ||QᵀQ - I||₁
lower leakage       = ||tril(QᵀA, -1)||₁ / ||A||₁

For each matrix, the required limits are:

factor residual     ≤ 20  n eps32
orthogonality error ≤ 100 n eps32

Lower-triangular leakage is reported as a diagnostic and is already covered by the factor residual against triangular R.

The stress generator includes column-scaled dense, rank-deficient, near-rank-deficient, banded, row-scaled, near-collinear, upper-triangular, clustered-scale, and heterogeneous mixed batches. The cond parameter is a deterministic scaling knob, not a request for an exact condition number.

The runner warms up each case, times with CUDA events, and synchronizes outside the measured region. Peak memory includes the input and returned factors. GPU clocks, thermals, PyTorch/CUDA versions, and concurrent workloads still affect results, so record the full JSON metadata when comparing runs.

The default shape list is intentionally modest enough for local iteration. Use challenge-exact batch sizes and conditioning values when producing leaderboard numbers.