arXiv subjects

Julian Walker

Publications and source records attributed to Julian Walker.

4 recordsLinked to original sources

arXiv2026

MaxKernel: Agentic Kernel Generation for TPUs

Designing and authoring high-performance custom kernels for accelerators is a complex task that requires deep hardware-level expertise. Large Language Models (LLM) can be leveraged together with real-time compiler feedback to build agentic systems for kernel generation. In this work, we present MaxKernel, a multi-agent system that implements three distinct paradigms for TPU kernel development: (1) a Human-in-the-Loop (HITL) agent for collaborative, step-by-step design; (2) an Autonomous (Auto) agent that executes a fully automated, metric/trace-driven optimization loop; and (3) a Graph-Based Autonomous Search that scales the Auto agent for global exploration of the design space. All three paradigms leverage a shared pool of specialized sub-agents to handle planning, implementation, self-debugging, testing, and hardware profiling. We evaluate MaxKernel on JaxBench, a comprehensive suite of 50 diverse kernel tasks for TPUs, alongside complex, real-world workloads from state-of-the-art open-source models. We demonstrate that MaxKernel consistently generates highly optimized implementations, matching expert hand-tuned baselines and delivering significant performance across the benchmark. Our agent is open-sourced and available https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel.

cs.AI↗

arXiv2026

JAXBench: Benchmarking Autonomous TPU Kernel Optimization

Rigorous benchmarks have driven progress in autonomous GPU kernel performance optimization by establishing a shared target to hillclimb on, but no equivalent exists for TPUs. We present JAXBench, a TPU-native benchmark suite for AI-generated kernel optimization on Google Cloud TPUs. JAXBench comprises 50 JAX workloads that are both relevant and provide headroom for optimization. We extract 17 production ML operators from architectures in the public MaxText library such as Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, and AlphaFold2, and translate 33 operators from KernelBench that are validated for correctness and set with new problem sizes that achieve high TPU v6e MXU utilization. Eight of the 17 production operators ship with hand-optimized Pallas kernels from the public Tokamax library and block-size tuned to establish an expert upper-bound baseline. We evaluate four feedback-driven methods on generating candidate Pallas kernels for JAXBench. Across the full suite with Gemini 3 Flash, we find that target-specific context matters more than model scale on a sparsely-documented DSL like Pallas. Conditioning on curated TPU documentation raises per-sample correctness from 5.8% to 37.3% and solves 48 of 50 benchmarks at a 1.28x geomean speedup. Search structure yields significant gains once correctness is achieved, with Autocomp's beam-search pipeline reaching a 1.36x geomean speedup over XLA. On the 8 hand-tuned kernels, Autocomp reaches 1.60x geomean over XLA, recovering most of the 2.08x Tokamax upper bound but trailing on the specialized paged and ragged attention operators. High-quality TPU kernel optimization remains a challenging task, and we release the JAXBench benchmark, evaluation harness, and baseline results to support open source contributions.

cs.AI↗

arXiv2026

Adapting AlphaEvolve to Optimize Fully Homomorphic Encryption on TPUs

The deployment of Fully Homomorphic Encryption (FHE) at scale is hindered due to its heavy computational overhead. While specialized hardware accelerators like Google Tensor Processing Units (TPUs) can help, mapping complex cryptographic kernels onto such architectures remains a challenge. Efficient execution requires co-optimization between the systolic array-based Matrix Multiplication Unit (MXU) and Vector Processing Units (VPUs), as well as the orchestration of data movement across the vector register files. Existing compiler stacks often abstract low-level hardware utilization, requiring developers to adopt a manual trial-and-error process that often results in fragmented execution and underutilized resources. To accelerate this development process, we use AlphaEvolve to automate the exploration of hardware-aware cryptographic-kernel optimizations. We frame optimization as an evolutionary search problem, utilizing the closed-loop system provided by AlphaEvolve, that leverages LLM-driven code generation. We use real-world feedback from hardware execution and rigorous correctness testing to guide the evolution process. We evaluate AlphaEvolve optimization on primitives for both the TFHE (Jaxite) and CKKS (CROSS) FHE schemes on Google Cloud TPUv5e, a contemporary TPU architecture. Within 24 hours of automated exploration, AlphaEvolve discovered implementation-level optimizations that improve TFHE bootstrap latency by 2.5x and CKKS rotation and multiplication latency by 1.31x and 1.18x, respectively, relative to human-engineered state of the art. These results demonstrate that AlphaEvolve can be used to enable researchers to navigate the optimization trade-offs between cryptography, compilers, and hardware accelerators.

cs.CR↗

arXiv2023

Electric-field-induced non-ergodic relaxor to ferroelectric transition in BiFeO3-xSrTiO3 ceramics

While BiFeO3-based solid solutions show great promise for applications in energy conversion and storage, realizing this promise necessitates understanding the structure-property relationship in particular pertaining to the relaxor-like characteristics often exhibited by solid solutions with polar-to-non-polar morphotropic phase boundaries. To this end, we investigated the role of the compositionally-driven relaxor state in (100-x)BiFeO3-xSrTiO3 [BFO-xSTO], via in situ synchrotron X-ray diffraction under bipolar electric-field cycling. The electric-field induced changes to the crystal structure, phase fraction and domain textures were monitored via the {111}pc, {200}pc, and 1/2{311}pc Bragg peaks. The dynamics of the intensities and positions of the (111) and (11-1) reflections reveal an initial non-ergodic regime followed by long-range ferroelectric ordering after extended poling cycles. The increased degree of random multi-site occupation in BFO-42STO compared to BFO-35STO is correlated with an increase of the critical electric field needed to induce the non-ergodic-to-ferroelectric transition, and a decrease in the degree of domain reorientation. Although both compositions show an irreversible transition to a long-range ferroelectric state, our results suggest that the weaker ferroelectric response in BFO-42STO is related to an increase in ergodicity. This, in turn, serves to guide the development of BFO-based systems into promising platform for further property engineering towards specific capacitor applications.

cond-mat.mtrl-sci↗