AutoSAGE: Input-Aware CUDA Scheduling for Sparse GNN Aggregation (SpMM/SDDMM) and CSR Attention

AI summary

Overview

Research area: GPU kernel scheduling for sparse linear algebra in graph neural networks (CSR SpMM, SDDMM, and a composed CSR-attention pipeline).

Technical level: Advanced. The paper assumes familiarity with CUDA execution models, CSR sparse formats, tiling/mapping strategies, and GNN aggregation operators.

Scope (one sentence): AutoSAGE is an input-aware CUDA scheduler that picks per-input tiling and row-mapping strategies for CSR SpMM and SDDMM using cheap estimates refined by on-device micro-probes, with a guardrail that falls back to vendor kernels and a persistent cache for deterministic replay.

What This Paper Is About

Sparse GNN aggregation kernels such as CSR SpMM and SDDMM vary widely in performance depending on degree skew, feature width, and GPU micro-architecture, so a single static kernel configuration underperforms across datasets and devices. The paper's goal is operator-local adaptivity: choose the kernel variant per input at runtime, measure candidates on a small induced subgraph, and accept a change only when it beats the baseline under a guardrail. AutoSAGE covers both SpMM and SDDMM and composes them into a CSR-attention pipeline (SDDMM → row-softmax → SpMM).

Key Contributions

1. Per-input micro-probes with guardrail fallback. Kernel-granularity decisions are made by timing shortlisted candidates on an induced subgraph, with a guardrail that the authors claim never regresses versus the baseline under identical inputs.
2. Hub-aware split and op-aware caching. Unified handling of SpMM and SDDMM with CTA-per-hub mapping for heavy rows and persistent replay across the tuple (device, graph signature, F, op).
3. Deployment toggles and telemetry. One-line controls (probe budget, thresholds, vectorization) plus CSV and JSON logs for reproducibility.
4. Released artifacts. CUDA sources, Python bindings, a reproducible harness, and replayable cache logs are released, with a code repository listed as github.com/SV25-22/AutoSAGE.

Main Findings

