NUMA-Torch: Compiler-Driven NUMA-Aware LLM Prefill on Multi-Socket CPUs

NUMA-Torch: Compiler-Driven NUMA-Aware LLM Prefill on Multi-Socket CPUs Thisara Rathnayaka Weerakoon ∗ , Ishan Dewpura ∗ , Niman Ransindu ∗ , Shavin Anjitha ∗ , Tharindu R. Patabandi † , Sanath Jayasena ∗ ∗ Dept. of Computer Science & Engineering, University of Moratuwa, Sri Lanka

arXiv:2611.07105v1 [cs.DC] 4 Nov 2026

{thisara.21, ishandewpura.21, niman.21, shavin.21, sanath}@cse.mrt.ac.lk † Independent Researcher, San Jose, CA, USA tharindu@protonmail.com

Abstract—For sufficiently large transformer-based Large Language Models (LLMs), whose parameters cannot be retained in on-chip caches, prefill throughput on multi-socket CPUs is often limited by sustainable DRAM bandwidth rather than peak floating-point throughput. In the standard PyTorch compilation (i.e., TorchInductor) path, model weights are allocated without considering the Non-Uniform Memory Access (NUMA) topology—specifically, the affinity between memory banks and compute cores is disregarded. This causes a substantial fraction of weight accesses to traverse the inter-socket interconnect, severely degrading the effective memory bandwidth available to each socket. We address this bottleneck through a compiler-level solution that makes TorchInductor NUMA-aware. Treating each NUMA socket as an independent tensor-parallel partition, we introduce NUMA-aware Tensor Parallelism (NUMA-TP). This approach transparently shards weights across sockets and schedules compute kernels such that thread pools and memory allocations remain strictly local to their respective sockets using CPU affinity. We evaluate the approach on GPT-2 [1] (117M, 345M, 762M) and BERT [2] (Small, Large), which exercise standard multi-head attention and MLP, as well as Gemma 3 270M [3], which combines Multi-Query Attention with a GeGLU-gated MLP. By strictly enforcing data locality, our approach successfully mitigates crosssocket traffic, achieving prefill speedups of up to 2.63× on a dualsocket Haswell platform, up to 1.87× on a dual-socket Cascade Lake platform, and up to 2.27× on a four-socket Haswell-EX platform. Index Terms—Compilers, Data Locality, NUMA, Tensor Parallelism, Large Language Models, CPU inference

remote DRAM access (high cost)

Socket 0 / NUMA Node 0

DRAM 0 (local) local access

Socket 1 / NUMA Node 1

Core 0

Core 1

Core 0

Core 1

L1 / L2

L1 / L2

QPI / UPI

L1 / L2

L1 / L2

Shared L3 Cache

cross-socket interconnect

Shared L3 Cache

Mem. Controller

Mem. Controller

DRAM 1 (local) local access

Fig. 1. Logical view of a two-socket NUMA CPU treated as a distributed system; the dashed arrow marks a high-cost remote DRAM access across the cross-socket interconnect.

deployment environments. The prefill phase of LLM inference which processes the full input prompt in one forward pass is traditionally assumed compute-bound owing to the high arithmetic intensity of dense matrix multiplications over long sequences. Yet on multi-socket CPUs under the standard TorchInductor compilation pipeline, prefill unexpectedly becomes severely DRAM-bound as NUMA-oblivious memory allocations force massive cross-socket traffic. As illustrated in Figure 1, we view each socket of a multisocket NUMA CPU as an independent node owning its locally attached DRAM and memory controller, together with cores that have private L1/L2 caches and share an L3. Sockets are linked by a QPI/UPI interconnect serving as the inter-rank network; accesses to another socket’s memory must cross it, I. I NTRODUCTION incurring substantially higher latency and lower bandwidth Driven by the necessity of accommodating many cores, than local accesses, and are therefore treated as explicit modern high-performance servers have universally adopted communication in our compilation model. In standard execution, compute threads are distributed across Non-Uniform Memory Access (NUMA) architectures [4]. all cores without regard to data locality, causing roughly half There is consequently a tremendous, yet underutilized, potential of all weight accesses to incur remote DRAM latency and to deploy Large Language Model (LLM) inference on these saturating the inter-socket link. To empirically validate this, multi-socket CPU platforms, given their cost-effectiveness, we conducted a Roofline analysis using Intel Advisor [5] broad availability, and the frequent absence of GPUs in on the kernels that dominate the prefill phase of a GPT-2 © 2026 IEEE. Personal use of this material is permitted. Permission from 345M model [1] (batch size 1, sequence length 4096) on a IEEE must be obtained for all other uses, in any current or future media, including reprinting/republishing this material for advertising or promotional dual-socket Haswell platform (Figure 2). For TorchInductorpurposes, creating new collective works, for resale or redistribution to servers compiled unmodified Triton-CPU baseline (circles), these or lists, or reuse of any copyrighted component of this work in other works. kernels exhibit poor data locality and sit near the main-DRAM Accepted for publication in the 2026 IEEE 33rd International Conference on High Performance Computing, Data, and Analytics (HiPC), Bengaluru, bandwidth ceiling, confirming that the prefill operates firmly in India, December 16–19, 2026. This is the authors’ accepted version; the final the DRAM-bound regime. Our NUMA-Torch path (triangles) published version will appear in IEEE Xplore.

@torch.compile

Torch Dynamo

Decompositions (PrimTorch)

NUMA-Torch NUMA MLP pass

FX Passes (Graph-Level Opt.)

NUMA Attention pass NUMA Normalisation pass emitted by NUMA passes

Lowering and scheduling (Inductor IR)

NUMA partitioning schema (per-block metadata)

Fig. 2. Roofline analysis on dual-socket Haswell: baseline Triton CPU kernels (circles) versus the NUMA-Torch path (triangles) for the dominant GPT-2 345M prefill kernels (batch size 1, sequence length 4096). Arrows indicate the per-kernel trajectory from baseline to NUMA-Torch.

block type: MLP / Attn / Norm sockets: N ; placement policy: 1. column-parallel 2. row-parallel 3. replicated

Triton Kernel Codegen

consumed at wrapper codegen NUMA-aware memory allocation

