How Gemini 4 Argon Upleveled Agents on Antigravity: Trillion-Edge Graphs and Custom TPU Kernels

Recently, we introduced Gemini 4 Argon in Google Antigravity, showing how Argon’s 1-million-token context window, tunable reasoning, and native sub-agent orchestration upleveled Teamwork—our structured multi-agent framework—to optimize datacenter memory and rewrite legacy C++ libraries into safe Rust. Today, we are sharing how Google’s Graph Mining and Graph Learning teams used Gemini 4 Argon with Teamwork to solve a different class of internal engineering bottleneck: compressing months of low-level RDMA networking and custom TPU kernel development into hours.

In just days of agentic collaboration, a small team of researchers and engineers used Teamwork powered by Gemini 4 Argon to build and deploy two foundational capabilities inside our internal Distributed Graph Flow stack:

• Planet-Scale Distributed Graph Mining (FarArray): Compressing an estimated one-year distributed systems roadmap into a 5-hour gated engineering sprint, the team built FarArray, a hardware-accelerated Partitioned Global Address Space (PGAS) over RDMA. FarArray scales lock-free shared-memory graph algorithms to trillion-edge graphs, delivering 60–90× higher throughput and 40–180× faster runtimes than TCP key-value baselines, over 100× lower cluster DRAM than distributed shuffle pipelines, and 11–27× faster Personalized PageRank embeddings (InstantEmbedding) with up to 56× lower cumulative memory.
• Custom O(N)-Memory TPU Kernels for Sparse GNNs: Bridging the gap between dense-matrix TPUs and sparse graph message passing, Teamwork authored and verified custom Pallas Sparse Matrix Multiplication (SpMM) TPU kernels that reduce High-Bandwidth Memory (HBM) usage by up to 185× (111.5 MiB flat vs. 20.2 GiB) and accelerate forward and backward GNN passes by 8–15× on TPU v6e with exact numerical parity.

How Gemini 4 Argon Upleveled Teamwork for Low-Level Systems Co-Design

Modern graph workloads sit at an intersection of distributed systems and hardware acceleration. Scaling a graph algorithm past a single machine’s RAM requires not only careful algorithmic design but also deep distributed systems engineering. Accelerating a Graph Neural Network (GNN) on a TPU requires hand-writing Pallas kernels that manage on-chip scratchpad memory (VMEM/SMEM), vector register alignment, and asynchronous DMA transfers. Because both tasks demand rare, cross-stack expertise, promising research ideas routinely wait months or years to reach production scale.

With Gemini 4 Argon, Teamwork crossed a capability threshold on these deep systems tasks: specialized sub-agents—spanning C++ RDMA verbs, Linux io_uring, lock-free graph algorithms, and Pallas/XLA compiler lowering—reliably executed end-to-end implementation, testing, and verification across strict quality gates.

Crucially, human researchers stayed in control of the architecture and acceptance criteria throughout: humans defined the system interfaces, memory layouts, and zero-mock verification harnesses, while Teamwork’s proposer and adversarial verifier agents handled the low-level implementation, debugging, and performance tuning.

1. Scale: From Single-Machine Graph Mining to 1.14 Trillion Edges with FarArray

The Bottleneck: The Distributed Graph Tax

Shared-memory multicore libraries like our open-source Graph Based Benchmark Suite (GBBS) process billions of edges in seconds using lock-free atomic operations on dense arrays—until the graph outgrows a single server. When scaling past a single host, pipelines traditionally rely on distributed batch frameworks like MapReduce, often paired with low-latency Remote Key-Value (RKV) systems like CliqueMap, or TCP key-value stores to accelerate lookups.

However, generic key-value stores surround each vertex state with heavy per-key metadata, hash table structures, and serialization framing. For billion-scale graphs, this multiplies cluster memory requirements and saturates server CPUs with kernel network interrupts and RPC overhead. Achieving single-machine efficiency across a cluster requires executing atomic operations directly in remote physical memory at wire speed—bypassing the OS kernel so host CPUs remain dedicated to graph computation.

The Teamwork Sprint: Building FarArray in Hours

To eliminate that tax, the team designed FarArray: a low-latency, kernel-bypass Partitioned Global Address Space (PGAS) inspired by the SNAP microkernel that exposes pinned DRAM across a cluster as a flat, hardware-addressable array.

Using Teamwork in Google Antigravity, the team specified a milestone gated pipeline (Transport Abstraction → Core Pinned Arena → Batched Doorbells → Distributed Algorithms → Live Audit). Specialized Gemini 4 Argon agents implemented pluggable backends across SmartNIC RDMA, Linux io_uring, and POSIX shared memory:

• Zero-Mock Adversarial Verification: An independent auditor agent enforced a strict “Zero-Facade” gate—forbidding mocks or simplified stubs. Within hours, the agent team delivered reviewable change lists passing end-to-end tests across functional, scale, stress, and failure-recovery tiers.
• Zero Remote CPU Overhead: When backed by RDMA-capable SmartNICs, FarArray executes remote 64-bit reads, writes, and atomic CompareAndSwap (CAS) operations directly in NIC hardware via PCIe AtomicOps in microseconds — leaving target host application CPUs idle while TCP key-value server CPUs saturate at 100%.
• Closing the Loop on Tail Latency: When cluster benchmarks surfaced tail-latency spikes under heavy fan-in, Teamwork analyzed hardware telemetry, diagnosed and tuned FarArray—cutting maximum read latency spikes by 89%.

