Shmoosh update:
This is the most encouraging speed result so far.
For context, Shmoosh is my current research project around making diffusion models run faster by compressing parts of attention in ways that preserve the behavior of the model. The goal is to find techniques that actually matter in local creative tools like Underpaint and Comfy, where people are running image generation and editing workflows on consumer GPUs.
And today, the kernel finally moved.
I pushed the next slice:
1f0ee4e Add TF32 packed attention fast mode.
The short version is that I added opt-in `dot_precision` support for fused packed attention.
The default remains `ieee`, because fidelity matters and I don't want to silently change the high-quality policy. But there is now an explicit `tf32` fast mode for cases where you want speed and can tolerate a very small quality tradeoff.
The measured changes
On synthetic 1024px K7/no-QJL packed_t testing on the RTX 4070, packed attention moved from:
0.6232ms to 0.2831ms
That's the packed attention kernel phase getting more than 2x faster.
The total encode+attention path moved from:
0.8177ms to 0.4761ms
And the relative RMSE barely changed:
0.023998 to 0.024101
That is the kind of result that makes you stop squinting at the profiler and start paying attention.
The 1024px image suite also held up well:
Mean speedup: 1.068x
Mean PSNR: 54.17 dB
Min PSNR: 51.70 dB
The packed attention phase in the image suite moved from:
1.2806ms to 0.7293ms
Synthetic results are useful, but they can lie. Seeing the packed attention phase move meaningfully in the actual image generation is more interesting.
An explicit fast-mode policy
I didn't promote TF32 into the high-fidelity default policy, because one of the test cases, reading-nook, paid a small quality tax:
52.07 dB to 51.70 dB
Still very high, but it is a real difference, and it deserves to be treated as a policy choice.
So I added a separate fast-mode config:
configs/underpaint-juggernaut-sdxl-up0-self-attn1-firstblocks-gated70pct-k7-noqjl-tf32-policy.json
The validation is clean:
pytest: 87 passed
compileall: passed
jq: passed
git diff --check: passed
repo clean and synced
The next shape of the system
I'm slowly narrowing in on the right shape.
K-only compression has been more promising than touching V.
Packed representation beat the direct rotated-K path, because density matters. Late-gated self-attention continues to look like the sane lane.
K7/no-QJL is holding up well enough to be interesting.
And now TF32 gives the packed attention kernel an actual shove.
Before this, Shmoosh was in the “scientifically promising, modestly faster” zone. But, a local creative tool needs wins that survive contact with actual image workflows, real GPUs, and quality expectations.
Still working on it.
This is the first result where the packed-attention path became meaningfully faster.
I'm excited.