PaperlayerView exact revision on arXiv

Source-linked AI summary

One Megakernel is Not Enough: Predicting Where to Cut for Streaming Inference

Aditya Chatterjee, Anand Joseph X M, Sanjay Singh Chauhan

arXiv:2610.35336v1cs.PFcs.DCcs.LG

TL;DR

Choosing where to cut a streaming model for megakernel fusion is costly to evaluate through exhaustive implementation and timing. The paper uses profiles, benchmarks and implementation contracts to rank supported compositions without timing candidates end to end. The planner selected the measured fastest design in all 13 held-out cases within the tested validation pools.

  • Problem

    Exhaustively implementing and timing candidate compositions is costly; selecting boundaries requires accounting for first publication, persistent resource lifetimes and streaming dependencies.

    Source: §I · p. 2

  • Method

    A planner combines conventional-graph profiles, operator benchmarks, implementation contracts and streaming dependencies, then simulates tile and event graphs to rank supported compositions without end-to-end candidate measurements.

    Source: §I · p. 2, D. Megakernel Designs and the Boundary… · p. 2, B. Tile Graph (tGraph) and Service Table · p. 4

  • Results

    The planner selected the measured fastest design in all 13 held-out cases within the declared validation pools, across the tested workloads on H100 NVL.

    Source: B. Selection accuracy on held-out cases · p. 7

  • Takeaways & Limitations

    The best composition ranges from zero to four megakernels across models and batch sizes, so fusion is a question of where to cut rather than whether to fuse.

    Source: §IX · p. 8

    Selection accuracy is established within the declared validation pools on H100 NVL, not over unmeasured compositions or other GPU architectures.

    Source: B. Selection accuracy on held-out cases · p. 7

Abstract

from arXiv · show

We address a fundamental question: How far should kernel fusion extend and which boundaries are worth preserving? Unlike the general convention of using one monolith megakernel at lower batch sizes or supporting only a selective model component as megakernel and rest of the model as conventional graph, we show that the best configuration is a composition of multiple megakernels with no overlapping conventional graph component, or a fully conventional graph. We present a framework that selects megakernel regions in a streaming inference model from conventional-graph profiles and standalone kernel benchmarks, without running any candidate design end to end. It simulates each candidate's tile and event graphs to predict time to first output and selects the measured best design in all 13 combinations of model and batch size we tested. Across NH-TTS, Qwen3-TTS and Llama-1B, the best composition ranges from two to four megakernels to fully conventional execution and none is a single megakernel. We introduce four model categories (megakernel-stable/ receding/ advancing/ averse) that capture how a model's best megakernel composition changes with batch size. Using NH-TTS model on H100 NVL GPU, we demonstrate that preserving three specialized megakernels can improve both latency and throughput despite introducing additional kernel boundaries. At batch size 32, this design delivers 21.4% lower mean batch-p95 time to first audio (TTFA) and 28% higher throughput than the evaluated monolith megakernel. Against conventional TensorRT graph execution, it reduces engine-level mean batch TTFA p95 by 4.05x and increases throughput by 19.7%.

I. INTRODUCTION

The introduction frames fusion as a boundary-selection problem because streaming phases have different execution needs. It presents a planner that selects supported compositions and reports that its choices match measured best designs across 13 model-batch cases.

I. INTRODUCTION

  • When RTFx exceeds 1, response latency is dictated by TTFA or TTFT, making those measures critical for bidirectional streaming.Real-time serving also requires supporting the maximum batch size while keeping key metrics within acceptable limits.Source: §I · p. 1
  • TTS phases run at different rates and prefer different thread-block shapes, making each fusion cut a choice about launch configuration, resource lifetimes and GPU sharing.The paper describes prefill, autoregressive production, local-transformer work and codec decoding as distinct phases.Source: §I · p. 1
  • At NH-TTS batch 32, three megakernels publish the first audio 4.05× sooner than conventional TensorRT and 21.4% sooner than one megakernel.Figure 2 presents TTFA and RTFx for the selected three-megakernel design.Source: Figure 2 · p. 1, §I · p. 1
  • The boundary planner combines conventional-graph profiles, operator benchmarks, implementation contracts and streaming dependencies to rank supported compositions without timing candidates end to end.Exhaustively implementing and timing NH-TTS’s 81 candidates is costly.Source: §I · p. 2

