ESC
其他 13 分钟阅读

Auto-research with codex: How I achieved a 232x Faster Kernel

Auto-research with codex: How I achieved a 232x Faster Kernel

来源:Hacker News

So a Householder reflection is about finding the perpendicular part and subtracting it twice.

A Householder vector is the mirror, stored compactly. In code, we don’t carry around the whole mirror plane. We store one vector v sticking straight out of it. The mirror is everything perpendicular to v, and the reflection moves along v.

The perpendicular part is just the shadow of x along v, which is v⊤xv⊤v copies of v. Plug that into the subtraction above:

ℋx=x−τv(v⊤x),τ=2v⊤v

So tau is just 2v⊤v: the factor of 2 and the length of v bundled into one precomputed number. v picks the mirror, tau scales the update. (I’ll write the mathematical reflector as ℋj and reserve H for the compact output matrix.)

x reflected x v x_parallel x_perp Mirror rule: keep x_parallel, flip x_perp. So reflected x = x - 2x_perp.

A simplified 2D Householder step. Drag the orange vector: the purple mirror changes so the green reflected vector lands on the horizontal axis. In higher dimensions, landing on the axis is exactly what makes the below-diagonal entries become zero.

(function(){ var svg = document.getElementById(‘hhref-svg’); if(!svg) return; var readout = document.getElementById(‘hhref-readout’); var grid = document.getElementById(‘hhref-grid’); var xaxis = document.getElementById(‘hhref-xaxis’); var yaxis = document.getElementById(‘hhref-yaxis’); var input = document.getElementById(‘hhref-input’); var reflected = document.getElementById(‘hhref-reflected’); var normal = document.getElementById(‘hhref-normal’); var plane = document.getElementById(‘hhref-plane’); var bridge = document.getElementById(‘hhref-bridge’); var parallel = document.getElementById(‘hhref-parallel’); var perpPos = document.getElementById(‘hhref-perp-pos’); var perpNeg = document.getElementById(‘hhref-perp-neg’); var handle = document.getElementById(‘hhref-handle’); var lx = document.getElementById(‘hhref-label-x’); var lhx = document.getElementById(‘hhref-label-hx’); var lv = document.getElementById(‘hhref-label-v’); var lpar = document.getElementById(‘hhref-label-parallel’); var lperp = document.getElementById(‘hhref-label-perp’); var NS = ‘http://www.w3.org/2000/svg'; var C = {x:360,y:220}, S = 64; var vec = {x:2.45,y:1.35}; var dragging = false;

function pt(x,y){return {x:C.x+xS,y:C.y-yS};} function setLine(el,a,b){ el.setAttribute(‘x1’,a.x); el.setAttribute(‘y1’,a.y); el.setAttribute(‘x2’,b.x); el.setAttribute(‘y2’,b.y); } function setText(el,p,txt,dx,dy){ el.setAttribute(‘x’,p.x+(dx||0)); el.setAttribute(‘y’,p.y+(dy||0)); el.textContent = txt; } function dot(a,b){return a.xb.x+a.yb.y;} function norm(a){return Math.sqrt(dot(a,a));} function fmt(n){return (Math.round(n*100)/100).toFixed(2);} function clamp(n,a,b){return Math.max(a,Math.min(b,n));}

function makeGrid(){ grid.innerHTML = ‘’; for(var i=-5;i=-2 && i= 0 ? -r : r; var v = {x:x.x-alpha,y:x.y}; var vv = Math.max(dot(v,v),0.001); var tau = 2 / vv; var vx = dot(v,x); var hx = {x:x.x - tauv.xvx, y:x.y - tauv.yvx}; var vn = norm(v); var u = {x:v.x/vn,y:v.y/vn}; var d = {x:-u.y,y:u.x}; var along = dot(x,d); var xp = {x:d.xalong,y:d.yalong}; var origin = pt(0,0); var px = pt(x.x,x.y); var phx = pt(hx.x,hx.y); var pxp = pt(xp.x,xp.y); var pvu = pt(u.x1.25,u.y1.25); var pa = pt(d.x*-4.5,d.y*-4.5); var pb = pt(d.x4.5,d.y4.5);

setLine(xaxis,pt(-5.1,0),pt(5.1,0)); setLine(yaxis,pt(0,-2.7),pt(0,3.0)); setLine(input,origin,px); setLine(reflected,origin,phx); setLine(normal,origin,pvu); setLine(plane,pa,pb); setLine(bridge,px,phx); setLine(parallel,origin,pxp); setLine(perpPos,pxp,px); setLine(perpNeg,pxp,phx); handle.setAttribute(‘cx’,px.x); handle.setAttribute(‘cy’,px.y); setText(lx,px,‘x’,10,-10); setText(lhx,phx,‘reflected x’,10,18); setText(lv,pvu,‘v’,8,-8); setText(lpar,pxp,‘x_parallel’,8,16); setText(lperp,{x:(pxp.x+px.x)/2,y:(pxp.y+px.y)/2},‘x_perp’,8,-6); readout.textContent = ’tau = ’ + fmt(tau) + ’ · reflected x = (’ + fmt(hx.x) + ‘, ’ + fmt(hx.y) + ‘)’; }

function eventToVec(ev){ var r = svg.getBoundingClientRect(); var sx = (ev.clientX-r.left)/r.width720; var sy = (ev.clientY-r.top)/r.height420; var x = clamp((sx-C.x)/S,-4.25,4.25); var y = clamp((C.y-sy)/S,-2.35,2.65); if(Math.sqrt(xx+yy)

Why QR cares about mirrors

QR wants to turn A into an upper-triangular matrix R. Column 1 should become something like (*, 0, 0), column 2 should have zeros below row 2, and so on.

A Householder mirror is useful because it can do that to a column in one shot. Take the first column of the 3×3 example above: (12, 6, -4). We want to send it to the x-axis so the lower entries become zero. A reflection can only change direction, not length, so the target must also have length 14. One valid target is (-14, 0, 0). After that reflection, the 6 and -4 entries are gone, which is exactly what we wanted.

How do we find the mirror? It sits halfway between the column and its target, so v, the vector poking through the mirror, is just the column minus its target:

v = (12, 6, -4) - (-14, 0, 0) = (26, 6, -4)

Then compute tau = 2 / (vᵀv). The reflector itself is:

ℋ=I−τvv⊤,τ=2v⊤v

ℋx=(I−τvv⊤)x

ℋx=x−τv(v⊤x)

ℋAactive=(I−τvv⊤)Aactive

ℋAactive=Aactive−τv(v⊤Aactive)

This turns the current column into (-14, 0, 0) and rewrites the other columns consistently, so the next reflector is built from the We keep doing this once per column. In rough notation, the repeated updates look like:

A(1)=A(0)−τ1v1(v1⊤A(0))

A(2)=A(1)−τ2v2(v2⊤A(1))

A(3)=A(2)−τ3v3(v3⊤A(2))

Each line uses the matrix produced by the previous line. Each reflector zeroes out everything below the diagonal of its column without disturbing the columns already finished. After the last one, A has walked down to an upper-triangular R:

Mirrors don’t change lengths or angles, so each ℋj is orthogonal, and so is their product Q. That’s where the orthogonality the checker verifies comes from for free.

Once column j is processed, everything below its diagonal is dead space. geqrf reuses those slots to stash the tail of v_j (the leading 1 is implicit). On and above the diagonal you’re looking at R; below it, the reflectors; and tau rides along as a separate vector. That’s why the checker needs both H and tau to rebuild Q.

Householder QR zeroes out A below the diagonal one column at a time. Each step builds a reflector from the current column and applies it to everything on the right. The problem is reflector j+1 is built from the matrix after reflector j has already hit it. So you can’t reorder the steps and you can’t fuse them. It’s serial, and the serial matrix-vector work runs in the slow vector lanes of the SM while the tensor cores just sit there idle.

The classic fix is the blocked algorithm. You pick a narrow panel of b columns (say 32 or 64) and do all the serial work inside it. That’s fine, because the panel is only b columns wide, so it stays cheap. Then, instead of applying the panel’s b reflectors to the rest of the matrix one at a time, you compress them into a single rank-b update (the “WY representation”) and hit the entire trailing block in one shot with three back-to-back matrix multiplies. The serial work stays confined to the panel, and everything else turns into GEMMs, which is exactly the shape the tensor cores want.

Concretely, the WY representation collapses a panel’s b reflectors into a single rank-b update. Stack the panel’s Householder vectors as columns of V=[v1,v2,…,vb], build a small b×b upper-triangular T, then:

and the trailing-block update becomes three GEMM-shaped steps:

The panel is where the serial matrix-vector work is stuck, but it is only b columns wide. Everything to its right is the big trailing block, and that is pure GEMM. As the panel walks down the diagonal, the trailing block shrinks.

If you want to understand the math, I recommend brainstorming with Claude. Also check out Mike’s writeup where he has covered math in a more descriptive and visual way than me. He placed 5th in the contest and shared his learnings focusing on the problem.

Two more things proved to be challenging: reliably using low precision internally (especially for ill-conditioned inputs), and coping with the wide spread of shapes (n = 32, 176, 352, 512, 1024, 2048, 4096) and batch sizes, where the largest matrices had too few batches to fill the tensor cores while n = 32 was so small we had to pack many matrices into one kernel launch.

I participated with ChatGPT Pro (200 USD subscription) and Claude Pro (20 USD subscription). I also used Modal for profiling (they provide 30 USD worth of credits free every month btw).

After I got a basic understanding, I asked Codex to do the basic setup: add problem_statement.md, mention basic details in AGENTS.md on how to submit and use popcorn CLI, and maintain a log.md where we did bookkeeping of the submissions and their status (accept/reject along with shape-wise timings). If we have to contrast with Karpathy’s auto-research, my AGENTS.md and problem_statement.md were initially my program.md.

Logs serve as the evidence of the ideas that worked and didn’t work. Future agent sessions could read the logs and quickly check if an idea had been tried or not. I put more investment into logging after the 3000 µs mark as things started to get harder.

The cool thing about Codex is you can just tell it to do stuff and it will actually do the stuff. You can make it work for hours if you give it a detailed enough prompt with targets. It knows all the math and code, and it has all the feedback it needs. I had this intuition, but I was still surprised by how much GPT-5.5 could push the performance beyond the baseline solution (which was the torch.geqrf/cuSolver function).

My initial few sessions were manual prompts to implement a solution in Triton and optimize for the n = 512 and n = 1024 shapes, as they had the most weight in the final geometric mean.

However, if you want to make the model loop until a certain objective is achieved, use /goal. You can give specific, achievable, quantitative goals. I found that giving a good numeric goal followed by specific criteria worked well.

Example: “Use only Triton or CUDA and beat our active best’s n = 512 timings. Try several ideas either by submitting directly to the leaderboard or using Modal profiling. Remove cuSolver altogether in the new set of experiments. We will only use it as a fallback.” At one point, this goal ran for over a day.

I gave inputs every 2-3 hours to move the model in the direction I wanted. I also let it run without supervision overnight on some days. In the initial few days, my instructions were mostly around steering the model to try out different ideas (more on this later) for different shapes, shouting at it to use more Triton, less PyTorch, and fewer fallbacks! A lesson I learned here: I often had the itch to check on my agents frequently, but you gotta trust it and let it do its work. Get out of the agent’s way you must; steer it back only when it gets stuck.

When using /goal, you can ask the model questions without pausing the loop by using /btw or /side. This creates a temporary thread with context from your main conversation. I liked this way of checking on the agent and providing oversight.

I would ask questions like: Are you winning son? What algorithm/changes are there in the current active submission? Explain this concept to me. What are the ideas you are currently working on? What’s the progress? What are some recent breakthroughs? Can we transfer it to another shape? What are your next best ideas? Based on the responses, I would consult Claude to improve my understanding of the concepts and bottlenecks involved and then provide my thoughts to Codex. After questioning, I would just go back to the main thread and dump thoughts to steer the model.

if a /goal or some loop type thing is running, then either you can queue up the instruction in codex to ask stuff or you can do /btw. (img me checking on codex to ask what’s it doing, what next ideas are etc.). if you do esc to interrupt the loop, then the /goal pauses https://t.co/Yx6EpH5PKo pic.twitter.com/imqZc6RPkz

The baseline torch.geqrf path was around 419 ms (419,000 µs) overall. I was able to reach 5000 µs within a day on the n = 512 shape after implementing the blocked Householder route for that shape. This was the shape weighted most heavily in the geomean.

The 232x number above comes from comparing the rough 419,000 µs baseline to the final 1,805 µs tracked result. The lineage chart below starts from the first recovered point in my tracked submission history, so it shows the later 108,803 to 1,805 µs arc rather than the full baseline-to-final ratio.

When optimizing kernels, making the work more matrix-shaped so the tensor cores stop being idle is your life’s purpose.

Optimizations were much harder after the 3000 µs point. I had to get more involved in the loop in terms of learning the concepts and steering the model. I gave Modal profiling access to Codex and let it run torch profiling / nsys profiling to test out different ideas, compare implementations, and sweep parameters faster. (Later on, the organizers also provided a way to do NCU profiling.)

The rough structural evolution of the QR kernel, from baseline/library-heavy paths toward custom panel work, fused assembly, and GEMM-shaped trailing updates:

There was a lot of back and forth on profiling the entire shape and then identifying the bottlenecks. For this problem, launch overhead and panel overhead dominated. We were rarely memory or compute bound. Most of the effort was spent finding ways to reduce panel overhead and make the WY-update more GEMM-able.

Questions I found myself repeatedly asking:

What is panel overhead/launch overhead? How do we solve it?

What info are we missing and how can we get it?

What’s the latest profile that you have done? What’s your take on it?

Look for fusion candidates. Are there Triton fusion-level candidates?

Look for possible compiler-based optimizations. What can we convert from a runtime signal to a static signal? The idea was to give compiler hints. I once had a 200 µs jump because of this

Asked it a few times to look at Triton-generated compiled artifacts

What are some numerical tricks to exploit the precision slack?

Can you deploy sub-agents to do some math and find possible optimizations?

Deploy agents to search for bugs that may be bogging us down

Are there reduction fusions possible? I saw Codex discover one and then I started hammering it often. Reductions are operations like doing a sum or finding max. Since work needs to be done to iterate through the entire sequence, we can perform multiple of these in a single go.

A major challenge I started facing in the 3000 -> 1800 µs range was the model getting stuck in local maxima. This looked like endless hand-tuning of parameters and small variants of the same idea.

A couple of years ago, agents got stuck in (doom) loops because they were not smart enough or just didn’t have the knowledge (or the verification loop was not robust enough). People used to experiment with sampling, temperature variation and different decoding strategies for this.

Then over time, models got smart enough, pretrained on newer knowledge, and we got tons of RL scaling and inference-time compute (training the model to think for a larger number of tokens).

Now the models struggle with finding new ideas and “research taste” - what next best idea/experiment should we do given the verifier feedback, prior evidence, and our evals (profiler feedback in our case). Good idea generation is the next great adventure. I recently wrote an article too on this theme.

I used the following strategies to help the model get out of local maxima:

Beam of candidates: For a long time, I did a dumb thing. I kept a single best candidate against which new candidates were tested. If the agent tries a new structural idea or a significantly big change, then it’s highly likely that it would score less than our current best submission. However, after a few iterations, that change may outperform our best candidate. With this observation, I introduced some instructions to maintain a beam of 3-5 candidates.

Human in the loop: I would act as the secret sauce to steer the model when it was stuck for long times without improvement

Encourage the model to take more risks and try ambitious ideas. You will not believe it but this worked. Also, Claude would often give up after a few rounds, with excuses like “we have exhausted all optimizations to reach x geomean”; Codex, on the other hand, is more persistent.

Use a stronger advisor model that produces more varied ideas. In my harness, this would look like an instruction in AGENTS.md to encourage the model to use headless calls claude -p to get ideas, provide profiling data, etc.

Instruct the model to frequently use sub-agents to try out ambitious ideas, search the web for blogs and papers, go through the list I mentioned above, and find micro-optimizations

Profile using NCU and Modal, compare the results, and work on bottlenecks that both Modal and NCU profiling pointed out

A multi-agent swarm approach is something I thought of but didn’t try.

I think the strong advisor strategy (like GPT-5.6 Sol or Fable) is going to be a standard strategy in auto-research flows. Think very big models with state-of-the-art training. Claude Code provides an /advisor command for this too.

We’re bringing the advisor strategy to the Claude Platform.

Pair Opus as an advisor with Sonnet or Haiku as an executor, and get near Opus-level intelligence in your agents at a fraction of the cost. pic.twitter.com/fRkegyMs5t