Will Bui

27x faster Attention Residuals!!! 🚀 We implemented Block AttnRes as a pip-installable package. !pip install flash-attn-res No annoying kernel nonsense. No compile/autograd plumbing. Call it like a regular PyTorch op. It just works. Methodology: 🔹 fused triton kernels 🔹 batched attention over residual blocks 🔹 online-softmax merge 🔹 flash attention-style split-KV reduction Thanks @LLMenjoyer and @cartesia for the support and guidance✌️

@flappyairplanes I don't get it but I see kittens I like

(4/5) One thing we’ve built is a “kittens” virtual machine that takes over the whole GPU and allows new kinds of co-optimization. We can go past the traditional sequential kernel model – for example, fusing entire training runs into a single kernel and even weirder stuff.

(1/5) Great to be at @sequoia to give a sneak peek of one of our research directions! TL;DR one path to data-efficiency may be to “abuse GPUs like they’ve never been abused before”

@m_sirovatka as a vllm contributor, i approve this post

I have now officially became one of the vLLM contributors, after long months of hard work. It has been a long and hard journey, I would like to thank my family, friends and my company for supporting me along the way. nah jk my 3 line PR just got merged 🫡

@retr0sushi_ Be around those with better taste than you

how does one develop a taste for research? explain to me like i am a guy entering college with unlimited energy and enthusiasm

@anneouyang @Norapom04 cattention

TIL the Huawei Ascend linear algebra kernels library is called "CATLASS" 🐱

@danofer @LLMenjoyer Great to know!!!

@will_ea @LLMenjoyer Dayum. (My implementation really got about ~6-8 sec/it, and hence my dropping it in favour of going to 1.6 it/sec). On a tiny mlm.

27x faster Attention Residuals!!! 🚀 We implemented Block AttnRes as a pip-installable package. !pip install flash-attn-res No annoying kernel nonsense. No compile/autograd plumbing. Call it like a regular PyTorch op. It just works. Methodology: 🔹 fused triton kernels 🔹 batched attention over residual blocks 🔹 online-softmax merge 🔹 flash attention-style split-KV reduction Thanks @LLMenjoyer and @cartesia for the support and guidance✌️

@Birchlabs @LLMenjoyer okay tbh jvp support isn't planned at the moment but PRs accepted! :)

@will_ea @LLMenjoyer I am just a guy I don’t have a budget

@Birchlabs @LLMenjoyer if you hire me to do so, yes.

@will_ea @LLMenjoyer will you add JVP support?

@Birchlabs @LLMenjoyer The benchmarking script and output here verified gradient correctness github.com/catswe/flash-a…. I also did some private internal tests verifying output correctness. More thorough tests will be added in the future.

@will_ea @LLMenjoyer the kind of test I'm asking for is something like what I did in this gist.github.com/Birch-san/3e0b… it also shows you how to add (and test) JVP support to your triton attention kernel, which is crucial for supporting MeanFlow models.

@slLuxia @LLMenjoyer Gotcha. I’m keeping that in mind for future development

@will_ea @LLMenjoyer yeah it's orchestration level more than the raw ops; but it is important because that's double the layer routing overhead in practice and where the chunk of the slowdown is reported from if i remember right in the original paper (and why they do blocks vs all layers)

@slLuxia @LLMenjoyer The package exposes lower-level phase 1/phase 2 ops that users can use to compose the routing however they want. I believe what you’re referring to is the experimental API/example, which is still rough around the edges and mainly meant for research/prototyping.

@will_ea @LLMenjoyer hmm; it seems like you drop some routing; in the original paper they do 2x per block for pre-MLP & pre-attention but your repo does each block as a single layer. can see in fig 2 and section 2 of the moonshot paper. is this intended?

@slLuxia @LLMenjoyer Yes it allows arbitrary blocks.

@will_ea @LLMenjoyer should be agnostic to block size/allow arbitrary blocks yeah? i find maintaining predicted interlayer circuitry gives the best perf for attnres

@ricklamers indeed. im personally very excited to contribute to the ai infra space

software stack matters, a lot

@Kimi_Moonshot I loved this paper so much I built the first public two-phase kernels for AttnRes ;) x.com/will_ea/status…

Introducing 𝑨𝒕𝒕𝒆𝒏𝒕𝒊𝒐𝒏 𝑹𝒆𝒔𝒊𝒅𝒖𝒂𝒍𝒔: Rethinking depth-wise aggregation. Residual connections have long relied on fixed, uniform accumulation. Inspired by the duality of time and depth, we introduce Attention Residuals, replacing standard depth-wise recurrence with learned, input-dependent attention over preceding layers. 🔹 Enables networks to selectively retrieve past representations, naturally mitigating dilution and hidden-state growth. 🔹 Introduces Block AttnRes, partitioning layers into compressed blocks to make cross-layer attention practical at scale. 🔹 Serves as an efficient drop-in replacement, demonstrating a 1.25x compute advantage with negligible (<2%) inference latency overhead. 🔹 Validated on the Kimi Linear architecture (48B total, 3B activated parameters), delivering consistent downstream performance gains. 🔗Full report: github.com/MoonshotAI/Att…

@LLMenjoyer this is actually pretty fire now that i think more about it...

this mfer is dropping speedful attention residuals kernels today i tried to tell bro to use the good diagrams, but bro didn't listen...... sad!

Check it out here: github.com/catswe/flash-a…

@MainzOnX @A_K_Nain i love writing kernels... unless i'm racing against a deadline

@A_K_Nain Yeah no one tells you that it’s actually not the fun part 😂 I’m writing backwards ops for 25% of my time last week and probably next week it’s no fun

I hate writing kernels. We shouldn't be doing this at least frequently. I hope the next gen compilers are much better and smart.

@blelbach I spent a couple weeks just writing kernels

I have not worked hours this long since the start of my career. So much to do these days. So many possibilities that are now unlocked.

@henrylhtsang I mean you need a decent amount of prerequisites knowledge to start with any substantial kernel work. The actual kernels may not seem like much but to get to that point requires narrowing down a bunch of ways to do it optimally that having experience would help.

one impression of working a new grad: they can seem... slow? like slow to finish a task and turn around in general is that what my TL thought of me when I first started

@Alexintosh @ritakozlov How is this AI slop, Alex?

@ritakozlov ai slop spotted

i've picked up the pen so many times to write about being a woman in tech and every time i chicken out because there's this catch-22: to talk about being a woman in tech, you need to have credibility. and once you start talking about it as a woman, you lose said credibility so i'm going to mortgage some of my credibility to get this off my chest, as someone who has both had a pretty successful career in tech, and leads a team with a lot of women on it: every woman you work with has had the most insane shit happen to her — on an almost daily basis. shit that makes you look at the camera and go "how did i end up here". from wild remarks about appearance to stalking and trauma dumping, and just constant dismissal from so many directions (employees, customers...). shit that you never tell anyone because they wouldn't believe you... i recently learned that like 97% of my followers on here are men. so my challenge to you is just to sit with that for a moment. you don't need to do anything about it (other than try not to be that person). but you should be aware that that's what every woman you work with deals with