improves core-to-memory affinity and, as the trajectory arrows indicate, shifts the kernels toward the higher-bandwidth L3 and L2 cache tiers, increasing overall throughput (GFLOPS). Despite this architectural reality, existing CPU optimization efforts frequently overlook the NUMA penalty, focusing on intra-core compute optimizations. Addressing this NUMAinduced DRAM bandwidth wall must instead be the primary objective, and the most direct solution is strict enforcement of data locality: explicitly placing weights in socket-local DRAM via NUMA-aware memory allocation, and binding computations to the corresponding cores via NUMA-aware thread scheduling. While mainstream frameworks for CPU inference, such as llama.cpp [6], deliver high baseline performance, they are not explicitly optimized for NUMA data locality. Recent work such as ArcLight [7] addresses this memory bottleneck through cross-NUMA tensor parallelism. However, because ArcLight is a custom inference engine written entirely in C++, it relies heavily on manual engineering and requires the model code to be meticulously rewritten to support different LLM architectures. In this work, we focus on lifting cross-NUMA Tensor Parallelism into PyTorch 2.0’s compiler ecosystem [8] to provide a PyTorch-native implementation. By extending the standard TorchInductor pipeline, our transformations remain entirely transparent to the user, requiring zero modifications to the model code, optimizer, or the inference loop. We present the first compiler-based, PyTorch-native implementation of NUMA-aware tensor parallelism targeting multisocket CPUs. We make the following contributions: 1) A compiler-driven solution to the NUMA DRAM bottleneck that leverages post-grad FX graph transformations to partition operations across sockets, minimizing manual engineering overhead. 2) An extended Triton-CPU backend [9], [10] that emits per-socket buffer allocations and dispatches kernels via a persistent, NUMA-aware thread pool. 3) An evaluation of the compiler approach across standard models (GPT-2 [1], BERT [2]) and modern LLMs with complex attention and MLP variants (Gemma 3 [3]),

NUMA cache

Wrapper Codegen

NUMA extended clone TorchInductor Triton CPU Launcher

NUMA-aware thread scheduling

Triton CPU Runtime

Fig. 3. Overview of NUMA-Torch: NUMA-Torch extends TorchInductor and Triton-CPU compilation stacks with NUMA-aware tensor parallel graph transformations, per-socket buffer placement, and a NUMA-aware runtime that dispatches kernels to socket-pinned worker groups. The NUMA partitioning schema (folded-corner note) is not a pipeline stage: it is per-block metadata emitted by the NUMA passes at FX-graph level and consumed by the wrapper code generator when materializing socket-local buffers (dashed arrows).

showing up to 2.63× prefill speedup on dual-socket platforms, comparing against process-level NUMA policies (numactl), and assessing scalability from dual-socket to four-socket platforms. This work targets the prefill phase at fixed shapes. Throughout, “end-to-end prefill” denotes a single prefill forward pass at a fixed shape, not end-to-end serving: we measure prefill in isolation, so decode, continuous batching, KV-cache placement, and serving-level metrics are outside the evaluation. Figure 3 illustrates the overall design of our approach, which we refer to as NUMA-Torch. NUMA-Torch is built on top of TorchInductor and Triton-CPU, extending the compiler pipeline with NUMA-aware graph transformations, socket-local buffer management, and affinity-aware kernel dispatch. Section II describes each component of this pipeline in detail. Section III provides a comprehensive evaluation of NUMA-Torch, Section IV discusses related literature, and we present concluding remarks in Section V. II. NUMA-T ORCH A. Explicit NUMA Placement via Extended Clone Semantics Effective NUMA-aware tensor parallelism requires precise control over where a tensor’s physical memory resides. In standard PyTorch, the clone operator creates a semantically

independent copy. However, it does not let callers specify which activation and A, B are the weight matrices. As shown in socket’s DRAM wiill hold the copy. Instead, cloned shards Algorithm 1, we partition the up-projection A column-wise follow the calling thread’s first-touch allocation policy [11], and the down-projection B row-wise. Each socket i computes so physical pages land in the memory bank of whichever its shard of the intermediate activation Y i = GeLU(XAi ) and socket’s thread first writes each virtual address. This typically the matching partial output Z i = Y i B i entirely from socketconcentrates storage in one socket’s DRAM, causing severe local memory. Because matrix multiplication distributes over cross-interconnect traffic during subsequent parallel computa- addition, the partial outputs P are combined with a single crosstion. socket all-reduce, Z = i Z i , confining inter-socket traffic to To resolve this limitation, we extend the native PyTorch one activation-sized reduction per block. clone operator by introducing an explicit NUMA placement Algorithm 1: NUMA-TP MLP Block parameter. This modification empowers the runtime to bypass Input: activation X; up-projection A; down-projection B; socket default allocation policies [11] and assign each tensor replica count N to a specific DRAM socket at the exact moment of creation, Output: output activation Z leveraging the Linux mbind system call [12]. 1 Partition A column-wise and B row-wise across the N sockets This extended functionality provides a straightforward, 2 foreach socket i ∈ {0, . . . , N − 1} in parallel do Y i ← GeLU(XAi ) // column-parallel up-proj. Pythonic interface for tensor sharding across socket DRAMs. 3 4 Z i ← Y iBi // row-parallel down-proj. For example, a global weight matrix can be easily partitioned PN −1 i 5 Z ← Z // cross-socket all-reduce i=0 along its rows and distributed across multiple sockets using 6 return Z the new numa_node argument: # Global weight matrix; defaults to first-touch allocation weight = torch.randn(4096, 4096) # Place tensor partitions across distinct DRAM sockets shard_0 = weight[:2048, :].clone(numa_node=0) shard_1 = weight[2048:, :].clone(numa_node=1)

B. NUMA-Aware Tensor Parallelism

2) Multi-Head Attention Block: The same column-then-row split extends to multi-head attention, following the MegatronLM [13] decomposition. The query, key, and value projections WQ , WK , WV are partitioned column-wise so that each socket owns a disjoint subset of attention heads, while the output projection WO is partitioned row-wise. Because heads are independent within scaled dot-product attention, every socket evaluates its local attention without exchanging queries, keys, values, or scores across sockets. The block then closes with a single all-reduce over the per-socket partial outputs, as detailed in Algorithm 2.