II. BACKGROUND

  • In codec-based autoregressive TTS, prefill initializes the request context, while the producer emits acoustic codes and the consumer converts buffered frames into waveform samples.Producer and consumer name scheduling roles relative to the acoustic-code buffer.Source: §II · p. 2

A. Streaming Inference and First Publication

  • First publication is when the first audio chunk or token becomes available, and the planner minimizes a deterministic first-publication score.For TTS, publication can wait for enough frames to decode a chunk and for a further producer step confirming the stream continues.Source: A. Streaming Inference and First… · p. 2

B. Conventional Graph Execution

The paper contrasts conventional graph execution with persistent megakernels, then frames boundary selection as a choice of partition, backend, and GPU-sharing policy. In NH-TTS, three overlapping megakernels publish audio sooner than the evaluated conventional and monolithic designs.

B. Conventional Graph Execution

  • Conventional frameworks execute models as kernel sequences, while CUDA graphs replay recorded launches with one host call to remove host-launch overhead.At small batch sizes, host launches, barriers, and host synchronization take a large share because kernels are short.Source: B. Conventional Graph Execution · p. 2

C. Megakernels

  • Megakernels keep resident thread blocks executing operators and can enable prefetch and cross-kernel tile execution, but shared launch configurations and register pressure can cause spilling.The cited mechanisms and costs distinguish potential utilization gains from fusion overhead.Source: C. Megakernels · p. 2

D. Megakernel Designs and the Boundary Selection Problem

  • A design chooses contiguous model regions, a conventional or megakernel backend for each, and a policy for sharing the GPU.The planner ranks supported compositions using a deterministic first-publication score, a proxy rather than a direct latency-distribution or p95 prediction.Source: D. Megakernel Designs and the Boundary… · p. 2

III. COMPOSING MEGAKERNELS: A NEMOTRON-H BASED TTS CASE STUDY

  • NH-TTS combines prefill, acoustic-token generation, and codec decoding; first publication waits for both a decoded chunk and one further generation step.Its phases run at different rates and use different tensor shapes and thread-block sizes.Source: §III · p. 2

A. Multiple Combinations

  • At batch size 32, three overlapping megakernels publish first audio 4.05× sooner than TensorRT CUDA graphs and 21.4% sooner than one megakernel.The ordering also holds at B64 (47.4 ms versus 61.0 ms and 221.5 ms) and B1 (27.5 ms versus 32.7 ms and 73.0 ms), respectively.Source: A. Multiple Combinations · p. 2

IV. AUTOMATED FIRST-ORDER BOUNDARY DISCOVERY FROM FIRST PRINCIPLES

The first-order planner uses conventional-graph and standalone measurements, implementation contracts, and explicit tile, launch, and capacity models to rank candidate boundaries without timing candidate compositions end to end.

IV. AUTOMATED FIRST-ORDER BOUNDARY DISCOVERY FROM FIRST PRINCIPLES

  • The planner builds region, tile, and event graphs, simulates each supported composition, and ranks candidates by first-publication score without measuring them end to end.Its inputs include conventional-graph profiles, standalone operator or step benchmarks, and source-declared implementation contracts.Source: §IV · pp. 2–3

A. Region graph (rGraph) and candidates

  • The region graph represents each group as a node and records data, state, and publication dependencies between groups.Groups may run conventionally as one kernel per operator or as a megakernel in one persistent launch.Source: A. Region graph (rGraph) and candidates · p. 3
  • A cover partitions ordered groups into contiguous regions and assigns each a backend; each candidate pairs a supported cover with an applicable execution policy.The policies include serial execution, fixed producer-consumer CTA splits, and dynamic splits that require separate launches and eligible megakernel regions.Source: A. Region graph (rGraph) and candidates · p. 3
  • The streaming schedule specifies execution frequency, delay before output, first-chunk length, and lookahead that reveals whether the stream continues.These parameters determine how much producer work precedes first publication.Source: A. Region graph (rGraph) and candidates · p. 3
  • For NH-TTS, four ordered groups yield 54 serial covers, with supported fixed and dynamic splits bringing the candidate total to 81.The groups are the Nemotron-H backbone, local transformer, and codec, with the passage also describing the producer and consumer roles.Source: A. Region graph (rGraph) and candidates · p. 3

B. Tile Graph (tGraph) and Service Table

The tile graph models supported output tiles and their dependencies, then schedules them under resource and implementation constraints to estimate service costs for candidate compositions.

