COS568 Assignment 1: Performance Profiling due TBD

Introduction

You cannot make a system faster until you know what limits it. In this assignment you will take one training step of a large language model apart on a single GPU, find out where its time and its memory go, and learn the tools for finding that out: GPU timers, the GPU's own counters, and NVIDIA's profiler, Nsight Systems.

A GPU offers a program three resources:

compute: arithmetic operations per second (TFLOP/s);
memory bandwidth: bytes per second moved between memory and the cores (GB/s);
memory capacity: bytes that fit at once (GiB).

Throughout this assignment the question is the same: which resource is the limit here?

The model is OLMoE-1B-7B, a mixture-of-experts (MoE) language model from the Allen Institute for AI. It has 6.9 billion parameters, but any one token uses only 1.3 billion of them. That gap shapes almost everything about how the model performs, and you will meet it in all three assignments of this course:

Assignment	You will	GPUs
1	Measure where time and memory go in a training step	1
2	Add data and expert parallelism, fine-tune the whole model, and compete on throughput	1–8
3	Serve your fine-tuned model with vLLM and SGLang, and study inference configurations	1–8

You will not write any model code. The file olmoe/model.py is a complete and reasonably efficient OLMoE. Your job is to build a few small measuring tools, to measure, and to explain what you measure.

The assignment has four parts:

One operation. You learn to time GPU code correctly, to bound how fast an operation can run with the roofline model, and to read the GPU's own counters.
One forward pass. You examine a forward pass of the real model with Nsight Systems.
One training step. You take a step apart: its forward pass, its backward pass and its optimizer, the memory they need, and two ways of trading time for memory, activation recomputation and micro batches.
Predicting performance. You write a small model that predicts the time and memory of training steps you have never run.

Parts 1 and 2 are mostly a guided tour, in which you run things and study what they print. Parts 3 and 4 ask more of you. Expect to spend 10 to 20 hours in all.

The text is a walkthrough, and what you have to do is set apart from it in numbered boxes of three kinds. Each kind is numbered on its own, and is graded in its own way:

Coding Exercise (write code · autograded)

You write a function. make test checks it. There are 9 of them.

Experiment (run it and hand in the results · checked automatically)

You run scripts that we provide and hand in the files they write. make test checks that the files are complete and plausible. There are 7.

Written Question (answer in a1/answers.md · graded by hand)

You explain what you measured, in a few sentences of your own. Every answer must cite numbers from your own runs. There are 7, and they are the only part that a person grades.

There are also optional Challenges at the end. They earn no points: the best solutions are presented in class.

Software Setup

You will work on the neuronic cluster: 32 nodes, each with 8 NVIDIA L40 GPUs. The machine you log in to is shared by everyone, and it is only for editing files and submitting jobs. Everything else runs on a GPU node, through the Slurm scheduler. The Makefile takes care of that for you.

The files you will need for this and the later assignments are distributed using the Git version control system. Clone the course repository on the cluster:

$ git clone https://github.com/als244/cos568_spr27.git
$ cd cos568_spr27

Your clone is your own. Commit to it as often as you like; nobody else can see it. The course repository is public, so do not publish your solutions: no public forks, and no pushing your work anywhere others can read it. Corrections to this assignment, and Assignments 2 and 3, will arrive as new commits in the course repository. To get them, commit your work and run git pull --no-rebase --no-edit, which merges them into it.

There is nothing to install, and nothing to source when you log in. Every make target loads a shared, read-only environment (Python 3.12 and PyTorch 2.13) and knows where the OLMoE checkpoint is. Do not pip install anything, and do not download the model: your home directory holds 16 GB, and the checkpoint alone is 14 GB.

To check that everything is in place, run make check. It submits itself to a GPU node, so it can take a moment to start. You should see something like this:

$ make check
COS568 setup check
  [ok] running inside a Slurm job
  [ok] Python 3.12.16
  [ok] PyTorch 2.13.0+cu130
  [ok] CUDA is available
  [ok] GPU: NVIDIA L40, 44.4 GiB
  [ok] NVIDIA Nsight Systems version 2026.2.1.210-262137639646v0
  [ok] OLMoE checkpoint: /n/fs/cos568/shared/weights/OLMoE-1B-7B-0924 (13.8 GB)
  [ok] your course directory: /n/fs/cos568/users/netid
All good.

If a line says FAIL, read the hint under it, and tell the course staff if it is not something you can fix.

Getting Started

Take a look at what you have:

olmoe/            the model and training code, shared by all three assignments
  model.py        OLMoE: attention, router, experts
  data.py         micro batches: flat token tensors, and where the sequences are
  train.py        the training step
  prof.py         timing, NVTX ranges, GPU counters
  load.py         load and save HuggingFace-format checkpoints
a1/
  perfmodel.py    your performance model
  answers.md      your written answers
  *.py, nsys.sh   the experiments: you run them, and need not edit them
grade/            the public tests

In this assignment you edit five files: olmoe/prof.py, olmoe/train.py, olmoe/model.py (one line), a1/perfmodel.py and a1/answers.md. Only those are graded, so do not put code that your solution needs anywhere else.

Read olmoe/data.py and olmoe/model.py now, top to bottom. Together they are about 280 lines, and everything later assumes you know them. One convention deserves a word, because it differs from most code you have seen: nothing here has a [batch, seqlen] shape. The sequences that go through the model in one pass are a micro batch, and a micro batch is all of its sequences concatenated. Tokens are [T] and activations are [T, D], where T is the total number of tokens in the micro batch and D = 2048 is the hidden size. A small Seqs object says where each sequence starts and ends. For a micro batch of three sequences with 3, 5 and 2 tokens:

tokens      [44, 39, 33,   63, 79, 27,  3, 97,    1, 66]     all three sequences: T = 10
targets     [39, 33, 60,   79, 27,  3, 97, 83,   66, 56]     for each token, the token that follows it
cu_seqlens  [0, 3, 8, 10]
max_seqlen  5

Sequence i is rows cu_seqlens[i] to cu_seqlens[i + 1]. Sequences in one micro batch may have different lengths, and nothing is ever padded. Attention uses FlashAttention's variable-length kernel (torch.nn.attention.varlen.varlen_attn), which takes cu_seqlens directly. The MoE layers do not care about sequences at all: routing is per token. The experiment scripts take --micro-batch-size B --seqlen S, which means B sequences of S tokens each, so T = B * S.

Now run the tests:

