PaperlayerView exact revision on arXiv

Source-linked AI summary

CREDIT: Cost-guided Reduction-reuse with Efficient DSMEM Inter-CTA Tiling

Zhengxiong Li, Tsung-Wei Huang, Umit Ogras

arXiv:2609.01864v1cs.DC

TL;DR

CREDIT addresses the challenge of determining when DSMEM improves performance despite remote-access, synchronization, and resource costs. It profiles DSMEM, transforms reduction-reuse workloads, and uses a cost model to select profitable cases. Across six workloads, it achieves 91.7% profitability prediction accuracy and outperforms the fastest baselines on RTX 5090 and H100.

  • Problem

    DSMEM introduces remote-access, synchronization, allocation, and occupancy costs, making it difficult to determine when clustered execution is profitable.

  • Method

    CREDIT combines DSMEM performance profiling, a reduction-reuse transformation using distributed local staging and compact reduction, and a profitability cost model.

  • Results

    CREDIT achieves 91.7% profitability prediction accuracy and beats the fastest baseline on all six workloads, with geometric-mean speedups of 1.466× on RTX 5090 and 1.318× on H100.

  • Takeaways & Limitations

    CREDIT supports selective DSMEM use by generating clustered variants when recovered source time amortizes distributed staging and compact inter-CTA reduction.

Abstract

from arXiv · show

NVIDIA distributed shared memory (DSMEM) enables direct shared-memory access within a thread block cluster. However, cluster synchronization, remote access, and resource costs make it difficult to determine when DSMEM improves performance. To fill this gap, we propose CREDIT, a cost-guided framework that identifies DSMEM-profitable workload patterns, predicts their profitability range, and delivers consistent speedups across diverse workloads. CREDIT combines three innovations: (1) a profiling-driven characterization that identifies workload patterns likely to benefit from DSMEM; (2) a transformation that applies DSMEM to reduction-reuse workloads; (3) a cost model based on profiling data, to determine its profitability range. Evaluations on diverse workloads show CREDIT achieves 91.7% prediction accuracy on profitability. CREDIT beats torch.compile, Triton, and optimized non-DSMEM CUDA baselines on all six workloads, with geometric-mean speedups of 1.466x on RTX 5090 and 1.318x on H100. CREDIT's source code is publicly available at https://github.com/zhengxiongli08/CREDIT.

I. INTRODUCTION

CREDIT addresses the difficulty of using DSMEM profitably as GPU compute-to-bandwidth ratios rise and DSMEM introduces synchronization, communication, and occupancy costs. It combines profiling, a reduction-reuse transformation, and a cost model, achieving broad speedups and 91.7% profitability prediction accuracy.

  • Motivation: The dense FP16 compute-to-bandwidth ratio rises from 139 FLOPs/byte on V100 to about 295 FLOPs/byte on H100, increasing pressure to reduce memory traffic.Many operators fail to reach peak performance unless they reduce memory footprint and increase arithmetic intensity.
  • Motivation: DSMEM enables cluster-wide shared-memory access, but remote accesses, synchronization, larger allocations, and fixed cluster sizes can reduce performance.These costs make DSMEM beneficial only for suitable workloads and configurations.
  • Approach: CREDIT identifies DSMEM-beneficial workload patterns, predicts their profitability range, and applies the transformation only when it pays off.Its approach targets wide reduction-reuse kernels that can avoid rereading data from global memory between reduction passes.
  • Results: CREDIT outperforms the fastest baseline on all six workloads, with geometric-mean speedups of 1.466× on RTX 5090 and 1.318× on H100.The evaluation compares against torch.compile, Triton, and optimized no-DSMEM CUDA baselines and measures 33–60% off-chip traffic reductions.
  • Approach: CREDIT quantitatively characterizes DSMEM read/write latency, throughput, and synchronization overhead to identify promising workload patterns.These measurements support workload selection and profitability analysis.