With NUMA-aware data allocation established, the subsequent challenge is NUMA-aware computation scheduling. Rather than relying on manual engineering, we use the compiler to automatically identify and schedule computations, which Algorithm 2: NUMA-TP Multi-Head Attention Block requires a robust underlying scheduling model. To formulate Input: activation X; projections WQ , WK , WV , WO ; output bias this model, a multi-socket CPU can be abstracted as a bO ; per-head dimension dk ; socket count N distributed system. Each socket acts as an independent node Output: output activation Y consisting of its compute cores and local DRAM. The sockets 1 Partition WQ , WK , WV column-wise and WO row-wise across the N sockets are connected by an interconnect such as Intel’s QuickPath 2 foreach socket i ∈ {0, . . . , N − 1} in parallel do Interconnect (QPI) or Ultra Path Interconnect (UPI). i , XW i , XW i 3 Qi , K i , V i ← XWQ K
V √ i i In distributed multi-GPU cluster environments, a fundamen- 4 Z ← softmax Q (K i )⊤ / dk V i // owned heads i tal goal is to reduce communication overhead to accelerate 5 Y i ← Z i WO // row-parallel out-proj. both training and inference. To mitigate these bottlenecks, 6 Y ← PN −1 Y i + bO // all-reduce, add bias i=0 works such as Megatron-LM [13] utilize tensor parallelism 7 return Y to minimize synchronization points within the Multi-Layer Perceptron (MLP) and attention blocks [14]. Similarly, the 3) Normalization Block: Layer normalization and RMS primary performance bottleneck in our multi-socket CPU normalization require a different strategy. Instead of computing topology is the latency and bandwidth limitation of cross-socket the normalization on a single socket and broadcasting the communication. Therefore, we adapt this tensor parallelism result, NUMA-TP replicates this computation on every socket, strategy to our single multi-socket CPU environment to as shown in Algorithm 3. Each socket maintains its own minimize cross-NUMA socket communication. normalized copy and feeds it directly to the next local kernel. We now detail how NUMA-aware Tensor Parallelism This completely eliminates the need for cross-NUMA data (NUMA-TP) is applied to the MLP, multi-head attention, and movement, such as an all-reduce operation. We selected this normalization sub-blocks that comprise a transformer layer. strategy empirically on multi-socket platforms, where this redundant computation proved cheaper than the eliminated 1) MLP Block: Consider the MLP block of GPT-2, which inter-socket communication. consists of two GEMMs separated by a GeLU activation, Y = GeLU(XA) followed by Z = Y B, where X is the input

Algorithm 3: NUMA-TP Normalization Block (LayerNorm / RMSNorm) Input: activation X; normalization parameters γ, β; socket count N Output: per-socket normalized activations {Y 0 , . . . , Y N −1 } 1 Replicate γ, β to all sockets // weights not sharded 2 foreach socket i ∈ {0, . . . , N − 1} in parallel do 3 Xi ← X // local input copy 4 Y i ← Norm(X i ; γ, β) // replicated full norm 5

return {Y 0 , . . . , Y N −1 }

// no all-reduce

def forward(self, x): # Input activation ’x’ has shape [s, h] # 1. Up-projection: Y = XA # - linear_1 applies weight A of shape [h, d] # - Output y expands to intermediate shape [s,d] y = self.linear_1(x) # 2. Element-wise activation # - Maintains the shape [s, d] y = torch.nn.functional.gelu(y) # 3. Down-projection: Z = YB # - linear_2 applies weight B of shape [d,h_out] # - Output z is projected to the terminal shape [s, h_out] z = self.linear_2(y) return z Fig. 4. Standard PyTorch forward definition of the MLP block.

C. FX Graph Transformation for Tensor Parallelism

where i ∈ {0, . . . , N − 1} denotes the target NUMA socket index for that specific compute shard. Similarly, for the down-projection, the weight matrix B ∈ Rd×hout is sliced along its rows (the d axis) to produce shards of shape [d/N, hout ]. The corresponding operation is:   d d slice B, dim = 0, start = i · , end = (i + 1) · (2) N N As established in Section II-B, yielding the correct final output Z requires an element-wise reduction (addition) of the partial results computed by each socket. However, the N -way reduction node is absent from the transformed FX graph in Figure 5b. This omission is a deliberate architectural choice, rooted in the fact that the FX graph maintains a “Single-Program, Multiple-Data” (SPMD) representation. The graph does not instantiate N separate matrix multiplication (mm) nodes; instead, there is exactly one logical mm node whose dimension metadata has been scaled down from the global [s, d] to the local [s, d/N ]. Because the graph encodes the execution logic for a single generic shard, there are no N distinct partial-sum nodes in the intermediate representation to combine. Consequently, the physical parallelization and the N -way reduction are deferred entirely to the code-generation phase (discussed in Section II-E). The actual existence of N separate partial outputs (Z 0 . . . Z N −1 ) only manifests at runtime when the wrapper codegen allocates N independent buffers and launches the identical kernel N times across the hardware domains.

Since our approach focuses on applying tensor parallelism D. The NUMA Partitioning Schema: Bridging Graph Transwithin the PyTorch compilation pipeline, the standard PyTorch formations and Code Generation definition of the MLP block, shown in Figure 4, serves as our The NUMA partitioning schema is the metadata bridge starting point. Here, s denotes the sequence length, h the input hidden dimension, d the MLP intermediate dimension, and hout between the mathematical rewriting performed in the FX graph and the physical multi-socket allocation performed by the the output dimension. As illustrated in Figure 5a, TorchDynamo [8] traces this backend code generator. As depicted in Figure 3, it records for module (Figure 4) into an FX graph that sequentially projects each tensor its placement policy (column-parallel, row-parallel, an [s, h] input to an [s, d] intermediate activation, applies a or replicated), which the wrapper code generator later uses decomposed pointwise GeLU subgraph, and projects it to a to materialize socket-local buffers. In the MLP block above, the up-projection A is column-parallel, the down-projection B final [s, hout ] output. To apply tensor parallelism at FX graph level, we rely on a row-parallel, and the input activation X replicated across all custom post-grad FX pass that captures the canonical Linear sockets. → Activation → Linear pattern and transforms its E. Wrapper Code Generation and Buffer Materialization execution topology, as illustrated in Figure 5b. To illustrate this So far the partitioning exists only as metadata, with no buffer transformation, consider a dual-socket CPU (N = 2) similar allocated and no page bound to a socket. Both happen in the to the depiction in Figure 1. The parallelization pass injects code TorchInductor generates, which is Triton-CPU kernel aten.slice.Tensor nodes into the fx graph (Figure 5b). code [10] together with the Python wrapper code that drives it. Fundamentally, aten.slice.Tensor is an Aten library The wrapper function call allocates each buffer, places it, and operation [15] that extracts a contiguous sub-region of a tensor only then launches the kernels that read it. For every kernel along a specified dimension. argument it consults the schema, classifies the tensor as a As discussed in Section II-B, the up-projection weight matrix column shard, a row shard, or a full replica, and emits the h×d A∈R is sliced along its columns (the d axis) to yield matching slice with a clone(numa_node=i). That explicit local shards of shape [h, d/N ]. For a generalized N -socket numa_node binds shard i’s pages to socket i’s DRAM at system, this is achieved via the operation: creation instead of leaving them to first touch. Because each   d d slice A, dim = 1, start = i · , end = (i + 1) · (1) placement finishes before the kernel that reads it is launched, N N every socket’s kernel finds its operands already local. The FX