$ make test
...
==================================== score =====================================
  Setup         (no points; these must pass before you start)
  Coding 1      FAIL   0/13
  Coding 2      FAIL   0/4
  Coding 3      FAIL   0/11
  Coding 4      FAIL   0/8
  Coding 5      FAIL   0/12
  Coding 6      FAIL   0/8
  Coding 7      FAIL   0/20
  Coding 8      FAIL   0/4
  Coding 9      FAIL   0/20
  Experiment 2  FAIL   0/5
  Experiment 3  FAIL   0/5
  Experiment 4  FAIL   0/5
  Experiment 5  FAIL   0/5
  Experiment 6  FAIL   0/5
  Experiment 7  FAIL   0/5
  Total                0/130
28 failed, 8 passed in 21.85s

The 8 tests that pass are the Setup tests: they compare our model with HuggingFace's reference implementation, on the real checkpoint and on sequences of different lengths. Everything else fails, because you have not done it yet: the coding exercises with NotImplementedError, and the checks of your experiments because you have not run any. The tests for all the code you will write are already there. make test C=1 runs those of Coding Exercise 1 alone, and make test E=2 runs the checks of Experiment 2.

For experiments you need a shell on a GPU node. make gpu-shell gives you one, with the environment loaded, for at most two hours:

$ make gpu-shell
(cos568) $ python a1/hello.py "The capital of France is"
loaded 6.92B parameters in 3 s; 12.9 GiB on the GPU
The capital of France is Paris.

The capital of the United States is Washington, D.C.

The
(cos568) $ exit

Exit when you are done, so that someone else can use the GPU. Anything that takes more than a few minutes is better as a batch job, which needs no open terminal:

$ sbatch slurm/job.sbatch python a1/sweep.py
Submitted batch job 4221588
$ squeue -u $USER                  # is it running yet?
$ cat slurm-4221588.out

make test, make check and make gpu-shell use the cluster's interactive partition, which starts jobs quickly and is meant for short things. The cluster is busy, and GPUs are scarce before deadlines: start early.

Experiment 1 (run it · not graded)

Run make check and make test, and confirm that your setup is complete and that the Setup tests pass. Then open a GPU shell and run a1/hello.py with a prompt of your own.

Hand-In Procedure

You hand in three things, one for each kind of box: your code, the results of your experiments, and your written answers.

Coding Exercise	File	You implement	Points
1	`olmoe/prof.py`	`cuda_time`	13
2	`a1/perfmodel.py`	`gemm_flops`, `gemm_bytes`, `roofline_tflops`	4
3	`a1/perfmodel.py`	`param_count`, `active_params`, `flops_per_token`, `rows_per_expert`, `forward_floor_ms`	11
4	`olmoe/prof.py`	`annotate`	8
5	`olmoe/model.py`	the recomputation branch of `Olmoe.forward`	12
6	`olmoe/train.py`	`accumulated_step`	8
7	`olmoe/train.py`	`stage_times`	20
8	`a1/perfmodel.py`	`state_bytes`, `optimizer_floor_ms`	4
9	`a1/perfmodel.py`	`predict_step_ms`, `predict_peak_gib`	20
		public tests, and again for hidden tests	100 + 100

Experiment	You run	Points
1	`make check`, `make test`, `a1/hello.py`	none
2	`a1/bench_gemm.py`, `a1/regimes.py`, `a1/bench_ops.py`	5
3	three recordings of a forward pass	5
4	`a1/micro_batches.py`	5
5	`a1/stages.py`, without and with recomputation	5
6	three recordings of a training step	5
7	`a1/sweep.py`, over configurations of your choice	5
		30

Written Question	About	Points
1	one operation: the three regimes	10
2	one forward pass, small and large	10
3	micro batches and the global batch	10
4	where a training step spends its time	10
5	what the profiler shows of a step	10
6	a slow backward pass	10
7	shapes, and your predictor	10
		70

That is 300 points in all. make test reports 130 of them: the public tests of your code and the checks of your experiments. The hidden tests check the same properties as the public ones, with other shapes and seeds, so code that passes the public tests honestly will pass them too.

You answer the written questions in the file a1/answers.md. Put your name there, your answers, and a sentence on how you used AI tools. A few sentences and a small table per answer are enough: we look for a correct explanation, not for length.

When you are ready to hand in, or whenever you want to be safe, run make handin:

$ make handin
Handed in 92 files (4.9 MB) as /n/fs/cos568/users/netid/handin/a1/20261111-174959.tar.gz
You can hand in again at any time; the last hand-in before the deadline is graded.

It checks that the results of your experiments are there and that every written question has an answer, and copies your work to your course directory, where the staff collect it. You can hand in as often as you like. Your last hand-in before the deadline is the one we grade.

This assignment is done alone. You may use AI tools for anything in it, and three things remain your responsibility. The numbers must come from your own runs on the cluster. The files you hand in must agree with your answers. And you must be able to explain anything you hand in if we ask. If you let a coding agent run commands on the cluster, it must follow AGENTS.md; above all, it must never run GPU work on the login node.

Part 1: One Operation

Before profiling a model, you need to be able to measure one operation and to say what limits it. The operation in this part is the one that language models spend nearly all of their arithmetic on: a matrix multiply.

Timing GPU Code

When Python calls torch.matmul(a, b) on GPU tensors, the call returns almost at once. It has only enqueued a kernel, and the GPU runs the kernel later. Timing the call with a clock on the CPU therefore measures the time to enqueue, not the time to multiply. Here are three ways of timing the same multiply of two 8192 × 8192 matrices.

The first reads the CPU's clock before and after the call:

t0 = time.time()
c = torch.matmul(a, b)
t1 = time.time()

The second waits for the GPU to finish before it reads the clock again:

t0 = time.time()
c = torch.matmul(a, b)
torch.cuda.synchronize()
t1 = time.time()

The third does not use the CPU's clock at all. It puts a CUDA event before the multiply and another after it, and asks for the time between the two:

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
c = torch.matmul(a, b)
end.record()
torch.cuda.synchronize()
ms = start.elapsed_time(end)

On an L40 the three report:

1: no synchronize       0.043 ms
2: with synchronize     9.954 ms
3: CUDA events          9.647 ms

The first number is wrong by a factor of several hundred: it is the time Python took to hand the work to the GPU. The second is right, because torch.cuda.synchronize() blocks until the GPU has finished everything it was given. The third is right too, and it is the one we will use. A CUDA event is a marker placed in the GPU's queue. The GPU stamps it with the time at which it reaches it, so the time between two events is the time the GPU spent on the work between them, and nobody had to wait in the middle. That matters once you want to time the parts of something larger without disturbing it, as you will in Part 3.