• Parity on real graphs at large feature widths: On Reddit, at F = 128 the baseline was chosen (38.115 ms baseline vs 38.153 ms chosen, speedup 0.999) and at F = 256 the baseline was chosen (89.801 ms vs 89.778 ms, speedup 1.000).
• Wins at small feature widths on real graphs: On Reddit at F = 64, AutoSAGE was chosen with 16.241 ms baseline vs 15.093 ms chosen (speedup 1.076). On OGBN-Products at F = 64, AutoSAGE was chosen with 27.534 ms baseline vs 26.385 ms chosen (speedup 1.044); at F = 128 the baseline was chosen (56.152 ms vs 56.161 ms, 1.000) and at F = 256 the baseline was chosen (114.550 ms vs 114.385 ms, 1.001).
• Large synthetic speedups: On an Erdős–Rényi graph with N = 200,000 and p = 2×10⁻⁵, AutoSAGE was chosen at every width: F = 64 gave 1.523 ms → 0.323 ms (4.720x), F = 128 gave 1.720 ms → 0.563 ms (3.055x), and F = 256 gave 2.231 ms → 1.067 ms (2.091x).
• Hub-skew wins, with the guardrail declining at wide F: On a hub-skew synthetic with N = 200,000, k = 4, h = 0.15: F = 64 gave 10.190 ms → 8.976 ms (1.135x), F = 128 gave 19.267 ms → 18.074 ms (1.066x), and at F = 256 the baseline was chosen (37.362 ms vs 37.351 ms, 1.000).
• Guardrail sensitivity changes selection: On Reddit with guardrail = 0.98, AutoSAGE was chosen at F = 64 (16.222 ms vs 15.031 ms, 1.079x) and at F = 128 (38.179 ms vs 36.951 ms, 1.033x), while F = 256 stayed on baseline (89.706 ms vs 89.735 ms, 1.000).
• Feature-width sweep on Reddit: AutoSAGE chosen at F = 32 (7.871 ms → 6.614 ms, 1.190x) and F = 64 (16.257 ms → 15.070 ms, 1.079x); baseline chosen at F = 96 (26.486 ms vs 26.584 ms, 0.996), F = 128 (38.482 ms vs 38.074 ms, 1.011), F = 192 (63.256 ms vs 63.328 ms, 0.999), F = 256 (89.855 ms vs 89.540 ms, 1.004), and F = 512 (199.971 ms vs 200.245 ms, 0.999).
• Feature-width sweep on OGBN-Products: AutoSAGE chosen only at F = 32 (13.714 ms → 12.540 ms, 1.094x); baseline chosen at F = 64 (27.572 ms vs 27.551 ms, 1.001), F = 96 (41.806 ms vs 41.802 ms, 1.000), F = 128 (56.170 ms vs 56.202 ms, 0.999), F = 192 (85.265 ms vs 85.204 ms, 1.001), F = 256 (114.585 ms vs 114.441 ms, 1.001), and F = 512 (232.336 ms vs 232.254 ms, 1.000).
• vec4 helps on synthetics, hurts on one real case: On ER (N = 200k, p = 2×10⁻⁵), vec4 improved latency by +4.6%, +20.2%, and +17.1% at F ∈ {64, 128, 256}, reported as speedups of 1.046, 1.202, and 1.171. On Reddit at F = 64, vec4 was slower (−27.5%, speedup 0.725); for larger F the baseline is selected so vec4 is moot.
• CTA-per-hub split under skew: On hub-skewed synthetics at F = 128, split outperformed baseline by 2.2x–3.4x with bests up to approximately 4.6x in sweeps. Specific rows: N = 20k, hub = 5k, other = 64 gave 1.603 ms → 0.717 ms (2.236x); N = 20k, hub = 12k, other = 32 gave 3.205 ms → 0.934 ms (3.432x); sweep bests in two regimes were 2.985 and 4.563.
• Probe overhead is bounded and tunable: With AUTOSAGE_PROBE_FRAC = 0.03 and a 1.0 ms cap, probe overhead was approximately 6–9% of a full-graph iteration at Reddit F = 64; a low-overhead setting (0.02, 0.5 ms) yielded approximately 3–4% with mildly higher variance.
• CSR attention pipeline: With sddmm_csr_auto and a row-softmax forming csr_attention_forward, uncached mode is dominated by probe costs; once cached or replayed the pipeline runs at approximately 200 ms with both sub-operations selecting AutoSAGE on ogbn-products. Per-op probe logs indicate SDDMM at approximately 1.5–1.8x and SpMM at approximately 5.2–5.5x over their baselines in the chosen regimes.
• Non-regression claim: Proposition 1 states that with α ≤ 1 the chosen runtime satisfies t_chosen ≤ t_b, because the candidate is accepted only if t* ≤ α·t_b (default α = 0.95) and otherwise the baseline is used.

Methodology in Plain English

The authors build a set of CUDA kernel templates for sparse matrix operations on CSR data: a warp-per-row SpMM baseline with feature tiling (tile widths starting at 32, 64, 128, and so on), a CTA-per-hub SpMM variant that assigns one thread block to each heavy row (with a split threshold hubT), and a row-wise SDDMM that computes CSR dot products. Both SpMM variants have scalar and vec4 versions, where vec4 requires F mod 4 = 0 and 16-byte alignment. A numerically stable CSR row-softmax is included so that SDDMM, softmax, and SpMM can be chained into CSR attention.

Scheduling happens in three stages. First, the scheduler extracts features (row count, nonzero count, degree quantiles, feature width F, device capabilities) and uses a roofline-style estimate to shortlist candidates. Second, it times the top-k candidates on an induced subgraph — by default 2–3% of rows with a minimum of 512 rows — for n iterations under a wall-time cap. Third, a guardrail compares the best candidate time t* against the baseline time t_b and accepts the candidate only if t* ≤ α·t_b (default α = 0.95); otherwise it falls back to the baseline. Decisions are keyed by (device signature, graph signature, F, op) and stored in a persistent cache, with a replay-only mode for deterministic reproduction.

