Tsinghua’s Chen Wenguang: Building an RTX 5090 Attention Kernel with Claude Code, and How AI Is Redrawing the Boundaries of Low-Level Programming
How far has AI-assisted programming come? A professor in Tsinghua University's Department of Computer Science and Technology is already using Claude Code to develop high-performance GPU kernels, with enormous efficiency gains.
Recently, at the "Agent Innovation Workshop" hosted by the Beijing Key Laboratory of On-Device Intelligence and co-organized by ModelBest and the OpenBMB open-source community, Professor Chen Wenguang of Tsinghua's Department of Computer Science and Technology gave a distinctive talk.
Using a hardcore engineering project as his case — "developing a Flash Attention kernel on the consumer-grade RTX 5090 that exploits FP8 to run 20% faster than existing solutions" — he walked through his real month of "pair programming" with Claude Code.
This is not just a technical story about squeezing out maximum performance; it is a deep exploration of the future software development paradigm. Professor Chen points out that technology is moving from conversational tools like ChatGPT, which stay within the information space, to agentic tools like Claude Code that can actually operate a computer, call tools, and execute complex tasks — and the way software is developed is undergoing a fundamental change. We are moving from "AI-assisted programming" to a new era of "AI-led programming."
This raises an industry-wide question: in the AI era, what are the core skills of a programmer? Do you still need to grind through C++ syntax details? Do you still need to know assembly? Professor Chen ended his talk with a "question" rather than an "answer." Below is an edited transcript of his talk:
Why Bring In Claude Code?
What I'm sharing today has nothing to do with academia; it is basically engineering. We wanted to build a somewhat complex program to see just how far existing AI could help us improve it. So we spent about three weeks building a GPU kernel on a low-end consumer graphics card and improved its efficiency by about 20%.
Zhiyuan (Professor Liu Zhiyuan of Tsinghua's Department of Computer Science and Technology) just said that he thinks the arrival of OpenClaw is no less influential than ChatGPT. I largely agree. Conversational models like ChatGPT basically stay in the information space — looking things up, answering, generating reports. But starting with Claude Code, the capability became tool use, and especially the ability to operate a machine improved dramatically.
Especially last December, when Anthropic released a new Claude model, and later when OpenAI released new models, I think the whole way software is developed changed dramatically. People used to use Cursor for AI-assisted programming, but in this period of developing algorithms, it was no longer AI-assisted programming. I basically conversed with it in natural language, stating requirements fairly precisely, and Claude Code could implement almost all of them — far more efficiently than I could alone.
I think tools like Claude Code, with strong model and tool-calling capabilities, operate your computer: with your various authorizations, it can use shell scripts, drive your browser, and do a great many things. It doesn't just show you results in a chat box; it can actually send emails and handle all kinds of scheduled tasks.
OpenClaw, building on that, greatly simplified the user interface. Claude Code is still aimed at professionals like me — people who work at a computer, writing programs, writing documents, sending emails. OpenClaw completely opens up the interface: you can access AI through instant messaging apps. The underlying logic is the same: pick useful tools and call them. Tasks can be very complex and require many steps of tool calls, but they share the same core. Recent AI progress made this combination feasible, which is why OpenClaw became very popular starting in January and February this year.
But there are still differences between OpenClaw and Claude Code. Claude Code is more about a project: source files go into the project directory, and it first searches the code and documents in that folder. OpenClaw's interface is more like a collection of functions — I tell you five things or eight things today, and it organizes them well. AI has a major context or memory problem, and OpenClaw needs better management there.
Overall, I am still a light user of OpenClaw, so most of what follows is about Claude Code.
Developing high-performance GPU kernels with Claude Code: In short, what we did was develop a kernel called Flash Attention on the consumer-grade RTX 5090. Why tackle this specifically? If you follow AI chips, you know that NVIDIA's ecosystem is strong, with CUDA compatibility that is better than domestic alternatives. But look closer and you'll find that each generation of NVIDIA architecture differs enormously, so "CUDA compatibility" is not really that compatible. With a new architecture, essentially all high-performance kernels must be rewritten. A lot of programmers and a lot of money today go into writing a faster kernel on a new architecture. NVIDIA's RTX 5090 is a generation beyond the H series, more advanced, but it is built for consumers. Although it is also a B-series part, its underlying architecture differs greatly from B200 and B300 cards, which are also B-series. Still, they share some things: both support Blackwell FP8 and a new data format in which a group of parameters shares one exponent. The H series did not support this format. So the 5090 and B200 both have a data format we believe is very helpful for training and inference. B200 has indeed exploited it. If you have been following, Flash Attention 4 was just released with many new techniques, including excellent use of FP8 and strong results, but there is no matching kernel for the 5090 — that part has not been done well.
Transformer Engine now implements some functionality on the 5090. It wraps a layer on top of cuDNN to make AI training easier. It supports an MX FP8 training strategy, split into parts: some linear layers use MX FP8, but the attention part still uses BF16.
My project asked: since the hardware has this capability and existing libraries don't provide it, can I use the hardware's FP8 capability to implement a faster Flash Attention? Then we could use cheap consumer cards like the 5090 for model training, providing a low-cost solution.
In reality this is very complex. The 5090 looks like a B-series card, but its structure is very different from the B200. First, it has no Tensor Memory and no separate asynchronous compute instructions; even its most important memory access and compute capabilities fall short. Data can only be loaded each time from the card's HBM (global memory) into Shared Memory that compute units can access quickly, then into registers for computation.
Here a problem emerges: Shared Memory is actually very small — only 128 KB per Stream Processor — which heavily constrains how algorithms are written.
Precisely because the 5090 lacks Tensor Memory, the optimized Tensor Memory-based methods used on the B200 are unavailable. So how do you compute? You must load data from Shared Memory into registers and perform FP8 MMA (matrix multiply-accumulate) in registers. Now register capacity becomes a problem: only 255 registers per thread. Exceed that and you get so-called "spill" — writing currently unused values to Shared Memory and reading them back when needed, which is much slower.
Due to architectural limits — no Tensor Memory, no so-called WGMMA instructions, memory access and compute that must be written rigidly by hand, plus very small Shared Memory and registers — writing the kernel became a huge obstacle.
Attention has a forward and a backward pass. The forward pass computes through the Attention architecture; the backward pass computes derivatives and propagates back. The forward pass was done in two or three days and performed well: FP8 forward is 1.7× the performance of BF16 forward, which is quite good, with good register usage and no spill. But the backward pass ran into big problems, because it involves five matrix multiplies — how to connect them, how to pass data between them, all while keeping Shared Memory and registers within limits — which took a great deal of effort.
This work is not like ordinary front-end programming; it is essentially a process of logic flow with many decision points and optimizations along the way: should I go this way or that way? Which direction to explore first, which later? And any single optimization may raise compute efficiency while increasing register pressure, so that in the end there is no real speedup.
This is our performance optimization process. The figure below shows kernel execution time, from 8-something milliseconds at the start down to 3-something milliseconds later — a big drop. On the right is the comparison with BF16; the work implemented in TransformerEngine reached up to 1.25×. But that 1.25× version was not quite right and had bugs; after fixing them, some shapes reach a 1.2× speedup.
The most effective optimizations we found over the whole process were these:
The first was in the "compute, then multiply" stage, where both were consuming time; you should iterate so that one side does Load while the other computes. We also did an "unroll": unrolling reduces branches and improves fusion, generally increasing performance, but it also increases register pressure. As I said, registers are very limited, so unroll 1 means no unrolling, which lowers register pressure and cut spill from over a thousand bytes to 272 bytes — a big performance gain.
The second optimization was an in-place matrix transpose in Shared Memory, which avoids writing back to global memory and reading it in again — there is no room in Shared Memory, so without this in-place transpose there would be nowhere to put it and you would have to dump it out and load it back.
Both optimizations worked very well.
The AI is not dumb: it doesn't do only the one thing you ask. When I say optimize, it automatically brings up the SASS framework, analyzes your instructions, the distribution of each instruction type, and what performance it should reach. AI has a systematic methodology for optimizing performance, and I don't even need to tell it. But I haven't given it things like Nsight that require superuser privileges; analysis done in user space might yield even higher performance.
I won't go into too much detail. Let me talk briefly about what AI does well after working with Claude Code for a while.
The first is architecture exploration. It can analyze existing code — how many of which instructions, how many spills — including spreading out data reads to get the most performance. These things are tiring and mechanical for humans, but AI does them well and tirelessly. For it these are routine, and it mostly gets them right without many mistakes.
The second is systematic optimization. We humans have to find many hotspots and then check whether compute is redundant, memory access is redundant, compute and memory access overlap poorly, or thread loads are imbalanced, and so on. Claude Code can do these relatively automatically. It is as if a performance engineer's basic skills are built in, presumably from its training data. Every time we run correctness tests: what passes we keep, what fails we feed back, and this iteration works very well.
Also, methods like "pingpong" and FA4 that split threads into two classes — one for compute, one for memory access — we tried 4+4, 6+6, 8+8, and other configurations. If I had to write it myself, I'd have been stuck long ago. But letting AI write it, reverting when wrong, is just experimentation for it; failure doesn't matter, and a machine doesn't get tired. That includes whether to unroll and exploring the parameter space. Because we were modifying Flash Attention, we had to integrate with low-level pieces like cuBLAS and our training scripts, and the system handled those relatively automatically too.
Also, for our 5090 (SM120) architecture kernel, with my prompts AI could read papers and reports on the Hopper (SM90) and B200 (SM100) architectures to see which techniques could be applied here. These things — including generating reports at the end — are relatively easy for it.
Finally, let me briefly cover what we did not succeed with; it doesn't matter much, since our own work also has failed directions.
One direction was 4+4 pingpong. Before doing it we estimated how much it could improve: overlapping compute and memory access should give better performance. But when applied to the real, full Attention, register pressure and other constraints imposed many limits, and in the end we couldn't make it work.
Not all exploration was useless, though; these attempts gradually deepened our understanding of the problem. With 8 warps, 128-byte rows, and 255 registers occupied, this is the stable result we ended up with, because other attempts always broke some of these factors and made overall performance worse.
What Will Programming Look Like in the AI Era?
As I worked, I kept thinking: if someone else did this, could they do it better? Or in what ways could it be done better?
First, as I mentioned, I did not open up low-level details like Nsight to it. That is not a systemic problem, but a detail.
Second — and this is actually important — memory and context need managing. Complex tasks require reading a lot of code, searching many things, and reading many experiment results, so I used a model with a one-million-token context window. But as work progressed, issues still arose later on. You have to keep watching this, give it hierarchical memory organization, and compress or clear it from time to time.
Common skills — anyone familiar with OpenClaw knows how to use them — but for complex things like plan mode, you can't let it take on too much clutter. For a hard problem, first let it think and set a detailed plan, then implement that plan. There is also a brainstorm skill in the system; installing it improves results.
A problem we often hit in practice is thinking we have reduced register pressure when we actually haven't. The reason is that NVIDIA only exposes the PTX instruction set, a virtual instruction set with no control over the register allocation layer beneath it; the compiler responsible for register allocation is still a black box to AI. Improving register allocation from outside is very hard, because you don't know whether the allocation algorithm has already eliminated temporary variables that aren't needed.
The model has no good way to perfect the coordination of these two aspects yet. If I change only one frequently edited file, compilation takes about a minute; if I change more, it takes about an hour. My early tests were relatively simple; later I found loss problems, so I switched to longer tests of about ten to twenty minutes.
There were other failures too — for example, AI modifying things without permission and causing errors. But we put good software engineering safeguards in place: after every change we run relatively complete correctness and functional tests, keep what helps, feed back what doesn't, and use version control to make the whole process safer.
Finally, although on some tensor shapes our kernel is 20% faster than the previous BF16 Flash Attention, when we brought it into model training tests, actual training speed was comparable — not much faster. Whether FP8 can be used well and its potential truly realized still requires creative work. Today's AI can look at what others have done and see whether it applies to my scenario, but it cannot yet do truly complete innovation. I don't think AI can't do it — it is essentially search, reinforcement learning, and collaboration with humans — but at least for now it hasn't done this.
After finishing the kernel, with decent performance, I sent it to a student to try. But we soon found something off: a convergence problem. By a few dozen iterations — after 60 — a problem became fairly obvious. Our student thought there were errors and wrote some test scripts.
I've been fixing this for the past couple of days, trying various methods, finding and fixing some bugs. Around iteration 80, it previously converged to 5.96, while BF16 should reach around 5.5, so now at 80 iterations it also converges to around 5.5. It is basically working.
As for GPU kernels, I have long worked on performance training for domestic platforms. There used to be a view that the main problem with domestic platforms is the lack of high-performance kernels, so how to write high-performance kernels for them is the biggest issue.
After doing this work, I found that the NVIDIA platform is not simple either. It looks CUDA-compatible, but underneath things differ enormously; much cannot be reused and must basically be rewritten. Hence many new programming abstractions have appeared, like Triton and TileLang, recently built at Peking University.
Everyone hopes for a layer of abstraction friendlier to programmers that still controls the low level, plus a very general compiler that handles the kernels underneath. After this period of practice, I now somewhat doubt this path, because very low-level things — many architectural features of the 5090 — are simply not supported. The underlying Triton compiler, Cutlass, and others do not support the 5090 well; to develop high-performance kernels you can only write the lowest-level CUDA, or even PTX assembly, to extract its performance.
If that is the case, what is the value of such mid-level abstractions? I now have some doubts. But this work is still preliminary, so this is not a firm view.
When we use AI for all kinds of work, memory use matters because you are doing different tasks; even within one task we split into many subtasks, some succeeding and some failing, each time with major code refactoring. How do you manage memory well and put what's needed into context? The existing "skill mechanism" and "memory mechanism" are fundamentally about hierarchical memory: don't cram everything in at once, but insert what was known before as needed. So far, management is not yet fully transparent and automatic.
So I believe that the more you understand the problem, the more you can still amplify your own abilities and do better. I know the basic concepts of compiler and performance optimization well, and I know the fundamentals of CUDA programming. But for an architecture like the 5090, whose details I am not yet familiar with, I could still get to a certain point together with AI.
Some numbers for the whole project: about 5,000 lines of CUDA/C++ code, 30 files changed, roughly 19 successful optimizations and 15 unsuccessful ones, with support for multiple GQA shapes. Forward reaches up to 347 TFLOPS, backward over 240. Broadly, AI can help systems experts amplify their abilities. AI can extend what we know and help us do better, but it is not a complete replacement.
In the AI Era, Do We Still Need to Learn C++?
I'd like to close with a post from Weibo:
Teacher Yu Shiqi is a very dedicated teacher at SUSTech who teaches C++. He said C++ will return to "ancient-method coding": folding interviews into the kernel, because homework has lost its meaning — whatever the assignment, students can hand it to a model and it will surely be done. And assignments are ten thousand times simpler than what I talked about today, doable at will. There is no way to assess how much students have really learned through homework.
The first reply was mine. I was thinking: what do students really need to learn? If they truly need to learn it, no exam format is a problem. But in the current situation I am a bit less confident. Many people are no longer familiar with assembly, including those of us who write operating systems. Will today's high-level languages become like assembly — needed by very few people? That is, we only need to talk to models in natural language, the model produces high-level code, and the compiler compiles it, and we can still do a great deal. But that is still an open question.
Two other people who reposted the post are also influential in China's developer community. One said that if assembly is useless, don't offer the course; if you offer it, teach it properly. Zou Xin said the curriculum requirements for training researchers in core CS technology and for training application talent are very different. I think his current view is probably similar to mine half a year ago — that AI can only make one-off personal websites and the like. After this period of practice, I somewhat doubt his current statement. Doesn't someone who writes kernels count as a core CS researcher? I think they do, yet it seems you can do a lot without understanding the low-level details.
So the question on the right: in the AI era, do researchers in core CS technology need to know C++ syntax details? I'm honestly not sure. The direction I sense is that collaboration with AI and big-picture knowledge matter more — what you ask it to do, and which direction to optimize toward.
You can imagine us pruning a large search space so AI spends less effort. Endless investment in one thing is not acceptable when the investment exceeds the value produced; it is only interesting when investment is less than the value you create. So pruning still matters.
I can't say micro-level knowledge is definitely unnecessary; people who write assembly and compilers are still needed. It may be like the CSAPP course. People asked me at the time: what is the point of taking Computer Systems: A Programmer's Perspective? The course says you don't necessarily need to write assembly yourself, but you need to read it, know what the compiler generates, and know what a high-level language you write will be compiled into. Symmetrically: you may know what kind of code a model will turn a piece of natural language into.