Two more things can spoil a measurement. The first calls of anything on a GPU are slow, because libraries are loaded and memory is allocated, so throw some calls away before you measure. And, as you are about to find out, a GPU that has just been idle runs faster than one that has been busy for a second.

Coding Exercise 1 (write code · autograded)

In the file olmoe/prof.py, you must implement the function listed below:

cuda_time

It runs a function a few times without timing it, then a few times with, and returns the GPU time of each timed call. Every later experiment uses it. Read its docstring carefully for what exactly it must do.

You can run make test C=1 to test your solution:

$ make test C=1
..                                                                       [100%]
==================================== score =====================================
  Coding 1      OK    13/13
  Total               13/13
2 passed, 34 deselected in 1.93s

Make sure your code passes both tests before you go on.

The Roofline

Every kernel does arithmetic, and it moves bytes between memory and the cores. An L40 can do at most 181 TFLOP/s of bf16 matrix arithmetic, and it can move at most 864 GB/s. So a kernel that performs F floating-point operations and must move B bytes takes at least F / 181e12 seconds, and also at least B / 864e9 seconds. Whichever is larger is its bound:

if the first is larger, the kernel is compute-bound;
if the second is larger, it is memory-bound.

The roofline puts both bounds on one plot. A kernel's arithmetic intensity is F / B, in FLOPs per byte, and the most TFLOP/s it can reach is

min(peak compute, intensity × peak bandwidth)

The two meet at the ridge point, 181e12 / 864e9 ≈ 210 FLOPs per byte. Below it a kernel is memory-bound, and above it compute-bound.

Take the matrix multiply C = torch.matmul(A, B) in bf16, two bytes per number, where A is [M, k], B is [k, n], and C is [M, n]. It performs M * n * k multiply-adds, which is 2 * M * n * k FLOPs. It must at least read A and B and write C: 2 * (M * k + k * n + M * n) bytes. As an example, take one of the attention projections of our model, a 2048 × 2048 weight matrix, applied to M = 16 tokens:

FLOPs      = 2 * 16 * 2048 * 2048                        = 134 million
bytes      = 2 * (16 * 2048 + 2048 * 2048 + 16 * 2048)   = 8.5 million
intensity  = 134e6 / 8.5e6                               = 15.8 FLOPs per byte

That is far below the ridge point, so this multiply is memory-bound: its 8 megabytes of weights have to be read however few tokens there are. The best it can do is 8.5e6 / 864e9 = 10 microseconds, at 15.8 × 864 GB/s = 13.6 TFLOP/s, a thirteenth of the GPU's peak. Give the same matrix 4096 tokens and the intensity is 819: now it is compute-bound, and could run at the full 181 TFLOP/s.

Notice what happened. M is a number of tokens and B is a weight matrix. With few tokens a layer is limited by the bytes of its weights, and with many it is limited by its arithmetic. This one observation will explain a great deal of what follows.

Coding Exercise 2 (write code · autograded)

In the file a1/perfmodel.py, you must implement all the functions listed below:

gemm_flops

gemm_bytes

roofline_tflops

Measuring Matrix Multiplies

Now measure. In a GPU shell, run

(cos568) $ python a1/bench_gemm.py

For each GEMM shape that appears in OLMoE, it sweeps M from 1 to 32768 and prints one line per measurement:

expert gate_proj       M=    64     0.0266 ms    10.08 TFLOP/s    172.3 GB/s  intensity    58.5
expert gate_proj       M=  1024     0.0348 ms   123.36 TFLOP/s    301.2 GB/s  intensity   409.6
...
lm_head                M=    64     0.3103 ms    42.50 TFLOP/s    685.7 GB/s  intensity    62.0
lm_head                M=  4096     7.6170 ms   110.80 TFLOP/s     83.4 GB/s  intensity  1329.3

The TFLOP/s and GB/s columns divide the FLOPs and bytes of your functions from Coding Exercise 2 by the measured time. The script writes the table to results/a1/gemm.csv, and plots it over the roofline in results/a1/roofline.png. Copy the plot to your laptop with scp and look at it for a while. Some shapes follow the rising part of the roofline at a respectful distance and then flatten out. Others stay far below it until M is large.

At the end, the script runs one large GEMM back to back for six seconds, starting from a GPU that has been idle, and writes what happens to results/a1/gemm_sustained.csv. Keep it for Written Question 1.

Three Regimes

There is a third regime, which the roofline does not show. Launching a kernel has a fixed cost of some microseconds, on the CPU and on the GPU. When the work is small enough, that cost is all there is: the kernel is overhead-bound, and neither resource is anywhere near its limit.

You can tell the three regimes apart without a model, because the GPU counts what it does. a1/regimes.py runs one workload of each kind for a few seconds, and samples the GPU's own counters while it runs:

(cos568) $ python a1/regimes.py
large GEMM, 8192^3            10.8910 ms  kernel running 100%  memory busy  16%  1411 MHz  280 W
add two 1 GiB tensors          4.8808 ms  kernel running  97%  memory busy  97%  2490 MHz  207 W
small GEMM, 64 x 2048 x 1024   0.0153 ms  kernel running  51%  memory busy   0%  2490 MHz  141 W

The first number on each line is the time of one call. From this output you can conclude a few things:

The large GEMM is compute-bound. A kernel is running all of the time, but memory is being read or written only 16% of the time: the arithmetic works out of the caches.
Adding two large tensors is memory-bound. There is almost no arithmetic to do, and memory is busy 97% of the time.
The small GEMM is overhead-bound. Half of the time no kernel is running at all, although the loop launches them as fast as Python can. The GPU is waiting for the CPU.

The counters come from NVML, the library behind nvidia-smi, through the GpuMonitor class in olmoe/prof.py. They average over up to a second, so they are only meaningful for a workload that runs steadily for a few seconds. Look at the last two columns as well, the clock frequency of the GPU's cores and its power draw. They are not the same on every line.

Finally, run the same kind of experiment for operations with almost no arithmetic:

(cos568) $ python a1/bench_ops.py

It times a copy, an addition, silu(x) * y and a normalization on tensors from a tenth of a megabyte to several gigabytes, and writes results/a1/ops.csv and results/a1/ops.png: achieved bandwidth against the number of bytes moved.

Experiment 2 (run it and hand in the results · checked automatically)

