SBUF / PSUM Bank Geometry

All addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310; cp310/11/12 are byte-identical). The geometry constants live in neuronxcc/starfish/lib/libwalrus.so; the engine handle is in libBIR.so. For .text/.rodata the virtual address equals the file offset. Other wheels differ — treat every address as version-pinned.

Abstract

Every on-chip address the backend computes — every SBUF partition offset, every PSUM bank index — is bounded by a handful of compile-time constants baked into libwalrus.so. They describe two physically distinct memories: SBUF (the State Buffer, the 128-partition scratchpad the engines read and write) and PSUM (the Partial-SUM accumulator, a small array of fixed 2 KiB banks the PE array drains into). This page recovers those constants for all four hardware generations, the struct that holds each one, and the accessor chain the allocators walk to read them. The bar is that a reimplementer can reconstruct the SBUF and PSUM address spaces — partition count, per-partition byte budget, bank count, bank size — for Inferentia, Trainium (gen2), gen3, and gen4.

The geometry is not carried on the BIR bir::EngineInfo object. That object is an 8-byte {EngineType, id} handle with no size fields (1.01); the numbers live one indirection deeper, in a static singleton object tree Board → Device → Core built once at library load. The Core holds six engine sub-objects, and two of them — Statebuf (SBUF) and Psumbuf (PSUM) — are pure geometry records. An allocator reads a field by calling getArchModel(codename) and walking Board+0x8 → Device+0x10 → Core+0x{20,28} → field. This page documents that walk, the per-generation immediates the per-arch constructors stamp into those records, and the one structural surprise: Inferentia is a half-width part — 64 in the engines' second dimension where gen2+ are uniformly 128.

The numbers in this build are corroborated two ways. The constructor immediates are read directly off the <Arch>Statebuf / <Arch>Psumbuf ctor bodies (InferentiaPsumbuf, SundaStatebuf, …), and the field names come from a real C++ header shipped in the wheel, data/include/hwm/ctm/ctm.hpp, which declares class Statebuf, class Psumbuf, class Core with their member layout and doc comments. The two sources agree field-for-field, so the geometry here is CONFIRMED at the strongest available level: observed immediate plus named declaration.

For reimplementation, the contract is:

• The Board → Device → Core singleton tree and the two accessor offsets (Core+0x20 = Psumbuf*, Core+0x28 = Statebuf*) that key every geometry read.
• The Statebuf and Psumbuf record layouts and the four per-generation immediate sets that fill them.
• Partition-axis vs free-axis addressing: why SBUF is a 2-D (partition × byte) space and PSUM is a 1-D (bank) space, and what an address means in each.
• The half-width Inferentia case, so a reimplementer does not hard-code 128 into the engine second dimension.

The geometry object tree

Purpose

The hardware model is a static singleton tree, constructed once when libwalrus.so loads and shared by every pass that needs a hardware constant. A consumer never instantiates it — it asks getArchModel for the Board belonging to a codename and walks down to the field it wants. The tree exists so that the dozens of allocators, schedulers, and legalizers that depend on geometry all read one authoritative copy keyed by the compile target, rather than each carrying its own constants.

Entry Point

getArchModel(codename)           @0x17344c0   ── string-compare dispatch over 10 codenames → Board*
  └─ returns one of 4 static Board singletons (one per generation)
       Board   @ +0x08 → Device*
       Device  @ +0x10 → Core*
       Core    @ +0x20 → Psumbuf*   (PSUM geometry)
       Core    @ +0x28 → Statebuf*  (SBUF geometry)
       Core    @ +0x00..0x18 → Pe*, Pool*, Act*, Dve*  (compute-engine geometry)
       Core    @ +0x30 → dramSizeGb (HBM window — see 1.06)

The four Board objects are selected by codename. getArchModel clusters ten device aliases onto four generations; core_v5 (gen5, ArchLevel 50) has no Board and aborts with assert("0 && Unknown architecture") at ctm.cpp:0x5F — gen5 is a forward stub in this build. The clustering is CONFIRMED on both binaries (the libBIR side via getArchModel @0x478f90, the libwalrus side here):

CORRECTION — the alias count is ten, not eleven. The compare chain has exactly ten arms: tonga/inf1/inferentia (3) + sunda/trainium/trn1/inf2 (4) + cayman/gen3 (2) + core_v4 (1) = 10, as the table below enumerates. core_v5 is not a compare arm (its literal sits past the __assert_fail string in .rodata).

