Welcome to my blog! đź‘‹

Hi there, This is Li.

FlashAttention-3 on Hopper

Contents 1 Introduction and Motivation 1.1 Why attention is the bottleneck 1.2 The Hopper problem: FA2 leaves most of the H100 idle 1.3 The three ideas of FlashAttention-3 1.4 The code base: from CUTLASS C++ to the CuTe DSL 1.5 Organization of this post 2 The Hopper GPU Architecture: Features Used by FA3 2.1 Chip-level anatomy of the H100 2.2 The SM: warps, warpgroups, and the asynchronous datapaths 2.3 TMA: the Tensor Memory Accelerator 2.3.1 Multicast and clusters 2.4 mbarriers: arrive/wait synchronization in shared memory 2.5 WGMMA: asynchronous warpgroup matrix multiply 2.6 The async proxy and memory fences 2.7 Register reallocation: setmaxnreg 2.8 Threadblock clusters, persistent kernels, and occupancy 2.9 The SFU: why softmax is expensive 2.10 Putting it together: the FA3 hardware checklist 3 Background: Attention, Online Softmax, FA1/FA2 3.1 Attention and the quadratic intermediate 3.2 Safe softmax and its streaming state 3.2.1 The block-combine interpretation 3.2.2 Consequences for implementation 3.3 FlashAttention-1: IO-aware exact attention 3.3.1 Backward recomputation 3.3.2 First-generation work partitioning 3.4 FlashAttention-2: parallelism and work partitioning 3.4.1 Fewer non-matrix operations 3.4.2 A query block becomes an independent thread-block task 3.4.3 Split-Q warp partitioning 3.4.4 Performance envelope and the remaining Hopper gap 3.5 FA1, FA2, and FA3 in one comparison 3.6 Kernel anatomy inherited by FA3 4 The FA3 Forward Pass: Algorithm 4.1 Work decomposition 4.2 Online-softmax state and invariant 4.3 Producer–consumer warp specialization 4.3.1 Pipeline-stage ownership 4.4 Pingpong scheduling: hiding softmax between warpgroups 4.5 Intra-warpgroup overlapping: the 2-stage GEMM–softmax pipeline 4.5.1 Prologue, steady state, and drain 4.6 Softmax details: rescaling, correction, and LSE 4.7 Causal masking and the split loop 4.8 Per-tile computation and data movement 4.9 Tile sizes, head-dimension variants, and register budgets 4.10 Epilogue 5 The FA3 Forward Pass: CuTe DSL Implementation 5.1 A note on the CuTe DSL 5.2 Implementation map 5.3 Kernel configuration 5.3.1 Compile-time variant selection 5.4 Shared memory plan 5.4.1 Tensor views and WGMMA partitions 5.5 Kernel entry: pipelines and role dispatch 5.5.1 Pipeline cursors and phase changes 5.6 The producer: load 5.6.1 From a scheduled tile to copy coordinates 5.7 The consumer: mma 5.7.1 Accumulator ownership and initialization 5.8 Core operation: one \(n\)-block with intra-warpgroup overlap 5.8.1 Why each wait and release is safe 5.9 Pingpong in code: the scheduler barriers 5.10 Softmax module 5.11 Epilogue in code 5.12 Differences from the paper’s CUTLASS C++ kernels 6 FP8 Forward: Layouts, Block Quantization, Incoherent Processing 6.1 FP8 formats and accumulation 6.2 The layout problem: K-major operands and the in-kernel V transpose 6.3 Block quantization 6.4 Incoherent processing: the Hadamard transform 6.5 FP8 kernel schedule and performance shape 7 The FA3 Backward Pass: Algorithm 7.1 Gradient math 7.1.1 The softmax Jacobian, step by step 7.1.2 Why \(D\) can be computed without \(\mathbf{P}\) 7.1.3 Why LSE alone suffices to recompute \(\mathbf{P}\) 7.1.4 Gradients through LSE itself 7.1.5 Where the scale \(\alpha\) is applied 7.2 Tiling and the dQ problem 7.2.1 Why KV-parallel and not Q-parallel 7.2.2 Causal masking: trapezoidal iteration ranges 7.2.3 Three ways to resolve the dQ race 7.3 The preprocessing and postprocessing kernels 7.3.1 Preprocess: D, LSElog2, and zeroing dQaccum 7.3.2 Postprocess: dQaccum \(\to\) dQ 7.4 Warp specialization in the backward 7.4.1 Named barriers of the backward 7.4.2 SMEM budget 7.5 Scheduling the five GEMMs 7.5.1 Operand layouts: who is SS, who is RS 7.5.2 swapAB: transposing the whole GEMM instead of the data 7.6 dQ accumulation across thread blocks 7.7 Backward tile configurations 7.8 Causal masking, varlen, and other mainloop features 7.9 Backward performance 8 The FA3 Backward Pass: CuTe DSL Implementation 8.1 Configuration surface 8.2 SMEM layouts supporting two access orientations 8.3 Host-side: TMA atoms, schedulers, semaphores, PDL 8.4 Role dispatch: two producer warps, two consumer warpgroups 8.5 The producer load loop 8.6 The consumer inner loop: five GEMMs and two pointwise stages 8.7 The dQaccum store agent 8.8 Epilogue: dK/dV stores, and the GQA accumulation path 8.9 The postprocessing kernel 8.10 Preprocess kernel code 9 Tile Scheduling, Variable-Length Sequences, and Masking 9.1 What the tile scheduler controls 9.1.1 Where scheduling appears in the kernel 9.2 The scheduling problem: causal masking skews tile cost 9.2.1 Static, persistent, and dynamic scheduling are not synonyms 9.3 Variable-length sequences (varlen) 9.3.1 How the varlen coordinate map works 9.4 Scheduling and masking: division of responsibility 9.4.1 Correctness invariants for scheduler changes 9.5 GQA and PackGQA 9.6 Masking machinery 9.7 Paged KV and other serving features 10 Flash-Decoding: Split-KV and the Combine Kernel 10.1 Why decoding underutilizes the GPU 10.2 Parallelization over the KV sequence 10.2.1 The combine math 10.3 The CuTe DSL combine kernel 10.4 How the mainloop produces splits 10.5 Choosing num_splits 10.6 Interaction with GQA, paged KV, and varlen 10.7 Performance character and caveats 11 PackGQA: Packing Query Heads into the Query Tile 11.1 Motivation: filling the M tile 11.1.1 What packing changes, and what it preserves 11.2 The packed layout: a zero-copy CuTe transformation 11.2.1 TMA with a packed mode 11.3 The gather/scatter path: PackGQA 11.4 Masking and per-row head indices 11.5 Enablement criteria for PackGQA 11.6 Composition with split-KV, and provenance 12 Performance Analysis and Benchmarks 12.1 Forward pass, BF16 12.2 Backward pass 12.3 FP8 forward 12.4 An issue-time consistency check 12.5 Interpreting throughput curves 12.6 End-to-end implications of attention speedups 13 Caveats and Pitfalls 13.1 Numerical caveats 13.2 Resource and performance caveats 13.3 Minimum validation matrix for kernel changes 14 Sources and Further Reading How the sources were used ...