II. BACKGROUND AND RELATED WORK

The background introduces thread block clusters and DSMEM as an intermediate cooperation scope for reducing data movement, while related work spans IO-aware algorithms, compiler optimization, hardware characterization, and DSMEM-specific systems.

  • DSMEM Background: Thread block clusters schedule cooperative thread arrays concurrently and provide cluster-wide synchronization, while DSMEM permits direct loads, stores, and atomics on peer CTAs’ shared memory.This extends cooperation beyond the conventional single-block shared-memory scope.
  • Prior Approaches: IO-aware algorithms and compiler systems reduce off-chip traffic through tiling, fusion, scheduling, and performance models without requiring DSMEM.Examples include FlashAttention, TVM, PyTorch, DNNFusion, Ansor, Roller, AStitch, and Welder.
  • Hardware Characterization: Microarchitectural studies show that DSMEM latency and bandwidth depend on cluster configuration and network topology, so peak specifications are insufficient.CREDIT uses calibrated primitive measurements for the protocol it evaluates.
  • DSMEM Background: DSMEM accesses peer-owned allocations through the inter-SM interconnect, although each allocation remains physically local to its owner.The figure distinguishes logical addressability from physical ownership.
  • DSMEM Systems: Recent systems use DSMEM for communication abstractions, fusion planning, collective-traffic modeling, and memory-bound kernel optimization.These systems target compute-intensive operator chains, LLM decoding stages, or Hopper and Blackwell workloads.

C. Positioning of CREDIT

CREDIT is positioned as a workload-level framework for deciding whether clustered staging is profitable for a given operator shape, rather than as the first DSMEM use or a new reduction collective.

  • Positioning of CREDIT: CREDIT asks whether a given operator shape should use clustered staging relative to a competitive non-clustered implementation.It targets selection across both profitable and unprofitable cases on H100 and RTX 5090.
  • Positioning of CREDIT: The framework uses matched primitive and control measurements together with one no-DSMEM timing to predict profitability and crossover points.This distinguishes CREDIT from approaches focused only on introducing DSMEM functionality.
  • CREDIT Transformation: CREDIT transforms reduction-reuse kernels by retaining owner-local slices and exchanging compact scalar partials within a cluster.Each CTA can push partials to peer-owned buffers so final combining reads are local.
  • Primitive Characterization: CREDIT profiles DSMEM using latency, throughput, and repeated cluster-synchronization microbenchmarks that place two CTAs on different SMs.The measurements report the median of 15 runs.

B. Reduction-Reuse Transformation

CREDIT targets rows whose per-element data is reduced into scalar statistics and then reused across a full-size output. It partitions each row across CTAs, retains slices, and compares clustered staging against global rereads under SMEM capacity constraints.

  • Reduction-Reuse Workload: CREDIT models each row as per-element inputs contributing to staged scalar reductions, followed by an elementwise function that produces every output element.At stage ℓ, each statistic combines per-element contributions with earlier-stage statistics; after L stages, h computes y_i.
  • Row Partitioning: Rows are partitioned into P nonoverlapping CTA-owned slices, with each CTA retaining only its assigned elements and local partial statistics.The transformation distributes input ownership across CTAs before cluster-wide exchange.
  • Resource Constraint: Staging is legal only when retained bytes and per-CTA scratch fit within the configured dynamic-SMEM limit.The legality condition depends on per-element retained storage, fixed scratch, and the maximum dynamic-SMEM allocation.
  • Traffic Model: The clustered schedule eliminates reread bytes per element while preserving fixed traffic, such as output writes, that clustering does not affect.This defines the traffic difference used to estimate the memory-bound speedup ceiling.

C. Cluster Reduction and Reuse Protocol