Algorithm

The walk is the same for every geometry field; only the terminal offset differs. This is the SBUF-size read, lifted verbatim from the loop-aware SB_Allocator constructor (the PSUM read is identical but ends at Core+0x20, Psumbuf+0):

// neuronxcc::backend::ColoringAllocatorWithLoop::Rep::SB_Allocator::SB_Allocator @0xa97750
function read_SB_SIZE(module, allocator):
    codename = module.getArch()                       // 0xa97b7a → ArchLevel2string
    board    = getArchModel(codename)                 // 0xa97b82, rbx = Board*
    device   = *(Device**)(board  + 0x08)             // 0xa97baa, Board+0x8
    core     = *(Core**)  (device + 0x10)             // 0xa97bb3, Device+0x10
    statebuf = *(Statebuf**)(core + 0x28)             // 0xa97bb7, Core+0x28
    sb_size  = *(uint32*)(statebuf + 0x00)            // 0xa97bbb, Statebuf+0 = partitionSize
    allocator[0x2b8] = sb_size                        // 0xa97bbd, cached at allocator+696
    // 0xa97bc3: a follow-on `cmp $0xa,%r14d` (ArchLevel 10 = Inferentia) forks the
    //           gen1 byte budget; the stored value itself is the Statebuf+0 immediate.

NOTE — the cache slot allocator+0x2b8 (= 696 decimal) is the field earlier reports called "a1+696". It holds SB_SIZE bytes (Statebuf+0), not a partition count; the 128-partition figure is the separate Statebuf+0x8 field. A reimplementer must keep the two distinct — they are two fields of one record, never one overloaded slot.

Considerations

The geometry is a baked compile-time constant per codename, not device-config or runtime-probe driven. The compiler decides the SBUF/PSUM budget purely from the target string. The class family is defined once by the shipped ctm hardware-model library (ctm.hpp header + ctm.cpython-3xx.so); the per-arch constructors that hold the immediates are compiled into both libwalrus.so and libBIR.so, which carry the geometry tree byte-identically — a cross-binary consistency check (two compilations of one ctm source agree), not an independent derivation (1.01 makes that explicit). libBIR additionally exposes the engine-multiplicity counts (getEngineCount), a different facet. The genuinely independent corroborator of specific values — the 2048-byte bank size, the clock frequencies — is the cost model (bir::Hwm, e.g. getPsumBankSize), a separate code path documented in 1.04. The geometry on this page is read from the libwalrus copy.

The Statebuf record — SBUF geometry

Purpose

Statebuf describes the State Buffer: the engines' primary on-chip scratchpad, organized as 128 independent partitions, each a flat per-partition byte budget. It is a 2-D address space — (partition, byte-offset) — and the SBUF allocator (SB_Allocator, Part 8) colors tensors into partition × byte rectangles. The record holds the two numbers that bound that rectangle (per-partition size, partition count) plus alignment and a reserved-region size.

Record layout

The base constructor Statebuf::Statebuf (@0x17341e0) takes six arguments. The names are CONFIRMED from ctm.hpp:129's class Statebuf declaration; the offsets are read off the constructor's store sequence:

CORRECTION (D-M12 → ctm.hpp) — an earlier decode labeled Statebuf+0xC a "SBUF sub-bank granule" and Statebuf+0x18 a "partition alignment quantum". The shipped ctm.hpp names them midPartition and reservedSize respectively, with reservedSize defaulting to 16384 in the declaration. The 0x40 (=64) at +0xC is midPartition (the partition-quadrant midpoint), and the 0x4000/0x4008 pushed argument is reservedSize.

Per-generation immediates

Each <Arch>Statebuf constructor sets up the six arguments inline and tail-calls the base Statebuf ctor. The immediates below were read off the ctor bodies and independently re-disassembled; all four sets match byte-for-byte:

The total SBUF capacity per generation is SB_SIZE × 128 partitions: 12 MiB (gen1), 24 MiB (gen2), 28 MiB (gen3), 32 MiB (gen4). A reimplementer rarely needs the total — the allocator's budget is per partition, because every engine accesses all 128 partitions in lockstep and a tensor's footprint is bounded by its largest per-partition slice.