B. Tile Graph (tGraph) and Service Table

  • Each supported output tile becomes a task with compute, loads, writes, and staging storage; edges capture tile inputs while original kernel barriers remain separate.Tile dependencies draw on source-declared contracts, profiling kernels, grids, and resource footprints.Source: B. Tile Graph (tGraph) and Service Table · p. 3
  • An event-driven list scheduler runs tiles on a fixed worker pool, starting tasks when dependencies and staging space permit and routing transfers through one FIFO memory queue.Opaque operators remain tasks in the worker pool, with reduced-budget costs estimated by a CTA-wave approximation.Source: B. Tile Graph (tGraph) and Service Table · pp. 3–4
  • Megakernel scheduling removes original barrier edges so ready tiles can fill idle workers, while weight prefetch lets computing workers stage weights for upcoming tiles.The resulting service table estimates one execution time per group, backend, and CTA budget.Source: B. Tile Graph (tGraph) and Service Table · p. 4
  • Service costs depend on the candidate region’s implementation context, and incompatible timing evidence is excluded rather than transferred across configurations.Contracts capture geometry, storage, compiled-entry membership, synchronization, and other configuration details; compatible measurements must match their declared keys.Source: B. Tile Graph (tGraph) and Service Table · p. 4

C. Event Graph (evGraph)

The event graph models launches and dependencies through first publication, whose simulated time is determined by the later of first-chunk completion and producer-step completion.

C. Event Graph (evGraph)

  • The event graph assigns launch durations from the service table, models dependencies and resource-holding launches, and publishes only after both the first consumer chunk and producer step h complete.Its publication time is Tpub = max(Fchunk1, Fproducerh), where F denotes simulated completion time.Source: C. Event Graph (evGraph) · p. 4

D. Megakernel Impact Simulation

The simulator scores supported compositions by modeling tile service costs, launch allocations, shared bandwidth, and streaming publication dependencies. Boundary selection exhaustively returns the feasible candidate or candidates with the lowest first-publication score.

D. Megakernel Impact Simulation

E. Boundary Selection

  • At NH-TTS batch 32, the planner selects separate megakernels for prefill, the fused backbone/local transformer, and codec, with an 88/44-CTA dynamic producer/consumer split.The selected composition is P/(B + L)/C.Source: E. Boundary Selection · p. 4
  • In an illustrative case, an 88/44 split publishes at 28.2 ms versus 30.1 ms serially, despite slower individual producer and codec work under overlap.The example assumes chunk readiness at 20 ms, transition costs of 0.1 ms serially and 0.2 ms for overlap, and no bandwidth contention.Source: E. Boundary Selection · p. 4
  • The planner forms candidates from supported contiguous-region covers and applicable serial, fixed-CTA, or dynamic-sharing policies, then scores their first-publication time.A candidate combines a cover with an applicable policy and implementation contract.Source: E. Boundary Selection · p. 5

V. WHAT THE BOUNDARY MODEL REVEALS ABOUT COMPOSITION

Tile and event graphs predict and explain multiple performance observations for megakernels and kernels in a model execution environment.

  • Tile and event graphs predict and explain multiple megakernel and kernel performance observations in a model execution environment.Source: §V · p. 5

A. Observation 1: A faster kernel can make the pipeline slower

The observations show that pipeline performance depends on co-running implementations and resource contracts, not only isolated kernel speed or boundary count. Megakernel fixed costs can also make conventional execution preferable when its work-saving lead shrinks.

A. Observation 1: A faster kernel can make the pipeline slower

  • With 88 rather than 80 CTAs, the producer is 2.2% faster alone but the producer–conventional-codec pair takes 30.3% longer; beside the codec megakernel, the pair is 2.0% faster.The conventional codec loses eight CTAs in the first pairing.Source: A. Observation 1: A faster kernel can… · p. 6

B. Observation 2: Colocation cost depends on the peer implementation

C. Observation 3: Serial cuts can be neutral under fixed contracts

  • With implementation and resource contracts fixed and dispatch overhead excluded, a serial cut leaves the modeled schedule unchanged; the three-megakernel design's advantage over the monolith comes from Stage 0 specialization.The three-megakernel design uses 256-thread Stage 0 blocks, versus 512-thread blocks in the monolith.Source: C. Observation 3: Serial cuts can be… · p. 6

