Show HN: Phobos – A tiny scale-free kernel language with tile-DAG support

A developer built Phobos, a tiny scale-free kernel language that compiles to PTX and runs on NVIDIA GPUs, achieving 76% of cuBLAS SGEMM performance on a 2080 SUPER. The language supports tile-DAG design and was created as a personal research project to understand low-level GPU concepts.

TLDR:I’ve built Phobos , a tiny kernel language inspired by Triton. It lowers to PTX and runs on NVIDIA GPUs. It achieves acceptable performance at 76% of cuBLAS SGEMM GFLOP/s on a 2080 SUPER or 74% of the theoretical GFLOP/s peak . Phobos maps naturally to a distributed tile-DAG design. I only validated the cluster prototype on a single machine; there are no multi-node benchmarks here. This was a personal research project to gain a better understanding of low-level GPU concepts. I only realized the potential for distributed computing when implementing local tile optimizations. Ever since I was a kid and started to learn about artificial neural networks, I was fascinated by the human brain and AI to the point where I built object recognition software in my bedroom, using shitty webcams and a 1.8 GHz CPU for training. When I finished school I enrolled in Cognitive Informatics at Bielefeld University. Though, after a mere four days, I dropped out because I got a job offer from a startup in Paris and wanted to see the world. In the following years, I worked a lot on optimization, compilers and other machine learning related topics. I still follow the research, but I am an outsider to frontier AI research and tech. Here I am, back again with a shitty 2080 SUPER. Eager to learn. I strongly believe, in order to fully understand something, you really need to do the work and get your hands dirty. And because I have experience with compilers, optimization and LLVM, I thought writing a tiny language and compiler to target GPUs is approachable for me. I also have enough experience with software engineering, and compilers, to know that I must tread carefully. I want to keep this a learning exercise. It could very well snowball into a project that would consume too much of my available time if I take a wrong turn. Compilers are particularly tricky. Without being able to compile a program from start to end, you won’t see anything. Compiling a program is not trivial per se, so I am going to take some shortcuts. There won’t be any sort of linking or a phase model. The goal : create a small language that compiles and runs on NVIDIA GPUs. In doing so, I want to pick a model well suited for AI computations. It should make optimizations possible without having to reconstruct the semantics of typical AI kernels. I know that I want to support first class tensor types / tiles. In fact, the entire language is oriented around tiles. Here are two example programs I want to be able to compile. I’ll explain in more detail when we get to the language. Vector Addition This is \ \vec{c} = \vec{a} + \vec{b}\ . @autotune BLOCK in 1024 kernel add a: tensor<f32 N , b: tensor<f32 N , c: tensor<f32 N { let base = program id 0 BLOCK c base :+ BLOCK = a base :+ BLOCK + b base :+ BLOCK } Matrix Multiplication This is \ \mathbf{C} {M \times N} = \alpha \mathbf{A} {M \times K} \mathbf{B} {K \times N} + \beta \mathbf{C} {M \times N}\ . @autotune TILE M in 32, 128 , TILE N in 64, 256 , TILE K in 16, 64 kernel sgemm A: tensor<f32 M, K , B: tensor<f32 K, N , C: tensor<f32 M, N , alpha: f32, beta: f32 { let pm = program id 0 let pn = program id 1 var acc: tile<f32 TILE M, TILE N = 0.0 for kt in range 0, K, TILE K { var a = A pm TILE M :+ TILE M, kt :+ TILE K var b = B kt :+ TILE K, pn TILE N :+ TILE N acc += dot a, b } let c old = C pm TILE M :+ TILE M, pn TILE N :+ TILE N C pm TILE M :+ TILE M, pn TILE N :+ TILE N = alpha acc + beta c old } Side Note:If you are familiar with kernel development, then yes: it’s basically a mini-Triton. Unfamiliar Territory Let’s start to get to know CUDA and its target format. If we look at resources such as CUDA Refresher: Getting started with CUDA / CUDA kernel device code - CUDA Sample Codes Computes the vector addition of A and B into C. The three vectors have the same number of elements as numElements. / global void vectorAdd float A, float B, float C, int numElements { int i = blockDim.x blockIdx.x + threadIdx.x; if i < numElements { C i = A i + B i ; } } We are not interested in that. This is high-level code. We want the x64 of the GPU, not the C++. I first looked at the NVIDIA CUDA Compiler NVCC docs 1. However this is very convoluted as nvcc is what links entire executables. The CUDA docsthemselves have a section that’s more in our domain: 2 fn:2 PTX and low-level programming . Low-level programming is what we want. PTX provides a stable programming model and instruction set for general purpose parallel programming. It is designed to be efficient on NVIDIA GPUs supporting the computation features defined by the NVIDIA Tesla architecture. High level language compilers for languages such as CUDA and C/C++ generate PTX instructions, which are optimized for and translated to native target-architecture instructions. 3 There’s an existing LLVM backend we can use, called NVPTX Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ LLVM IR ⟶ NVPTX ⟶ PTX ⟶ GPU PTX is not the machine code the GPU runs but it is as close as we can get. SASS is the actual machine code 4, but undocumented. But there’s more. There always is. The feature creep. I can feel it. MLIR 5. Let’s use it instead of producing LLVM IR. In fact, I do not plan to use MLIR features, like a custom dialect. But it gives me the option to reach for CUTLASS-like performance for free…a man can dream. Seriously though, MLIR exposes many instructions that I’ll need. Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ PTX ⟶ GPU For this project I am also going to use Rust because it has bindings for LLVM with inkwell , MLIR with the melior 100x more pleasant. I had hopes that using MLIR would be as “simple” as using inkwell . But alas, it took me many hours of yak-shaving, patching the melior , mlir-sys and tblgen-rs crates, and other horrors, until I was able to compile a simple program. When I started this little adventure, I generated some PTX straight from LLVM IR just to get to know everything. That was the obvious path forward for me. But I quickly learned about MLIR. Getting started with inkwell took maybe a couple of hours. Getting melior and MLIR to work took me two days. Right-to-Left Because many things can break here, I am going from right-to-left. I am going to start with the most isolated task: run PTX code on the GPU. It is inevitably something I will have to do later. Let’s write the vectorAdd from above in PTX and run it with the cust First we need to understand the compute cap of our GPU. In my case, it’s a 7.5 for the Turing family. Yes, I have a pretty old GPU with only 8GB VRAM. Sadly, it also lacks float32 tensor cores — but I do have 16-bit precision available. We can work with that. bash $ nvidia-smi --query-gpu=name,compute cap name, compute cap NVIDIA GeForce RTX 2080 SUPER, 7.5 We can now start to write PTX. This should feel familiar if you’ve ever written LLVM IR by hand. At least we don’t have to write all of this in SSA form. In any case, this was the only PTX code I planned to write manually. Running the assembler directly with ptxas add.ptx was a useful sanity check to isolate any issues from Rust or the CUDA crate I am using. We can also compile kernels to PTX via nvcc -ptx in.cu -o out.ptx which is very helpful. We can write a kernel, compile it with the official toolchain, and inspect the generated PTX output. .version 9.0 // 9.3 should work, but does not? .target sm 75 // target arch compute cap .address size 64 .visible .entry add .param .u64 a, .param .u64 b, .param .u64 c, .param .u32 numElements { .reg .pred %p; .reg .f32 %f<4 ; .reg .b32 %r<5 ; .reg .b64 %rd<11 ; ld.param.u64 %rd1, a ; ld.param.u64 %rd2, b ; ld.param.u64 %rd3, c ; ld.param.u32 %r1, numElements ; // i = blockIdx.x blockDim.x + threadIdx.x mov.u32 %r2, %ctaid.x; mov.u32 %r3, %ntid.x; mov.u32 %r4, %tid.x; mad.lo.s32 %r2, %r2, %r3, %r4; // %r2 is i // if i = numElements, return setp.ge.s32 %p, %r2, %r1; @%p bra DONE; // byte offset = i 4 mul.wide.s32 %rd4, %r2, 4; // %f1 = a i 4 cvta.to.global.u64 %rd5, %rd1; add.s64 %rd6, %rd5, %rd4; ld.global.f32 %f1, %rd6 ; // %f2 = b i 4 cvta.to.global.u64 %rd7, %rd2; add.s64 %rd8, %rd7, %rd4; ld.global.f32 %f2, %rd8 ; // %f3 = %f1 + %f2 add.f32 %f3, %f1, %f2; // c i 4 = %f3 cvta.to.global.u64 %rd9, %rd3; add.s64 %rd10, %rd9, %rd4; st.global.f32 %rd10 , %f3; DONE: ret; } bash $ cargo run Compiling phobos v0.1.0 C:\Users\joaeb\code\phobos Finished dev profile unoptimized + debuginfo target s in 0.63s Running target\debug\phobos.exe 1048576 elements correct I include the Rust code here as we will come back to it later for testing purposes. Executing PTX on the GPU with Rust & cust js use cust::prelude:: ; const PTX: &str = include str "../ptx/add.ptx" ; fn main - anyhow::Result< { let ctx = cust::quick init ?; // init driver let module = Module::from ptx PTX, & ?; // compile the code let stream = Stream::new StreamFlags::NON BLOCKING, None ?; let func = module.get function "add" ?; let n: i32 = 1 << 20; // create some data let a: Vec<f32 = 0..n .map |i| i as f32 .collect ; let b: Vec<f32 = 0..n .map |i| 2 i as f32 .collect ; let mut c = vec 0.0f32; n as usize ; let a dev = a.as slice .as dbuf ?; // write to gpu let b dev = b.as slice .as dbuf ?; let c dev = c.as slice .as dbuf ?; let block: u32 = 256; // run our code let grid: u32 = n as u32 + block - 1 / block; unsafe { launch func<<<grid, block, 0, stream a dev.as device ptr , b dev.as device ptr , c dev.as device ptr , n ?; } // wait for gpu to complete stream.synchronize ?; // read from gpu c dev.copy to &mut c ?; // sanity check for i in 0..n as usize { assert eq c i , a i + b i ; } println "{n} elements correct" ; Ok } The source code we provide PTX is getting JIT compiled by the driver for the specific GPU to SASS. This makes sense as the same code should run on many different GPUs. If we introduce errors to the PTX, a mere “JIT error” message will be thrown at us. ptxas can help us identify what’s actually broken. Profiling our binary with ncu works as well. Nice. I have certainly a better understanding of the stack now. I can run low-level code on the GPU and verify it on the CPU. Let’s revisit the pipeline from earlier. Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ PTX ⟶ GPU We have started at the finish line. But that’s good. At least for me, this was the most uncertain part. Everything else is known territory from now on, mostly. The Devil You Knew I really like LLVM and I used to work a lot with it for a sea-of-nodes IR where I started with Java bytecode, parsed that to an IR of mine, did my transformations and then lowered to LLVM IR which I could then use to emit to x64 or WASM. MLIR did not exist back then. This was now more than 10 years ago. MLIR would have been a good fit to model Java’s high-level constructs. For example, property access throws a NullPointerException . The high-level IR I created contained nodes like null-check , to then emit many LLVM IR nodes out of these. Custom MLIR dialects would allow me to express this within the LLVM IR universe. 6 I thought initially that I was going with the same pattern … ⟶ Parser ⟶ IR ⟶ Optimizer ⟶ Codegen ⟶ LLVM IR ⟶ … but dropped the entire IR and Optimizer idea in favor of MLIR.Getting melior to work on Windows took some time. Lucky for you, I created PRs for melior mlir-sys tblgen-rs .instructions.md was committed to one of the repos and guess what. NTFS really does not like in file names. I’ll spare you the details of this entire saga, but fun fact: I compiled an x86 LLVM initially when Rust was targeting x64. If you know how long it takes to compile LLVM, you know what that means in terms of a time waste. With this out of the way, I can start to generate my first kernel via code. Creating PTX with MLIR We have a very simple c i = a i + b i kernel in PTX and we have Rust working with MLIR. The logical next step is to re-create the PTX with MLIR and run it. It’s not that simple though. melior can’t lower to PTX directly. I could use an external tool, mlir-opt , which lowers the MLIR to PTX for me but I want simple things. It’s LLVM all the way down. Funny enough, when I started this project, I used inkwell Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ LLVM IR ⟶ NVPTX Backend ⟶ PTX ⟶ GPU The code starts to get a lot more verbose so here are the basic ideas: // generic function to produce ptx code as a string. // it takes a closure to fill in the actual function body. pub fn gen ptx<'p, F ctx: &'p phobos base::context::Context, body: F - anyhow::Result<String where F: for<'c FnOnce &'p phobos base::context::Context, &'c Context, &Module<'c , - anyhow::Result< , { // MLIR setup let registry = DialectRegistry::new ; register all dialects &registry ; let context = Context::new ; //... let loc = Location::unknown &context ; let module = Module::new loc ; // fill in the function body body &ctx, &context, &module ?; // generate PTX string from MLIR lower mlir to ptx &ctx, &context, module } // recreating the ptx code for c i = a i + b i with MLIR fn build add ir<'a : &'a phobos base::context::Context, ctx: &'a Context, module: &Module<'a , - anyhow::Result< { let loc = Location::unknown ctx ; let i32t: Type = IntegerType::new ctx, 32 .into ; let f32t = Type::float32 ctx ; let idx t = Type::index ctx ; let mem space: Attribute = IntegerAttribute::new IntegerType::new ctx, 64 .into , 1 .into ; let memref t: Type = MemRefType::new f32t, & i64::MIN , None, Some mem space .into ; // ... this is where we add all the instructions ... let gpu module = OperationBuilder::new "gpu.module", loc .add attributes & Identifier::new ctx, "sym name" , StringAttribute::new ctx, "kernels" .into , .add regions gpu mod region .build ?; module.body .append operation gpu module ; Ok } run on gpu gen ptx ctx, build add ir ; Phase 1: The generated MLIR This could be achieved in many different ways. I wanted to use the gpu , and memref dialects early on. memref<?xf32, 1 states: we have a type with an unknown f32 dimension in address space 1 . We could actually say memref<1024xf32, 1 but, we don’t know the size of our vectors yet . module { gpu.module @kernels { gpu.func @add %arg0: memref<?xf32, 1 , %arg1: memref<?xf32, 1 , %arg2: memref<?xf32, 1 , %arg3: i32 kernel { // i = blockIdx.x blockDim.x + threadIdx.x %thread id x = gpu.thread id x %block id x = gpu.block id x %block dim x = gpu.block dim x %0 = arith.index cast %thread id x : index to i32 %1 = arith.index cast %block id x : index to i32 %2 = arith.index cast %block dim x : index to i32 %3 = arith.muli %1, %2 : i32 %4 = arith.addi %3, %0 : i32 // i index = i as usize %5 = arith.index cast %4 : i32 to index // i < n %6 = arith.cmpi slt, %4, %arg3 : i32 cf.cond br %6, ^bb1, ^bb2 ^bb1: // pred: ^bb0 // load a i index and b i index %7 = memref.load %arg0 %5 : memref<?xf32, 1 %8 = memref.load %arg1 %5 : memref<?xf32, 1 // a + b %9 = arith.addf %7, %8 : f32 // assign to c i index memref.store %9, %arg2 %5 : memref<?xf32, 1 cf.br ^bb2 ^bb2: // 2 preds: ^bb0, ^bb1 gpu.return } } } Phase 2: MLIR Lowering You can see that each memref argument expands to a total of five. The reason is that memref is overkill for this little example but still a good learning exercise. Instead of a single float pointer ptr , we now have ptr, align ptr, offset, size, stride . For example, memref.load %arg0 %5 : memref<?xf32, 1 expands to pseudo code of load element ptr with element ptr = align ptr + offset + stride INDEX sizeof element . In our case, we have sizeof element = 4 because the type is f32 . We have offset=0, stride=1 which means element ptr = align ptr + offset + stride INDEX sizeof element = align ptr + 0 + 1 INDEX 4 ergo element ptr = align ptr INDEX 4 . The layout of our memref type actually specifies offset=0, stride=1 . Since these are compile time constants, the additional arguments are unused but part of our ABI. module { gpu.module @kernels nvvm.target<O = 3, chip = "sm 75", features = "+ptx90" { llvm.func @add %arg0: llvm.ptr<1 , %arg1: llvm.ptr<1 , %arg2: i64, %arg3: i64, %arg4: i64, %arg5: llvm.ptr<1 , %arg6: llvm.ptr<1 , %arg7: i64, %arg8: i64, %arg9: i64, %arg10: llvm.ptr<1 , %arg11: llvm.ptr<1 , %arg12: i64, %arg13: i64, %arg14: i64, %arg15: i32 attributes {gpu.kernel, nvvm.kernel} { %0 = nvvm.read.ptx.sreg.tid.x : i32 %1 = llvm.sext %0 : i32 to i64 %2 = nvvm.read.ptx.sreg.ctaid.x : i32 %3 = llvm.sext %2 : i32 to i64 %4 = nvvm.read.ptx.sreg.ntid.x : i32 %5 = llvm.sext %4 : i32 to i64 %6 = llvm.trunc %1 : i64 to i32 %7 = llvm.trunc %3 : i64 to i32 %8 = llvm.trunc %5 : i64 to i32 %9 = llvm.mul %7, %8 : i32 %10 = llvm.add %9, %6 : i32 %11 = llvm.sext %10 : i32 to i64 %12 = llvm.icmp "slt" %10, %arg15 : i32 llvm.cond br %12, ^bb1, ^bb2 ^bb1: // pred: ^bb0 %13 = llvm.getelementptr inbounds|nuw %arg1 %11 : llvm.ptr<1 , i64 - llvm.ptr<1 , f32 %14 = llvm.load %13 : llvm.ptr<1 - f32 %15 = llvm.getelementptr inbounds|nuw %arg6 %11 : llvm.ptr<1 , i64 - llvm.ptr<1 , f32 %16 = llvm.load %15 : llvm.ptr<1 - f32 %17 = llvm.fadd %14, %16 : f32 %18 = llvm.getelementptr inbounds|nuw %arg11 %11 : llvm.ptr<1 , i64 - llvm.ptr<1 , f32 llvm.store %17, %18 : f32, llvm.ptr<1 llvm.br ^bb2 ^bb2: // 2 preds: ^bb0, ^bb1 llvm.return } } } Phase 3: MLIR to LLVM IR This is what mlir sys::mlirTranslateModuleToLLVMIRToString gives us. We are back in LLVM territory at this point. ; ModuleID = 'LLVMDialectModule' source filename = "LLVMDialectModule" define ptx kernel void @add ptr addrspace 1 %0, ptr addrspace 1 %1, i64 %2, i64 %3, i64 %4, ptr addrspace 1 %5, ptr addrspace 1 %6, i64 %7, i64 %8, i64 %9, ptr addrspace 1 %10, ptr addrspace 1 %11, i64 %12, i64 %13, i64 %14, i32 %15 { %17 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x %18 = sext i32 %17 to i64 %19 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x %20 = sext i32 %19 to i64 %21 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x %22 = sext i32 %21 to i64 %23 = trunc i64 %18 to i32 %24 = trunc i64 %20 to i32 %25 = trunc i64 %22 to i32 %26 = mul i32 %24, %25 %27 = add i32 %26, %23 %28 = sext i32 %27 to i64 %29 = icmp slt i32 %27, %15 br i1 %29, label %30, label %37 30: ; preds = %16 %31 = getelementptr inbounds nuw float, ptr addrspace 1 %1, i64 %28 %32 = load float, ptr addrspace 1 %31, align 4 %33 = getelementptr inbounds nuw float, ptr addrspace 1 %6, i64 %28 %34 = load float, ptr addrspace 1 %33, align 4 %35 = fadd float %32, %34 %36 = getelementptr inbounds nuw float, ptr addrspace 1 %11, i64 %28 store float %35, ptr addrspace 1 %36, align 4 br label %37 37: ; preds = %30, %16 ret void } ; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory none declare noundef range i32 0, 1024 i32 @llvm.nvvm.read.ptx.sreg.tid.x 0 ; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory none declare noundef range i32 0, 2147483647 i32 @llvm.nvvm.read.ptx.sreg.ctaid.x 0 ; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory none declare noundef range i32 1, 1025 i32 @llvm.nvvm.read.ptx.sreg.ntid.x 0 attributes 0 = { nocallback nofree nosync nounwind speculatable willreturn memory none } llvm.module.flags = { 0} 0 = {i32 2, "Debug Info Version", i32 3} Phase 4: The generated PTX This looks very close to the PTX code we generated manually. Success It is the LLVM IR lowered with the NVPTX backend. // // Generated by LLVM NVPTX Back-End // .version 9.0 .target sm 75 .address size 64 // .globl add .visible .entry add .param .u64 .ptr .global .align 1 add param 0, .param .u64 .ptr .global .align 1 add param 1, .param .u64 add param 2, .param .u64 add param 3, .param .u64 add param 4, .param .u64 .ptr .global .align 1 add param 5, .param .u64 .ptr .global .align 1 add param 6, .param .u64 add param 7, .param .u64 add param 8, .param .u64 add param 9, .param .u64 .ptr .global .align 1 add param 10, .param .u64 .ptr .global .align 1 add param 11, .param .u64 add param 12, .param .u64 add param 13, .param .u64 add param 14, .param .u32 add param 15 { .reg .pred %p<2 ; .reg .b32 %r<9 ; .reg .b64 %rd<8 ; mov.u32 %r1, %tid.x; mov.u32 %r2, %ctaid.x; mov.u32 %r3, %ntid.x; mad.lo.s32 %r4, %r2, %r3, %r1; ld.param.b32 %r5, add param 15 ; setp.ge.s32 %p1, %r4, %r5; @%p1 bra $L BB0 2; ld.param.b64 %rd4, add param 1 ; ld.param.b64 %rd5, add param 6 ; mul.wide.s32 %rd6, %r4, 4; add.s64 %rd1, %rd4, %rd6; add.s64 %rd2, %rd5, %rd6; ld.param.b64 %rd7, add param 11 ; add.s64 %rd3, %rd7, %rd6; ld.global.b32 %r6, %rd1 ; ld.global.b32 %r7, %rd2 ; add.rn.f32 %r8, %r6, %r7; st.global.b32 %rd3 , %r8; $L BB0 2: ret; } Running the PTX is as simple as copy-and-paste of the cust example from above now with the added parameters. Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ LLVM IR ⟶ NVPTX Backend ⟶ PTX ⟶ GPU I really like going from right-to-left here instead of the other way around. We get to see and validate things early. Getting everything working with MLIR and LLVM using Rust on Windows was challenging . But now, we really have to start with the source code. Introducing the Phobos Language What got me started on this project was the blog post “How to Optimize a CUDA Matmul Kernel” GEMM stands for General Matrix Multiplication \ C = \alpha A B + \beta C\ . SGEMM stands for Single-Precision General Matrix Multiplication. All of the types are 32-bit floats. Simon gives this example for a naive SGEMM implementation: js global void sgemm naive int M, int N, int K, float alpha, const float A, const float B, float beta, float C { // compute position in C that this thread is responsible for const uint x = blockIdx.x blockDim.x + threadIdx.x; const uint y = blockIdx.y blockDim.y + threadIdx.y; // if condition is necessary for when M or N aren't multiples of 32. if x < M && y < N { float tmp = 0.0; for int i = 0; i < K; ++i { tmp += A x K + i B i N + y ; } // C = α A@B +β C C x N + y = alpha tmp + beta C x N + y ; } } No worries. I am not going to recreate the entire blog post here. It’s excellent. Go read it. What matters is this: I want to enable some of the optimizations he is going through in my compiler. At least to some degree. I could of course write a little CUDA clone with a different syntax and call it a day. We have the pipeline set up and it would work. But this form of code is hostile to the optimizations we want to perform. I always thought about deep learning basically as map-reduce, but with matrices. I thought it would make sense to treat computations at the tile level. Tensor Cores are \ M \times N \times K\ processing units anyways. If you think of an \ M \times N \times K\ GEMM, then many of the optimizations are about computing tiny \ M' \times K'\ and \ K' \times N'\ matrices where \ M'\ , \ N'\ and \ K'\ are numbers with \ M' \le M \land N' \le N \land K' \le K\ . Choosing the right values depends on the actual hardware specification. Hence, the same optimizations that apply locally keeping tensor cores as busy as possible scale to the datacenter as well keeping nodes as saturated as possible . This is not my idea. There exists the Triton language Side-quest: Comparing Triton and Phobos Triton walks the AST of Python code, uses custom MLIR dialects and is able to emit highly optimized code. Vector add in Triton 7 is shown as python @triton.jit def add kernel x ptr, Pointer to first input vector. y ptr, Pointer to second input vector. output ptr, Pointer to output vector. n elements, Size of the vector. BLOCK SIZE: tl.constexpr, Number of elements each program should process. NOTE: constexpr so it can be used as a shape value. : There are multiple 'programs' processing different data. We identify which program we are here: pid = tl.program id axis=0 We use a 1D launch grid so axis is 0. This program will process inputs that are offset from the initial data. For instance, if you had a vector of length 256 and block size of 64, the programs would each access the elements 0:64, 64:128, 128:192, 192:256 . Note that offsets is a list of pointers: block start = pid BLOCK SIZE offsets = block start + tl.arange 0, BLOCK SIZE Create a mask to guard memory operations against out-of-bounds accesses. mask = offsets < n elements Load x and y from DRAM, masking out any extra elements in case the input is not a multiple of the block size. x = tl.load x ptr + offsets, mask=mask y = tl.load y ptr + offsets, mask=mask output = x + y Write x + y back to DRAM. tl.store output ptr + offsets, output, mask=mask The Phobos version for vector addition looks mostly the same if you ignore the comments: @autotune BLOCK in 1024 kernel add a: tensor<f32 N , b: tensor<f32 N , c: tensor<f32 N { let base = program id 0 BLOCK c base :+ BLOCK = a base :+ BLOCK + b base :+ BLOCK } Phobos is really not Triton in the sense that it is much simpler, less feature rich and will break in ways you would not expect. For example, tensors must be a multiple of four, the block size must be divisible by 32 etc. It’s a learning exercise, not a production-grade language and there’s no point in replicating an existing language 1:1. What Triton does not do by itself is clustering. GPU Performance We all know GPUs are fast, thanks to their parallel nature. But it’s like with an F1 car: you need a skilled driver to win the race. You have to be able to use the full potential of the hardware. If you spend thousands of GPUs on training, you want to use the resources as efficiently as possible. Let’s start with a GPU primer and understand the concepts of memory levels and execution in a cooperative thread array CTA on Streaming Multiprocessors SM . The CTA is a group of threads 8 in which threads can synchronize. The GPU schedules a CTA to an SM, and will not move to a different SM . SMs have on-chip shared memory. 9 fn:9 Within a CTA a subset of threads, called a warp, will execute the same instruction classically in lockstep SIMT . 10 fn:10 Now, we must understand the different memory regions: Global memory: This is where we upload the weights. Reading global memory is slow at roughly 300-600 cycles. Physically it’s DRAM. GDDR on consumer cards like mine, or HBM High-Bandwidth Memory on datacenter parts. Shared memory : This is memory on an SM per CTA and can be shared between threads. Much faster than global memory but still requires ~20-30 cycles. Physically, this is backed by SRAM Static Random-Access Memory . Registers : Registers are physically allocated to individual threads and are the absolute fastest memory ~1 cycle .They hold the exact data fragments that feed directly into the execution units, such as ALUs or Tensor Cores. 11 fn:11 If you want to compute \ C = \alpha \times A \times B + \beta \times C\ , you have to multiply every row of \ A\ with every column of \ B\ . The issue is that both of these live in global memory, but that’s slow. In fact, a naive kernel will be memory-bound because the execution units can compute math operations much faster than they can fetch the required data from global memory. What we want is to keep the GPU as busy as possible, which is achieved by making sure we can feed the ALUs and Tensor Cores constantly. But that’s not so simple. Here are the four primary techniques I am using in Phobos: Coalescing : If the 32 threads in a warp read memory in a random order, that results in up to 32 separate memory fetches. The GPU understands when the access is contiguous and can coalesce this into the minimum number of required transactions. Tiling: We can reduce fetches from global memory by prefetching the data into a shared memory tile per CTA. Warp Tiling: The shared memory tile is partitioned again among warps in micro-tiles to maximize register reuse. The memory access patterns within each warp must be structured so that we avoid bank conflicts, where multiple threads in a warp attempt to access the same memory bank simultaneously. 12 fn:12 Tensor Cores : Special hardware to compute \ M \times N \times K\ ; several times faster than normal SIMT instructions. In order to implement these techniques, a compiler would have to understand the blockIdx.x blockDim.x + threadIdx.x arithmetic and derive what each thread computes. Furthermore, one needs to prove the affine access to the data pointers x i + y , whether the pointers alias A == B and ultimately understand which loads are reused across threads. That’s extremely difficult. 13 fn:13 If this is however modeled by your language implicitly, you don’t have to reverse engineer anything. @autotune TILE M in 32, 128 , TILE N in 64, 256 , TILE K in 16, 64 @launch 256 @pipeline @tensorcore kernel sgemm A: tensor<f32 M, K , B: tensor<f32 K, N , C: tensor<f32 M, N , alpha: f32, beta: f32 { let pm = program id 0 let pn = program id 1 var acc: tile<f32 TILE M, TILE N = 0.0 for kt in range 0, K, TILE K { var a = A pm TILE M :+ TILE M, kt :+ TILE K var b = B kt :+ TILE K, pn TILE N :+ TILE N acc += dot a, b } let c old = C pm TILE M :+ TILE M, pn TILE N :+ TILE N C pm TILE M :+ TILE M, pn TILE N :+ TILE N = alpha acc + beta c old } - With @autotune , we can say that we want TILE M to be a constant in the range 32, 128 . - The first-class tensor<f32 M, K specifies we have a two-dimensional tensor with dimensions M and K of unknown size; I query these at runtime. - The thread-id arithmetic is abstracted away by a program id intrinsic that just states how we move in the x and y directions. tile<f32 TILE M, TILE N declares a CTA-wide tile. The codegen places it in shared memory, in registers, or split across both. for kt in range 0, K, TILE K advances kt from 0 to K by TILE K . A pm TILE M :+ TILE M, kt :+ TILE K selects the sub-tile of A with x: pm TILE M, y: kt, width: TILE M, height: TILE K . acc += dot a, b is the dot product of two tiles, accumulated into a tile. C pm TILE M :+ TILE M, pn TILE N :+ TILE N = acc assigns a complete tile. We do not compute individual values using threads. In fact, the whole concept of threads is abstracted away. Instead we specify arithmetic over 2D-tiles and the code generator will assign the computations to the threads. How tiles are read or written — to the metal — is up to the codegen. Furthermore, TILE M , TILE N and TILE K become extremely interesting optimization variables. For example, if we set TILE M=16 , TILE N=16 , and TILE K=16 , the codegen can map this tile<f32 16, 16 structure directly to native Tensor Core MMA instructions on my 2080 SUPER 14. Language I am allergic to semicolons so we are using Go-style semicolon insertion. A startOffset : endOffset is a range where endOffset is exclusive, A startOffset :+ numItems is the same as A startOffset : startOffset + numItems . Otherwise, there’s nothing really special here. It’s a small language. We have @attributes and that’s about it. The compiler does not have a phase model, compilation units or linking. It is a simple Lexer, feeding a parser, feeding the code generator. The code generator is the most complex piece as it performs analysis on the AST to choose the proper output. In a more advanced compiler one could use MLIR dialects and custom passes, or introduce another intermediate representation. Again, I want to keep things simple. Source code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ LLVM IR ⟶ NVPTX Backend ⟶ PTX ⟶ GPU Runtime This is not the whole picture. In fact you’ve seen the @autotune annotation. Kernels should be tuned for their specific use case. We have several parameters we can modify. For example, the tiling size that defines the shared memory region is something we can configure, or the padding for warp tiling. The idea is to specify these parameters and compile a different version of the kernel for the parameter set. If you specify @autotune M in 16, 64 I will generate versions of the kernel for M = 16 , M = 32 and M = 64 . This is intertwined with the code generator and a configuration that can be passed in. In order to find the best configuration we build the Cartesian product of all config values, compile and benchmark all of them, and then pick our winner. I do this in a brute-force fashion. I try all configurations once, select the best three and then do a longer eval on them. Many configurations do not compile on my GPU anyways. A TILE M=256 , TILE N=256 accumulator is 256 KB of live data, and it has to be materialized somewhere. Both options are bounded: staged in registers, it blows past the 255-registers-per-thread limit; staged in shared memory, it exceeds even a Blackwell datacenter GPU’s 228 KB per SM 15. With all of that out of the way, I can say that I am pretty happy with the performance. The more I worked on this project and the more optimizations I added, the larger the search space grew. I am not exploring it beyond the current @autotune parameters. For example, register pressure is another major axis as I learned: larger tiles or more aggressive unrolling can use more registers per thread, which may allow better local reuse, but it can also reduce the number of CTAs/warps resident on an SM. That occupancy tradeoff can have a profound impact on performance. Clustering Remember how I said I always thought of deep learning as map-reduce? What we currently achieve with first-class tiles is basically: Input at M x N ⟶ HBM at M x N ⟶ SRAM at M/m0 x N/n0 ⟶ Warp ⟶ Tensor Core ⟶ Warp ⟶ SRAM ⟶ HBM ⟶ Output at M x N The Tensor Core is our atom, the smallest possible unit. We can add another step and say the data is \ M \times N\ only for the purposes of the GPU, but it actually is much larger. The GPU receives just another sub tile of a much larger datacenter scale supertile. Previously, it was about keeping the tensor cores busy. And we dance around the latency of the SRAM and the HBM to achieve peak performance. Now it is about keeping HBM busy, and we have to dance around the latency of the disk and/or the network instead. The same source is compiled twice. The entire language is scale-free. The additional @cluster TILE M in ... , ... attribute just declares one more tiling level above @autotune ’s device tiles. It reuses the same symbol names. The source code compiles down via two different paths: AST ⟶ device codegen ⟶ MLIR ⟶ PTX : the per-node kernel AST ⟶ cluster codegen ⟶ cluster IR : a small dataflow program over supertiles @cluster TILE M in 4096, 16384 , TILE N in 4096, 16384 , TILE K in 4096, 16384 @autotune TILE M in 32, 256 , TILE N in 32, 256 , TILE K in 4, 32 kernel sgemm A: tensor<f32 M, K , B: tensor<f32 K, N , C: tensor<f32 M, N , alpha: f32, beta: f32 { // ... } A supertile is a large piece of a tensor, say 8192x8192 f32 256 MiB , and it’s the only thing that ever crosses the network if it has to . The kernel stays the same, just interpreted at different scale: program id now indexes the supertile grid, and a slice x y:z now refers to a potentially remote supertile. And I really mean the kernel stays the same. Because @cluster binds the same TILE M , TILE N , TILE K symbols as @autotune , the per-node kernel the cluster compiler distributes is almost the original body. The only thing it touches is the epilogue, or the final store, because each node now contributes one K -chunk of a much larger supertile instead of the whole result. 16 fn:16 For a pure reduction C .. = acc that’s literally one line: the = becomes += , and a tiny synthesized init leaf zeroes the supertile first. The full GEMM epilogue alpha acc + beta c old works in a similar way. The beta c old term can’t live inside a += body. It would be re-applied on every K -chunk. Therefore the compiler hoists it out of the loop. The init leaf seeds the supertile with beta C orig once, and the per-chunk store keeps the alpha acc + c old computation with an implicit beta of 1. Summed over the chunks it results in \ \beta C {\text{orig}}\ + \alpha \sum{A k \times B k} = \alpha \text{acc} + \beta C {\text{orig}}\ . Same math, and no accumulator ever crosses the network. That’s the split shown by the conceptual IR below. The body that does the actual math is interpreted twice: once at device scale CTAs on my GPU , once at cluster scale supertiles across nodes . All from the same AST. In order to achieve high throughput, one can compute the FLOP/byte and scale the supertiles accordingly for higher arithmetic intensity. This can also be tuned, the same way auto-tune works. The DAG Compiler The cluster codegen emits a parametric cluster-IR, because the problem shape M, K, N only arrives at submission time and doesn’t know the node topology. The scheduler then instantiates it into a concrete DAG of four primitive instructions: ALLOC , FETCH , COMPUTE , FREE plus LOAD / STORE at the storage boundary . What we then have to do is: - Classify each tile slice as supertile-aligned the cluster skeleton or intra-tile the leaf body; or the actual kernel . The alignment alone is used to split this information. - Loop-carried accumulators are promoted to read-modify-write RMW supertile parameters, and SSA over supertile versions \ t k \to t {k+1}\ makes each instruction’s read/write set exact. And that set is the DAG; nothing else carries dependency information. This is good old compiler def-use chain analysis. - Because the DAG makes every supertile’s last consumer explicit, lifetimes can be inferred by a backward pass that places the FREE s. There is no ownership or lifetime syntax in the language. It all falls out of the graph for free. This reads as: for each output supertile, loop over K and accumulate; then write // this is a conceptual idea of the cluster IR. // it's never materialized like this as it's directly consumed // by the scheduler to emit LOAD, FETCH, ..., ASSIGN // instructions to the actual nodes in the cluster cluster sgemm { params: M, K, N // bound at submission time super: SUPER M, SUPER N, SUPER K // bound by cluster autotuner tensors: A: f32 M, K in B: f32 K, N in C: f32 M, N out scalars: alpha: f32 beta: f32 grid pm, pn in M/SUPER M, N/SUPER N : t c old = ref C pm, pn // seed the accumulator with beta C t acc@0 = compute k init beta, t c old loop kt in 0 .. K/SUPER K: t a = ref A pm, kt t b = ref B kt, pn // accumulate alpha A k B k into the running C tile. t acc@kt+1 = compute k body rmw t acc@kt, t a, t b, alpha define C pm, pn = t acc@last } We can generate the scheduler’s plan for a hypothetical 4-node cluster computing Flash Attention 17. Every node LOAD s its own Q straight from storage 18 and writes its output back; the K and V tiles have been assigned to the first node and others FETCH them via DCN. These are four independent local programs. Each node owns four query blocks, so it loads K and V once and the COMPUTE s from the same buffers. Their FREE is deferred until the second block is done with them. That lifetime manifests via the backward pass over the graph. We insert a FREE at the tile’s last use. And because the second block’s 19 fn:19 LOAD s do not depend on the first block’s COMPUTE , the node prefetches them. That’s basically the same overlap I reached for with @pipeline at the device level, just one level up.Submission is done via a simple job description containing inputs, outputs and kernel parameters: source = examples/flash attention.ph dim Nq = 16384 dim Nk = 16384 dim D = 64 tensor Q = read f32 16384x64 file:///tmp/q.tensor tensor K = read f32 16384x64 file:///tmp/k.tensor tensor V = read f32 16384x64 file:///tmp/v.tensor tensor O = write f32 16384x64 file:///tmp/o.tensor scalar scale = f32 0.125 Where Does This Fit In? Some properties have emerged from the design choices: - There is no SYNCHRONIZE or barrier instruction. Dependencies are the synchronization mechanism. We have therefore an out-of-order task engine. - The scheduler is not in the critical path. A node’s dependencies are always on the same node; every cross-node edge is in tile space. A consumer’s FETCH blocks server-side on a node until the tile is available. That wait is the synchronization, without a scheduler round-trip. A producer’s FREE waits on a statically known serve-count. The scheduler only drains batched completion events; its load should be \ O \text{segments} + \text{failures} \ , not \ O \text{instructions} \ by construction. - Tile identifiers are deterministic and if a node dies, re-running the lost subgraph reproduces the same tiles lineage model . 20 fn:20 - When possible, accumulators stay on the same node and never travel across the network. This sits somewhere between the Single Program, Multiple Data SPMD and Asynchronous Many-Task AMT world. It lands mostly in the task-based distributed linear-algebra lineage, like PaRSEC, Legion, and most directly NVIDIA’s Legate. There is a tile-DAG, a central scheduler, and dumb nodes. It is not the SPMD+collectives model that the deep-learning stack converged on 21 fn:21 22 fn:22 23. Because a Phobos kernel is whole-program, with static shapes and affine subscripts, the compiler sees the entire dataflow and can build the DAG ahead of time and perform the hoisting trick . However, this design assumes static, dense DAGs and would need rethinking for data-dependent shapes or MoE-style routing. But for dense tensor algebra it provides global placement, inferred lifetimes, exact communication accounting, and lineage-based recovery — all supposedly painful under SPMD . 24 fn:24 The same Phobos source that runs on a single GPU could now run theoretically in a cluster of GPUs. Like with SPMD, the exact same code runs on every GPU simultaneously — though out-of-order — but given the dependencies, of course all nodes reach the same logical synchronization points. The dependencies, let’s say for a forward pass ⟶ reduce, are implicit and handled automatically. Though, they are not micro-managed by the cluster runtime. Nodes themselves will FETCH from each other when they need the data and FREE the data when it is known to be unused last dependency satisfied . So I basically have the implicit dependency management and out-of-order execution of AMT, without the scheduler in the critical path, but also the simplicity of SPMD, with the same code running everywhere as we just focus on the tiles. The tile size can also be autotuned given some heuristics and parameters like network speed against FLOP/byte. The Cluster Prototype I can run the scheduler and nodes on my single machine and constrain each node to use only a subset of the VRAM and make them all talk via gRPC cough Pathways cough . php SGEMM N=20480 M=N=K , super=10240 - 2x2 grid = 4 supertiles, 4 nodes on disk: 4.7 GiB of f32 A+B+C plan: 72 instrs, peak resident 1.2 GiB/node - arena 1.4 GiB/node the shared GPU must hold 4 x 1.4 GiB = 5.6 GiB at once operands replicate per node seeded A in 1.3s, B in 1.3s 0.000s INFO sched: node0 registered tile server 127.0.0.1:61121 0.000s INFO node0: attached to scheduler 127.0.0.1:61119, tile server 127.0.0.1:61121 0.006s INFO sched: node1 registered tile server 127.0.0.1:61123 0.006s INFO node1: attached to scheduler 127.0.0.1:61119, tile server 127.0.0.1:61123 0.011s INFO sched: node2 registered tile server 127.0.0.1:61125 0.011s INFO node2: attached to scheduler 127.0.0.1:61119, tile server 127.0.0.1:61125 0.017s INFO sched: node3 registered tile server 127.0.0.1:61127 0.017s INFO node3: attached to scheduler 127.0.0.1:61119, tile server 127.0.0.1:61127 0.173s INFO sched: planned 'matmul' supers= "TILE K", 10240 , "TILE M", 10240 , "TILE N", 10240 ingest=DirectLoad budget=1440MiB nodes=4 segments/node= 1, 1, 1, 1 peak=1200MiB fetch=0MiB 0.254s INFO sched: compiled 2 leaf kernel s 0.255s INFO sched: waiting for 4 node s to register 0.256s INFO sched: 4 node s registered: "node0=127.0.0.1:61121", "node1=127.0.0.1:61123", "node2=127.0.0.1:61125", "node3=127.0.0.1:61127" 0.257s INFO sched: - node0 issue segment 0 18 instrs: 5 ALLOC, 4 LOAD, 3 COMPUTE, 1 STORE, 5 FREE; incr 1200MiB 0.259s INFO sched: - node1 issue segment 1 18 instrs: 5 ALLOC, 4 LOAD, 3 COMPUTE, 1 STORE, 5 FREE; incr 1200MiB 0.260s INFO node0: received segment 0 18 instrs 0.260s INFO sched: - node2 issue segment 2 18 instrs: 5 ALLOC, 4 LOAD, 3 COMPUTE, 1 STORE, 5 FREE; incr 1200MiB 0.261s INFO node1: received segment 1 18 instrs 0.261s INFO node0: issue 1 ALLOC tile=0x20000000000000 shape= 10240, 10240 0.262s INFO sched: - node3 issue segment 3 18 instrs: 5 ALLOC, 4 LOAD, 3 COMPUTE, 1 STORE, 5 FREE; incr 1200MiB 0.262s INFO node2: received segment 2 18 instrs 0.263s INFO node1: issue 19 ALLOC tile=0x20000000000001 shape= 10240, 10240 0.266s INFO node1: issue 21 ALLOC tile=0x0 shape= 10240, 10240 ... 12.173s INFO node0: D2H tile=0x20000000000000 400.0 MiB read back from VRAM, writing to file://C:\Users\joaeb\AppData\Local\Temp\phobos cluster correctness\C.bin 12.173s INFO node2: issue 51 FREE tile=0x3 serves=0 12.175s INFO node2: issue 52 FREE tile=0x10000000000002 serves=0 12.176s INFO node2: issue 53 STORE tile=0x20000000000002 12.247s INFO node2: D2H tile=0x20000000000002 400.0 MiB read back from VRAM, writing to file://C:\Users\joaeb\AppData\Local\Temp\phobos cluster correctness\C.bin 12.447s INFO node0: STORE 17 wrote 400.0 MiB to file://C:\Users\joaeb\AppData\Local\Temp\phobos cluster correctness\C.bin 12.449s INFO node0: issue 18 FREE tile=0x20000000000000 serves=0 12.492s INFO node2: STORE 53 wrote 400.0 MiB to file://C:\Users\joaeb\AppData\Local\Temp\phobos cluster correctness\C.bin 12.494s INFO node2: issue 54 FREE tile=0x20000000000002 serves=0 12.495s INFO sched: all 72 instructions accounted 72 retired, 0 abandoned dispatched + ran in 12.3s check: 12 sampled entries, distributed vs CPU reference match max rel err 2.2e-6 Trust me bro: We validate 12 samples for a \ 20480 \times 20480\ matrix. I am somewhat in a bit of a weird position here. I like the idea of a potentially horizontally scaled scheduler that enables cooperative progress and recovery, yet is not micro managing the nodes and therefore less of a bottleneck. But it’s not something I can verify at scale. The idea is really “I am going to give you a blueprint of what you should do, plus here are your teammates that will do some part of the work; go talk to them if you need something and let me know when you’re done” versus “I am going to tell you exactly what you have to do for each step along the way; do let me know when you need something and I will coordinate with your teammate”. This may sound familiar when it comes to work. I hope for you that you’re in the first camp. Central scheduling is certainly not trivial as shown by Pathways 22. It is a huge effort and it would have been really interesting to learn more about PLAQUE , Google’s dataflow language. But that is closed-source at the time of writing. I also learned about DISTALwhen doing research for this blog post. It is almost exactly the same design and closest prior art. And with SpDISTAL 25 fn:25 there’s actually hope that a design like this could even work for sparse graphs / MoE-routing. DISTAL runs on Legion and the programming model is certainly different but the conceptual idea is similar. 26 fn:26 Are we cuBLAS yet? autotune: long probe of the top 3 TILE M=64 TILE N=128 TILE K=8: 16.69ms +/- 213.10µs TILE M=32 TILE N=256 TILE K=8: 17.48ms +/- 272.96µs TILE M=128 TILE N=64 TILE K=8: 18.43ms +/- 289.87µs autotune winner: TILE M=64 TILE N=128 TILE K=8 phobos sgemm fp32: 16.69ms +/- 183.19µs cuBLAS sgemm : 12.66ms +/- 181.79µs check: 16777216 elements, both correct phobos: 8235.9 GFLOP/s, cuBLAS: 10856.7 GFLOP/s phobos / cuBLAS: 1.32x / 75.86% Phobos is at the time of writing at ~76% of cublasSgemm v2 throughput for fp32 SGEMM at M=N=K=4096 . The 2080 SUPER I have can achieve 11 TFLOP/s on paper for fp32 27. Side-quest: Cheating autotune: long probe of the top 3 TILE M=128 TILE N=128 TILE K=16: 18.83ms +/- 12.92ms TILE M=64 TILE N=256 TILE K=16: 7.46ms +/- 316.64µs TILE M=128 TILE N=64 TILE K=16: 8.49ms +/- 127.78µs autotune winner: TILE M=64 TILE N=256 TILE K=16 phobos sgemm fp16tc fp32acc: 7.34ms +/- 71.74µs cuBLAS sgemm : 12.78ms +/- 143.52µs check: 16777216 elements, both correct phobos: 18723.8 GFLOP/s, cuBLAS: 10754.0 GFLOP/s phobos / cuBLAS: 0.57x / 174.11% How is this possible? I am using tensor cores, but my 2080 does only fp16 . I cast the data to fp16 before loading it into the tensor cores, perform the computation, load it and cast it back to fp32 . The accumulator stays fp32 but it’s of course a loss in precision. I am using Tensor Cores, do math with less precision, while cublasSgemm v2 is not using tensor cores and does everything in 32-bit. So, I include this here only for fun. I am not stating that I did something that’s faster than cuBLAS. It’s an apples-to-oranges comparison. cublasHgemm , the fp16 variant, is at 70 TFLOP/s : My pure fp16 GEMM hits 62% of real cuBLAS fp16 . As difficult as it is: I’ll have to stop here. At this point, we enter the territory of diminishing returns. There are still ideas to try but it’s not worth it from my perspective. I did burn some tokens on adding ping-pong buffering and pipelining for example, which gave me a modest performance gain, but this is snowballing in complexity. Furthermore, I read a lot about features of modern GPUs that are simply not supported by my ancient 2080 and I feel like I might chase the wrong ideas. Little aside:I wasnotable to stop and added more~~features~~complexity. Mainly XOR swizzling registers and moving to different tensor core instructions which was a nightmare. I’ve been working on this for months now and this blog post was somewhat of my development diary. The fp16 performance is where I spent a lot of time lately but once again I learned that a lot of the cool features are only available on Ampere+ GPUs. But now I am really going to stop. I can run the scheduler and node daemon on my machine, submit jobs and watch them run until completion. So “it’s working”. I just don’t have the GPUs to validate it at scale. It was a lot of fun for me to get my hands dirty again with compiler development, explore low-level GPU concepts and get a fresh look at LLVM with MLIR. I nerd-sniped myself and spent several nights again in front of the computer with typical “ah shit, it’s already 1AM” moments. I found it hard to navigate all of this and find an entry for a topic that is usually buried under massive corporate frameworks. I achieved my goal. It’s the same with training 28. You have to get the hours in. There’s no way around it. You just have to do the work. You can find the source code of Phobos at github.com/joa/phobos https://github.com/joa/phobos-lang . You can run ptxas add.ptx -o add.cubin and cuobjdump --dump-sass add.cubin to have a look at the actual machine code. ↩︎ fnref:4 Now may be relative for you; the last time I used LLVM seriously was around 2016 and MLIR was made public in 2019 Source https://en.wikipedia.org/wiki/MLIR software . ↩︎ fnref:5 This is more helpful than you might think. I did not want to rebuild inlining and use the LLVM optimization pass. But at that point, my high-level nodes were long gone. I could only hope that LLVM would be able to figure out certain optimizations, which would have been far easier to implement if the high-level nodes were simply inlined. ↩︎ fnref:6 Instructions like shfl sync allow you to actually share values across thread registers within a warp; so do not consider registers as “per thread” necessarily ↩︎ fnref:11 Shared memory is divided into equally sized banks typically 32 banks, each 4 bytes wide . They can be accessed simultaneously. However, if multiple threads within the same warp request different memory addresses that map to the same bank, the hardware cannot serve them in parallel. The requests are serialized. ↩︎ fnref:12 Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions https://arxiv.org/abs/1802.04730 and the PLUTO https://pluto-compiler.sourceforge.net/ language went down this route. The term is polyhedral compilation https://en.wikipedia.org/wiki/Polytope model . ↩︎ fnref:13 Note that there is no f32 tensor core instruction on the 2080 SUPER so I need to convert to f16 first, then light up the tensor cores and convert back to f32. That’s of course a lossy conversion. ↩︎ fnref:14 We now sum the K-chunks in a different order than the single-pass kernel, so floating-point non-associativity makes the result epsilon-correct against the single-GPU reference, not bit-identical. And that’s also interesting. Any LLM at temperature=0 is going to be non-deterministic if accumulators cross the network, since reduction order varies with parallelism. But, most kernels such as FlashAttention are already non-deterministic for high throughput. See also DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training — Qiang et al., 2026 https://arxiv.org/html/2601.21824v1 ↩︎ fnref:16 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness https://arxiv.org/abs/2205.14135 ↩︎ fnref:17 The IO load still exists; every node still needs to query storage / cache. ↩︎ fnref:18 If you are wondering how this is supposed to work in failure scenarios: the FREE will get retracted by the scheduler and a new, corrected FREE node, will be issued. ↩︎ fnref:19 Spark’s Resilient Distributed Datasets popularized lineage-based recovery. You re-derive lost partitions from their deterministic computation graph rather than from checkpoints: Resilient Distributed Datasets Zaharia et al., NSDI 2012 https://www.usenix.org/system/files/conference/nsdi12/nsdi12-final138.pdf . ↩︎ fnref:20 Very hard for me to verify. Accelerating large-scale distributed neural network training with SPMD parallelism https://dl.acm.org/doi/10.1145/3542929.3563487 suggests computation stalls due to SPMD data dependencies for example. ↩︎ fnref:24 Table 2. GeForce RTX 3080 vs GeForce RTX 2080 / 2080 Super; P.14 https://www.nvidia.com/content/PDF/nvidia-ampere-ga-102-gpu-architecture-whitepaper-v2.1.pdf ↩︎ fnref:27 ↩︎ fnref1:27 Physical : ↩︎ fnref:28