MegaUP · fusing computation and communication

We timed four operations around attention: the query/key/value (QKV) and output projections, plus their input-gradient calculations in the backward pass. Each pairs a matrix multiplication (GEMM) with an all-to-all exchange between GPUs. We compare separate Transformer Engine and NCCL operations with MegaUP’s fused version.

Bars show time saved in milliseconds: separate − fused. The top chart adds the four savings to estimate a net benefit; full attention and model execution were not measured.

Configuration: B200 GPUs · BF16 precision · B8 = 8 equal sequences · D128 = 128 features/head · TP1 = no tensor sharding.
CP = context-parallel GPU ranks. M = total tokens per rank; local tokens/sequence = M / 8, global tokens/sequence = M × CP / 8.

Net time saved across four operations → positive: potential overhead budget · ← negative: operation cost

Mean of initial/reversed operation medians, equally weighted. Each block has 10 warmups and 40 timed observations per operation, using the slowest GPU latency (rank maximum). Each bar sums four operations; click a row for its breakdown.

One linear ms scale across visible rows; scale changes with filters.

For differential integration overhead h = fused integration cost − native integration cost, estimated saving is Δ − h. A positive bar gives the break-even h budget, assuming remaining work is unchanged and operation times add without overlap or interactions. Benefit requires h < Δ.

These are sums of mean operation medians from the two run orders, not pooled raw-sample medians, a confidence interval, or a measured full attention step. Each original operation median uses synchronized rank-max timing. Native here is TE GEMM + one packed NCCL exchange; full MHA uses three pipelined QKV exchanges. Prepared operations already include recurring communication, layout/reorders, staging and required copies; setup/JIT, attention, wgrad and some full-integration costs are excluded. End-to-end percent needs an actual baseline or an explicit assumption for unchanged work. No 5% win/loss rule is used here.

56 measured configurations. Super/Ultra are shown only at CP2 because each configuration has two KV heads.