Run a1/bench_gemm.py, a1/regimes.py and a1/bench_ops.py in a GPU shell, and study the two plots and the three tables.

Written Question 1 (answer in a1/answers.md · graded by hand)

(a) At M = 64, the lm_head GEMM and an expert's gate_proj GEMM have almost the same arithmetic intensity, so the roofline gives them the same bound. How close does each come to it, and what limits each one?

(b) A large GEMM settles at a little over half of 181 TFLOP/s, although gemm_sustained.csv starts out much faster. Use that file and the output of regimes.py to explain why. What does that say about a benchmark that runs for a few milliseconds?

(c) In ops.png, some points lie above the 864 GB/s line. For which sizes, and how is that possible?

Part 2: One Forward Pass

Now for the whole model. In this part you will work out what a forward pass should cost, and then use a profiler to see what it does cost.

The Model in Numbers

Each of the 16 layers has an attention block, with four 2048 × 2048 projections, and an MoE block with 64 experts. An expert is three matrices: gate_proj and up_proj, from 2048 to 1024 dimensions, and down_proj, from 1024 back to 2048. A router sends every token to its 8 best experts and adds up their outputs.

Part of the model	Parameters
One layer: attention, 4 × 2048 × 2048	16.8 million
One layer: 64 experts, each 3 × 2048 × 1024	402.7 million
One layer: router and norms	0.1 million
16 layers	6.71 billion
Embedding and `lm_head`, each 50304 × 2048	0.21 billion
Whole model	6.92 billion

So 96% of a layer is experts, but a token is multiplied by only an eighth of them. That is the gap between the "7B" and the "1B" in the model's name. (Your active_params will come to 1.18 billion: the weights that a token is multiplied by. The 1.3 billion of the introduction also counts the embedding table, of which a token reads one row.)

Coding Exercise 3 (write code · autograded)

In the file a1/perfmodel.py, you must implement all the functions listed below:

param_count

active_params

flops_per_token

rows_per_expert

forward_floor_ms

The docstrings give the conventions, and the tests check them exactly.

forward_floor_ms returns two lower bounds on the time of a forward pass: its FLOPs at peak compute, and reading every weight once at peak bandwidth. This is Part 1's roofline, applied to the whole model. Evaluate it for one sequence of 512 tokens, and for 32 of them. You should find that the first pass is bounded by memory and the second by compute. Write both pairs of numbers down. You are about to compare them with reality.

Nsight Systems

Nsight Systems (nsys) records every kernel the GPU runs and every CUDA call the CPU makes, on one timeline. To find your own code in that timeline, you mark regions of it with NVTX ranges: torch.cuda.nvtx.range_push("name") opens one and range_pop() closes it. The model already marks two regions in each MoE layer, router and experts, and train_step marks forward, backward and optimizer.

One thing about ranges is easy to get wrong. A range marks a stretch of CPU time, the time during which your code launched kernels. The kernels themselves run later. Nsight can project a range onto the GPU, which means taking the span of the kernels that were launched from inside it. That is the number you usually want.

Coding Exercise 4 (write code · autograded)

In the file olmoe/prof.py, you must implement the function listed below:

annotate

It gives every submodule of a model an NVTX range named by its path, such as layers.3.attn, using forward hooks, so you do not have to edit model.py. A range must be closed even when the module's forward raises an exception: look at the always_call argument of register_forward_hook. You will need that in Part 3, where PyTorch sometimes ends a layer's forward early by raising one.

Make sure your code passes the tests of this exercise (make test C=4).

The script a1/nsys.sh records a run of a1/profile_step.py and summarizes the recording. It is short, and at its heart are two commands:

nsys profile --trace=cuda,nvtx --sample=none --cpuctxsw=none \
    --gpu-metrics-devices=all \
    --capture-range=cudaProfilerApi --capture-range-end=stop \
    --output=reports/a1/NAME python a1/profile_step.py ...

nsys stats --report nvtx_gpu_proj_trace --report cuda_gpu_kern_sum --report cuda_api_sum \
    --format csv --output results/a1/NAME reports/a1/NAME.nsys-rep

--trace=cuda,nvtx records CUDA calls and kernels, and NVTX ranges.
--sample=none --cpuctxsw=none turns off sampling of the CPU, which we do not need.
--gpu-metrics-devices=all samples the GPU's hardware counters, 10,000 times a second. More about them below.
--capture-range=cudaProfilerApi records only between torch.cuda.profiler.start() and stop(). profile_step.py calls them around two steps, after loading the model and warming up, so the recording is small and shows the steady state.
nsys stats turns a recording into tables. nvtx_gpu_proj_trace is the projection of your ranges onto the GPU that was just described.

Try the script on the first four layers and eight short sequences:

(cos568) $ a1/nsys.sh demo --layers 4 --micro-batch-size 8 --seqlen 256
note: this node does not let you read the GPU's performance counters; recording without GPU metrics.
...
one step: 23.9 ms before the recording started, 29.5 ms while recording
peak memory allocated: 3.73 GiB
...
NVIDIA L40: 142 SMs, 2490 MHz, memory bandwidth 864 GB/s, L2 cache 96 MiB
range                              GPU ms/step instances  GPU ops
step                                     29.35         1   1589.5
forward                                  29.35         1   1589.5
layers.*                                 26.32         4   1561.5
layers.*.moe                             22.34         4   1425.5
experts                                  21.61         4   1397.5
layers.*.attn                             3.21         4     96.0
lm_head                                   2.54         1      1.0
router                                    0.51         4     28.0
CCCL:cub::DeviceRadixSort                 0.44         4     80.0
layers.*.attn_norm                        0.20         4     16.0
layers.*.moe_norm                         0.17         4     16.0
CCCL:cub::DeviceScan::InclusiveSca        0.06         5     10.0
norm                                      0.04         1      4.0
embed                                     0.02         1      1.0
(no GPU metrics in this recording)
wrote results/a1/demo_ranges.csv
wrote reports/a1/demo.nsys-rep and results/a1/demo_*.csv

The line that starts with one step comes from profile_step.py itself, which times a pass before the recording starts and again while it is running. The line above the table describes the GPU as its driver reports it; you have met most of these numbers already. The table is a summary of your NVTX ranges, projected onto the GPU and added up over the layers:

GPU ms/step is the span of the kernels launched inside ranges of that name. Ranges nest, so the rows overlap: layers.* contains layers.*.moe, which contains experts.
instances is how often a range of that name was opened in a step: once per layer for anything inside a layer.
GPU ops is the number of kernels and memory operations launched inside.
The ranges whose names start with CCCL are not yours. NVIDIA's own libraries mark some of their internals in the same way.