August 29, 2026 Â· 162 min Â· 34331 words Â· Li Cao

CMU Graduation

I graduated from CMU today — a journey I’ll always be grateful for, having completed the 2-year MS in ECE (Advanced Study) in 3 semesters, earning A’s in all classes, and serving as a teaching assistant and research assistant. For my last semester, I took classes in Distributed Systems, Cloud Computing, and Optimization. They were all solid and valuable classes. Together, they round out the machine learning and systems foundation I’ve been building toward throughout my time at CMU. That intersection is where I’ve found my passion, and I’m excited to keep building on it—ML performance and accelerated computing in AI after graduation. ...

May 10, 2026 Â· 1 min Â· 210 words Â· Li Cao

Distributed Systems Project

A high-performance inference gateway in C++17 that routes client requests to a cluster of LLM serving replicas. The gateway provides prompt-prefix affinity via consistent hashing, weighted load balancing, fault tolerance with mid-stream failover and request hedging, circuit breaker for degraded replica detection, streaming token delivery, backpressure management, and zero-downtime rolling updates. Replicas participate in a SWIM gossip protocol for decentralized membership and failure detection. Final Project Report

April 10, 2026 Â· 1 min Â· 67 words Â· Li Cao

CMU Course Reviews - 2nd Semester

Here are my thoughts on the courses I took or audited during my second semester at CMU. On Grading at CMU I received straight A’s again this semester. However, I’ve noticed some peculiar aspects of CMU’s grading system. Graduate students in the Carnegie Institute of Technology (College of Engineering) cannot receive A+ grades in transcript even if they scored A+ in a class not offered by CIT, capping their GPA at 4.0. Since an A+ is worth 4.33 on the GPA scale, graduate students in other colleges can offset an A- with an A+ to maintain a 4.0, whereas engineering students cannot. This policy likely stems from CIT being CMU’s oldest and original college, preserving a traditional grading system without A+ grades. ...

December 17, 2025 Â· 3 min Â· 608 words Â· Li Cao

Parallel Binomial Option Pricing

