kernel_logical_visualizer

⚠️ Beta — tool under testing. Interfaces, parsed output, and the schedule model may change. Bug reports and sample kernels that fail to parse are very welcome.

Visualize the logical async pipeline of a CUTLASS warp-specialized GPU kernel (SM90+): producer/consumer warpgroups, multi-stage shared-memory buffers, TMA fills, WGMMA consumption, mbarrier hand-offs, and epilogue overlap — as an interactive GTKWave-style timeline in your browser.

No GPU, no profiler, no simulator required. The tool statically parses the CUTLASS source, extracts the pipeline structure, and simulates the logical schedule client-side in JavaScript. Timing is in abstract steps, not measured cycles — the point is to understand the shape of the pipeline: fill, steady state, drain, bubbles, and warpgroup overlap.

Try it now

Open examples/sm90_gemm_pingpong_waveform.html in a browser (no server needed). It visualizes the canonical CUTLASS SM90 warp-specialized GEMM mainloop (sm90_mma_tma_gmma_ss_warpspecialized.hpp) with warpgroup roles from the pingpong kernel layer (sm90_gemm_tma_warpspecialized_pingpong.hpp): 1 producer (load) warpgroup + 2 consumer (math) warpgroups.

Drag the sliders and watch the pipeline change. Click any block to follow one k-tile through TMA → buffer FULL → WGMMA → release.

What the timeline looks like

Rendered from the example HTML (SM90 warp-specialized GEMM, 4 stages, 4 output tiles × 5 k-tiles, pingpong with 2 consumer warpgroups):

Reading it: tile o.k (output tile o, k-tile k) is TMA-issued by the producer, fills a stage (blue), sits FULL (pink) until its consumer WG can take it, then is consumed by WGMMA (yellow = even output tile, orange = odd). The pingpong hand-off is visible in the two consumer lanes: WG1's orange mainloop for output tile 1 runs while WG0 is in its green EPI o0 epilogue, and vice versa — the math order barrier passes WGMMA ownership back and forth. Red WAIT blocks are bubbles; the verdict line (top right) totals them.

Legend:

block	meaning
purple o.k (producer lane)	producer_acquire + TMA issue for that tile
blue FILLING (stage lane)	TMA in flight → mbarrier arrive_and_expect_tx
pink FULL (stage lane)	buffer ready, waiting on consumer_wait
yellow/orange o.k	WGMMA consuming (color alternates by output tile — makes pingpong visually obvious)
green EPI	epilogue for an output tile
red WAIT (consumer lane)	bubble — tooltip says why (TMA latency exposed vs. math order barrier)

Usage

# waveform (default) — mainloop only
python3 viz.py path/to/sm90_mma_tma_gmma_ss_warpspecialized.hpp -o out.html

# with warpgroup roles from the kernel layer (recommended)
python3 viz.py \
  cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp \
  --kernel-layer cutlass/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp \
  -o out.html

# static dataflow node-graph instead of a timeline
python3 viz.py <mainloop.hpp> --mode graph -o graph.html

# dump the parsed structure as JSON
python3 parse_cutlass.py <mainloop.hpp> --pretty

Requires Python 3.9+, stdlib only. Output HTML is self-contained.

Which file do I point it at?

Point at the collective mainloop header, not the example driver .cu. Example drivers (e.g. CUTLASS examples/48_hopper_warp_specialized_gemm) use CollectiveBuilder + KernelScheduleAuto — the pipeline API lives in the mainloop the builder resolves to, e.g.:

cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp

If you point at a driver file by mistake, the tool detects it and prints mainloop candidates found in the adjacent CUTLASS include tree.

Interactive knobs

knob	meaning
stages	pipeline depth S (source has symbolic DispatchPolicy::Stages)
output tiles	how many output tiles the CTA processes
k-tiles/output	K-loop trip count per output tile
TMA fill	abstract latency of a TMA load
WGMMA	abstract latency of consuming one k-tile
epilogue	abstract epilogue latency per output tile
consumer WGs	1 or 2 math warpgroups
schedule	pingpong or cooperative

