FlashSpread: IO-Aware GPU Simulation of Non-Markovian Epidemic Dynamics via Kernel Fusion

Abstract: Non-Markovian (renewal) epidemic simulation on multi-million-node contact networks is essential for realistic forecasting under general age-dependent holding-time distributions (log-normal, Weibull, Erlang, and similar), but the age-dependent hazard forces dense per-step updates that render the sparse event-queue strategies of standard CPU methods ineffective. We present FlashSpread, a GPU framework that consolidates the per-step renewal pipeline (CSR traversal, numerically stable erfcx-based hazard evaluation, Bernoulli tau-leaping, state transition, and next-step infectivity write-back) into a single fused Triton kernel whose intermediates never leave streaming-multiprocessor registers, with block-scalar skips that preserve CUDA Graph capture and a degree-aware CSR dispatch (thread / warp / edge-merge) that keeps the peak throughput on scale-free graphs. On an NVIDIA A100 the fused CUDA-Graph engine reaches 8.09 Giga-NUPS at N = 10^6 on a uniform-degree graph, a 217x strict hardware speedup over optimised CPU tau-leaping at the same N; on a Barabasi-Albert graph of the same size the merge-based dispatch recovers 4.5x (0.45 to 2.0 Giga-NUPS) over the default kernel, and the framework scales to N = 10^8 on a single A100 (40 GB), with a mixed-precision storage path that extends the L2-reachable scale by roughly 3x and delivers a 1.15x throughput lift at the far bandwidth-bound end. Validation against an exact non-Markovian Gillespie reference shows a structural-bias floor of approximately 6% on peak infection and approximately 7% on final attack rate that does not detectably decrease as epsilon nears 0 across two decades of tolerance, comfortably within typical epidemiological parameter uncertainty. Code: https://github.com/Shakeri-Lab/FlashSpread.

• The paper introduces a novel GPU-optimized framework that combines Markovian and non-Markovian engines through kernel fusion for efficient epidemic simulation.
• It employs adaptive Bernoulli tau-leaping and degree-aware CSR dispatch to significantly boost throughput and scalability on heterogeneous graphs.
• Extensive validation shows that FlashSpread achieves remarkable speedups over CPU baselines while maintaining structural fidelity in simulating complex epidemic models.

FlashSpread: IO-Aware GPU Simulation of Non-Markovian Epidemic Dynamics via Kernel Fusion

Motivation and Context

Traditional epidemiological models often rely on memoryless Markovian assumptions where transitions (e.g., infection, recovery) between states occur with rates independent of the duration already spent in a given state. This assumption is mathematically convenient but biologically inaccurate, since empirical data overwhelmingly supports age-dependent (non-Markovian) distributions such as log-normal, Weibull, or Erlang for epidemic dwell times. Non-Markovian epidemic simulation on large-scale networks is thus essential for realistic forecasting yet computationally challenging: the hazard rates become continuously time-dependent, requiring dense per-step updates and prohibiting efficient sparse event-driven algorithms typically used in CPU implementations.

FlashSpread addresses this computational challenge by providing a unified, GPU-optimized framework capable of simulating both Markovian and non-Markovian spreading processes on networks, exploiting dense parallelism and advanced kernel fusion to achieve unmatched throughput.

Algorithmic Architecture

FlashSpread is architected around two computational engines: a Markovian (memoryless) engine optimized for sparse, event-driven updates and a Renewal (non-Markovian) engine optimized for dense, age-dependent hazard re-evaluation. The system automatically selects the optimal traversal and kernel launch strategy based on input graph topology, specifically leveraging the degree heterogeneity as measured by $D_{\max}/D_{\mathrm{avg}}$.

Kernel Fusion

Central to FlashSpread is the consolidation of the entire renewal simulation pipeline—CSR graph traversal, numerically stable $\mathrm{erfcx}$-based hazard evaluation, adaptive Bernoulli tau-leaping, state transition, and infectivity update—into a single fused Triton kernel. This design ensures that all intermediates reside in streaming-multiprocessor registers, minimizing off-chip memory traffic and maximizing arithmetic intensity. Block-scalar skips are used to exploit the compute sparsity inherent in epidemic dynamics while maintaining compatibility with CUDA Graph capture for efficient batched computation.

Degree-Aware CSR Dispatch