Now look at the first line of that transcript, and at the line under the table. Nsight can also record the GPU's hardware counters, the GPU metrics, but reading them needs a permission that the administrators of a machine must grant. When this was written, the cluster's nodes did not have it, and the script recorded without them. On a node that has it, the note does not appear, and the table has four more columns: the means of four counters over the time when the kernels of each range were running.

sm active is the share of that time in which the GPU's cores, its 142 streaming multiprocessors, had work to do. It is low when the GPU is waiting for kernels to be launched.
tensor active is the share in which the tensor cores were busy, the units inside each SM that do bf16 matrix arithmetic. It tells you how much of the GPU's compute a range is using.
dram read and dram write are the shares of the memory bandwidth that were in use for reading and for writing.

These are finer versions of the counters that regimes.py printed in Part 1: they are sampled every 100 microseconds, not once a second, so they can be attributed to a single layer, and they tell the tensor cores apart from the rest of the GPU. Without them nothing else changes, and the two written questions that use the GPU metrics say what to use in their place.

Every recording leaves these files:

reports/a1/NAME.nsys-rep                   the recording
results/a1/NAME_ranges.csv                 the table above
results/a1/NAME_nvtx_gpu_proj_trace.csv    every single NVTX range, projected onto the GPU
results/a1/NAME_cuda_gpu_kern_sum.csv      GPU time of every kind of kernel
results/a1/NAME_cuda_api_sum.csv           CPU time of every kind of CUDA call

In the kernel summary, each row is one kind of kernel, with the number of times it ran (Instances), its average duration, and its share of the GPU time. The names are long C++ names. Look for words in them: gemm is a matrix multiply, elementwise an operation applied to every element of a tensor, and index a gather or a scatter.

To see a recording as a timeline, install Nsight Systems on your laptop. It is free, and you need version 2026.2 or newer, the version on the GPU nodes. Copy the .nsys-rep file over and open it. Expand the Python process and its threads to find your NVTX ranges and the CUDA calls each thread made, and expand the GPU to find the kernels. Zoom into one layer. The written questions can be answered from the CSV summaries, but the timeline often shows the answer at a glance.

Profiling a Forward Pass

Experiment 3 (run it and hand in the results · checked automatically)

Record three forward passes of the full, pretrained model:
a1/nsys.sh fwd_small       --micro-batch-size 1  --seqlen 512
a1/nsys.sh fwd_large       --micro-batch-size 32 --seqlen 512
a1/nsys.sh fwd_small_naive --micro-batch-size 1  --seqlen 512 --moe-impl naive
The third one uses another implementation of the MoE layer, _experts_naive in model.py, which computes the same thing with a plain loop over the experts.

Written Question 2 (answer in a1/answers.md · graded by hand)

(a) For fwd_small and fwd_large, give the GPU time of one pass, split into attention, experts, router, norms and lm_head.

(b) Compare each pass with its two floors from forward_floor_ms. Which regime of Part 1 is each pass in? For the small one, the kernel summary and the GPU ops column tell you where the time goes. Give a rough cost per kernel.

(c) The small pass is much slower while it is being recorded than before, but the large one hardly differs. Why?

(d) How much slower is the naive MoE at this size? Which CUDA call, in cuda_api_sum, accounts for most of the difference, how often does it happen per layer, and which line of model.py causes it?

(e) Which resource, if any, limits experts in each of the two passes, and which limits lm_head? If your recordings have GPU metrics, base your answer on them: give the SM activity, the tensor-core activity and the memory bandwidth of both ranges in fwd_small and in fwd_large, and say which changes most between the two passes. If they do not, base it on the time per GPU operation of each range (GPU ms/step divided by GPU ops), on the FLOPs of its GEMMs, and on what regimes.py measured for a small and for a large GEMM. Either way, does your answer agree with the regimes you named in (b)?

Part 3: One Training Step

A training step has three phases. The forward pass computes the loss, and keeps the activations that the backward pass will need. The backward pass computes a gradient for every weight. The optimizer, AdamW, updates every weight. See train_step in olmoe/train.py: it is about ten lines.

In this assignment the weights, their gradients, and AdamW's two moments are all stored in bf16, two bytes each. That is 8 bytes per parameter. For all 16 layers it comes to 51.5 GiB, more than the L40's 44 GiB, before a single activation is stored. A model that does not fit on the GPU can still be trained there by offloading: part of that state is kept in the host's memory, and moved to the GPU only while it is needed. Offloading is a valid way to train large models, but we will not consider it in this assignment. Instead, in this part we train only the first four layers (--layers 4, 1.9 billion parameters). The layers are identical in shape, so what you learn about four of them carries over. The experiments in Parts 3 and 4 use random weights and random tokens, so that they need no checkpoint.

Recomputation

The forward pass keeps activations for the backward pass, and their size grows with the number of tokens and with the number of layers. With enough tokens, they are what fills the GPU's memory.

Activation recomputation, also called checkpointing, keeps only the input of each layer. When the backward pass reaches a layer, it first runs that layer's forward again, to rebuild the activations, and then computes the layer's gradients as usual. Only one layer's activations exist at a time, and the price is that every layer's forward runs twice in each step. It trades compute for memory capacity.

A step with recomputation therefore has four kinds of work, and the rest of this part accounts for each of them separately: forward, backward, recompute, and optimizer.

Coding Exercise 5 (write code · autograded)

In the file olmoe/model.py, fill in the recomputation branch of Olmoe.forward, using torch.utils.checkpoint.checkpoint with use_reentrant=False. It is one line.

The tests of this exercise check that the gradients are the same with and without recomputation, and that far fewer activations are kept. Make sure your code passes them (make test C=5).

Micro Batches and Gradient Accumulation

Recomputation is one way to make a step fit in memory. Here is a second one. It rests on a property of the loss.

Two different things are called "the batch" of a step, and from here on we keep them apart. The global batch is the set of sequences that one optimizer update is computed from, and the global batch size is how many sequences that is. A micro batch is the set of sequences that goes through the model in one forward and backward pass, and the micro batch size is how many sequences that is. So far every step has run its global batch as a single micro batch, so the two sizes were equal. They need not be.

The loss of a step is the average, over all tokens of its global batch, of each token's cross-entropy:

loss  =  (l_1 + l_2 + ... + l_N) / N