Warpgroup model (matches the CUTLASS SM90 kernel layer)

One CTA = 1 producer (load) warpgroup + N consumer (math) warpgroups, per WarpGroupRole { Producer, Consumer0, Consumer1 }:

• pingpong (sm90_gemm_tma_warpspecialized_pingpong.hpp, NumLoadWarpGroups=1, NumMmaWarpGroups=2): consumer WGs alternate output tiles. WGMMA ownership passes between them via the MathWarpGroupOrderBarrier, so one WG's epilogue overlaps the other WG's mainloop — the timeline makes this overlap directly visible.
• cooperative (sm90_gemm_tma_warpspecialized_cooperative.hpp): both consumer WGs split the same output tile along M and consume each k-tile buffer together; the epilogue sits on the critical path.

The insight to look for: in pingpong with epilogue > 0, the red WAIT blocks shrink and total time drops versus cooperative at the same knob settings — that's the epilogue hiding that motivates the schedule.

What the parser recognizes

Canonical CUTLASS 3.x pipeline API (not general C++):

• using MainloopPipeline = cutlass::PipelineTmaAsync<...> / PipelineAsync<...>
• struct SharedStorage { ... array_aligned<..., cosize_v<SmemLayoutX>> smem_X; ... }
• pipeline.producer_acquire / producer_commit / producer_get_barrier
• pipeline.consumer_wait / consumer_release
• copy(tma_load_X.with(barrier, ...), gmem, smem) → TMA op (buffer routing by A/B suffix)
• cute::gemm(tiled_mma, a, b, accum) → WGMMA op
• kernel layer: WarpGroupRole enum, NumLoadWarpGroups / NumMmaWarpGroups, pingpong/cooperative schedule detection

A function is classified producer if it calls producer_acquire/ producer_commit, consumer if it calls consumer_wait/ consumer_release. Every block in the timeline carries the source line numbers it came from (e.g. TMA at L384-385, consumer_wait at L487).

Tested on:

• sm90_mma_tma_gmma_ss_warpspecialized.hpp (canonical bf16)
• sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp
• sm90_mma_multistage_gmma_ss_warpspecialized.hpp (cp.async, PipelineAsync)

Files

file	purpose
viz.py	entry point: parse → render HTML
parse_cutlass.py	static parser for the CUTLASS pipeline API
render_waveform.py	GTKWave-style timeline renderer (default)
render_html.py	static dataflow node-graph renderer (--mode graph)
examples/	pre-generated HTML you can open directly

Limitations / roadmap

• Beta / under testing — expect rough edges; the logical-schedule model is a simplification of the real hardware behavior.
• Logical schedule, not measured timing. For real cycle data, the natural extension is feeding NVBit-instrumented traces or Accel-Sim replay records into the same lane layout.
• SM90 only. SM100 (Blackwell) tcgen05.mma / TMEM / PipelineTmemUmma patterns are not yet in the pattern set.
• Attention kernels (e.g. FlashAttention-3) parse but visualize approximately. The parser correctly extracts FA3's structure from flash-attention/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp — multiple pipelines (MainloopPipelineK/V/Vt), buffers (smem_q/k/v/vt), and all producer/consumer functions. However, the waveform renderer models a single pipeline feeding one consumer loop, which is GEMM-shaped. Attention actually has: separate K and V pipelines with independent timing; a two-phase consumer (GEMM0 Q·Kᵀ → softmax → GEMM1 P·V) where K stages release after GEMM0 but V stages after GEMM1; and a load warp that is simultaneously a producer (K/V) and a consumer (Vt transpose pipeline). Until per-pipeline lanes and a two-phase consumer land, treat the FA3 timeline as a rough approximation of the real schedule.
• Regex-based extraction — heavily macroed or non-canonical pipeline code may not parse. --print-json shows exactly what was extracted.