Fig. 5. FX graph transformation for NUMA-aware tensor parallelism: (a) The baseline AOTAutograd FX graph for the MLP block contains an up-projection, decomposed GeLU activation, and down-projection. (b) The NUMA-TP post-grad pass rewrites the graph by inserting slice nodes that encode per-socket weight shards and shard-plan metadata, allowing physical placement and per-socket execution to be deferred to code generation.

pass therefore states only how a block is partitioned; wrapper the wrapper inserts the all-reduce loop that sums the percode generation is where its shards first acquire physical pages. socket partial outputs into the final tensor Z. This stage also decides which socket-local buffers can be F. NUMA-Aware Kernel Dispatch and Deferred Reassembly reused across adjacent kernels. When an activation produced on socket i is consumed immediately by another NUMA(b) Terminal partial sums (a) Intermediate activations parallel kernel on the same socket, the wrapper records it in a Socket 0 Socket 1 Socket 0 Socket 1 per-socket cache instead of constructing a global tensor and Column-parallel Column-parallel Row-parallel Row-parallel slicing it again. In the MLP example, the column-parallel upkernel kernel kernel kernel projection output Y i = GeLU(XAi ) stays local and feeds Z 0 (partial) Z 1 (partial) Y0 Y1 the row-parallel down-projection consuming B i , so the only cross-socket step is the reduction that closes the block. All-reduce Next local Next local P cross-socket Z = i Zi kernel kernel Because the final reduction is not explicit in the FX graph (Figure 5b), the wrapper must know whether a kernel produces Stay socket-local: cached and final Z forwarded, no reduction a still-local intermediate or a value closing the NUMA-parallel region. The distinction matters most after fusion, where one lowered kernel may contain both matrix multiplication and Fig. 6. The two buffer-reconciliation patterns chosen by the wrapper from pointwise operations. We therefore attach an output tag to each each value’s semantic output tag. (a) Outputs of column-parallel kernels and ops are tagged as intermediate activations: each per-socket generated value, and the wrapper uses it to choose one of fused pointwise Y i stays in socket-local DRAM, is held in the per-socket cache, and two reconciliation paths, shown for a dual-socket system in shard is forwarded to the next local kernel with no cross-socket communication. Figure 6: (b) Outputs of row-parallel kernels are tagged as terminal partial sums Z i ; before the value leaves the NUMA-parallel block, the wrapper emits a single • Intermediate activations (Figure 6a): Outputs from cross-socket all-reduce that sums the partials into the final tensor Z. column-parallel projections and fused pointwise operations remain socket-local. The wrapper stores each shard in the After the wrapper has placed each tensor shard in its socketper-socket cache and forwards it to the next local kernel local memory, execution is handed to the Triton-CPU launcher. without reduction. In the unmodified backend, this launcher runs a kernel as a • Terminal partial sums (Figure 6b): Outputs from row- single program replicated across many parallel work items, parallel projections, such as the MLP down-projection or but it treats the machine as one flat pool of cores with no the attention output projection, represent partial values notion of NUMA locality. As a result, even when a shard lives Z i . Before such a value leaves the NUMA-parallel block, in a particular socket’s DRAM, the default launcher gives no

Algorithm 4: NUMA-Aware Kernel Launch Input: kernel K; its set of parallel work items G; when NUMA-tagged, the per-socket local operands {P0 , . . . , PN −1 } and the socket each shard belongs to; thread budget T Output: K executed, with every socket writing its output shard into local memory 1 if K is tagged for NUMA execution then 2 N ← number of sockets in the per-socket operands 3 T ← max(T, N ) 4 Split the T threads into N groups of Ti ≈ T /N , one per socket 5 foreach socket i ∈ {0, . . . , N − 1} in parallel do 6 Ci ← cores belonging to socket i 7 foreach worker j in socket i’s group of Ti threads do 8 bind worker j to a core in Ci 9 run K on worker j’s 1/Ti share of G, reading Pi wait for all groups to finish

10 11 12

else run K through the default, NUMA-oblivious launch

guarantee that the threads using that shard run on cores of the same socket. Many instead run on a remote socket, paying the interconnect penalty and undoing the placement the wrapper performed. We therefore extend this launch path with a NUMA-aware dispatch layer, summarized in Algorithm 4. Each time a kernel is launched, the layer checks whether the compiler has marked it for NUMA execution (i.e., whether the wrapper has attached the per-socket operands together with a note of which socket each shard belongs to). If this annotation is absent, the kernel is launched through the original NUMA-oblivious path, and hence, ordinary non-sharded operations are left untouched. If it is present, the launcher takes the NUMA-aware path described next. On this path, the launcher turns the single-program, multipledata view of the FX graph into execution along two levels of parallelism. Across sockets, every socket runs its own copy of the kernel over the full set of work items, reading only that socket’s local shards, so no socket reaches across the interconnect for data. Within a socket, the work is shared among a group of threads. The launcher first fixes a total thread budget, taking either the parallelism the kernel requests or the available hardware threads, but never fewer than one thread per socket. It then splits that budget into N roughly equal groups, one for each socket, and binds every group’s threads to cores of its own socket as reported by the OS topology. No thread therefore drifts onto another socket. The launcher waits for every group to finish, so the per-socket kernels complete together. With this binding in place, each thread runs on the socket that physically holds the shard it touches. This is what closes the locality gap left by the default launcher. This dispatch policy completes the path from graph rewriting to execution: the FX pass marks an operation as shardable, the schema records how each of its tensors should be placed, the wrapper prepares the socket-local operands, and the launcher runs one kernel instance per socket. Reassembly is left to the wrapper rather than the launcher and is deferred until a value actually needs to be combined (Figure 6). A value consumed

