Programmable switches can already sum data streams in flight, but only for distributed training's neatly aligned gradient tensors — real key-value streams (MapReduce shuffles, stream analytics) are asynchronous, with unforeseeable keys and unknowable completion, which breaks every existing in-network design's reliability and memory management — so ASK turns the switch into a best-effort key-value aggregation service, packing 32 tuples per packet across a two-dimensional aggregator layout, deduplicating retransmitted partially-aggregated packets with a per-tuple bitmap window, and periodically swapping shadow copies of switch memory so hot keys can evict squatting cold ones — up to 155× faster aggregation than Spark, 3–5× faster big-data jobs at 88.7% less CPU, while staying compatible with training workloads.
The argument in one glance
-
Problem
Merging key-value streams eats clusters alive
Summing values by key across machines is everywhere: reduce() in MapReduce, aggregations in stream processing, SUM() in databases, AllReduce in training. It is also brutally expensive — the ReduceByKey step takes 94.67% of a WordCount job's time, and gradient aggregation up to 79% of distributed training. Every tuple crosses the network to a receiver whose CPU then does nothing but add numbers.
-
What people did
Let the switch do the adding
Programmable switches process packets at terabits per second and can run small user programs on every packet. In-network aggregation (INA) systems — SwitchML, ATP, SHARP, NetReduce — have the switch sum gradient streams as they pass through, so one merged stream exits instead of N entering. For distributed training this works beautifully.
-
Why it fell short
Training streams are the easy special case
Gradient tensors are VALUE streams: every sender's data is a vector with linear, contiguous, aligned indices, and senders move in lockstep. That makes aggregation SYNCHRONOUS — the switch knows exactly when slot i has heard from all senders, so it can emit the result and recycle the memory cell immediately. Key-value streams offer none of this: keys arrive unordered and unforeseeable, senders are unsynchronized, and the switch can never know a key's last appearance. Consequences cascade — memory cells can never be recycled mid-task, the 1-bit-per-aggregator deduplication of training INA needs unbounded state, one-tuple-per-packet strawmen cap useful throughput at 9.76 of 100 Gb/s, and under real Zipf-skewed workloads a rare key that arrives early squats on scarce switch memory for the entire job.
-
This paper's idea
A best-effort switch service, with the host as the safety net
ASK (from Tsinghua, Peking University, Intel's Barefoot division, Wisconsin, and Harvard) co-designs switch and host so the switch aggregates what it can and forwards the rest — correctness never depends on the switch succeeding. Three mechanisms: multi-key packets vectorized across a two-dimensional aggregator-array layout (32 tuples per packet, with sender-side key-space partitioning so each key always lands in the same array); a lightweight reliability window where the switch tracks per-packet, per-tuple aggregation state in a compact bitmap so retransmitted partially-aggregated packets are deduplicated tuple by tuple; and shadow-copy swapping that periodically resets half the switch memory so hot keys statistically reclaim cells from squatting cold keys. Results on a Tofino testbed: up to 155× Spark's aggregation throughput, big-data jobs 3–5× faster with 88.7% less CPU, per-sender throughput flat as senders scale to 8, and training workloads supported at ATP/SwitchML speed.
Why this problem is worth solving
A distributed job’s aggregation phase is a strange kind of
waste: the network hauls every tuple to one machine, and that
machine’s expensive general-purpose cores then perform the
world’s simplest computation — addition — as fast as they can,
which is not very fast. The paper’s strawman experiment makes
the imbalance vivid: vanilla Spark needs 56 cores to hit its
peak aggregation throughput, while a
a datacenter switch (here an
Intel Tofino) whose packet pipeline is programmable in the P4
language. Packets flow through a sequence of match-action
stages, each with its own small SRAM; a program can read
header fields, do arithmetic, and update per-stage state at
full line rate — 25.6 Tb/s across ports on Tofino3. The
constraints are severe: roughly 15 MB of SRAM per pipeline,
memory isolated per stage, one sequential pass per packet, and
each register array readable and writable exactly once per
pass
using the switch to merge N
inbound data streams into one — each memory cell (an
“aggregator”) accumulates values for one key or index as
packets traverse, and only the merged result travels onward.
The network’s traffic drops by nearly the fan-in factor, and
the receiving host’s CPU is freed from the merge
entirely
What people did before, and where it broke
- SwitchML, ATP, SHARP, NetReduce (the last covered elsewhere on this site). In-network aggregation for training’s gradient tensors. Their designs lean on synchrony everywhere: aggregators recycle immediately on completion, and reliability needs just one bit per aggregator because each slot’s set of expected packets is known. Asynchronous key-value streams void both — completion is unknowable, and the per-key appearance state is unbounded.
- Key-value showcases (P4COM and workshop demos). Proofs that a switch can match and sum keys, without the system around it: no application interface, no correctness story under packet loss, no answer for keys that exceed switch memory or packets that carry a single 8-byte tuple against 78 bytes of headers.
- Middlebox/server-based aggregation (NetAgg, Camdoop, parameter servers). Aggregation on dedicated boxes along the path — real machines with real memory, but a fraction of a switch ASIC’s throughput, and more hardware to operate.
- Faster plumbing (SparkRDMA, shared-memory shuffle). The paper’s own baseline study lands the punch: after mappers pre-aggregate, intermediate data is small, so accelerating network and disk I/O barely moves job completion time. The bottleneck is the aggregation computation — which is exactly what these approaches do not touch.
How ASK works
A switch-host co-design, best-effort by construction
-
Service architecture new
Each server runs an ASK daemon (DPDK, kernel-bypass) with one control channel and several persistent data channels that many applications and tasks multiplex; apps exchange tuples via shared memory. The switch holds a pool of aggregator arrays; a controller leases array regions to tasks. Aggregation is BEST-EFFORT: the switch aggregates what it can and forwards the remainder to the receiver host, which finishes the job — correctness never depends on switch success.
-
Multi-key packets, vectorized (goodput fix) new
One tuple per packet means 8 useful bytes under 78 bytes of overhead — 9.76 Gb/s of goodput on a 100 Gb/s wire. ASK packs 32 tuples per packet. PISA allows each register array one read-write per pass but four arrays per stage, so ASK builds a TWO-DIMENSIONAL layout: 32 aggregator arrays (AAs), the i-th AA processing the i-th tuple slot of every packet. A header bitmap tracks which tuples remain unaggregated.
-
Sender-assisted addressing new
If the same key landed in different slots across packets it would reserve cells in several AAs. So senders partition the key space into 32 subspaces by hash and always place a key's tuples in its subspace's slot — one key, one AA, forever. Variable-length keys get coalesced placement: adjacent AAs are grouped and indexed by ONE hash of the whole key (hashing segments independently produces false matches when two long keys' fragments recombine); very long keys skip the switch entirely.
-
Reliability for partial aggregation (correctness fix) new
A flow now has three endpoints — sender, switch, receiver — and a retransmitted partially-aggregated packet must not re-aggregate its already-consumed tuples. Sender: sliding window of 256 packets, 100-microsecond retransmit timeout (ACKs arrive from both switch and receiver, so out-of-order ACKs cannot signal loss). Switch: a per-channel "seen" bitmap detects duplicates — compressed to W bits via an even/odd-segment parity trick where one atomic set/clear instruction simultaneously records, reports, and pre-initializes state — plus a PktState array storing each in-flight packet's per-tuple bitmap, restored onto retransmissions so only genuinely unaggregated tuples proceed. Total: 1056 bytes per data channel; 264 KB serves 64 servers.
-
Shadow-copy hot-key prioritization (memory fix) new
Aggregator cells go first-come-first-served, and Zipf-skewed reality means an early cold key can squat on a cell for the whole task. ASK splits each AA into two copies: the switch writes into the active copy while the receiver drains and resets the other; a one-bit indicator flips on a receiver-initiated notification (atomic because PISA processes one packet per stage at a time). Each swap gives every key a fresh chance at an empty cell — and statistically, frequent keys win the race. LRU eviction was rejected: a switch cannot retry a failed eviction to an unreliable network.
-
Application plugins
Spark plugin (1800 lines of Java) and BytePS plugin (500 lines of C++) convert application data to ASK tuples — big data and distributed training ride the same service. Switch program: ~5000 lines of P4; host stack: ~4500 lines of DPDK C.
The idea
Asynchrony is the real problem; best-effort is the answer that fits a switch. Training INA could promise complete aggregation because synchrony told the switch when it was done. ASK abandons the promise instead of faking the synchrony: the switch aggregates what its memory and hashing happen to catch, marks each tuple’s fate in a header bitmap, and forwards leftovers to the receiver host, whose merge is always correct. This one design choice absorbs every ugly case — hash collisions, full aggregator arrays, oversized keys, memory pressure — as a graceful performance loss rather than a correctness failure. And on production traces the best effort turns out to be excellent: the switch absorbs 85.7–94.3% of tuples and 72.0–90.4% of packets, which never even reach the receiver’s link.
Let the sender do what the switch cannot. PISA’s restrictions — one read-write per register array per pass, no cross-stage memory — make runtime key steering inside the switch impossible. ASK’s response is to move the intelligence to the endpoints: senders partition the key space and build packets so that each tuple arrives pre-addressed to the one aggregator array that will ever hold its key; receivers initiate the shadow-copy swaps and the eviction that the switch’s programming model cannot express (an eviction needs retries on an unreliable network, and a switch cannot suspend a packet to retry). The switch is left doing only what it does at line rate: hash, compare, add, and flip bits.
Bound state by the window, not by the data. The fatal flaw of porting training-INA reliability was per-key state for unbounded keys. ASK’s insight is that reliability never needed to track keys at all — only packets in flight, of which the sliding window permits exactly 256 per channel. The per-tuple aggregation bitmaps (PktState) plus the duplicate-detection bits (seen, halved to W bits by a parity trick where set-bit on even window laps and clear-bit on odd laps lets one atomic instruction record the current lap and pre-arm the next) total 1056 bytes per channel — fixed, tiny, and independent of how many keys or tuples the stream carries.
Two packets through the switch
A WordCount shuffle, tuple by tuple
AA = aggregator array (one per packet slot; 32,768 cells each). Bitmap = per-packet header bits marking unaggregated tuples. seen / PktState = the switch's per-channel duplicate and aggregation-state windows. Cell layout: half the register holds the key (kPart), half the running sum (vPart).
Step 1 of 6
The sender builds a multi-key packet
The mapper's tuples ("hello",1) and ("world",1) hash into key subspaces 0 and 1, so the daemon places them in slots 0 and 1 of one packet, sets both bitmap bits, stamps the task ID and sequence number, and sends. Every future "hello" tuple, from any packet, will ride slot 0 — one key, one aggregator array, always.
goodput scales almost linearly up to 32 tuples/packet even the most skewed production trace (yelp) averages 16.91 filled slots
Why it works
- The switch was already touching every packet at line rate → adding hash-compare-add logic costs no wall-clock time → aggregation becomes a side effect of forwarding, and each fully-absorbed packet is traffic the receiver’s link and CPU never see (72–90% of packets on production traces).
- Best-effort semantics push all hard cases to the receiver host → collisions, overflow, and long keys degrade throughput but never correctness → the switch program stays simple enough to fit PISA’s one-access-per-array, stage-isolated constraints.
- Sender-side key-space partitioning fixes each key’s slot and array before the packet is built → no in-switch steering state, no single-key-multiple-cells waste, and a whole-key hash for coalesced long keys eliminates fragment-recombination false matches.
- Reliability state scales with the window, not the stream → 256 in-flight packets × per-tuple bitmaps = 1056 bytes per channel regardless of key count → the unbounded-state problem that voided synchronous-INA reliability vanishes by accounting for packets instead of keys.
- Periodic shadow swaps re-run the race for cells → frequency wins statistically, so hot keys accumulate in-switch without the switch ever knowing which keys are hot → aggregator memory 16× smaller than the key set still captures 95.85% of tuples.
What it costs
- Per-packet overhead and host cores. 32 small tuples per packet still carry ASK headers: peak goodput is 73.96 Gb/s versus 91.75 Gb/s for plain MTU-sized transfer, and saturating the NIC takes 4 cores versus 2. The claim is that the aggregation offload repays this — and the CPU numbers (7.14% versus up to 100%) back it.
- Switch memory is leased, not infinite. Aggregation tasks reserve aggregator-array regions; concurrent tenants share ~15 MB of SRAM. Long keys bypass the switch entirely, and the medium-key group sizing (m=2, k=8) is an empirical fit to studied datasets.
- A three-endpoint transport of its own. ASK replaces TCP with a custom windowed protocol over DPDK — fine-grained timeouts, switch-generated ACKs — which must coexist with datacenter congestion control (the paper argues compatibility with ECN/loss-based INA schemes, bounded by the reliability window).
- Single-rack scope for the strong version. Extending through multiple switch hops would multiply per-channel state; the proposed deployment keeps ASK on top-of-rack switches, with cross-rack traffic bypassing to the receiver host.
- Trust in specific hardware tricks. The compact seen window leans on Tofino’s atomic set/clear-and-return instructions, and shadow-swap atomicity leans on PISA’s one-packet-per-stage execution — portable in spirit, but verified on one ASIC family.
The evaluation
Real hardware end to end: a 32-port Tofino switch and nine 100 Gb/s servers, production key-value traces (Yelp reviews, newsgroup text, blog corpus, movie reviews) alongside synthetic uniform and Zipf streams, and three integration levels — microbenchmarks, Spark WordCount via HiBench, and BytePS distributed training against ATP and SwitchML.
Where the aggregation happens
| Property | Host-only (Spark / PreAggr, plus RDMA and shared-memory variants) | Training-oriented INA (SwitchML / ATP / NetReduce) | ASK |
|---|---|---|---|
| What it can aggregate | |||
| Handling unreliable networks | TCP end-to-end | ||
| Use of scarce switch memory under key skew | Not applicable | Not applicable — value-stream slots recycle on completion | |
| Aggregation throughput / job impact | Large training speedups (in their domain) | ||
| Scalability with senders | Scales for training topologies |
Check yourself
Check your understanding
1. A packet carrying ("goodbye",1) and ("frost",1) is partially aggregated — the switch consumed "goodbye" but a collision bounced "frost" — and then the packet is lost before reaching the receiver. The sender retransmits. Predict what must NOT happen, and how ASK prevents it.
2. Training-INA systems handle packet loss with one bit per aggregator. Why does that mechanism collapse for key-value streams, and what does ASK track instead?
3. For keys longer than one register, a naive scheme hashes each key segment independently into its own array and aggregates when all segments match. ASK rejects this. What goes wrong, and what is the fix?
4. Under Zipf skew, an early-arriving cold key can hold an aggregator cell for a task's entire lifetime. The obvious fix is an LRU-style eviction policy in the switch. Why is that impossible, and how does the shadow copy achieve the same end?
Sources
- ACM DL: A Generic Service to Provide In-Network Aggregation for Key-Value Streams — ASPLOS 2023 (open access)