Here l_i is the cross-entropy of token i, and N is the number of trainable tokens in the global batch, the tokens that have a target. Every one of them weighs 1/N. The gradient of a sum is the sum of the gradients, so the gradient that the optimizer needs is

gradient  =  (gradient of l_1) / N  +  (gradient of l_2) / N  +  ...  +  (gradient of l_N) / N

That is a sum of N terms, and it does not matter in what groups they are computed. So split the global batch into micro batches, each a few whole sequences, and run the forward and the backward pass of one micro batch at a time. Each backward pass produces the terms of its own tokens, and PyTorch adds them to whatever is already in .grad. After the last micro batch, .grad holds the gradient of the whole global batch, and the optimizer steps once. This is gradient accumulation.

The argument needs two things, and you should be able to say why:

A token's loss depends only on the tokens of its own sequence. Nothing in our model connects one sequence to another: attention never crosses a sequence boundary, the router looks at one token at a time, and there is no batch normalization. A micro batch therefore computes l_i for its tokens exactly as the whole global batch would in a single pass.
Every token's loss is divided by N, the number of trainable tokens in the global batch, and not by the number in its own micro batch. This is why lm_loss in olmoe/train.py adds the cross-entropies up (reduction="sum") and divides by a count that you pass in, instead of taking the usual mean. With the mean, a token in a micro batch of few tokens would weigh more than a token in a micro batch of many, and when sequences have different lengths, micro batches of the same size rarely have the same number of tokens.

The update is then the same whether a global batch of 16 sequences runs as one micro batch of 16 or as sixteen micro batches of 1. It is the same up to rounding: floating-point addition gives slightly different results when the same numbers are added in a different order. The two gradients agree to several digits, not bit for bit, and the tests compare them with a small tolerance.

What does change is the cost. Only one micro batch's activations are alive at a time, so the memory that a step needs depends on its micro batch size, while the update depends only on its global batch. In exchange, the model makes several small passes instead of one large one. The optimizer still runs once per step.

One more observation, on which Assignment 2 is built. Nothing in the argument says that the micro batches must run one after another, or on the same GPU. Give each of several GPUs some of the global batch's sequences as its micro batch, let each compute the terms of its own tokens with the same 1/N, and add the results up: the update is the same again. That is data parallelism. Gradient accumulation spreads a global batch over time, data parallelism spreads it over GPUs, and the two can be combined.

Coding Exercise 6 (write code · autograded)

In the file olmoe/train.py, you must implement the function listed below:

accumulated_step

It makes one optimizer step on a global batch that is given as a list of micro batches. train_step, just above it, does the same for a global batch that is a single micro batch. Yours must make the same update on the same sequences, however they are split. Read lm_loss first.

The tests of this exercise compare the update of accumulated_step with the update of train_step on the same sequences, for micro batches of equal and of different sizes. They also check that the optimizer stepped once, and that the step needed less memory. Make sure your code passes them (make test C=6).

a1/micro_batches.py runs one global batch with micro batches of several sizes, without and with recomputation. Here it is on a small case, two layers and a global batch of eight sequences of 1024 tokens:

(cos568) $ python a1/micro_batches.py --layers 2 --global-batch-size 8 --seqlen 1024 \
               --micro-batch-sizes 8 4 1
global batch size 8, 1024 tokens per sequence: 8,192 tokens per step, 2 layers
micro batch size  micro batches  recompute   step ms  tokens/s  peak GiB
               8              1         no     174.4    46,965     14.22
               4              2         no     186.9    43,827     11.99
               1              8         no     316.2    25,910      8.85
               8              1        yes     207.2    39,537     10.65
               4              2        yes     222.3    36,858     10.20
               1              8        yes     400.4    20,460      8.40
wrote results/a1/micro_batches.csv
wrote results/a1/micro_batches.png

Every row is the same step: the same sequences, and the same update. When a row does not fit in the GPU's memory, the script says out of memory and goes on to the next one. PyTorch prints a warning of its own about the failed allocation, which you can ignore.

The script also draws results/a1/micro_batches.png: the memory allocated over the course of one step, for a global batch of half the size, which still fits in memory as a single micro batch. It is run in four ways: as one micro batch or as four, and without or with recomputation. The curves come from PyTorch's own record of every allocation and every free, the record behind PyTorch's memory visualizer, which you will use in Profiling a Step. The two dotted lines mark the training state: the weights with AdamW's two moments, which are allocated all the time, and the same plus one gradient per weight.

Experiment 4 (run it and hand in the results · checked automatically)

Run python a1/micro_batches.py. With no arguments it takes the first four layers and a global batch size of 16, with sequences of 2048 tokens. That is four times the step that the rest of this part uses. It runs the step with micro batch sizes 16, 8, 4, 2 and 1, without and with recomputation, and writes results/a1/micro_batches.csv and results/a1/micro_batches.png. Look at the figure.

Written Question 3 (answer in a1/answers.md · graded by hand)

(a) Explain why the update of a step does not depend on the micro batch size. Which property of the model does the argument need, and what kind of layer would break it? Why does lm_loss divide by the number of tokens in the global batch, and what would go wrong with the mean of each micro batch? Why are the two gradients equal only up to rounding?

(b) In micro_batches.csv, which micro batch sizes fit in memory, and which gives the fastest step? Recomputation and a smaller micro batch both save memory. What does each pay for it, and which is the better deal here? Why does a step get slower as its micro batch size goes down, although it processes the same tokens?

(c) Look at micro_batches.png. For each of its four panels, say where the peak is and what is alive at that moment. What do micro batches take out of the picture, what does recomputation take out, and what is left when both are on? With four micro batches, memory does not return to where the step began until the very end. What stays allocated in between?

(d) Suppose the four micro batches of the row with micro batch size 4 ran at the same time on four GPUs, one each, instead of one after another on one GPU. What would each GPU have to exchange with the others before the optimizer can step, and how many bytes is that for this model? If that exchange cost nothing, about how long would a step take, and what would the throughput be?

Timing the Stages of a Step

How long does each layer take, in each of the four kinds of work? You can find out with CUDA events and two kinds of hooks.

A forward hook is a function that you attach to a module. PyTorch calls it every time the module's forward has finished, and passes it the output. A tensor hook is a function that you attach to a tensor. PyTorch calls it during the backward pass, when the gradient of that tensor has been computed.

Here is a small example that uses both. Its model has two stages, the linear layers first and second, and x goes through one and then the other. Each stage gets a forward hook, and that hook attaches a tensor hook to the stage's output:

