I recently ran an agent loop to find if it could improve on pytorch’s flash attention (F.scaled_dot_product_attention) on a rtx a5000. It ended up finding a kernel that is ~1.5-2% faster and uses 28% less DRAM traffic, for Wan self attention. Although the perf gain is not significant, in this blog I will cover the agent loop flow and the key changes it made to achieve those numbers.

Kernel Stats


Measurement Methodology

A 1.5-2% margin is smaller than most benchmarking noise and thus we need to be careful while measuring this. Every experiment in this blog ran on an rtx a5000 (sm_86, 64 SMs, 24 GB).

Clock vs Temp


For the measurements, I didn’t pin the clocks to a low freq (or use NCU’s timings) as I wanted to get close to real world perf. I noted that as the temperature increased (>60 degrees), the clock frequencies automatically went down to ~1700 MHz and ~7600 MHz. As a result of this fluctuation, no two runs executed at the same fixed clock. I did the following to handle this discrepancy:

  • Warmup: there were 3 rounds of warmup before doing 12 rounds of measurement.
  • Interleave and rotate: In each round all kernels were run. There were other variants that are not covered in this blog. And in each round their execution order was shuffled.
  • Data randomize: On each run Q, K and V were calculated from a different seed randomly.
  • L2 scrub: Before each run the L2 cache (6 MB for a5000) was flushed by doing a dummy write of 256 MB.
  • Hot & Cold runs: This entire 12 rounds experiment was captured twice. Once when the GPU was cold (at the very beginning), ~58-62 degrees and again when it reached higher temperatures, ~68-73 degrees.

Agent Loop

As mentioned earlier, the kernel was found through an automated search loop. I ran deepseek v4 flash for about 1hr that cost me $2. The loop worked similar to breeding/evolution. The loop starts with a working base and every round LLM proposes few changes in it. Every proposal that clears the validation checks is measured for speed. All the past kernels (along with failures) are stored in history and used for all subsequent generations. History contains a short description of the previous tests such as “bigger tiles than this collapses occupancy, much slower”, which helps inform future decisions.

Still a lot of improvements are required in this loop, like having an IR (similar to the newly released CAKE) would make this process much better.

Agent Loop


Implementation Details

Attention is a weighted lookup, given by the formula out = softmax(Q @ K^T / sqrt(d)) @ V. If we were to compute it naively we would end up materializing the score matrix tensors into memory. FA’s core proposition is to never materialize the score tensor (QK) completely in the memory. It computes it in small blocks and then discards them after use. In our case we are running Wan2.1 and at every self attention, the compute runs over the entire latent. Q, K and V are all 21504 long. For the curious, we got 21504 by multiplying the VAE output of 21 latents by the token count of 32^2 = 1024 (512/8 VAE x 2 patch). The attention runs on 21504 tokens x 12 heads x 128 dims x 2 batch (pos and neg).

FA2 vs Ours


Most of the stuff in our kernel is the standard FA2 algo like 128 row query tiles, 32 row KV chunks, online softmax with scale and log2e folded into Q. The key difference is that instead of running one query tile per program, we run two query tiles. The effect of this is that the online softmax state - running max, running sum and the output accumulator now have two copies, one for each tile, which halves the kv fetch request. This also increases the register pressure (our kernel reaches the max 255 registers), but as seen in my previous blog, it’s not always a bad metric. One other thing that this kernel does differently is to use fixed contiguous tensors and tags the offset vectors with tl.multiple_of / tl.max_contiguous hints, similarly while loading Q and KV, hints are provided to ‘evict_last’ Q while ‘evict_first’ for KV. Using the .cg modifier KV is directly streamed from L2 into shared mem and it doesn’t touch L1, saving a wasteful stop along the path.

This kernel drops the memory traffic by 28% as we use the same kv for both the query tiles. But how does that only give us only ~2% perf increase ? Because the kernel didn’t speedup the compute and math pipes only bumped a little in their utilization. Also the L2 hit rate was already 97.6%, so most of the KV was streamed directly from L2, cutting mem traffic didn’t have much effect on this. Another interesting point to note, that directly applies from Volkov’s principles, is that as the occupancy fell the kernel got faster. In our case the same kernel is doing more work, requiring more registers.

FA2 vs our stats


More perf can likely be squeezed out if we drop down to CUDA or PTX. I’ll try to work on it in a future blog. You can find the code for this kernel here.