QUIRK — numPartitions = 128 is constant across all four generations, including Inferentia. The half-width nature of Inferentia (next section) does not touch SBUF — its 128 partitions are full-width. The "64" on Inferentia appears only in the compute engines and PSUM. A reimplementer who assumes "Inferentia is 64 everywhere" will mis-size SBUF by 2×.

Considerations

The function map below ties the SBUF symbols to addresses. The <Arch>Statebuf ctors are full bodies (~0x4d bytes), not thunks — they construct the Align descriptor inline and reach the base ctor via the PLT stub _ZN8StatebufC2EjjjjR5Align_j@plt (0x5f3870). The in-binary Statebuf::Statebuf @0x17341e0 is the inlinable copy.

The Psumbuf record — PSUM geometry

Purpose

Psumbuf describes the Partial-SUM buffer: the matmul accumulator space the PE systolic array drains into (MemoryType::PSUM = 32, 1.01). Physically it is a small array of fixed-size banks, each 2 KiB. Unlike SBUF, PSUM has no addressable byte axis the allocator searches — a tensor occupies a whole number of banks, and the PSUM allocator is a 1-D bank-interval colorer (Part 8). The record holds the bank count, the partition count, and the bank size.

Record layout

The base constructor Psumbuf::Psumbuf (@0x17341b0) takes three arguments. Names are CONFIRMED from ctm.hpp:111's class Psumbuf:

bufLen64/bufLen32 are derived element-count helpers the header marks "unnecessary - deprecate these"; the allocator uses only numBanks and partSize.

Per-generation immediates

Each <Arch>Psumbuf ctor is a 5-instruction thunk: load the three args, tail-jmp to the base ctor via PLT (_ZN7PsumbufC2Ejjj@plt, 0x5ff360). Immediates re-disassembled and independently confirmed:

NOTE — the previously-open bank count is now closed. The PSUM allocator caches numBanks at its this+0 from a getArchModel walk, and an earlier study of that allocator (D-K10) could only infer the value = 8 from the 16 KiB PSUM window — the per-arch integer was "arch-table-resident, not a literal" in the colorer. The immediates above are that literal: read directly off the four <Arch>Psumbuf ctor bodies, numBanks = {4, 8, 8, 8} for gen1/2/3/4. The bank count is CONFIRMED, no longer INFERRED. Inferentia is the only 4-bank part; gen2+ are uniformly 8.

The bank size 2048 is corroborated three independent ways and is CONFIRMED: (1) the Psumbuf+0x08 immediate above; (2) three per-arch getter bodies that each compile to a literal return —

; TrainiumHwm::getPsumBankSize @0x1851900   (Gen3Hwm @0x1857cf0, CoreV4Hwm @0x185e570 identical)
1851900:  b8 00 08 00 00   mov    $0x800,%eax        ; 2048
1851905:  c3               ret

— and (3) getPsumBankNum (@0x1088bd0) computing a tensor's bank footprint as ceil(tensorBytes / getPsumBankSize()), i.e. ceil(bytes / 2048).

The PSUM read-chain

The PSUM allocator's constructor reads numBanks with the same getArchModel walk as SBUF, ending one offset earlier (Core+0x20 = Psumbuf* rather than Core+0x28 = Statebuf*):

// neuronxcc::backend::ColoringAllocatorWithLoop::Rep::PSUM_Allocator::PSUM_Allocator @0xad9970
function read_NUM_PSUM_BANKS(module, allocator):
    board    = getArchModel(module.getArch())         // 0xada120
    device   = *(Device**)(board  + 0x08)             // 0xada125, Board+0x8
    core     = *(Core**)  (device + 0x10)             // 0xada131, Device+0x10
    psumbuf  = *(Psumbuf**)(core  + 0x20)             // 0xada135, Core+0x20
    allocator[0x00] = *(uint32*)(psumbuf + 0x00)      // 0xada139→0xada13b, Psumbuf+0 = numBanks

The colorer then guards every tensor: if numBanks < tensor.liveBanks, it aborts with "<name> is too big for PSUM, requires <N> banks with PSUM size <M>" (Part 8). That M is the cached numBanks, so the constant on this page is exactly the budget a reimplementer's PSUM allocator must enforce.

Considerations

Partition-axis vs free-axis addressing