first = torch.nn.Linear(8, 8)     # a model with two stages:
second = torch.nn.Linear(8, 8)    #     x -> first -> second -> loss
x = torch.randn(4, 8)

def watch(name, module):
    def forward_hook(module, args, output):   # called when the forward of `module` has finished
        print(f"forward:  {name} has finished")
        output.register_hook(tensor_hook)

    def tensor_hook(grad):                    # called when the gradient of `output` is known
        print(f"backward: about to enter {name}")

    module.register_forward_hook(forward_hook)

watch("first", first)
watch("second", second)

loss = second(first(x)).sum()
print("calling loss.backward()")
loss.backward()

This prints

forward:  first has finished
forward:  second has finished
calling loss.backward()
backward: about to enter second
backward: about to enter first

The forward hooks fire in the order of the stages. The tensor hooks fire in the reverse order, during loss.backward(): the backward pass goes through the stages from last to first, and the gradient of a stage's output becomes known at exactly the moment the backward pass is about to enter that stage. So the hooks give you the boundaries between the stages in both directions, and a CUDA event recorded at each boundary gives you their times.

Recomputation needs one more boundary. A recomputed forward is an ordinary forward of the layer, so its hooks fire again, but in the middle of the backward pass. olmoe/prof.py has a function for telling the two apart: in_backward() is true while a backward pass is running. A forward pre-hook, which PyTorch calls when a module's forward is about to start, and a forward hook then give you the start and the end of the recomputation. Register the forward hook with always_call=True, as in Coding Exercise 4: PyTorch stops a recomputed forward as soon as it has rebuilt what it needs, by raising an exception.

Coding Exercise 7 (write code · autograded)

In the file olmoe/train.py, you must implement the function listed below:

stage_times

For every stage of a training step it returns three GPU times, forward, backward and recompute, and it returns the time of the optimizer. The stages are the embedding, each layer, and the head, which is the final norm, lm_head and the loss. Some advice:

Do not synchronize between stages. That would change what you measure. Record events, and read them after the last step.

Recompute time is not backward time. A layer's backward time is what is left of its share of the backward pass once the recomputation is taken out.

Remove your hooks before you return.

Testing Your Code

How do you test a timer? The tests of this exercise make one layer artificially slow, by 30 ms going forward and 50 ms going backward, and check that your function finds the extra time in that layer and in the right place. With recomputation on, the right answer is 30 ms more forward time, 50 ms more backward time, and 30 ms more recompute time, since the slow forward runs twice. Run the tests with make test C=7, and make sure your code passes all of them.

Reading a Step

Coding Exercise 8 (write code · autograded)

In the file a1/perfmodel.py, you must implement all the functions listed below:

state_bytes

optimizer_floor_ms

a1/stages.py runs your stage_times and prints the result. For two layers and eight sequences of 2048 tokens, with recomputation, you should see something like this:

(cos568) $ python a1/stages.py --layers 2 --micro-batch-size 8 --seqlen 2048 --recompute
stage         forward ms  backward ms  recompute ms
embed               0.19         1.15          0.00
layers.0           34.19        64.81         32.58
layers.1           34.78        67.36         38.03
head               42.24        79.33          0.00
total             111.40       212.66         70.62
optimizer 23.24 ms   (floor 16.94 ms)
step 417.9 ms = 39,204 tokens/s, MFU 31.9%
peak memory 15.45 GiB   (weights + gradients + Adam: 7.79 GiB)
wrote results/a1/stages_rc.csv
counters: a kernel was running 99% of the time, memory was busy 65%, clock 1923 MHz, 298 W

Look at the peak memory. This step needs twice as much as its weights, gradients and optimizer state together, and that is with recomputation: the rest is activations. Without --recompute the same step peaks at 22.6 GiB. The last line shows the same counters as regimes.py in Part 1, over a few seconds of these steps.

Experiment 5 (run it and hand in the results · checked automatically)

Run the configuration that the rest of this part uses, four layers and four sequences of 2048 tokens, without and with recomputation:
python a1/stages.py --layers 4 --micro-batch-size 4 --seqlen 2048
python a1/stages.py --layers 4 --micro-batch-size 4 --seqlen 2048 --recompute
The first run should report a step of roughly 300 ms. The two runs write results/a1/stages.csv and results/a1/stages_rc.csv.

Written Question 4 (answer in a1/answers.md · graded by hand)

(a) What is the ratio of backward to forward time for a layer, and for the head? Explain the ratio. How does the head compare with a layer, and does flops_per_token agree?

(b) Compare the optimizer's time with optimizer_floor_ms. What share of the step is it here? Estimate its share when the step has 512 tokens instead of 8192. Why does an MoE spend a larger share of each step in the optimizer than a dense model with the same FLOPs per token?

(c) How much memory does recomputation save in this configuration, and how much time does it cost? How does a layer's recompute time compare with its forward time, and what happened to its backward time? (Written Question 5 comes back to the first of these.)

Profiling a Step

Nsight Systems shows the same step from the outside, with two additions to what you saw in Part 2. Ranges that are opened during a backward pass belong to a recomputation, so olmoe/prof.py puts recompute. in front of their names: the recomputed forward of layer 3 appears as recompute.layers.3, with recompute.layers.3.attn and the others inside it. And the table gets one row that is not a range: (backward) is what is left of each step outside forward, optimizer and the recomputations.

Experiment 6 (run it and hand in the results · checked automatically)

Record three training steps:
step="--mode train --layers 4 --micro-batch-size 4 --seqlen 2048"
a1/nsys.sh train       $step --memory-snapshot train
a1/nsys.sh train_rc    $step --memory-snapshot train_rc --recompute
a1/nsys.sh train_naive $step --moe-impl naive --steps 1
--memory-snapshot also writes reports/a1/NAME.pickle, a record of every allocation during one more step. Drop that file on https://pytorch.org/memory_viz to see memory over time.

Written Question 5 (answer in a1/answers.md · graded by hand)

(a) In train_ranges.csv and train_rc_ranges.csv, the backward range has almost no GPU time, although the backward pass is most of the step. Where did its kernels go? (Which thread launches them? The timeline shows it.) What does that tell you about using NVTX ranges to time a backward pass?

(b) Compare the recompute. rows of train_rc_ranges.csv with the recompute times that your stage_times reported. Do the two tools agree? A recomputation does not have to rebuild the layer's output, because that is the next layer's input, which was kept. How many GPU operations does each recomputed layer skip, and which are they? Is its recompute time smaller than its forward time by what you would expect?