For highly heterogeneous graphs (e.g., scale-free networks), traditional one-thread-per-node traversal loses efficiency due to hub-induced warp divergence. FlashSpread implements warp-per-node and edge-partitioned merge-based kernels, auto-selecting the optimal strategy based on degree distribution thresholds, recovering up to $4.5\times$ throughput in adverse scenarios (Figure 1).

Figure 1: Active-node compaction: NUPS$(t)$ stratified into ten temporal windows of the $\mathrm{TF}=50$ benchmark, illustrating topology-dependent speedup.

Parallel Tau-Leaping and Fidelity

FlashSpread employs Bernoulli tau-leaping, an integrator where each node can transition at most once per step, naturally matching compartmental epidemic dynamics. Step size is dynamically adapted to bound transition probabilities by a user-selected tolerance parameter $\varepsilon$. Validation against exact non-Markovian Gillespie methods demonstrates a structural bias floor of $\sim$6\% in peak infection and $\sim$7\% in final attack rate that saturates as $\varepsilon \to 0$, well within typical epidemiological parameter uncertainties.

Figure 2: Fidelity validation. Fused GPU kernel (solid), fused CUDA Graph (dashed), and exact non-Markovian Gillespie (dotted) benchmark trajectories.

Performance Evaluation

Extensive benchmarks on NVIDIA A100 GPUs using regular and scale-free graphs show that FlashSpread's fused renewal engine achieves:

• $8.09$ Giga-NUPS at $\mathrm{erfcx}$0, representing a $\mathrm{erfcx}$1 speedup over optimized CPU tau-leaping baselines using identical tau-leaping algorithms and vectorized PyTorch operations.
• Scalability up to $\mathrm{erfcx}$2 nodes on a single GPU using mixed-precision storage (state, age, infectivity, weights), extending the memory-limited tractable scale by $\mathrm{erfcx}$3 (Figure 3).

  Figure 3: Wall-clock–fidelity Pareto, plotting computational efficiency versus peak infection and attack rate bias relative to exact reference.

On scale-free topologies, edge-partitioned merge kernels recover up to $\mathrm{erfcx}$4 Giga-NUPS from $\mathrm{erfcx}$5 Giga-NUPS under naive traversals, mitigated by the auto-dispatch mechanism.

Structural Fidelity and Error Analysis

Empirical sweeps of tau-leaping tolerance reveal a flat error profile: peak infection and final attack rate discrepancies relative to exact reference are unchanged across two decades of $\mathrm{erfcx}$6, confirming a structural bias inherent to synchronous tau-leaping updates. Discretization error is negligible for $\mathrm{erfcx}$7, with self-consistency errors decaying monotonically but remaining invisible in biologically relevant metrics.

Figure 4: Error decomposition for Bernoulli tau-leaping across tolerance values, showing convergence of discretization error and structural bias floor saturation.

Practical Implications

FlashSpread enables the simulation of non-Markovian spreading processes at previously unattainable scales for ensemble forecasting, policy evaluation, and reinforcement learning in epidemic control. The framework is validated across SIS, SIR, and SEIR models (Figure 5, Figure 6), demonstrating tight agreement with exact references for both Markovian and renewal regimes.

Figure 5: FlashSpread Markovian engine versus exact Doob–Gillespie on SIS/SIR, demonstrating tau-leaping fidelity across sparse and dense event regimes.

Figure 6: Ensemble-mean SEIR trajectories overlaid on exact non-Markovian Gillespie reference with tau-leaping tolerance variation.

Future Directions

Scaling beyond $\mathrm{erfcx}$8 will require multi-GPU domain decomposition. The current implementation assumes static network topology; temporal networks and dose-dependent dynamics will necessitate explicit per-edge age tracking and memory optimizations to retain tractability. Further hardware-aware optimizations, such as intra-warp segmented reductions for high-contention hubs, are noted as promising but currently omitted due to portability constraints.

Conclusion

FlashSpread demonstrates that IO-aware GPU kernel fusion, coupled with degree-sensitive traversal and adaptive tau-leaping, enables biologically realistic non-Markovian simulation at extreme scales. The framework’s validation against exact references and robust handling of complex network topologies positions it as a rigorous foundation for high-performance computational epidemiology. The modular, user-controllable architecture offers both practical throughput and theoretical fidelity for advanced researchers in network epidemiology and stochastic modeling.