Planet-Scale Results: Connected Components and InstantEmbedding

By combining FarArray’s hardware atomics with thread-local Lowest Common Ancestor (LCA) early-exit pruning—which resolves 97%–99.9% of redundant graph edges in local worker RAM before touching the network—the team scaled shared-memory graph algorithms to planet-scale datasets:

Workload & Benchmark	Graph Scale	Key Result with FarArray
Connected Components (ClueWeb12 & HyperLink 2012)	Up to 3.56B vertices, 113.6B–128.7B edges	60× and 90× higher throughput, 40× (ClueWeb12) and 180× (HyperLink 2012) faster wall-clock time, and 6×–22× less RAM than TCP key-value baselines (>100× less DRAM than distributed shuffle pipelines: 28.5 GB vs. 4.0 TB), with 100% exact partition parity against single-machine GBBS.
Planet-Scale Connected Components	18.0B vertices, 1.14T unique (2.27T symmetric) edges	Completed using only 148 GB of total pinned DRAM across 64 targets (2.31 GB/node, 0 OOMs where >1 TiB TCP KV and >20 TB shuffle baselines are infeasible).
Distributed InstantEmbedding (Personalized PageRank)	955M–3.56B vertices, up to 98.0B directed edges	Storing quantized embeddings as globally addressable mutable state in FarArray reduced peak concurrent RAM by 8.7–13.9×, cumulative memory by up to 55.9×, and 8-hop iteration wall time by 11.3–26.5× compared to MapReduce + Remote KV pipelines.

2. Speed: Breaking the TPU Memory Wall for Sparse GNNs

The Bottleneck: Sparse Message Passing on Dense Accelerators

Google Cloud TPUs are custom-engineered for the dense matrix multiplications (MXU) and contiguous memory access patterns of Transformers. GNN message passing (Y = A @ X), by contrast, aggregates node embeddings across sparse, irregular edge lists (row_ids, col_ids). In standard JAX (jax.ops.segment_sum), gathering features onto edges materializes an O(|E| × D) intermediate tensor in High-Bandwidth Memory (HBM)—consuming 20.2 GiB of temporary HBM at 20 million edges and exhausting a single TPU v6e device.

The Teamwork Sprint: Custom O(N)-Memory Pallas TPU Kernels

Eliminating that O(|E| × D) intermediate tensor requires custom low-level Pallas TPU kernels that never materialize edge-sized feature tensors in HBM. Using Teamwork with Gemini 4 Argon, researchers dispatched an iterative kernel-design and verification loop that co-designed a drop-in message passing engine for GIN and GCN layers:

• Streaming COO Edges via SMEM with Scoped VMEM: Rather than converting sparse graphs to CSR on the host, the Teamwork-authored kernel streams raw Coordinate (COO) (row_ids, col_ids) blocks (4,096 edges/block) directly through TPU Scalar Memory (SMEM) while expanding XLA’s scoped Vector Memory (VMEM) allocation from 16 MiB to 122 MiB—holding dense node feature and accumulator tables entirely on-chip.
• Hardware-Aligned 3D VREG Tiling & Forwarding: To maximize vector unit throughput on TPU v6e, the kernel reshapes 2D node features into hardware-aligned (8, 128) 3D vector register tiles and uses 4-way loop unrolling with software store-to-load register forwarding to avoid VMEM pipeline stalls on duplicate destination nodes.
• Zero-Float-Residual Backward Pass (custom_vjp): For backpropagation (dX = A^T @ dY), the custom VJP saves zero floating-point activation residuals from the forward pass—retaining only integer edge indices and swapping row_ids and col_ids to preserve strict O(N) HBM scaling during training.

The Results: Up to 185× Less HBM and 8–15× Faster GNN Passes

Evaluated on TPU v6e across progressive edge-scaling sweeps (50K nodes, 256-dimensional features, 214K to 20.97M edges) and multi-layer GNN architectures (GIN with SUM and MEAN aggregators, plus 4-layer GNN stacks), the custom Pallas kernels delivered:

• Up to 185.8× HBM reduction on 20.97-million-edge batches (111.5 MiB constant HLO temporary memory vs. 20.2 GiB in native JAX)—turning workloads that previously hit single-chip HBM limits into fast, constant-O(N)-memory operations.
• 8–15× speedups across both forward passes (10.5×–14.7×) and backward / end-to-end training passes (8×–12×) when plugged directly into GNN layers, with strict floating-point numerical parity against native JAX.

Human in Control: A New Cadence for Systems and ML Research

Across both FarArray and our Pallas TPU GNN kernels, Gemini 4 Argon didn’t replace human systems intuition—it amplified it. Engineers and researchers set the architectural boundaries, chose the mathematical invariants, and defined uncompromising end-to-end verification suites. Teamwork’s coordinated Gemini 4 Argon agents then turned specifications that once required a year of specialized driver and kernel engineering into verified production code in a matter of hours.

Try Teamwork via /teamwork-preview in Google Antigravity today.

Contributors to the Graph Mining and GNN acceleration work featured in this post include Brandon Mayer, Kuba Lacki, Laxman Dhulipala, Michael Galkin, Simon Meierhans, Ben Finkelshtein, Yogesh Tewari, Mathieu Guillame-Bert, Bryan Perozzi, and Vahab Mirrokni, in collaboration with the Google Antigravity and Teamwork teams.