HBM DMA 对齐契约

本页中的所有地址、结构体偏移、行号和魔数常量均适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，build libtpu_lts_20260413_b_RC00，clang/LLVM trunk）。其他版本会有所不同。

摘要

libtpu 中的每一次 HBM DMA 都受同一个对齐量子约束：jf_driver::kHbmMinimumDmaAlignment = 1024 B。这不是建议值，而是在两个独立层级、通过两种不同机制被执行两次：

1. DMA 发起时（tpu::JfDmaIssuer::WritePremappedHbm 及其同族函数），字节偏移、缓冲区大小和最小传输长度会通过可恢复的 RetCheck 校验。未对齐请求不会使进程崩溃；它会短路进入调用方的错误回调，并携带非 OK 的 Status。
2. 描述符构建时（asic_sw::driver::deepsea::jxc::HbmWriteDescriptor::SetHbmAddress），绝对 HBM 地址会用 0x3FF 掩码检查，并通过致命的 LogMessageFatal CHECK 校验。地址到达硬件描述符时已经太晚，无法恢复；这里的未对齐地址是编程错误，会导致 abort。

1024 B 下限就是 HBM BestFit 分配器把 1024 B 用作其 HBM Config.alignment_in_bytes_ 量子的原因：如果分配器发放更细粒度的 HBM 偏移，就会把 DMA 发起器拒绝的地址交给它。本页负责说明 DMA 对齐下限、WritePremappedHbm 预映射路径，以及对齐 CHECK。它不负责分配器的空闲链表数学（那是 hbm-allocator.md）或片内描述符字段布局（那是 ../dma/intra-chip-descriptor.md）；本页只把它们链接起来。

对于重新实现，契约如下：

• 下限。 kHbmMinimumDmaAlignment = 1024 = 0x400；向下取整掩码是 ~0x3FF = 0xFFFFFFFFFFFFFC00；取模测试是 x & 0x3FF == 0。同一个常量也作为 kMinimumDmaLengthBytes = 1024，即最小合法传输长度。
• 三个发起时 RetCheck（可恢复）：byte_offset % 1024 == 0、size % 1024 == 0、size >= 1024。每个失败都会经由 ScheduleCallbackOnError，用户回调会收到错误，不会 abort。
• 两个描述符时 CHECK（致命）：address < kAddressOffsetMaxBytes（2^50）和 (address & 1023) == 0。它们保护位打包到 JXC HbmWriteDescriptor 中的 64 位 HBM 地址字段。
• 分配器量子镜像该下限。 HBM 的 Config.alignment_in_bytes_ 就是 1024 B，因此分配器发出的每个偏移已经是该下限的倍数。
• 16 KiB 编译期程序对齐（FLAGS_xla_jf_program_hbm_alignment_in_kib = 16）是一个严格得多的、上游量子，在 MSA 放置前应用；它与 DMA 下限无关，也不会放宽该下限。

DMA 下限常量	jf_driver::kHbmMinimumDmaAlignment = 1024 B (0x400)
向下取整掩码	~0x3FF = 0xFFFFFFFFFFFFFC00
取模测试	x & 0x3FF == 0
最小传输长度	jf_driver::kMinimumDmaLengthBytes = 1024 B（同一常量）
发起时保护	JfDmaIssuer::WritePremappedHbm @ 0xe73db80（3 个 RetCheck，可恢复）
单块变体	JfDmaIssuer::WritePremappedHbmSingleChunk @ 0xe73eec0
读取变体	JfDmaIssuer::ReadPremappedHbm @ 0xe73c880
描述符时保护	HbmWriteDescriptor::SetHbmAddress @ 0xe7ce7e0（2 个 CHECK，致命）
最大 HBM 地址	kAddressOffsetMaxBytes = 0x4000000000000 (2^50)
发起源码标签	learning/45eac/tpu/runtime/hal/internal/jxc/jf_dma_issuer.cc
描述符源码标签	platforms/asic_sw/driver/deepsea/jxc/common/hbm_write_descriptor.h
分配器量子	HBM Config.alignment_in_bytes_ = 1024 B（见 hbm-allocator.md）
置信度	CONFIRMED（字节锚定），除非某行或标注另有说明

---

范围和边界