D. Observation 4: A megakernel’s fixed cost decides when conventional execution wins

  • Qwen3-TTS's four-region megakernel lead falls from 15.9 ms at B16 to 0.65 ms at B64 against fixed costs of 0.7–1.1 ms, so conventional execution wins at B48 and B64.For NH-TTS, the megakernel lead keeps increasing with batch size, and conventional execution never wins at tested sizes.Source: D. Observation 4: A megakernel’s fixed… · p. 6

VI. EVALUATION ACROSS MODELS AND FRAMEWORKS

On held-out model-batch cases, the frozen planner selected the measured fastest design in all 13 validation pools. Its event simulation improves selection over simpler selectors on overlap-sensitive cases, while evaluation remains bounded by the implemented pools and tested H100 NVL workloads.

A. Setup

  • The evaluation freezes scoring rules, candidate generation, and permitted profiles before end-to-end timing, using evaluation fixtures excluded from development.Qualified designs are checked against an unfused reference before measurement.Source: A. Setup · pp. 6–7

B. Selection accuracy on held-out cases

  • Predicted absolute times can err despite matching measured winners: selected NH-TTS B32 is predicted at 42.52 ms and measured at 35.99 ms, while closest choices are within their reported 95% intervals.The planner ranks designs rather than calibrating absolute time.Source: B. Selection accuracy on held-out cases · p. 7
  • The full planner matches 13/13 measured winners, while additive scoring and removal of overlap each match 11/13, with maximum regret of 4.50%.Both simpler-selector misses are NH-TTS B32/B64 cases where they choose serial execution; event simulation accounts for chunk readiness, CTA allocations, and concurrent completion.Source: B. Selection accuracy on held-out cases · p. 7
  • 13/13 held-out model-batch cases had a uniquely selected measured-fastest design, establishing accuracy within the declared validation pools on H100 NVL.This does not establish optimality over unmeasured compositions or transfer to other GPU architectures.Source: Figure 4 · p. 7, B. Selection accuracy on held-out cases · p. 7

VII. HOW FAR SHOULD KERNEL FUSION EXTEND?

There is no universal fusion extent: the best choice depends on model architecture, batch size, and objective. Megakernel regions pay when their savings exceed fixed costs, while the four categories describe where those regions are expected to help.

VII. HOW FAR SHOULD KERNEL FUSION EXTEND?

  • Advancing models are predicted where work-proportional traffic or serialized compute, memory, and communication can be overlapped by a persistent schedule.The advancing and averse categories are predicted, not observed.Source: §VII · pp. 7–8
  • Averse models are predicted for regions already compute-bound, although a region with little isolated gain may benefit when colocated, as in the Nemotron experiment.The averse category is predicted, not observed.Source: §VII · pp. 7–8
  • Megakernel boundaries are worth preserving when their crossing costs less than the specialization and publication control they protect.The choice also depends on model architecture, batch size, and objective.Source: §VII · p. 7
  • Stable kernels account for most of the time on the path to the objective; Llama-1B is stable for first-token latency but receding for decode.For first-token latency, the relevant path is only prefill.Source: §VII · p. 8
  • Form contiguous regions only when their savings exceed fixed costs at every batch size; stable kernels interleaved with others may not pay for the surrounding cuts.Co-located conventional components can also erase megakernel gains.Source: §VII · p. 8

VIII. RELATED WORK

Across 13 cases, the selected design matched the measured best, and NH-TTS achieved a 4.05× speedup in engine-level mean batch-p95 TTFA at batch 32. The related work covers persistent-program boundaries, pipeline tuning, forecasting, and interference-aware sharing; the conclusion leaves model superiority and software-level interference prevention open.

VIII. RELATED WORK

  • Prior work selects persistent-program boundaries or tunes pipeline groups and resource assignments; this work focuses on first-publication-aware composition using conventional and operator-level evidence.Performance forecasting and interference-aware sharing are described as complementary approaches.Source: §VIII · p. 8

IX. CONCLUSION

  • The planner selected the measured best design in all 13 cases, and NH-TTS achieved a 4.05× speedup in engine-level mean batch-p95 TTFA at batch 32.The conclusion describes the selection as prediction without end-to-end timing of candidate compositions.Source: §IX · p. 8
  • The conclusion asks whether megakernel-stable architectures outperform comparably accurate models and whether software scheduling can prevent interference between co-located conventional graphs and megakernels.Source: §IX · p. 8

Related papers

Loading 2610.35336v1…