(c) The two snapshots have the shapes of the left half of your figure from Experiment 4, for a smaller step and in more detail: in the memory visualizer every block is one tensor. Describe the shape of memory over one step in each snapshot. Which tensors are allocated at the peak, and why does the peak move when recomputation is on?

(d) Which of the three resources limits each phase of the step: forward, (backward) and optimizer? If your recordings have GPU metrics, base your answer on those of the three rows in train_ranges.csv: SM activity, tensor-core activity and memory bandwidth. If they do not, base it on the counters that stages.py printed for the whole step, and on how close the optimizer comes to optimizer_floor_ms and the forward pass to forward_floor_ms.

Written Question 6 (answer in a1/answers.md · graded by hand)

The naive MoE is slower going forward, as you saw in Part 2. Going backward it is far slower still. How much? Use the kernel summary of train_naive to find the kernels responsible, and work out from their number and their sizes what they are doing. Which expression in _experts_naive causes them, and why does _experts_grouped not have the problem?

Part 4: Predicting Performance

A performance model is the difference between having measured a system and understanding it. If you can predict the time and memory of a configuration you have never run, you know where the costs come from.

A Sweep of Your Own

a1/sweep.py measures training steps over a grid of configurations. For each one it records the step time, the four kinds of work in it, throughput, MFU (model FLOPs utilization, the share of peak compute spent on the model's arithmetic), peak memory, and the GPU's counters:

(cos568) $ python a1/sweep.py --layers 2 --micro-batch-sizes 1 8 --seqlens 512
 2 layers   1 x 512     47.7 ms  10,726 tok/s  MFU  8.5%   7.8 GiB  kernel  83%  memory  77%  2490 MHz
 2 layers   8 x 512    100.0 ms  40,972 tok/s  MFU 32.5%  10.0 GiB  kernel  98%  memory  76%  2298 MHz
appended to results/a1/sweep.csv

The last three columns are the counters of Part 1: the share of time a kernel was running, the share of time memory was busy, and the clock. Rows accumulate in results/a1/sweep.csv from one run to the next, so you can build up a table over several runs. Delete the file to start over.

Experiment 7 (run it and hand in the results · checked automatically)

Use a1/sweep.py to measure the configurations you think you need: numbers of layers, micro batch sizes, sequence lengths, with and without --recompute. The default arguments give one grid for four layers, which takes about two minutes. That is a start, and not enough. Use a batch job for long sweeps.

The Predictor

Coding Exercise 9 (write code · autograded)

In the file a1/perfmodel.py, you must implement all the functions listed below:

predict_step_ms

predict_peak_gib

They predict the GPU time and the peak memory of train_step, as sweep.py measures them, for a model with cfg.n_layers layers and a step whose global batch is a single micro batch of micro_batch_size sequences of seqlen tokens, with or without recomputation. Both must be plain arithmetic: no GPU, and no files.

Testing Your Predictor

The public tests compare your predictions with measurements by the staff, in grade/a1_truth.csv. The hidden tests use measurements of other configurations: other numbers of layers, up to 8, and other shapes, up to 16,384 tokens per step in sequences of up to 4096. A time prediction earns part of the credit if it is within 30% for four configurations out of five, and all of it if it is within 15%. For memory the two thresholds are 20% and 8%.

Pasting the measured table into your function will not work, because the hidden configurations are different. Build each model from what Parts 1 to 3 taught you about where time and memory go, and use your sweep to find its constants. Run make test C=9 to see how you are doing.

Written Question 7 (answer in a1/answers.md · graded by hand)

(a) In your sweep, how do throughput, MFU and the GPU's counters change as a step gets more tokens? At a fixed number of tokens, how much does the sequence length matter, and why so little? Which resource, if any, limits a step of 256 tokens, and which a step of 16,384?

(b) Describe your two models: their terms, what each term stands for, and where the constants come from. Where is each model least accurate, and what would it take to fix that?

This completes the assignment. Make sure that make test passes everything, coding exercises and experiments, and do not forget your written answers in a1/answers.md. Then run make handin.

One look ahead, before you go. Run python a1/routing.py. It counts how many tokens the router sends to each expert, for a passage of real text. In Assignment 2 the experts will live on different GPUs. Think about what the busiest expert will do to that.

Challenges (optional)

The challenges are not graded and earn no points. They are here for the fun of it, and we will present the best solutions in class. If you take one up, say in a1/answers.md what you tried and what you measured before and after.

Challenge!

Fix the naive MoE's backward pass by changing one or two lines, and show the kernel summary before and after.

Challenge!

At small micro batch sizes a forward pass is nowhere near either of its floors. Get closer. Some ideas: torch.compile; padding each expert's rows to a common length and using one batched matmul; fusing the norm. Report the times, and explain what your change removed.

Challenge!

The head keeps several [T, vocab] tensors alive. Compute the same loss with less memory, and show the memory snapshot before and after.

Challenge!

Build the most accurate predictor in the class. We will rank the predictors of Coding Exercise 9 by their median error on the hidden configurations, and announce the three best.

Challenge!

Predict a machine that you have never run a training step on. Your predictor knows one GPU, the L40. The directory a1/rtx5090/ holds what the experiments of Part 1 wrote on a different machine, a workstation with one NVIDIA GeForce RTX 5090 and 32 GB of GPU memory: regimes.csv, gemm.csv, gemm_sustained.csv and ops.csv, in the same format as your own files in results/a1/. There, regimes.py printed
large GEMM, 8192^3             4.9836 ms  kernel running 100%  memory busy  31%  2655 MHz  513 W
add two 1 GiB tensors          2.0478 ms  kernel running  92%  memory busy  92%  2856 MHz  347 W
small GEMM, 64 x 2048 x 1024   0.0062 ms  kernel running  82%  memory busy   0%  2899 MHz  269 W
and Nsight describes the GPU as 170 SMs, 2550 MHz, memory bandwidth 1792 GB/s, L2 cache 96 MiB.

That is all you get. Add a function predict_step_ms_rtx5090(cfg, micro_batch_size, seqlen, recompute=False) to a1/perfmodel.py. We have measured the configurations of the public and the hidden tests on that machine, and will rank the predictions by their median error. Your L40 predictor is off by almost a factor of two there. Each term of your model is bound by one of the three resources: which one? Two of those configurations do not run on that machine at all. Say which, and why, in a1/answers.md.