本页负责 HBM DMA 对齐下限，以及执行该下限的发起时/描述符时检查。相邻关注点放在别处；不要在这里重复：

关注点	负责页面
BestFitAllocator 空闲链表、best-fit 搜索、合并、拆分策略，以及分配器的对齐量子	hbm-allocator.md
片内 DMA 描述符字段布局（slice offset/size、stride、ring index）	../dma/intra-chip-descriptor.md
五层片上内存映射以及 HBM 所在位置	overview.md
设备缓冲区的偏移如何映射到设备端 tile 布局	tpu-buffer-layout.md
Host→device / device→host 传输路径和预映射 staging 池	overview.md, ../dma/host-device-dma.md

本页负责：1024 B 下限常量、WritePremappedHbm 预映射路径和对齐 CHECK，也就是每个 HBM DMA 必须满足的边界条件以及验证它们的两个代码位置。

NOTE — 这里的 "premapped" 指的是某次 DMA 的 HBM 侧端点，其主机缓冲区已经向驱动注册（pinned 且 DMA-mapped），因此发起路径可以跳过 bounce-buffer 复制并直接流式传输。预映射的主机池（PremappedMemoryManager 对 posix_memalign 分区的轮转）是分配器侧的独立对象；见 overview.md。本页中的函数是消费这些预映射缓冲区的设备 HBM 侧发起路径。

---

下限常量

jf_driver::kHbmMinimumDmaAlignment 是 1024（0x400）。虽然二进制未 strip（nm -C 能解析每个方法），但该常量自身没有符号（nm 找不到任何 kHbmMinimumDmaAlignment 条目）：它是编译期常量，优化器将其折叠进每个使用点的立即数，而不是作为 .rodata 数据发出。因此它是从反编译中三个独立且相互一致的编码恢复出来的：

编码	出现位置	反编译证据
向下取整掩码	WritePremappedHbm、ReadPremappedHbm、single-chunk	v8 = (a2 + 1023 if a2<0 else a2) & 0xFFFFFFFFFFFFFC00
取模（向上归零）掩码	所有发起位置的 size 检查	v6 & 0x3FF (0x3FF = 1024 − 1)
字面量	最小长度比较	*(_QWORD *)&v50 = 1024; if (size < 0x400) RetCheck
kAddressOffsetMaxBytes 字面量	SetHbmAddress	MakeCheckOpString(a2, 0x4000000000000, "address < kAddressOffsetMaxBytes")

掩码 0xFFFFFFFFFFFFFC00 正是 ~0x3FF，即标准的向下取整到 1024 操作。由 if (a2 >= 0) 保护的 + 1023 调整，是编译器在掩码前实现有符号向零取整的分支形式；它只对负偏移有意义（负偏移本身也会在后续 RetCheck 中失败）。对实际出现的非负偏移来说，v8 = a2 & ~0x3FF 就是向下取整后的值，而 a2 != v8 恰好表示 a2 未对齐。

kMinimumDmaLengthBytes 也是 1024。反编译在每个发起点都把字面量 1024 加载为比较操作数，因此最小合法 HBM 传输是一个完整对齐量子；不存在小于 1024 字节的 DMA。

NOTE — 该下限是一个字节量子，而不是 granule 数。每个发起例程后续的 DMA 分块算术会除以 a1[22]（发起器偏移 +0xB0 处的每发起器最大块大小），以把大传输拆成 ring 大小的块；该块大小是另一个更大的量子，并在 ../dma/intra-chip-descriptor.md 与描述符一同记录。1024 B 下限是每个块边界继承的对齐，不是块大小本身。

---

WritePremappedHbm — 发起时保护

tpu::JfDmaIssuer::WritePremappedHbm（0xe73db80）是规范的预映射 HBM 写入。其签名（demangled）是：

void JfDmaIssuer::WritePremappedHbm(
        int64_t                          byte_offset,   // HBM offset to write at
        SlicedDmaBuffer                  buffer,        // premapped source view (ptr+size at a3[1])
        absl::AnyInvocable<void(const Status&)> on_done,
        const std::optional<asic_sw::deepsea::SyncFlag>& sync)

该函数做的第一件事——在接触任何队列、描述符或硬件状态之前——是在 byte_offset（a2）和缓冲区的 size（a3[1]，捕获为 v6）上运行三个对齐保护：