only by another NUMA-aware kernel stays split, so the wrapper keeps the per-socket buffers and forwards them as they are. A value that is a final partial sum, or that is about to leave the parallel region, is closed with the cross-socket all-reduce described earlier, after which the reassembled tensor re-enters the ordinary PyTorch program. III. E XPERIMENTAL E VALUATION A. Experimental Setup We evaluate NUMA-Torch, on three Intel Xeon platforms: two dual-socket machines (CascadeLake and Haswell) and one four-socket machine (Haswell-EX), the last sharing the Haswell microarchitecture with the dual-socket Haswell server so that comparing the two isolates the number of NUMA domains. Table I summarizes the hardware and Table II the software stack. TABLE I H ARDWARE SPECIFICATIONS OF THE THREE EVALUATION PLATFORMS Specification

CascadeLake

Haswell

Haswell-EX

Processor Microarchitecture Total vCPUs Cores per socket Sockets / NUMA nodes Clock speed L3 cache Environment

Xeon Gold 6253CL Cascade Lake 60 15 2 3.10 GHz 49.5 MiB GCP c2-standard-60

Xeon E5-2620 v3 Haswell 24 6 2 2.40 GHz 30 MiB Server (bare metal)

Xeon E7-4820 v3 Haswell 80 10 4 1.90 GHz 100 MiB Server (bare metal)

TABLE II S OFTWARE CONFIGURATION Component

Dual-socket

Four-socket

Operating system Linux kernel Python PyTorch Triton-CPU LLVM (Triton-CPU backend)

Ubuntu 20.04 LTS Ubuntu 24.04 LTS 5.15 7.0 3.8 3.9 2.6.0a0 (torch.compile / TorchInductor) commit 7e394470 19 (commit adba14a)

B. Models and Configurations We evaluate three model families, summarized in Table III. GPT-2 [1] (decoder-only) and BERT [2] (encoder-only) exercise the core column- and row-parallel FX passes with standard multi-head attention, while Gemma 3 270M [3] additionally tests the KV-replication path for Multi-Query Attention (MQA) [16] (nkv = 1) and the GeGLU-gated MLP pass. TABLE III A RCHITECTURAL PARAMETERS OF EVALUATED MODELS Family

Model

Params

Layers

dmodel

Heads (nh )

nkv

MLP type

GPT-2 [1]

117M 345M 762M

117M 345M 762M

12 24 36

768 1024 1280

12 16 20

12 16 20

GeLU GeLU GeLU

BERT [2]

Small Large

29M 340M

4 24

512 1024

8 16

8 16

GeLU GeLU

Gemma 3 [3]

270M

270M

18

640

4

1

GeGLU

For each model we measure the mean prefill latency (ms) across a representative sweep of batch-size and sequence-length

configurations. Each latency data point discards 10 warmup iterations and averages the wall-clock time of 50 timed iterations, ensuring that compilation artifacts and cold-cache effects are excluded from reported figures. On CPU, TorchInductor generates C++/OpenMP code by default [8]. We instead select TritonCPU [10] as the CPU code-generation backend through torch._inductor.config.cpu_backend, because NUMA-TP needs control of the kernel launch path in order to bind thread groups to sockets. Both configurations below use it, so reported speedups are relative to Triton-CPU and do not transfer as absolute latencies to the default C++/OpenMP backend. The baseline is torch.compile with unmodified TorchInductor and the unmodified Triton-CPU backend. NUMA-Torch modifies both, as described in Section II.

(a) Haswell (Intel Xeon E5-2620 v3, 6 cores/socket)

C. Performance Results The mean prefill execution times of NUMA-Torch are presented as per-configuration heatmaps in Figure 7, spanning every evaluated model across the full sweep of batch sizes and sequence lengths on both the Haswell and Cascade Lake architectures. We evaluate the end-to-end prefill speedup of the NUMA-Torch compiler extension relative to the baseline. As established by our initial roofline analysis, the baseline prefill phase is bottlenecked by main memory bandwidth. By strictly enforcing data locality via our NUMA-aware tensor parallelism passes, NUMA-Torch successfully mitigates crosssocket interconnect traffic, unlocking speedups of up to 2.63× on Haswell and 1.87× on Cascade Lake for all evaluated (b) Cascade Lake (Intel Xeon Gold 6253CL, 15 cores/socket) transformer architectures (Figure 8). Fig. 7. Mean prefill execution time of NUMA-Torch, for every evaluated model 1) Architectural Impact (Haswell vs. Cascade Lake): The across the full sweep of batch sizes (rows) and sequence lengths (columns), gain tracks the number of cores per socket. On Haswell on the dual-socket (a) Haswell and (b) Cascade Lake platforms. Each cell (Figure 8a), with just 6 cores per socket, speedups frequently is annotated with the mean wall-clock time of one prefill call in seconds; logarithmic color scale, shown above each heatmap, is set independently exceed 2.0× at sequence lengths of 1024 and above: with the per platform. Over all measured configurations, the coefficient of variation few cores, NUMA-oblivious scheduling drives a high volume of prefill time has mean 1.58% (max 6.74%) on Haswell and mean 4.43% of remote traffic through each core, so threads stall on QPI (max 26.23%) on Cascade Lake. fetches and the baseline is highly inefficient. Cascade Lake (Figure 8b), with 15 cores per socket, spreads that traffic across context workloads intermediate activations dominate the prefill more threads, dampening the baseline’s cross-socket bottleneck; footprint. Figures 7 and 8 together report every evaluated configuration speedups are accordingly more moderate but highly consistent, in absolute and relative terms, where the savings at the longer ranging from 1.10× to 1.87×. 2) Impact of Model Scale and Architecture: Larger hidden sequence lengths reach multiple seconds per prefill call. dimensions and weight matrices yield higher relative speedups, D. Comparison with Process-Level NUMA Policies placing a heavier burden on the memory sub-system: on Do the operating system’s own placement knobs already Haswell at batch size 8, GPT-2 [1] (345M) peaks at 2.63× recover this benefit? We re-ran the full 144-configuration (sequence length 256) against 2.10× for the smaller BERT [2] Cascade Lake sweep under three numactl policies [17], each (Small) under identical conditions. Gemma 3 [3] (270M) measured against default first-touch allocation (Table IV). combines Multi-Query Attention with a GeGLU-gated MLP. None of the three comes close, and each falls short for It improves robustly as well, reaching 1.96× on Haswell a different reason. Interleaving reaches a geometric-mean and 1.76× on Cascade Lake, which confirms that the FX speedup of only 1.08× and a median speedup of 1.02×. Roundtransformations generalize to heterogeneous sub-layers without robin placement equalizes which memory controller serves a model code rewrites. request, but no access becomes local. It removes bandwidth 3) Batch Size and Sequence Length Scaling: The optiimbalance and leaves the interconnect volume that Figure 2 mizations scale efficiently as workload size grows. For long-