CREDIT compresses each CTA’s local partial statistics, pushes them into peer-owned shared-memory slots, and uses barriers so every CTA can reconstruct the full statistics locally. Replicated all-gather avoids a root bottleneck and later remote reads at the cost of replicated scalar stores.

  • Partial Reduction: Each CTA reduces its retained slice to compact partial statistics before performing cluster-wide communication.Thread accumulation, warp shuffles, and SMEM reduction produce the per-CTA statistic vector.
  • Replicated Push: Source CTAs write their partials into corresponding remote slots owned by destination CTAs.The protocol reserves a local-SMEM slot for each source CTA and statistic at every destination.
  • Synchronization: Two cluster-wide barriers expose pushed partials and protect scratch reuse while each CTA combines only its local slots.The first barrier makes writes visible; the second separates reconstruction from subsequent scratch reuse.
  • All-Gather Trade-off: Replicated all-gather lets every CTA reconstruct the full statistic independently, avoiding a serial root bottleneck and a second round of remote reads.The trade-off is replicated scalar stores across peer-owned buffers.
  • Output Generation: After reconstruction, each CTA applies h to its retained slice and writes its output segment.Subsequent statistic reads are local because each destination owns the slots containing all partials.

D. Profitability and Candidate Selection

CREDIT estimates DSMEM profitability using baseline timing, measured control overheads, and resource-dependent communication costs across legal cluster sizes. It predicts crossovers while rejecting workload classes whose reuse or state-growth assumptions do not hold.

  • Baseline Timing: CREDIT uses one timing of the selected non-DSMEM CUDA path as a cheap first screen before clustered variants are generated or tuned.This avoids unreliable peak-bandwidth assumptions caused by cache residency and architecture-dependent baseline efficiency.
  • Control Measurement: A work-free control kernel with matching grid, cluster size, dynamic-SMEM footprint, and K = 2L barriers measures cluster launch, scheduling, occupancy, and synchronization costs.The excess control time relative to a one-CTA control supplies Tctrl.
  • Occupancy: When one CTA per SM limits residency, a row cluster occupies P SMs and executes in at most ⌊S/P⌋ concurrent waves.The resulting wave count W(M, P) contributes to the modeled execution time.
  • Communication Costs: The model accounts for local replay, remote-store issue, and only the non-overlapped deposit residual because source CTAs issue stores concurrently.Deposit time is charged as [T_wSMEM − Tinput]+, which is zero on both measured devices because local-store issue is faster than the per-SM input stream.
  • Candidate Selection: CREDIT defines N⋆ as the smallest tested width where the profitability condition holds for at least one legal cluster size, while final tuning may select another legal size.The predictor identifies a pair and crossover rather than a cycle-exact configuration.
  • Rejected Cases: The framework rejects elementwise or scalar-output reductions without rereads, scans with growing ordered dependencies, and kernels whose remote state grows with N rather than P.These cases lack amortizable reread traffic or violate the compact-partial assumption.

A. Experimental Setup and Baselines

CREDIT is evaluated on six reduction–reuse workloads across RTX 5090 and H100, using three non-DSMEM baselines and sweeps over row widths and legal cluster sizes.

  • Workloads and Hardware: The benchmark covers six FP32 reduction–reuse workloads from five categories, with 2,048 rows for LayerNorm and 4,096 for the others.Widths are swept from 4K to 64K, and legal cluster sizes are P ∈ {2, 4, 8}.
  • Workloads and Hardware: Experiments run on an NVIDIA RTX 5090 and H100 SXM using CUDA 13.0, PyTorch 2.11.0, and Triton 3.6.0.Compilation and autotuning costs are measured once and excluded from the timed region.
  • Baselines: CREDIT is compared against torch.compile, Triton, and an optimized CUDA baseline selecting local-SMEM staging or global rereads.The CUDA baseline provides an ablation mechanism against non-DSMEM alternatives.
  • Metrics and Timing: Speedup is defined as SB = TB/TCREDIT, so values above 1 indicate CREDIT is faster than baseline B.“Best” denotes the fastest of the three baselines at that shape.