Experiments run on a single node with 8 NVIDIA A800-SXM4-40GB GPUs (compute capability 8.0, 108 SMs per GPU, 40 GB HBM each, 42.48 GiB reported by CUDA), 2 Intel Xeon Platinum 8378A processors (2 sockets × 32 cores, 128 hardware threads), and 1.0 TiB RAM, running Ubuntu 22.04.5 LTS, Python 3.12.2, PyTorch 2.8.0+cu128 (CUDA build 12.8), cuDNN 9.1.2, and CUDA toolkit 12.4 (nvcc V12.4.131). Timings are on a single GPU unless stated, reported as medians over 10–15 iterations after warm-up. Datasets are Reddit and OGBN-Products, plus synthetic Erdős–Rényi and hub-skew graphs. The SpMM baseline is cuSPARSE via PyTorch and the SDDMM baseline is a gather–dot implementation.

Why This Matters

Impact on research. The paper argues that sparse GNN aggregation performance is input-sensitive and cannot be captured by one static kernel, and it offers a reusable pattern — estimate, micro-probe on an induced subgraph, then gate the choice behind a guardrail — for making runtime kernel decisions without risking regression. It positions itself as complementary to compiler and autotuning work (AutoTVM, Ansor, GraphIt/G2, Triton) and to GNN system work (GNNAdvisor, AdaptGear, TC-GNN, Graphiler), focusing on runtime kernel-level choices over sparse primitives. The released cache logs and replay mode also address reproducibility of scheduling decisions.

Real-world applications (from the operator coverage described in the paper):

• GNN training pipelines where message passing is bounded by sparse aggregation throughput.
• Attention-style operator stacks built from a sparsity pattern, via the CSR-attention pipeline (SDDMM → row-softmax → SpMM).
• Graph workloads with heavy-tailed degree distributions, including recommender-style data, which the paper names explicitly as a setting its external validity may not fully cover.
• Sparse kernels with extreme sparsity, where the synthetic stressors show the largest reported gains.

Industry relevance. The guardrail design means deployments can adopt AutoSAGE without accepting a systematic slowdown relative to vendor baselines under identical inputs and devices, and the persistent cache means the per-input probing cost is paid once and then replayed. The toggles (AUTOSAGE_FTILE, AUTOSAGE_WPB, AUTOSAGE_HUB_T, AUTOSAGE_PROBE_*, AUTOSAGE_CACHE, AUTOSAGE_REPLAY_ONLY) and CSV/JSON telemetry are aimed at practical integration rather than research-only use.

Future Directions

• End-to-end gains. The paper states that wins are largest on synthetics and small-F regimes, and that more end-to-end gains remain future work; the reported results are kernel-level and pipeline-level.
• Wider vectorization and precision. Extensions named in the discussion include vec8 and mixed precision (FP16/BF16 reads with FP32 accumulators).
• Multi-GPU partitioning. The discussion lists multi-GPU partitioning as an extension; all timings reported are on a single GPU unless stated.
• Generalization beyond the tested graphs. Threats to validity note that results on Reddit and OGBN-Products may not generalize to bipartite recommenders, and that a vendor-tuned SDDMM baseline instead of the gather–dot implementation could change the measured gaps.

Target Audience

This paper benefits GPU kernel engineers, sparse linear algebra and GNN systems researchers, and performance engineers who deploy GNN training on CUDA hardware and want per-input scheduling of SpMM and SDDMM. It is also relevant to reproducibility-focused researchers, since the harness, cache logs, and replay mode are part of the release. Readers without background in CUDA execution models, CSR formats, and GNN aggregation will find the implementation-level details hard to follow, as the paper does not report end-to-end model training results or full end-to-end speedups.

Authors’ abstract

Sparse GNN aggregations (CSR SpMM/SDDMM) vary widely in performance with degree skew, feature width, and GPU micro-architecture. We present AutoSAGE, an input-aware CUDA scheduler that chooses tiling and mapping per input using a lightweight estimate refined by on-device micro-probes, with a guardrail that safely falls back to vendor kernels and a persistent cache for deterministic replay. AutoSAGE covers SpMM and SDDMM and composes into a CSR attention pipeline (SDDMM -> row-softmax -> SpMM). On Reddit and OGBN-Products, it matches vendor baselines at bandwidth-bound feature widths and finds gains at small widths; on synthetic sparsity and skew stress tests it achieves up to 4.7x kernel-level speedups. We release CUDA sources, Python bindings, a reproducible harness, and replayable cache logs.