(a) Haswell (Intel Xeon E5-2620 v3, 6 cores/socket)

(b) Cascade Lake (Intel Xeon Gold 6253CL, 15 cores/socket) Fig. 8. End-to-end prefill speedup of NUMA-Torch over the unmodified Triton-CPU baseline for every evaluated configuration, grouped by model, batch size, and sequence length, on the dual-socket (a) Haswell and (b) Cascade Lake platforms.

identifies as binding. It is also the least reliable policy, running slower than default in 44 of 144 configurations. localalloc is the best of the three by geometric-mean speedup, at 1.12×, and still recovers little. Two limits explain why. No process-level policy can shard a tensor, so each weight matrix lands whole on one socket and the other socket reads all of it remotely. First-touch locality also fails to survive execution, because TorchInductor’s NUMA-oblivious thread pool means the thread that first touches a buffer is rarely the one that reads it. NUMA-Torch removes both limits, pairing clone(numa_node=i) placement with thread-group binding. Single-socket execution is the most informative control. Its locality is perfect by construction, yet its geometric-mean speedup is only 0.59×, making it the worst option by a wide

margin; it runs slower than default in 143 of 144 configurations. It trades all remote traffic for half the cores and half the bandwidth. NUMA-Torch reaches a geometric-mean speedup of 1.38× and a median speedup of 1.44×, and it is alone in never running slower than default, doing so in 0 of 144 configurations. The goal is not locality alone, but locality without sacrificing the machine’s full compute and memory bandwidth. TABLE IV P ROCESS - LEVEL NUMA POLICIES ON C ASCADE L AKE Policy

Geo. mean

Median

Worse than default

--interleave=all --localalloc Single socket (--cpunodebind=0 --membind=0)

1.08× 1.12× 0.59×

1.02× 1.06× 0.59×

44/144 19/144 143/144

NUMA-TP (this work)

1.38×

1.44×

0/144

Fig. 9. End-to-end prefill speedup of NUMA-Torch over the unmodified Triton-CPU baseline on the four-socket Haswell-EX platform, as the geometric mean over all models and batch sizes at each sequence length. The dashed line marks parity; in the shaded region the geometric mean exceeds 1.0×.

E. Scaling to Four NUMA Domains We repeat the evaluation on the four-socket Haswell-EX platform of Table I, sweeping the same six models over same batch size and sequence length configurations. Figure 9 reports the resulting speedups. Speedup rises monotonically with sequence length, from a geometric mean of 0.48× at s=128 to 1.56× at s=4096, crossing 1.0× between s=1024 and s=2048. Every measured configuration at s ≥ 2048 is faster than the baseline, with a geometric mean of 1.38× and a peak of 2.27× for BERT Small at B=2 and s=4096, where the mean prefill time falls from Fig. 10. End-to-end prefill speedup of NUMA-Torch over the unmodified 25.8 s to 11.3 s. No configuration at s ≤ 512 is faster than the Triton-CPU baseline on the four-socket Haswell-EX platform, for every baseline. Lower speedups at the shorter sequence lengths are evaluated model, batch size and sequence length. Cells are colored about 1.0×, so green cells are faster than the baseline and red cells slower. expected on this platform, for two compounding reasons. Both reasons are already present on the dual-socket platforms and IV. R ELATED W ORK simply combine unfavourably here. The first is amortization. A. LLM Inference Optimizations on CPUs Thread-group binding, the barrier (Algorithm 4), and the allA substantial body of work has sought to alleviate the reduce cost about the same per kernel whatever the kernel computes, while the work each socket does shrinks as domains compute and memory bottlenecks of executing large language are added. A socket here takes a quarter of the partitioned models (LLMs) on CPU architectures. Foundational software intermediate width rather than a half, and every reassembly efforts like llama.cpp deliver high baseline performance for point closes with a four-way all-reduce rather than a two-way CPU inference through highly optimized, manually engineered one. At short sequence lengths the per-socket GEMM is then C/C++ backends [6]. Similarly, specialized frameworks such too small to cover that fixed cost. The second reason is how as ArcLight address the memory bottleneck via cross-NUMA much there was to recover in the first place. The Haswell and tensor parallelism on many-core CPUs [7]. However, because Cascade Lake comparison already showed the gain shrinking as these are custom inference engines written entirely in C++, cores per socket rise, because more threads spread the baseline’s they rely heavily on manual engineering and require meticulous remote traffic (Figure 8b). With 10 cores per socket against rewriting of model code to support different LLM architecdual-socket Haswell’s 6, the baseline here is less stressed to tures [7]. In contrast, our compiler-driven approach integrates seamlessly into the standard TorchInductor pipeline, remaining begin with. Figure 10 reports every individual configuration. Of the 144, entirely transparent to the user without requiring modifications the 120 shown are measured; the remaining 24 could not be to the model code or inference loop. Algorithmic and hardware-aware optimizations have also measured on this 64 GB host, as the per-socket sharing of activations places their working set beyond available memory. been extensively explored to maximize CPU efficiency. Shen et al. proposed an automated INT4 weight-only quantization flow paired with tailored runtime kernels to accelerate CPU deployment [18]. He et al. introduced KV cache size reduction techniques and distributed inference optimizations relying on