// WritePremappedHbm @ 0xe73db80 — alignment guards, byte-confirmed
v6 = a3[1];                                   // buffer.size()
v7 = a2 + 1023; if (a2 >= 0) v7 = a2;         // signed round-toward-zero
v8 = v7 & 0xFFFFFFFFFFFFFC00LL;               // floor to 1024  (~0x3FF)
*(int64_t*)&v43 = a2 - v8;                    // remainder = byte_offset % 1024

// CHECK 1 (line 440): byte_offset is 1024-aligned
if (a2 != v8) {                               // remainder != 0
    v32 = MakeCheckOpString<long,int>(&v43, &v50,
              "byte_offset % jf_driver::kHbmMinimumDmaAlignment == 0");
    if (v32) { v33 = 440; goto fail; }        // RetCheckFailSlowPath -> error callback
}

// CHECK 2 (line 442): buffer size is a multiple of 1024
*(int64_t*)&v43 = v6 & 0x3FF;
if ((v6 & 0x3FF) != 0) {
    v32 = MakeCheckOpString<unsigned long,int>(&v43, &v50,
              "size % jf_driver::kHbmMinimumDmaAlignment == 0");
    if (v32) { v33 = 442; goto fail; }
}

// CHECK 3 (line 444): transfer is at least one quantum long
*(int64_t*)&v43 = v6;
*(int64_t*)&v50 = 1024;
if ((unsigned long)v6 < 0x400) {              // size < kMinimumDmaLengthBytes
    v32 = MakeCheckOpString<unsigned long,unsigned long>(&v43, &v50,
              "size >= jf_driver::kMinimumDmaLengthBytes");
    if (v32) { v33 = 444; goto fail; }
}

fail 标签是针对源码 learning/45eac/tpu/runtime/hal/internal/jxc/jf_dma_issuer.cc 及记录行号的 RetCheckFailSlowPath，随后调用 tpu::ScheduleCallbackOnError，它用失败的 Status 调用调用方的 on_done（a4），然后销毁两个 StatusBuilder。这是可恢复的：未对齐的预映射写不会 abort 进程，而是像任何其他 DMA 错误一样，以非 OK 状态异步完成。

为什么是三个检查，且按这个顺序

这三个检查构成后续分块算术的完整前置条件。一旦 byte_offset 和 size 都是 1024 的倍数且 size >= 1024：

• 快速路径（size <= a1[22]，即发起器偏移 +0xB0 处的最大块大小）会把单个 {offset, MaybeOwningDmaBuffer, optional<SyncFlag>} 元组入队到发起器的 BufferedQueue（a1[5] + 0x100）；
• 慢速路径会在一个 AsyncTaskGroup 下把传输拆分成 ceil(size / max_chunk) 个块，每个块从已对齐的 byte_offset 和已对齐的每块大小继承对齐，并且每个块都有一个 slice_offset + slice_size <= base.size() 边界 CHECK（dma_buffer_utils.h:40 处的致命 LogMessageFatal）。

由于输入已经预先取整，每个块边界自动是 1024 对齐的，而最终到达 SetHbmAddress 的每块描述符地址都保证能通过下面的致命掩码检查。发起时 RetCheck 是可恢复的第一道防线；描述符时 CHECK 是致命的后备防线。

函数	地址	角色
WritePremappedHbm	0xe73db80	预映射 HBM 写；3 个发起时对齐 RetCheck + 分块拆分
WritePremappedHbmSingleChunk	0xe73eec0	单块变体；相同检查，行 551/553/555
ReadPremappedHbm	0xe73c880	预映射 HBM 读；镜像检查位于行 344/346/348
RetCheckFailSlowPath	（发起器本地）	构建错误 Status，送入 ScheduleCallbackOnError
ScheduleCallbackOnError	（发起器本地）	用错误触发 on_done；不 abort

---

单块和读取变体

预映射家族中的检查序列完全相同；只有记录的源码行号和保存 size 的字段不同。

WritePremappedHbmSingleChunk（0xe73eec0）从 *(buffer + 8) 而不是 a3[1] 读取 size，其检查落在行 551 / 553 / 555，消息以 buffer.size() 为键：