B. Workload-by-Workload Scaling

CREDIT’s advantage depends on row width, architecture, and workload: small widths favor non-DSMEM baselines, while sufficiently wide rows can make CREDIT faster across workloads.

  • Overall Scaling: At 4K, CREDIT is slower than the one-CTA baseline, with geometric-mean speedups of 0.738× on RTX 5090 and 0.815× on H100.The baseline benefits from cache or local staging without cluster overhead.
  • Overall Scaling: At 64K, CREDIT beats the fastest baseline on all six workloads, achieving geometric-mean speedups of 1.466× on RTX 5090 and 1.318× on H100.These values summarize the cross-workload result at the widest tested shape.
  • Architecture Dependence: RTX 5090 first exceeds parity at 16K and wins every workload at 32K, whereas H100 wins every workload only at 64K.H100’s higher barrier latency, 851 versus 404 cycles, shifts its amortization point to wider rows.
  • Workload Variation: Weighted variance rises from 0.982× at 4K to 2.401× at 64K on RTX 5090, the largest reported gain.Pearson and softmax-logits show later moderate crossovers, while LARS reaches only 1.014× and 1.045× at 64K.
  • Workload Variation: Quantization crosses the best baseline at 16K on RTX 5090 but only at 64K on H100.The workload spread shows that width alone does not determine the crossover.

C. Cost-Model Validation

CREDIT’s profile-guided model screens workload–shape pairs for DSMEM profitability and narrows tuning without serving as a cycle-exact runtime predictor.

  • Model Inputs: The model predicts whether CREDIT beats matched non-DSMEM CUDA using control measurements, baseline timing, and statically analyzed traffic quantities.It uses measured primitive rates and one selected non-DSMEM timing rather than DSMEM workload timings.
  • Profitability Accuracy: 91.7% overall accuracy is achieved across 60 profitability pairs spanning five widths, six workloads, and two GPUs.Accuracy is 90.0% on RTX 5090 and 93.3% on H100.
  • Profitability Accuracy: All 28 points using a pure global-reread CUDA baseline are classified correctly.The five remaining errors occur near the local-SMEM staging boundary and remain within 2.2% of parity.
  • Crossover Prediction: The prediction is exact for seven of 12 device–workload pairs and within one factor-of-two tested width interval for the remaining five.No crossover is missed by more than one interval.
  • Structural Screening: Negative controls reject one-pass reductions, scans, and selection or stencil kernels because they lack the required recoverable reduction–reuse dataflow.The screening rules identify opportunities from dataflow rather than vector width alone.

D. Traffic Validation

CREDIT reduces DRAM traffic by retaining owner-local slices and replacing repeated global reads with compact inter-CTA statistics, but traffic savings do not alone determine speedup.

  • Measured Traffic: 33–60% DRAM traffic reductions are measured on RTX 5090 across the six workloads.LayerNorm drops from 24 to 12 B/element, while quantization drops from 9 to 5 B/element.
  • Mechanism: The reduction comes from retaining owner-local slices and replacing repeated global reads with compact statistics, while output traffic remains unchanged.This validates the traffic mechanism targeted by CREDIT’s reduction–reuse transformation.
  • Traffic Versus Speed: LARS removes 38% of traffic but gains only 1–5% over the best baseline at 64K.Its framework baseline is already substantially faster than matched CUDA, so traffic reduction does not directly determine practical speedup.
  • Cost Accounting: CREDIT explicitly subtracts cluster control, replay, and DSMEM transport costs instead of treating traffic reduction alone as the benefit.This cost accounting explains why the framework evaluates avoided source time at the baseline’s achieved rate.
  • Scope: CREDIT is presented as a selective reduction–reuse mechanism rather than a general fusion fabric.Clustered variants should be generated only when recovered source time amortizes distributed staging and compact inter-CTA reduction.
Loading 2609.01864v1…