In this blogpost I will discuss how I implemented a mega kernel on a B300, and we’ll go through the implementation details plus the mistakes I made by leaning too much on AI agents.
The idea of a mega kernel is to fuse all the operations of the model into one big kernel and do the scheduling ourselves, inside it. We can win on a few fronts here:
But of course it comes with a cost:
The idea is to generate a graph on the Python side, so that we also get the ability to express any type of computation. (The goal was a generalized framework for static graphs ,in the sense of no data-dependent operations, so that I can schedule every operation ahead of time.)
The whole pipeline is roughly these steps:
Let me go through them.
A model isn’t written as eager PyTorch. It’s composed by calling ops on a Graph object. Here is one decoder layer:
def _decoder_layer(g, refs, w, residual, ...):
# QKV gemv fused with RMS Norm
qkv = g.matmul(residual, w.qkv, gamma=w.attn_norm)
att = g.attention(qkv, kcache, vcache, ...)
# o projection with residual
h = g.matmul(att, w.o, residual=residual)
# Gate+Up with Norm + SiLU·mul fused
gu = g.matmul(h, w.gate_up, gamma=w.mlp_norm, silu_mul=True)
return g.matmul(gu, w.down, residual=h) # down-projSo we can control some things through the parameters of the ops. Here I can’t say I like it or dislike it. Ideally, in future work, I would rather have a pass over the generated graph that detects these opportunities automatically, since we can easily miss some of them. :P
The important detail: each op returns a tensor, and every tensor remembers which task wrote which range of its columns. That’s how the next op knows what it depends on.
This is the step where the high-level op becomes something the GPU can actually run. Each operator becomes one or more 64-byte “task descriptors”, which encode everything about that op.
The interesting part is the fan-out. A single matmul doesn’t become one task - it becomes one task per column tile. The output columns are split into N tiles (roughly one per SM), and each tile is an independent GEMV task that computes only its slice and writes it back. So qkv with 3072 output columns and 8 tiles becomes 8 GEMV tasks, each owning a 384-column range.
Attention is split differently: one ATTENTION_PARTIAL task per (sequence, KV-head group, KV chunk), followed by one ATTENTION_COMBINE task that merges the partials. This is plain flash-decoding.
A couple of things are not tasks. RMS-norm, the residual add, and the SiLU gate don’t get their own task - they become flags on the GEMV that consumes them. So a single GEMV does norm → dot → residual in one pass. The final RMS-norm before the lm_head is fused straight into the lm_head GEMV, which means the whole forward is literally one kernel with no separate norm step.
Here is the descriptor that all of this lowers into:
// csrc/tasks/task.cuh
struct TaskDescriptor {
uint16_t op_id; // 0 - OpId: which task (Gemv, Attention, ...)
uint16_t config_id; // 2 - which template instantiation to dispatch
uint16_t step; // 4 - draft step tag (speculative decode) // not yet used
uint16_t flags; // 6 - TaskFlags: fused epilogues
int32_t remaining_deps; // 8 - int32 so the device can atomicSub it
int32_t num_dependents; // 12 - how many edges leave the task
uint32_t dependents_offset; // 16 - index into the separate edge array
uint16_t in_buf[3]; // 20 - buffer-table indices (NOT pointers)
uint16_t out_buf; // 26 - output buffer index
int32_t m; // 28 - rows / batch
int32_t n; // 32 - cols / output dim
int32_t k; // 36 - inner dim (gemv)
int32_t kv_len; // 40 - KV length (attention, runtime-variable)
uint32_t reserved[5]; // 44 - per-op scalars (e.g. gemv: col_begin, col_count)
};
static_assert(sizeof(TaskDescriptor) == 64, "ABI: TaskDescriptor must be 64 B");
static_assert(std::is_standard_layout_v<TaskDescriptor>);
static_assert(std::is_trivially_copyable_v<TaskDescriptor>);
static_assert(offsetof(TaskDescriptor, remaining_deps) % 4 == 0,
"remaining_deps must be 4 B aligned for atomics");A few design decisions worth explaining:
Indices, not pointers. in_buf/out_buf are uint16 into the buffer table (more on that in step 5) - a separate array of device pointers, uploaded once. This keeps the descriptor a POD that is identical on both sides and never needs to be patched with addresses.
reserved[5] = overloaded per-op scalars. There is no per-task struct. The m/n/k/kv_len fields plus reserved[] are reinterpreted depending on op_id. For GEMV, reserved = (col_begin, col_count) - the column tile this task computes. For attention, they are (seq_idx, head_begin, head_count, ...). Split-KV attention doesn’t change the layout - it just reuses reserved/m/n.
remaining_deps is int32 on purpose. The old queue path did atomicSub on it from the device when a parent finished. The cooperative path doesn’t use it (it relies on the topological levels), but the field stays 4-byte aligned for atomics.
Flags are epilogues, not tasks. flags (uint16) encodes the fusions: bit0 = SILU_MUL, bit1 = ADD_RESIDUAL (gemv) / FUSED_RESIDUAL (rmsnorm), bit2 = RMS_NORM_INPUT. This is what lets a single GEMV do norm + dot + residual in one pass.
There is a lot of room to improve here, and I still have plenty of leftovers to clean up from earlier experiments, but so far this was an interesting solution. :)
While lowering, every task records its parents - the tasks that wrote the data it reads. At the end, I flip this around: for each task I compute its dependents (the children) and flatten them into a single edge array. Each descriptor then only stores dependents_offset and num_dependents, pointing into that flat array.
The tasks are emitted in topological order already (a task can never depend on one emitted after it), so the descriptor array is ready to ship as-is. The edge array is the compressed form of the whole dependency graph.
Because TaskDescriptor is a fixed 64-byte POD, the Python side builds a numpy array with the exact same byte layout and there is no real serialization - it’s a memcpy. The numpy dtype mirrors the C++ struct field by field, and a single assert guards that the layout stays 64 bytes. If the C++ struct and the Python dtype ever drift, the ABI version is bumped and things fail loudly instead of silently corrupting.
So at this point we have two byte blobs: the task array and the edge array.
This is the part I glossed over above. The descriptors reference buffers by uint16 index, not by pointer. The buffer table is where those indices resolve.
On the Python side, every tensor the graph touches - weights, per-step activations, the KV cache, the scratch for attention partials - is registered into an append-only table and gets an index. Weights and the KV cache are already on the GPU; activations and scratch are allocated on demand. When it’s time to upload, the table is flattened into a plain list of device pointers (tensor.data_ptr() for each), in index order.
So the contract is simple: descriptor says “read buffer 1, write buffer 3”, the kernel looks up entries 1 and 3 in the pointer table, and gets the real device addresses. The table is built entirely on the Python side and uploaded once.
The Python side hands three things to the C++ extension: the task bytes, the edge bytes, and the list of buffer pointers.
ext.upload_queue(tasks.tobytes(), edges.tobytes(), buffers.ptrs())On the C++ side, upload_queue lays out a single contiguous blob in device memory - a small 64-byte header, then the task array, then the edge array - and cudaMemcpys it over. It also makes a second “pristine” copy of the blob, so the same graph can be re-run later by copying the pristine version back (device-to-device) without re-uploading from the host. The pointer list becomes its own small device array: that’s the buffer table the descriptors index into.
Note that none of the model data moves here. The weights and KV cache are already on the GPU; we only copy descriptors, edges, and a handful of addresses.
This brings us to the level abstraction.
A level is a set of tasks with no data dependency between them, so they can all run at the same time. Formally, a task’s level is 1 + max(level of its parents), and tasks with no parents are level 0. Two tasks at the same level are guaranteed independent - neither is an ancestor of the other.
Right after the upload, the host walks the graph once (the tasks are topologically ordered, so a single forward pass is enough) and assigns every task a level, then buckets the task indices by level. The kernel only needs two arrays out of this: the task indices grouped by level, and the offsets where each level starts.
We launch one cooperative-grid kernel - one block per SM, all blocks co-resident - and it walks the levels:
for (int lvl = 0; lvl < num_levels; ++lvl) {
for (int i = lo + blockIdx.x; i < hi; i += gridDim.x) // this level's tasks
coop_run_task(tasks[level_order[i]], ...);
if (lvl + 1 < num_levels) grid.sync(); // one barrier
}We use the cooperative group to synchronize across the grid (not a solution I love, but it works and it’s something worth doing now and replacing later). Because every task in a level is independent, the single grid.sync() between levels is enough to enforce all the ordering - every parent of a level-(L+1) task lives in level L or earlier - and it also gives cross-SM memory visibility, so we don’t need per-task fences. This is the whole reason the levels exist: they let us replace fine-grained per-task synchronization with one cheap barrier per level.
Inside coop_run_task, the dispatch is a switch on op_id that jumps to the right task implementation. Since the descriptor is block-uniform, every thread in the block takes the same branch.
These are the 10 OpId values (task.cuh), of which a Llama layer effectively uses 4:
| OpId | Task | Role in a layer |
|---|---|---|
GEMV (2) | tiled linear projection | qkv, o_proj, gate_up, down, lm_head |
ATTENTION_PARTIAL (8) | online-softmax over one KV chunk → (m,l,o) scratch | split-KV attention |
ATTENTION_COMBINE (9) | LSE-merge of the partials → final O | split-KV attention |
RMSNORM (3) | standalone rms-norm | (usually fused into GEMV, rarely standalone) |
SAMPLE (6) | argmax over vocab | only at the end, end-to-end decode |
ATTENTION_APPEND (7) | RoPE-append current k,v | append new KV |
NOP(0) / MEMCPY(1) / ATTENTION(4) / MLP(5) | - | utilities / aliases |
We ran Llama-3.2-1B-Instruct with a single decode token (B=1) and a KV cache of 512, comparing the mega kernel against vLLM on the same setup.
| GPU | Speedup vs vLLM |
|---|---|
| B300 | 1.21× faster |
| RTX 4070 | 1.00× (parity) |
On the B300 we got a 1.21× speedup over vLLM. On the RTX 4070 we reached parity (1.00×), so no speedup was recorded there.
Note: we did not run exhaustive benchmarks yet.
I had a lot of problems working with Claude Code, because this is not a trivial problem and there are many possible solutions and a lot to research. I started by thinking through the design myself and delegating the implementation to agents (for example: “implement the gemv, attention, rmsnorm tasks”). It did them correctly, but there were no shared patterns between them - each one called the device kernel function in its own way. The benchmarking file, on the other hand, did follow the existing code patterns, which was good in a way.
So after I started delegating to Claude and reviewing the code, I noticed most of it could easily have been turned into reusable helpers (another mistake - maybe I should have started with the helpers and simpler code first), so I went and made them. First of all, the code quality was that of a student who spent time writing correct CUDA without prior C knowledge. :D I had to step in and ask for things like proper namespaces and templated parameters for the constants.
Then I wanted to push the capabilities of my implementation, since I wasn’t beating vLLM at that point. Here I think I got a good RNG roll from Claude, because it found some nice ideas around not copying data redundantly and some fusion opportunities (which we’ll discuss in another part). It went well, what can I say.
I could write a whole separate blogpost just on how I got scammed by the agents here. What happened: I didn’t check the Python benchmarking code, because I was sure it could handle something this simple (CUDA events, run vLLM, run my mega kernel). Well, it turned out we had some nasty problems. I told it “run the vLLM baseline” and “benchmark my mega kernel” and got two completely different things measured. The vLLM benchmark was the correct idea - generate 256 tokens with a KV cache of 512 and compute the time-between-tokens divided by the number of tokens. But the mega kernel benchmark was running only a single decode layer.
That doesn’t sound too bad on its own. On the first run I got an insane 1.44x and I was surprised - but then I checked and found out, for some reason, we weren’t even using the final logit head. What?
After fixing this and making both do a single full decode, I still got a nice 1.21x speedup over vLLM, which is a decent result. The problem is that adding --lm-head corrupted the rest of my code and I lost some neurons being angry at a machine.
I definitely learned a lot about using these agents. Mostly, they don’t make you that much more productive. Of course they code faster, that’s true, and I’ve kind of lost my muscle memory for the lower-level details (like getting a feel for the tiling of an op, or writing the code by hand - something I admit I’m not practicing enough these days, but it’s not lost, it just means I’ll need more time to get back to it). But the real problem is that I need to keep my eyes on it constantly. So whether it’s me coding or the agent coding, I still have to go through the review. It’s like having a buddy who codes while I’m the guy who approves the changes.
Also, “knowing how to ask” is something everyone keeps repeating these days, but feeling it on your own skin hits harder than ever.
It’s a good project. I got to learn more about mega kernels and how to actually build them, and also about the agentic workflow side.
I’d like to continue this project, but I need to learn more about a few aspects. :) I feel like the approaches you can take here are infinite - from how we tile the data, to how we emit the code, to how we test it, etc.
So another direction would be going the MPK style: having proper IR lowering passes where I run some optimizations (on the math side, to find equivalences) and also on the CUDA side, to generate the code. If you want to hear about my future ideas, I’m always happy to discuss this with folks.