// WritePremappedHbmSingleChunk @ 0xe73eec0 — same floor, different lines
v7 = (a2 + 1023 if a2<0 else a2) & 0xFFFFFFFFFFFFFC00;
if (a2 != v7)  RetCheck("byte_offset % jf_driver::kHbmMinimumDmaAlignment == 0");  // 551
v8 = *(uint64_t*)(a3 + 8);                                                          // buffer.size()
if (v8 & 0x3FF)  RetCheck("buffer.size() % jf_driver::kHbmMinimumDmaAlignment == 0");// 553
if (v8 < 0x400)  RetCheck("buffer.size() >= jf_driver::kMinimumDmaLengthBytes");     // 555
// then a fatal CHECK: tensor_node_ != nullptr (line 559) before enqueue

ReadPremappedHbm（0xe73c880）精确镜像写路径：byte_offset % 1024 == 0（行 344）、size % 1024 == 0（行 346）、size >= 1024（行 348），确认该下限在读写之间对称。随后读路径会按发起器偏移 +0xB0 处的每读最大大小（*(_QWORD*)(v50 + 176)）分块，并把每块分派给 NodeFabricTransferHbmToHostInternal。

NOTE — 单块变体额外增加了一个多块路径不会在顶部暴露的致命 CHECK：tensor_node_ != nullptr（行 559）。这是结构性前置条件（发起器必须绑定到 tensor node），不是对齐检查，但它说明单块路径期望不经 AsyncTaskGroup fan-out 直接入队。

---

SetHbmAddress — 致命描述符后备防线

当某个块的 HBM 地址被写入 JXC 硬件描述符时，asic_sw::driver::deepsea::jxc::HbmWriteDescriptor::SetHbmAddress（0xe7ce7e0）会重新校验它——这次是致命的：

// HbmWriteDescriptor::SetHbmAddress @ 0xe7ce7e0 — byte-confirmed
__int64 SetHbmAddress(HbmWriteDescriptor *this, uint64_t address) {
    uint64_t v7 = address;
    // CHECK A (line 37): address fits the 50-bit HBM address field
    if (address >> 50) {                                  // address >= 2^50
        v4 = MakeCheckOpString(address, 0x4000000000000LL,
                 "address < kAddressOffsetMaxBytes");
        v5 = 37;  goto fatal;                             // LogMessageFatal -> Flush -> ~ -> abort
    }
    // CHECK B (line 38): address is 1024-aligned
    if (address & 0x3FF) {                                // low 10 bits set
        v4 = MakeCheckOpString(address & 0x3FF, 0,
                 "(address & (kHbmMinimumDmaAlignment - 1)) == 0");
        v5 = 38;
fatal:
        LogMessageFatal("platforms/asic_sw/driver/deepsea/jxc/common/hbm_write_descriptor.h", v5, v4);
        LogMessage::Flush(...); ~LogMessageFatal(...);    // process abort
    }
    return BitCopy(this, 64, &v7, 0, 64);                 // pack the 64-bit address field
}

两个事实使它成为后备防线，而不是第一道防线：

• 它是致命的。 与发起时 RetCheck 不同，这里的失败是 LogMessageFatal → Flush → 析构函数 → abort。没有错误回调，也没有恢复。未对齐地址到达描述符会被视为 libtpu/编译器 bug，因为发起路径本应已经以可恢复方式拒绝它。
• 它用位掩码而不是取模表达下限。 (address & (kHbmMinimumDmaAlignment − 1)) == 0，其中 kHbmMinimumDmaAlignment − 1 = 0x3FF，与发起器的 byte_offset & 0x3FF 测试在代数上等价，确认两层共享同一个 1024 常量。地址最后通过 BitCopy(this, 64, &addr, 0, 64) 位打包进 64 位描述符字段。

CHECK A 还把地址限制到 kAddressOffsetMaxBytes = 0x4000000000000 = 2^50，也就是 JXC HBM 地址字段为 50 位宽（1 PiB 可寻址跨度）。该地址送入的描述符字段布局记录在 ../dma/intra-chip-descriptor.md。