A reimplementer reasoning about on-chip addresses needs the shape of each address space, not the full allocator. The two memories are shaped differently, and the difference is the single most useful thing on this page.

  SBUF — 2-D address space  (partition, byte)              PSUM — 1-D address space  (bank)
  ┌───────────────────────────────────────────┐           ┌─────┬─────┬─────┬─────┬───┐
  │ part 0  │ 0 .................... SB_SIZE-1 │           │bank0│bank1│bank2│ ... │ N │   N = numBanks-1
  │ part 1  │ 0 .................... SB_SIZE-1 │           └─────┴─────┴─────┴─────┴───┘
  │   ...   │ ...                              │            each bank = partSize (2048) bytes,
  │ part 127│ 0 .................... SB_SIZE-1 │            spanning ALL numPartitions at once
  └───────────────────────────────────────────┘
  free axis = byte offset within a partition                free axis = bank index
  partition axis = which of 128 partitions                  partition axis = implicit (a bank is full-width)

SBUF — partition axis × free axis. An SBUF address is a pair: a partition (0…numPartitions-1) on one axis and a byte offset (0…SB_SIZE-1) within that partition on the free axis. The engines operate partition-parallel — one instruction touches a strip across many partitions at the same free offset — so the allocator searches a 2-D rectangle: how many partitions tall (the partition dimension of the tensor) by how many bytes wide. The per-partition budget SB_SIZE bounds the free axis; numPartitions (128) bounds the partition axis. A reimplementer's SBUF allocator must place a tensor as a partitions × bytes rectangle inside the 128 × SB_SIZE grid.

PSUM — bank axis only. A PSUM address is a single bank index (0…numBanks-1). There is no byte axis the allocator searches: a bank is an indivisible 2048-byte unit that already spans all numPartitions partitions. A tensor occupies ceil(bytes / 2048) consecutive banks (getPsumBankNum), and "where it lives" is just its base bank. The partition axis is implicit — every bank is full-width — which is why the PSUM colorer is 1-D where the SBUF colorer is 2-D. A reimplementer's PSUM allocator is a bank-interval first-fit over numBanks slots, nothing more.

GOTCHA — numPartitions differs between the two records on Inferentia. Statebuf.numPartitions is 128 on every generation, but Psumbuf.numPartitions is 64 on Inferentia (128 on gen2+). They are not the same field and not the same number on gen1. The PSUM partition count matters for how many PE columns a bank serves, not for the bank-index arithmetic above — but a reimplementer copying "128 partitions" across both records will be wrong on Inferentia's PSUM.

The half-width Inferentia case

Across the four generations, gen2/3/4 are uniformly 128 in every "second dimension", but Inferentia (gen1) is half-width — 64 in the compute-engine and PSUM second dimension while keeping a full-width 128-partition SBUF. This is the one place a reimplementer cannot assume 128. The pattern, read off the per-arch compute-engine and PSUM ctors:

The reading is that Inferentia's systolic array is a 128×64 PE grid (half the columns of gen2+'s 128×128), its PSUM is half the banks and half the partitions, and its pooling/activation engines process 64 channels. Only the SBUF stays full-width at 128 partitions. Every other generation is the uniform 128 a naive reimplementation would assume.

CORRECTION (S2-08 → D-M12) — the "128 partition" constant earlier studies attributed to a field at EngineInfo+100 is not a bir::EngineInfo field — that object is the 8-byte {EngineType, id} handle with no geometry. The 128 is one of the Core sub-object geometry fields (Statebuf+0x8, Psumbuf+0x4, Pe+0x4, …), and on Inferentia several of those are 64, not 128. The geometry is keyed by EngineType via the Core sub-object table, not by an EngineInfo-indexed array.

At-a-glance: the full geometry matrix

The consolidated per-generation table. Every cell is CONFIRMED from a constructor immediate or a literal getter body, re-disassembled independently.

Related Components

Cross-References

• 1.01 Arch Object Model — the bir::EngineInfo handle and EngineType enum that key the Core sub-object table; PSUM = MemoryType 32, SBUF = 16.
• 1.04 Hardware Constant Matrix — the full per-generation constant table this page's SBUF/PSUM rows feed into.
• 1.06 DRAM / HBM Geometry — the off-chip side: Core.dramSizeGb read through the identical getArchModel walk, and the HBM budget check.
• SBUF allocator (Part 8) — consumes SB_SIZE and the 128-partition layout to 2-D color State-Buffer tensors.
• PSUM allocator (Part 8) — consumes numBanks and the 2048-byte bank size to 1-D color the partial-sum accumulator.