Stream-K++ Scheduling: Unified Adaptive Strategies

• Stream-K++ Scheduling is a unified concept that combines heterogeneous coded scheduling for distributed clusters with adaptive GPU kernel tuning for GEMM, optimizing job latency and throughput.
• The approach leverages convex optimization and Bloom filter–based policy pruning to partition and map computational workloads efficiently in high-dimensional resource spaces.
• Empirical results show up to 43% GFLOPS gains on GPUs and near-optimal delay performance in clusters, demonstrating the strategy’s practical scalability and efficiency.

Stream-K++ Scheduling is a unified term for two distinct, state-of-the-art resource management frameworks targeting high-throughput parallel computing: a joint coded scheduling system for heterogeneous distributed clusters (Cohen et al., 2021), and an adaptive kernel selection methodology for GPU-based general matrix multiplication (GEMM) (Sadasivan et al., 2024). Despite different architectural foci, both share the principle of adaptively partitioning and mapping computational workloads—leveraging heterogeneity-aware scheduling policies, modern coding techniques, and rapid selection/pruning mechanisms—to achieve near-optimal job latency or kernel throughput in highly variable, large-scale environments.

1. Distributed Coded Computing: Stream-K++ Scheduling Framework

The distributed Stream-K++ framework addresses the challenge of stream job execution across heterogeneous clusters, characterized by stochastic job arrivals and variable worker compute/communication capabilities (Cohen et al., 2021). Jobs arrive at the master node according to a Poisson process, with per-step arrival rate $\lambda$. Each job $j$ is a non-iterative function $f(X_1(j), ..., X_m(j))$, potentially encoded via an $(K\Omega,K)$ MDS code that expands into $K\Omega$ tasks, any $K$ of which suffice for decoding.

Worker nodes $p$ are associated with distinct mean job times $E[T_p]$, communication rates $c_p$, and maintain independent M/G/1 queues with stability per $\lambda \phi_p \le \mu_p = 1/E[T_p]$. The core decision variable is the job split $j$0 ($j$1, $j$2), denoting the proportion of work routed to each worker. The system enforces a "no‐bottleneck" constraint on per-worker communication and encoding/decoding throughput, relative to $j$3.

The scheduling-coding interaction is formalized as an optimization of average in-order execution delay,

$j$4

with $j$5 governed by Pollaczek–Khinchin queueing, and $j$6 reflecting aggregate communication times. The joint selection of the code family $j$7 (e.g., PolyDot) and optimal worker subset $j$8, coupled with closed-form convex optimization for $j$9, constitutes the Stream-K++ approach.

2. GPU Kernel Autotuning: Stream-K++ for GEMM Scheduling

On modern GPUs, Stream-K++ refers to an adaptive methodology for GEMM kernel scheduling and selection, implemented in the Open-sieve C++ library and integrated with the AMD Composable Kernel framework (Sadasivan et al., 2024). The fundamental objective is to maximize GFLOPS for a triple $f(X_1(j), ..., X_m(j))$0 of matrix sizes by adaptively choosing from a pool of tightly parameterized scheduling policies, each orchestrating the mix between Stream-K (persistent-kernel, atomic-enabled) and traditional data-parallel (DP) execution modes.

Stream-K++ extends the original three Stream-K schedules to seven, indexed by $f(X_1(j), ..., X_m(j))$1; in addition, the “all_StreamK” policy ($f(X_1(j), ..., X_m(j))$2) disables DP entirely. Each policy executes $f(X_1(j), ..., X_m(j))$3 batches of Stream-K iterations before transitioning to DP for the remainder, with concrete mapping determined at the granularity of workgroups. The best policy is non-trivial to predict for each tuple $f(X_1(j), ..., X_m(j))$4, necessitating rapid policy selection.

3. Policy Pruning and Selection via Bloom Filters