常量 / 检查	值 / 消息	严重性
kHbmMinimumDmaAlignment	1024 (0x400)；掩码 0x3FF	—
kMinimumDmaLengthBytes	1024（同一常量）	—
kAddressOffsetMaxBytes	0x4000000000000 (2^50)	—
byte_offset % 1024 == 0	发起时，WritePremappedHbm:440	可恢复 RetCheck
size % 1024 == 0	发起时，WritePremappedHbm:442	可恢复 RetCheck
size >= 1024	发起时，WritePremappedHbm:444	可恢复 RetCheck
(address & 1023) == 0	描述符时，SetHbmAddress:38	致命 CHECK
address < 2^50	描述符时，SetHbmAddress:37	致命 CHECK

---

分配器量子与下限的关系

HBM BestFit 分配器用 Config.alignment_in_bytes_ = 1024 B 构造其 HBM 实例，精确等于 DMA 下限；正是这一事实把两个子系统连在一起：

• 分配器发出的每个偏移已经是 1024 的倍数。 Allocate 通过 (size + align − (size != 0)) & −align 把每个请求向上取整到 alignment_in_bytes_，构造函数把区域末尾向下取整到同一量子。因此 base_offset + offset，也就是 Allocate 返回的地址，总是 1024 的倍数。当该地址稍后成为 DMA byte_offset 时，发起时 byte_offset % 1024 == 0 RetCheck 按构造即满足。
• 分配大小也是 1024 的倍数，因此覆盖整个缓冲区的 DMA 会自动满足 size % 1024 == 0。size >= 1024 下限意味着缓冲区必须至少为一个量子，分配器的向上取整也会为任何非零请求保证这一点。
• 这种关系是契约，而不是巧合。 如果分配器使用更小的 HBM 量子（例如 256 B），它可能发出类似 0x300 的 byte_offset，会被发起器的 RetCheck 拒绝；更糟的是，如果它漏过去，SetHbmAddress 会以致命方式拒绝。因此分配器量子被固定到 DMA 下限，作为保持每个发出地址 DMA 合法的最小值。分配器量子是最小值；编译期程序对齐是叠加在其上的更大、独立量子。

            compile time                         load / run time
   ┌──────────────────────────┐        ┌────────────────────────────────────┐
   │ MSA / ProgramMemory-     │        │ BestFitAllocator (HBM)               │
   │ Allocator                │ proto  │   Config.alignment_in_bytes_ = 1024  │
   │   round up to 16 KiB ────┼───────▶│   round up to 1024  (no-op on        │
   │   (program HBM align)    │offsets │     already-16KiB static offsets)    │
   └──────────────────────────┘        └───────────────┬──────────────────────┘
                                                        │ base+offset (×1024)
                                                        ▼
                                        ┌────────────────────────────────────┐
                                        │ JfDmaIssuer::WritePremappedHbm       │
                                        │   RetCheck byte_offset % 1024 == 0   │ recoverable
                                        │   RetCheck size       % 1024 == 0    │
                                        │   RetCheck size       >= 1024        │
                                        └───────────────┬──────────────────────┘
                                                        │ per-chunk address
                                                        ▼
                                        ┌────────────────────────────────────┐
                                        │ HbmWriteDescriptor::SetHbmAddress    │
                                        │   CHECK address < 2^50               │ FATAL
                                        │   CHECK (address & 1023) == 0        │
                                        └──────────────────────────────────────┘

完整的分配器量子机制（向上取整公式、区域末尾向下取整、max_aligned_size_，以及为什么 padding 是由分配拥有的内部碎片）位于 hbm-allocator.md § 对齐量子，本页不重复。

---

16 KiB 编译期程序对齐

还有一个第二个、粗得多的 HBM 量子，完全在编译期起作用，不能与 1024 B DMA 下限混淆。

FLAGS_xla_jf_program_hbm_alignment_in_kib（默认 16，.data 全局位于 0x223b4888）会在 MSA 放置之前把每个程序级 HBM tensor 向上取整到 16 KiB。这是 XLA 侧逻辑，在编译期间运行于主机，用于适配 XLA 的 stride / sub-tile 寻址方案和 slice-prefetch 边界；这些关注点与 DMA 引擎的对齐要求无关。

两个量子之间的关系：