the oneAPI Collective Communications Library (CCL) [19]. Our setting differs along three axes that make these frameAt the architectural level, Dice and Kogan detailed software- works a poor direct fit for intra-node CPU inference. Process level scaling optimizations for Transformer execution on CPUs model: these frameworks launch one OS process per device. without degrading model accuracy [20]. Furthermore, recent NUMA-TP instead runs in a single process and treats each innovations such as NoMAD-Attention replace traditional socket as a rank executed by an affinity-pinned thread group. multiply-add (MAD) computations with ultra-low-latency in- Communication: these frameworks invoke a collective library register lookups utilizing CPU Single-Instruction-Multiple-Data over a network or device interconnect. By contrast, our cross(SIMD) capabilities [21]. socket reduction directly sums socket-local buffers across the QPI/UPI link. Partitioning interface: these frameworks expose B. Efficient Attention Mechanisms partitioning through user-facing APIs or model wrapping, often Several of the CPU strategies above target memory directly. with architecture-specific configuration. NUMA-TP instead Weight-only quantization [18] and KV-cache compression [19] uses a compiler FX pass and requires no model-code changes. It shrink the memory footprint and reduce bandwidth use. These recasts distributed tensor parallelism as an intra-node, compilerapproaches reflect that memory bandwidth, rather than raw level transformation tailored to the low-latency, high-bandwidth compute, is the dominant constraint for LLM inference. IO- cross-socket interconnect. aware attention restructures computation around the memory hierarchy. On GPUs, FlashAttention uses aggressive tiling D. NUMA-Aware Execution Strategies to minimize intermediate transfers between HBM and onSeveral NUMA-aware deep learning optimizations are close chip SRAM [22]. FlashAttention-2 further optimizes thread- in spirit to our work. We describe their similarities and block parallelism and work partitioning [23]. Although these differences below. algorithms are tailored to the GPU memory hierarchy, they NUMA-Caffe restructures layer-level parallelism into a multiunderscore the importance of strict data locality. Our compiler- solver design to improve CNN training throughput on NUMA based methodology applies this principle to multi-socket CPU hardware [33]. Like our work, it treats NUMA locality as topologies by minimizing cross-socket interconnect traffic. important on multi-socket machines. However, its workload and mechanism differ from ours. NUMA-Caffe targets CNN C. Distributed Execution and Tensor Parallelism training with data parallelism, using per-node model replicas To mitigate synchronization overhead and manage massive over a partitioned batch. It implements this approach as a parameter scales, distributed paradigms like Megatron-LM framework-level Caffe modification. In contrast, we target introduced tensor parallelism to minimize synchronization Transformer inference with tensor parallelism applied as a points within both the Multi-Layer Perceptron (MLP) and atten- compiler transformation. A second line of work characterizes tion blocks [13]. Generalizing distributed tensor computations, the bottleneck rather than removing it. Studies of cache conMesh-TensorFlow established a flexible language for compiling tention and NUMA remote accesses for Transformer inference Single-Program-Multiple-Data (SPMD) graphs across high- on many-core processors [34] confirm that remote-DRAM performance supercomputing clusters [24]. Extending these traffic dominates. However, they stop at analysis and do not compiler abstractions to modern primitives, Triton-distributed provide a transparent mechanism to eliminate it. augmented the Triton compiler with OpenSHMEM. This enClosest to our setting is Sandwich [35], a hardware-centric ables native overlapping of computation and communication for LLM serving engine. Like NUMA-TP, Sandwich partitions distributed AI workloads [25]. Our methodology adapts large- models across NUMA nodes to reduce cross-socket traffic and scale tensor parallelism principles to the PyTorch FX graph. tailors execution to the prefill and decode phases. The key Specifically, it uses column- and row-parallel decompositions differences are integration and binding time. Sandwich is a within a single multi-socket CPU to localize DRAM access standalone engine that partitions models statically using aheadand avoid inter-socket bottlenecks. of-time, per-phase execution plans. NUMA-TP instead runs Beyond these foundational techniques, production distributed inside the standard TorchInductor pipeline as a post-grad FX frameworks expose tensor and data parallelism as first-class pass that reuses the Triton-CPU backend. Sandwich, like the runtime services. PyTorch’s native torch.distributed C++ engine ArcLight [7], requires adopting a separate runtime stack organizes execution as per-device processes that ex- and porting models to it. NUMA-TP instead dispatches kernels change shards through collective primitives. Its APIs include dynamically through a NUMA-aware thread pool and works DistributedDataParallel [26], Fully Sharded Data Parallel with unmodified PyTorch model code. (FSDP) [27], and a DTensor-based tensor-parallel interface. V. C ONCLUSION DeepSpeed combines ZeRO-style state partitioning across ranks [28], [29] with tensor-parallel inference kernels. These This paper introduces a compiler-driven, PyTorch-native systems delegate data movement to collective communication solution to the DRAM bandwidth bottleneck of the Large libraries such as NCCL [30], oneCCL [31], Gloo [32], or MPI. Language Model (LLM) prefill phase on multi-socket CPUs. Their all-reduce and all-gather routines are tuned for Extending the TorchInductor compilation path, it implements NVLink, PCIe, or InfiniBand and scheduled to overlap with NUMA-aware Tensor Parallelism (NUMA-TP) via post-grad computation.