Stream-K++ leverages a set of Bloom filters—one per scheduling policy—to rapidly eliminate configurations unsuitable for a given GEMM problem size prior to benchmarking or cost-model evaluation. Filters are constructed offline using empirical wins across $f(X_1(j), ..., X_m(j))$5 fp16 GEMM sizes, with a false-positive rate $f(X_1(j), ..., X_m(j))$6 and 100% true-negative guarantee. Runtime selection entails, for each policy $f(X_1(j), ..., X_m(j))$7:

• If $f(X_1(j), ..., X_m(j))$8 is false, skip $f(X_1(j), ..., X_m(j))$9 (safe negative).
• If true, perform empirical timing (e.g., using ckProfiler) and select $(K\Omega,K)$0 minimizing observed latency.

This mechanism prunes approximately $(K\Omega,K)$1 of policy candidates on average, resulting in $(K\Omega,K)$2 μs overhead per selection and a tuning speedup approaching $(K\Omega,K)$3.

4. Algorithmic Outline and Optimization

Distributed Cluster Scheduling (Cohen et al., 2021)

1. Worker and Code Selection
   • For each candidate code in $(K\Omega,K)$4, assemble a valid worker pool $(K\Omega,K)$5 meeting the redundancy criterion $(K\Omega,K)$6.
2. Optimal Load Split
   • Solve the convex program:

   $(K\Omega,K)$7

   where $(K\Omega,K)$8, $(K\Omega,K)$9. - The optimal split $K\Omega$0 is derived in closed form:

   $K\Omega$1

   where $K\Omega$2, $K\Omega$3 enforces $K\Omega$4.

GPU Scheduling Policy Selection (Sadasivan et al., 2024)

1. Host-Side Selection

   • For each unpruned policy $K\Omega$5, measure kernel execution time $K\Omega$6 (optional: use a cost model).
   • Select $K\Omega$7.
2. Device Execution
   • Use persistent kernel parameterized by $K\Omega$8 for the launch.
   • Each workgroup dynamically interleaves Stream-K and DP iterations as dictated by $K\Omega$9.

5. Performance and Empirical Insights

Distributed Clusters

• Simulation with $K$0, worker compute and communication uniformly distributed, and PolyDot codes at moderate redundancy $K$1 shows Stream-K++ achieves mean delay within $K$2 of a non-causal, idealized scheduler.
• Uniform splitting of load is substantially suboptimal, exhibiting 3–5$K$3 greater mean delay and instability (queue divergence) for $K$4 near unity.
• The optimal split exploits heterogeneity by heavily loading the fastest nodes.

GPU Kernel Scheduling

• Benchmarked on AMD Instinct MI250X, Stream-K++ attains up to $K$5 GFLOPS gain over the best pure-DP kernel for certain $K$6.
• Pure DP ($K$7) wins in $K$8 of cases at zero tolerance; permitting $K$9 slowdown, Stream-K policies are within that envelope for $p$0 of sizes.
• For cases where Stream-K policies prevail, median gains approach $p$1, with maximal outliers at $p$2.
• Bloom-filter selection ensures minimal policy search without compromising optimality for observed problem configurations.

6. Flexibility, Extensibility, and Future Directions

Both instantiations of Stream-K++ are architected for extension. In the GPU schedule-selector, incorporating new $p$3 sizes, precision types, or additional tuning parameters (such as tile sizes or vector widths) requires only rebuilding the corresponding Bloom filters and header-only library; additional policies (e.g., “7_batch_then_DP”) are accommodated by training and addition of new filters. The distributed framework similarly generalizes over code families and worker selection, adapting to new hardware capabilities or evolving job characteristics.

A plausible implication is that O(1) policy pruning via filter-based lookup will remain beneficial as tuning parameter spaces expand, supporting future GPU autotuning needs and distributed computing environments with more intricate heterogeneity profiles. For both domains, the Stream-K++ approach exemplifies unified, data-driven, and low-overhead adaptivity in high-dimensional resource allocation.

---

References:

• "Stream Distributed Coded Computing" (Cohen et al., 2021)
• "Stream-K++: Adaptive GPU GEMM Kernel Scheduling and Selection using Bloom Filters" (Sadasivan et al., 2024)