Porting An Existing Kernel#

Note

This chapter is part of the agent-targeted Agent BLAS Performance Guide and is also published as Markdown for agent consumption, in the markdown/ folder of the HTML documentation tree.

The goal of a port is the first implementation that is correct, uses the intended instruction path and measures at parity with the reference, reached in as few build-and-run cycles as possible.

The deliverable is a correct port at parity, and remaining wall-clock budget is not a reason to continue. Once the port is correct, on the intended instruction path and inside the parity band, deliver the result. The variant budget in step 6 is a total for the whole run, not an allowance per axis or per idea.

Follow these phases in order:

Identify constraints and tuning axes.
Set up a release build.
Replace the contraction in place.
Build one comparison harness.
Check geometry and instructions.
Measure against a fixed parity band.
Validate and deliver.

1. Identify Constraints And Tuning Axes#

Read the target kernel and its build script. Record the output format, epilogue contract, required precision, thread interface and scheduling constraints.
Distinguish the workload’s matrix dimensions from the per-block tile dimensions in Size<M, N, K>. A fixed workload does not imply a fixed tile.
If the workload dimensions are not multiples of the candidate tiles, a pipeline will refuse them: see Shapes That Do Not Divide The Tile for padding, masking and remainder descriptors before fixing the tile.
Tile M, N and K are usually sweep axes. Also consider block dimension and pipeline depth wherever the surrounding kernel permits them to vary.
An output layout or resource budget can restrict the legal candidates without fixing every tile dimension. Record each actual restriction and leave the remaining dimensions free.
Mark an axis as fixed only when the application requires that value; do not treat the reference’s existing tile or thread count as a requirement by default.

2. Set Up A Release Build#

Follow the Build Card and confirm in the emitted compile line that -O3 -DNDEBUG, the architecture-specific target and the library include path are present and that none of the card’s traps is.
A build that is not a release build produces measurements that mean nothing.

3. Replace The Contraction In Place#

Do not write standalone prototypes or probe programs to learn the API; use the card for your path (Hopper, Blackwell or supermma).
Keep the reference’s structure — its schedule, output format and host-side caching — and replace only the contraction and the data movement it needs.
Do not change the reference’s scheduling or specialization: no free-running or work-stealing tile schedulers, no space-filling (Hilbert, Morton) tile orders, no new warp-role assignments beyond those the pipeline API itself imposes, no second pipeline.
A port that keeps the reference’s structure can be measured against it; one that redesigns it reaches a different kernel, measured against nothing.
Before you overwrite the reference source, copy it aside; you will need it, byte for byte, as the comparison arm, and re-creating it from memory costs a step and an unverifiable file.

4. Build One Comparison Harness#

Build and state#

One script with modes (build, check, bench); both arms compiled into the same binary or extension, with the reference kernel moved into its own namespace so the two coexist; configuration exposed as macros so an experiment is one flag, not a new script.
Give every test case its own handles and state (a crash that appears only when a second case runs in the same process is object lifetime in the harness, not the kernel).

Correctness#

When the output format is lossy or capacity-limited, validate against the reference kernel’s own output rather than a dense closed form, and choose test-operand statistics that the reference’s epilogue can absorb: operands the reference itself cannot represent fail both arms.

Timing#

Give the bench mode one small representative shape that runs in seconds; it is the gauge for every iteration and for the band check in step 6.
Run the full-size sweep once, after the final build, to confirm the band on the representative shapes and report any gap it shows.

5. Check Geometry And Instructions#

Launch geometry#

A persistent grid must be sized from the measured device at runtime — cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, dev); an H100 PCIe has 114 SMs, an H100 SXM 132, a B200 148 — never from a constant inherited from the reference or the guide.
Then check wave quantization: tiles that do not divide into whole waves serialize a partial wave.
A uniform gap at large problem sizes is this before it is anything in the finetuning chapter.
Record the measured count and the wave arithmetic in the notes: launch geometry is the one gap explanation that needs no control, so do not build a copy of the reference re-gridded to the device to test it, and do not carry such a copy as a second control arm through the tables.

Generated instructions#

Then confirm the instruction path once, per kernel rather than per binary — step 4 put the reference into the same binary, and a reference that uses wgmma itself would satisfy a whole-binary count:

cuobjdump -sass <binary> | awk '/Function :/{f=$3} /[HIQ]GMMA/{n[f]++} END{for (f in n) print n[f], f}'

on Hopper (HGMMA for 16-bit and tf32 inputs, IGMMA for int8 / uint8, QGMMA for fp8), with /UTC[A-Z]*MMA/ in place of /[HIQ]GMMA/ on Blackwell (the integer form is UTCIMMA); the port’s own kernel must show a non-zero count (a kernel with no match is not listed).

Finding a PyTorch extension binary#

A PyTorch extension’s .so is in <build_directory> when load was given build_directory=; otherwise under $TORCH_EXTENSIONS_DIR/<name>/ when that variable is set, and otherwise under py<ver>_cu<ver>/<name>/ in the root torch.utils.cpp_extension.get_default_build_root() returns.
Take the file from the returned module’s __file__, not from the name: a second load of the same name with changed sources or flags in one process writes <name>_v1.so next to the stale <name>.so.
With is_python_module=False, load returns the path itself.

6. Measure Against A Fixed Parity Band#

Define the parity band#

Before the first timing, write down the parity band: within measurement noise of the reference on the representative shapes — take the noise from two identical binaries timed against each other, typically two to three percent.
When the port is inside the band on the small shape, deliver.

Manage the variant budget#

If it is not, allow at most three further rebuilt-and-measured variants in total, across every axis (tile M/N/K, block dimension, depth, grid, unroll pragma, epilogue), then deliver the best correct checkpoint and report the remaining gap honestly; keep the count as a column in a small ledger and never restart it by changing chapter, parameter or idea.
A diagnostic variant — an epilogue removed, a phase of the kernel switched off — counts toward the three and bounds what is possible, not what is deliverable: do not set its number as a target.
The one falsification control that step 7 may ask for is a build of the reference, not a variant, and does not count.

Choose the next measurement#

Sweep the axes identified as free in step 1. Tile M/N/K, block dimension and pipeline depth can interact; compare a small set of coupled candidates within the variant budget.
Do not re-scan axes after a change that cannot couple to them.
Profile only after step 5 is done and a gap above the parity band remains. Focus on the relevant speedup estimates, waves and register metrics.

7. Validate And Deliver#

Re-run the full correctness suite only when the compiled binary changed (compare an md5sum of the .so, or the SASS); host-side cleanups, comment edits and notes do not need re-measurement.
Run the repository’s own integration path once, on the final build; host-only edits after it that change nothing the kernel consumes or the caller receives (launch geometry, pointers, shapes and strides, tile counts, shared-memory size, caching, the returned tensors) need only the SASS comparison, not a repeat.
Run a falsification control, if the task needs one, once, on one shape, after the kernel is final.
Record the final configuration, correctness results, timing measurements and any remaining gap.
Keep build and profiler logs alongside the results so the comparison can be reproduced.