FX graph transformations, coupled with an extended TritonCPU backend for per-socket buffer allocation and NUMAaware kernel dispatch. The transformation strictly enforces data locality, remains entirely transparent to the user, and yields prefill speedups of up to 2.63× on dual-socket Haswell and 1.87× on Cascade Lake platforms, exceeding what processlevel numactl policies achieve on the same sweep. Three directions follow from this work. First, NUMA-Torch should extend beyond the transformer sub-layers treated here to other locality-critical compute blocks, so that any block whose cost is dominated by remote memory traffic can be partitioned by the same passes. Second, the transformation should be chosen rather than applied unconditionally, with the compiler determining the communication-computation tradeoff for each block from dataflow analysis and the characteristics of the target platform, such as its socket count and interconnect bandwidth. The four-socket result makes the case for the second directly, since the same transformation that pays at high arithmetic intensity does not pay at low. The third sits on a separate track within LLM inference optimization. This work targets prefill, and carrying NUMA-aware partitioning into decode needs its own treatment of KV-cache placement and reduction scheduling rather than a direct reuse of these passes. The TorchInductor NUMA-TP FX passes and the extended Triton-CPU backend are publicly available at https://github. com/LLM-Compiler/numa-torch. ACKNOWLEDGMENT The authors sincerely acknowledge the financial assistance provided by the University of Moratuwa to cover the publication fees. Furthermore, the authors extend their gratitude to the anonymous reviewers for their thorough and insightful feedback. R EFERENCES [1] A. Radford et al., “Language models are unsupervised multitask learners,” OpenAI, Tech. Rep., 2019. [2] J. Devlin et al., “BERT: Pre-training of deep bidirectional transformers for language understanding,” in Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2019, pp. 4171–4186. [3] Gemma Team, “Gemma: Open models based on gemini research and technology,” Google DeepMind, Tech. Rep., 2024. [4] I. Calciu et al., “Black-box concurrent data structures for NUMA architectures,” in Proceedings of the Twenty-Second International Conference on Architectural Support for Programming Languages and Operating Systems, 2017, pp. 207–221. [5] Intel Corporation, Intel® Advisor, Intel Corporation, Santa Clara, CA, USA, 2024, available at: https://www.intel.com/content/www/us/en/ developer/tools/oneapi/advisor.html. [6] G. Gerganov, “llama.cpp: Port of Facebook’s LLaMA model in C/C++,” https://github.com/ggml-org/llama.cpp, 2023. [7] Y. Xu et al., “ArcLight: A Lightweight LLM Inference Architecture for Many-Core CPUs,” in Proceedings of the Annual Meeting of the Association for Computational Linguistics: System Demonstration, 2026. [8] J. Ansel et al., “PyTorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation,” in Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024, pp. 929–947. [9] P. Tillet, H. T. Kung, and D. Cox, “Triton: An intermediate language and compiler for tiled neural network computations,” in Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019, pp. 10–19.

[10] Triton Contributors, “Triton CPU: Experimental CPU backend for the triton compiler,” GitHub repository, 2024, experimental compiler backend. [Online]. Available: https://github.com/triton-lang/triton-cpu [11] Linux Kernel Developers, “Numa memory policy,” The Linux Kernel Documentation, 2026. [Online]. Available: https://www.kernel.org/doc/ html/latest/admin-guide/mm/numa memory policy.html [12] Linux Programmer’s Manual, “mbind(2),” kernel.org, 2026. [Online]. Available: https://man7.org/linux/man-pages/man2/mbind.2.html [13] M. Shoeybi et al., “Megatron-LM: Training multi-billion parameter language models using model parallelism,” in arXiv preprint arXiv:1909.08053, 2019. [14] A. Vaswani et al., “Attention is all you need,” in Advances in Neural Information Processing Systems, vol. 30, 2017. [15] A. Paszke et al., “PyTorch: An imperative style, high-performance deep learning library,” in Advances in Neural Information Processing Systems, vol. 32, 2019. [16] N. Shazeer, “Fast transformer decoding: One write-head is all you need,” arXiv preprint arXiv:1911.02150, 2019. [17] Linux Programmer’s Manual, “numactl(8),” man7.org, 2026. [Online]. Available: https://man7.org/linux/man-pages/man8/numactl.8.html [18] H. Shen et al., “Efficient llm inference on cpus,” arXiv preprint arXiv:2311.00502, 2023. [19] P. He et al., “Inference performance optimization for large language models on cpus,” arXiv preprint arXiv:2407.07304, 2024. [20] D. Dice and A. Kogan, “Optimizing inference performance of transformers on cpus,” arXiv preprint arXiv:2102.06621, 2021. [21] T. Zhang et al., “Nomad-attention: Efficient llm inference on cpus through multiply-add-free attention,” 2024. [22] T. Dao, D. Simpson, and C. Ré, “Flashattention: Fast and memory-efficient exact attention with io-awareness,” arXiv preprint arXiv:2205.14135, 2022. [23] T. Dao, “Flashattention-2: Faster attention with better parallelism and work partitioning,” arXiv preprint arXiv:2307.08691, 2023. [24] N. Shazeer et al., “Mesh-tensorflow: Deep learning for supercomputers,” arXiv preprint arXiv:1811.02084, 2018. [25] S. Zheng et al., “Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler,” arXiv preprint arXiv:2504.19442, 2025. [26] S. Li et al., “PyTorch distributed: Experiences on accelerating data parallel training,” Proceedings of the VLDB Endowment, vol. 13, no. 12, 2020. [27] Y. Zhao et al., “PyTorch FSDP: Experiences on scaling fully sharded data parallel,” Proceedings of the VLDB Endowment, vol. 16, no. 12, 2023. [28] J. Rasley et al., “DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters,” in Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020. [29] S. Rajbhandari et al., “ZeRO: Memory optimizations toward training trillion parameter models,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020. [30] NVIDIA, “NVIDIA collective communications library (NCCL),” https: //developer.nvidia.com/nccl, 2024. [31] UXL Foundation, “oneAPI Collective Communications Library (oneCCL),” https://github.com/uxlfoundation/oneCCL, accessed: Oct. 31, 2026. [32] Meta Platforms, Inc., “Gloo: Collective Communications Library,” https: //github.com/pytorch/gloo, accessed: Oct. 31, 2026. [33] P. Roy et al., “Numa-caffe: Numa-aware deep learning neural networks,” ACM Transactions on Architecture and Code Optimization, vol. 15, no. 2, pp. 1–26, 2018. [34] J. Jiang et al., “Characterizing and optimizing transformer inference on arm many-core processor,” in Proceedings of the 51st International Conference on Parallel Processing, 2023. [35] J. Zhao et al., “Sandwich: Separating prefill-decode compilation for efficient cpu llm serving,” arXiv preprint arXiv:2507.18454, 2025.

Related documents

• PDScope: Breaking the Prefill-Decode Seesaw for Memory-Bound MoE Inference#1490286
• PDScope: Breaking the Prefill-Decode Seesaw for Memory-Bound MoE Inference#1482967
• DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression#1033751
• TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning#1161013
• TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning#1157547