This project implements a high-performance parallel pricing engine for American options using the Binomial Options Pricing Model (BOPM), engineered to scale efficiently from multi-core CPUs to GPU-accelerated clusters. By leveraging OpenMP, CUDA and MPI, it addresses the algorithm’s sequential bottlenecks through a diverse set of optimization strategies. More details at: https://github.com/l1-ca0/parallel-binomial-option-pricing

December 9, 2025 Â· 1 min Â· 51 words Â· Li Cao

Last Day of Fall Semester

Today is the last — and coldest — day of the semester. I ended the term by giving a recitation for 10‑703 on applying reinforcement learning to diffusion models.

December 5, 2025 Â· 1 min Â· 29 words Â· Li Cao

15618 Project Proposal - Parallel Option Pricing

November 17, 2025 Â· 0 min Â· 0 words Â· Li Cao

Lock-free Programming is Hard

Lock-free programming has this magical aura around it. If you’ve ever heard of lock-free programming, you’ve probably seen those neat little Compare-And-Swap (CAS) loops that seem to solve everything. I found a bug in a CAS loop that had been sitting quietly in CMU’s 15-418 Parallel Computer Architecture lecture slide for years. The “Simple” Example Here’s what the example in lecture slide looked like: // atomic compare and swap int atomicCAS(int* addr, int compare, int val) { int old = *addr; *addr = (old == compare) ? val : old; return old; } // build atomic max using CAS void atomic_max(int* addr, int x) { int old = *addr; int new = max(old, x); while (atomicCAS(addr, old, new) != old) { old = *addr; new = max(old, x); } } The idea is: ...

October 11, 2025 Â· 6 min Â· 1078 words Â· Li Cao

Birds can fly -- LLM Edition

Why LLMs Give Confusing True/False Answers Ask an AI “Birds can fly, true or false?” and an AI might initially say “True,” only to concede “False” after a bit more probing. What’s happening here? LLMs don’t “know” facts like humans do. They’re pattern-matching systems that predict the most statistically probable response based on their training data. When they see “birds can fly,” they recognize this phrase appears far more often than “birds cannot fly” in human text, so they lean toward “True.” ...

September 28, 2025 Â· 2 min Â· 323 words Â· Li Cao

What is the Full Stack Equivalent of Systems Programming?

Web development has popularized the concept of a “full stack developer”—someone who is comfortable working on every part of an application, from the user-facing front-end to the server-side back-end and the database it connects to, with a holistic view of the entire web stack. But what is the equivalent of “full stack” in the world of systems programming? Redefining the “Stack” To answer this, we need to first identify what constitutes the “stack” in systems programming: ...

August 1, 2025 Â· 2 min Â· 324 words Â· Li Cao

Finetuning LLM for Text-to-SQL generation

I just completed a project that lets people ask database questions in plain English and get back proper SQL queries using a fine-tuned large language model. For the base model, I chose Mistral-7B-v3 and fine-tuned it specifically for SQL generation. Using QLoRA for efficient training, I was able to train the 7-billion parameter model on a single consumer-grade GPU (Nvidia Tesla P100) in around 3 hours. The resulting model performs well on common SQL patterns like filtering, joins, and aggregations, effectively handling the majority of real-world database queries. That said, it can be less accurate for complex subqueries or really intricate nested queries due to the limitations of the Mistral-7B model —- a larger model would handle these cases better, but this was a tradeoff between performance and computational requirements. ...

July 24, 2025 Â· 2 min Â· 235 words Â· Li Cao

CMU Course Reviews - 1st Semester

Looking back on my first semester at CMU, I wanted to share my thoughts and experiences with the courses I took. This might be helpful for future students planning their schedules. This semester I took 4 courses. The workload was intense but manageable with good time management. Here’s my breakdown: 18-613: Introduction to Computer Systems While the famous CSAPP course has been extensively reviewed and is taught at many universities, here’s my personal perspective on the CMU experience. ...

July 18, 2025 Â· 6 min Â· 1259 words Â· Li Cao

Curiosity is (Almost) All You Need

The landscape of learning has been fundamentally transformed. In an era where Large Language Models can generate code and explain complex concepts, the traditional barriers to learning have largely disappeared. What remains—and what has become more important than ever—is curiosity. The Great Democratization Not too long ago, learning new technologies or skills required: Access to expensive courses or textbooks Mentorship from experienced practitioners Trial and error through countless hours of debugging Physical presence in classrooms or labs Today, anyone with internet access can have a conversation with an AI that knows more about programming, mathematics, science, and virtually any field than most human experts. The means of learning are no longer the bottleneck—curiosity and the drive to learn are. ...

June 15, 2025 Â· 4 min Â· 795 words Â· Li Cao