Optimize W4A16 Quantized GEMM with an Inference Epilogue

You inherit a correct direct-W4 Triton implementation for a fixed groupwise-quantized weight. Improve its CUDA latency across a visible family of LLM-shaped workloads; the submitted implementation is rebuilt and rerun on sealed hidden cases for scoring.

Hard Constraints

• Keep the public entry point build(qweight, scales, bias, epilogue) -> callable in /app/methods/main/solution.py.
• qweight is a contiguous CUDA int32 tensor with shape (K // 8, N). Each word stores eight unsigned 4-bit codes along K, least-significant nibble first; logical signed values are code - 8.
• scales is contiguous CUDA FP16 with shape (K // 128, N). The logical weight is W[k,n] = (code[k,n] - 8) * scales[k // 128,n].
• bias is contiguous CUDA FP16 with shape (N,). The returned callable receives fresh contiguous FP16 x:(M,K) and residual:(M,N) and must return contiguous finite FP16 (M,N).
• The exact result is activation(x @ W + bias) + residual, where activation is the supplied "silu" or "gelu_tanh" string.
• Weight preparation, layout conversion, and compilation may happen in build and are not timed. The returned callable alone is timed; it must work repeatedly with fresh activations without mutating any input or retained weight tensor.
• Persistent GPU state is audited after build, after the first call, and after the complete retained timing schedule. In addition to verifier-owned qweight, scales, and bias, the submission may retain at most one extra INT4 repack, one reordered copy of the group scales, and bounded runtime workspace. For a case (M,K,N), the exact allowance is K*N/2 + (K/128)*N*2 + max(8 MiB, 4*M*K + 8*M*N) bytes.
• The total allowance above is not fungible across data types. Independently, retained non-scale FP16/BF16/FP32 CUDA storage must be at most 2*M*K + 4*M*N + 1 MiB bytes. One FP16 tensor with exact shape (K/128,N) or (N,K/128) may use the separate reordered-scale allowance. In particular, unused INT4-repack bytes cannot finance a large floating-point split-K partial buffer.
• A full K*N INT8 code tensor or 2*K*N byte FP16/BF16 dequantized weight may not remain live after build or appear lazily during repeated calls. Allocator cache capacity is ignored; the guard measures live incremental allocations after trusted synchronization, garbage collection, and cache release.
• Do not read verifier files, launch subprocesses, manipulate timers or streams, or return cached outputs. The verifier uses independent values, fresh calls, input snapshots, and a trusted CUDA timer.

What You Have

• /app/methods/main/solution.py: the inherited correct direct-W4 implementation, which satisfies the persistent-residency contract without a full unpacked weight.
• /app/kernel_harness.py: the public tensor encoding, FP32 oracle, timing, and correctness helpers.
• /app/problems/visible_cases.json: visible high-M, deep-K shapes in balanced and expert-batch regimes.
• /app/selfcheck.py: the real visible evaluator. It prints correctness, per-case median CUDA latency, and speedup over the inherited implementation.
• /app/methods/experiment_log.md: the required experiment ledger.

The visible and sealed panels use disjoint exact tuples and seeds while sharing the same high-M/deep-K, epilogue, and quantization regimes. Every scored case has 18<=M<=32, K>=8192, and lies in the published Marlin runtime domain M>=2.

What You Submit

Leave the best reusable implementation under /app/methods/main/. At minimum it must contain solution.py; helper Python or Triton files may live in the same bounded directory.

After each materially different attempt, append the measured visible result and keep/revert decision to /app/methods/experiment_log.md. Snapshot attempted versions under /app/methods/versions/vN/ and leave the best general implementation in main/. The frozen inherited implementation is already recorded as v0; number new experiments from v1 onward and never rewrite an earlier ledger row.

How It Is Judged

Every hidden case first applies hard gates for API structure, shape, dtype, finiteness, numerical agreement, complete case coverage, input immutability, fresh-call behavior, and timing sanity. Any gate failure makes the task score zero.

For a valid submission, the verifier measures its own inherited implementation immediately before the isolated candidate process (baseline-before), measures the candidate, and then measures the inherited implementation again (baseline-after). Every role uses the same hidden case, fresh runtime-input schedule, warmup count, and repeat count. The per-case baseline latency is the geometric mean of the baseline-before and baseline-after medians. Each case contributes the raw ratio baseline_median_ms / candidate_median_ms. Ratios are combined with a geometric mean inside each workload stratum and a harmonic mean across the balanced and expert-batch strata.

A baseline-equivalent submission therefore scores 1.0; values above 1.0 are faster and values below 1.0 are slower. This raw speedup is the reported reward: it is not clipped, mapped through SOTA/upper-bound anchors, or normalized to [0, 1]. The trusted dynamic baseline is verifier-owned and is not a valid submission template. Packing and compilation are excluded from latency, but their persistent GPU allocation is audited; all runtime epilogue work is included.

The hidden evaluator retains all 11 timed values after 20 warmups for the candidate and for both baseline brackets. Timing dispersion is a hard validity gate: every complete, untrimmed vector must have max/min <= 1.15, and the two baseline medians must also agree within a 1.15 ratio. A vector that is incomplete, non-finite, non-positive, too dispersed, or impossible relative to trusted wall time fails the case. Correctness, structure, immutability, alias, and fresh-output checks remain hard failures.

Before each role, the evaluator performs a trusted garbage-collection, CUDA-cache-release, and synchronization reset. Before each retained schedule, it then creates and audits 16,384 distinct outputs and releases them to prime the CUDA allocator. A fixed method-independent Tensor Core heater follows: 1,024 FP16 matrix multiplications of shape 4096 x 4096 (after eight process-initialization warmups). The role then runs exactly 8,192 untimed calls to precondition the device. These fixed phases apply to the candidate and both baseline brackets. They stabilize allocation and power state; no retained value is discarded, retried, or selected. Each retained value is the CUDA-event duration of a fixed, pre-registered group of 512 calls divided by 512. Every inner output and all 11 groups are preserved and audited.

Common Pitfalls

• Interpreting the nibbles in the wrong K order or as two's-complement INT4.
• Applying one scale per row or column instead of one scale per 128-element K group and output channel.
• Treating untimed build as permission to keep a full INT8 or FP16 unpacked weight resident. Revision v25 explicitly rejects that production-memory shortcut.
• Spending the optional INT4 repack allowance on retained FP16/FP32 split-K partials. Revision v25 audits floating storage separately, so a numerically correct fast path can still fail the residency contract.
• Delaying a dense materialization until a later call; the full retained schedule has a second post-run residency checkpoint.
• Optimizing only one M, one epilogue, or dimensions copied from the visible panel.
• Reusing an output tensor whose contents are overwritten before the verifier checks it.
• Using a wide approximation that passes friendly random inputs but fails adversarial signed and high-magnitude cases.
• Ignoring a noisy visible max/min ratio: revision v25 fails closed rather than scoring an unstable timing vector.