• 16 KiB 是 1024 B 的倍数（16 × 1024 = 16384 = 16 × 1024），因此任何 16-KiB 对齐的静态偏移都会自动 1024 对齐。当 MSA 冻结的偏移在运行时分配器中重放时，分配器自己的 1024 B 向上取整对它们是 no-op。
• 两者彼此独立。 16 KiB 程序对齐比 DMA 下限严格得多，并且只应用于静态（MSA 放置）表面。动态运行时分配（transfer staging、async-copy scratch）直接通过分配器的 1024 B 量子，并且不是 16-KiB 对齐。两个表面都满足 DMA 下限；只有静态表面满足程序对齐。
• 调大该 flag 不会改变 DMA 下限。 提高 xla_jf_program_hbm_alignment_in_kib 会让静态 tensor 对齐更粗（因内部碎片浪费更多 HBM），但不能放宽 1024 B 发起时/描述符时检查，因为这些是二进制中的硬编码常量。

内存空间	最小 DMA / 访问对齐	编译期程序对齐（默认）	量子来源
HBM / kHbm	1024 B (jf_driver::kHbmMinimumDmaAlignment)	16 KiB (xla_jf_program_hbm_alignment_in_kib = 16)	DMA 下限（二进制常量）；程序 flag（XLA）
kPinnedHbm	1024 B（同一下限，加上主机侧 pinning 锁）	16 KiB（继承自 HBM）	相同
VMEM / kVmem	按 generation 而定（32 / 64 / 128 B；来自 chip_parts.binarypb）	按 codec bundle 宽度而定	Config.alignment_in_bytes_ — 见 vmem-allocator.md
SMEM / SFLAG / CMEM	按 generation 而定（通常 32 B；SFLAG 4 B）	按 tier 而定	Config.alignment_in_bytes_

NOTE — 上面的按 generation 的 VMEM / SMEM / CMEM 行只作为上下文；它们的对齐由该 tier 的 Config.alignment_in_bytes_ 管理，该值来自嵌入式 chip_parts.binarypb 资源，并由同级 tier 页面负责（vmem-allocator.md、cmem-pool.md）。这里仅负责 HBM 和 kPinnedHbm 行，即 1024 B DMA 下限与 16 KiB 程序对齐。

---

重新实现者检查清单

要复现 HBM DMA 对齐契约：

1. 定义 kHbmMinimumDmaAlignment = 1024 和 kMinimumDmaLengthBytes = 1024。 它们数值相同，但表达不同意图（对齐 vs. 最小长度）。派生掩码是 0x3FF；向下取整掩码是 ~0x3FF。
2. 在 DMA 发起时，以可恢复方式校验。 在任何分块之前，检查 byte_offset % 1024 == 0、buffer.size() % 1024 == 0 和 buffer.size() >= 1024。失败时，把非 OK 状态路由到完成回调；不要 abort。这些是返回 Status 的 RetCheck。
3. 在描述符构建时，以致命方式校验。 当把绝对 HBM 地址打包进硬件描述符时，CHECK address < 2^50 和 CHECK (address & 1023) == 0。这些是 LogMessageFatal；到这一层时未对齐地址是逻辑错误，因为步骤 2 本应已经捕获。
4. 把分配器量子固定到下限。 用 alignment_in_bytes_ = 1024 构造 HBM 分配器。绝不要更细；更细的偏移会让步骤 2/3 失败。更粗（例如匹配 16 KiB 程序对齐）是合法的，但会浪费 HBM。
5. 保持程序对齐独立。 仅对静态 MSA 放置 tensor，在放置上游应用 16 KiB 编译期向上取整；绝不要让它放宽运行时 1024 B 下限。

---

交叉引用

• hbm-allocator.md — 运行时 BestFitAllocator；其 Config.alignment_in_bytes_ = 1024 B HBM 量子是本页下限所固定的值
• ../dma/intra-chip-descriptor.md — JXC DMA 描述符字段布局，SetHbmAddress 会把已校验地址打包进去
• overview.md — 五个片上 tier + 主机内存；HBM 和预映射 staging 池所在位置
• vmem-allocator.md — 同级 VMEM tier，使用相同分配器类，但 Config 量子更小且不同
• tpu-buffer-layout.md — 设备缓冲区的偏移如何映射到设备端 tile 布局
• ../dma/host-device-dma.md — 消费预映射缓冲区的 host↔device 传输路径
• ../compiler/msa-overview.md — 在冻结静态偏移前应用 16 KiB 程序对齐的编译期放置 pass
• 返回索引