MoRI-IO XGMI backend: BatchRead ignores the base offset of a registered sub-region (silent data corruption; SIGSEGV in production)

[carlushuang]

Summary

When a buffer is registered as a sub-region of a larger device allocation
(i.e. register_memory(base + offset, size, ...)), the XGMI backend's
BatchRead / BatchReadWrite ignores that base offset and transfers from the
allocation base instead. The transfer reports StatusCode.SUCCESS, but the
wrong bytes are moved.

This is the normal paged-KV-cache layout for PD-disaggregation KV connectors:
one big [num_layers, num_blocks, block_size, head_dim] allocation, with each
layer registered as an offset view. Under that layout the XGMI backend silently
returns wrong data, and in a real inference engine it escalates to a SIGSEGV
inside hipMemcpyPeerAsync (details + evidence below).

Registering the whole allocation and addressing the sub-region via the
batch_read offset argument works correctly — only the registered region's
base offset is dropped, and only on the XGMI backend.

Environment

• mori 1.2.1.dev20260610 (also reproduces via the prebuilt wheel in rocm/atom-dev:sglang-latest)
• ROCm 7.2.4, libamdhip64.so.7
• 8 × AMD Instinct MI355X (gfx950), all XGMI-connected (hipDeviceCanAccessPeer == 1 for every pair)
• torch 2.10.0+rocm7.2.4
• Single node, intra-node XGMI (no RDMA NIC involved)

Minimal reproducer (pure mori + torch)

Two processes, two GPUs, XGMI backend. Same transfer run two ways:

• base — register the whole allocation, address the sub-region via the batch_read offset argument.
• offset — register the sub-region at base + OFF (size = sub-region), read at offset 0.

#!/usr/bin/env python3
# repro_xgmi_offset.py  —  pure mori + torch, no other deps
#   HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python repro_xgmi_offset.py prod base
#   HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python repro_xgmi_offset.py cons base      -> CORRECT
#   HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python repro_xgmi_offset.py prod offset
#   HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python repro_xgmi_offset.py cons offset    -> WRONG (val=0)
import os, sys, time, torch
from mori.io import (IOEngine, IOEngineConfig, EngineDesc, MemoryDesc,
                     MemoryLocationType, BackendType, XgmiBackendConfig)

role, mode = sys.argv[1], sys.argv[2]          # role: prod|cons ; mode: base|offset
PROD_DEV, CONS_DEV = 5, 1                       # two distinct GPUs on the node
BIG = 256 * 1024 * 1024                         # backing allocation (like a KV pool)
OFF = 64 * 1024 * 1024                          # sub-region offset (like a per-layer slice)
RGN = 1 * 1024 * 1024                           # sub-region size
N   = 9216                                      # bytes to transfer (one paged block)
D = f"/tmp/mori_repro_{mode}"; os.makedirs(D, exist_ok=True)

def wf(p, b):
    with open(p + ".tmp", "wb") as f: f.write(b)
    os.replace(p + ".tmp", p)
def rf(p):
    while not os.path.exists(p): time.sleep(0.05)
    time.sleep(0.1)
    with open(p, "rb") as f: return f.read()

dev = PROD_DEV if role == "prod" else CONS_DEV
torch.cuda.set_device(dev)
eng = IOEngine(role, IOEngineConfig(host="127.0.0.1", port=0))
eng.create_backend(BackendType.XGMI, XgmiBackendConfig())

big = torch.zeros(BIG, dtype=torch.uint8, device=f"cuda:{dev}")
if role == "prod":
    big[OFF:OFF + RGN] = 7                      # payload lives at [OFF, OFF+RGN)
torch.cuda.synchronize()

if mode == "base":
    reg_ptr, reg_size, rd = big.data_ptr(), BIG, OFF       # register whole alloc; address via read offset
else:  # mode == "offset"
    reg_ptr, reg_size, rd = big.data_ptr() + OFF, RGN, 0   # register the sub-region directly

desc = eng.register_memory(reg_ptr, reg_size, dev, MemoryLocationType.GPU)

if role == "prod":
    wf(f"{D}/eng", eng.get_engine_desc().pack())
    wf(f"{D}/mem", desc.pack())
    print(f"[prod {mode}] ready on cuda:{dev}", flush=True)
    rf(f"{D}/done")
else:
    eng.register_remote_engine(EngineDesc.unpack(rf(f"{D}/eng")))
    sess = eng.create_session(desc, MemoryDesc.unpack(rf(f"{D}/mem")))
    st = sess.batch_read([rd], [rd], [N], sess.allocate_transfer_uid())
    rc = eng.wait_all([st], 15000)
    torch.cuda.synchronize()
    got = big[OFF].item()
    verdict = "CORRECT" if got == 7 else "WRONG (data silently lost)"
    print(f"[cons {mode}] rc={rc}  big[{OFF}]={got} (expected 7)  ->  {verdict}", flush=True)
    wf(f"{D}/done", b"1")

Launcher:

#!/bin/bash
VIS=0,1,2,3,4,5,6,7
for mode in base offset; do
  rm -rf /tmp/mori_repro_$mode
  HIP_VISIBLE_DEVICES=$VIS python3 repro_xgmi_offset.py prod $mode > /tmp/rp_$mode.log 2>&1 &
  sleep 8
  HIP_VISIBLE_DEVICES=$VIS timeout 60 python3 repro_xgmi_offset.py cons $mode > /tmp/rc_$mode.log 2>&1
  grep -aE "cons" /tmp/rc_$mode.log
  wait 2>/dev/null
done

Observed output

[cons base]   rc=StatusCode.SUCCESS  big[67108864]=7 (expected 7)  ->  CORRECT
[cons offset] rc=StatusCode.SUCCESS  big[67108864]=0 (expected 7)  ->  WRONG (data silently lost)

The producer writes 7 into [OFF, OFF+RGN) and registers exactly that
sub-region. The consumer's BatchRead returns SUCCESS but reads from the
allocation base (which is 0), not from base + OFF.

What it does in a real engine (production crash)

In an actual PD-disaggregation deploy (paged MLA KV cache, TP=4 producer + TP=4
consumer, intra-node XGMI), each layer's KV slice is registered as an offset view
of one big allocation. The decode worker crashes with SIGSEGV on the first
BatchRead. Backtrace (rocgdb):

#0  libc __memmove_avx512_unaligned        <-- faulting insn: vmovdqu64 (%rsi),%zmm16
#1-#5 libamdhip64.so.7  (hipMemcpyPeerAsync internals)
#6  mori::io::XgmiBackendSession::BatchReadWrite(...)
#7  mori::io::IOEngineSession::BatchRead(...)

The fault is a CPU memcpy over a GPU device pointer. Intercepting
hipMemcpyPeerAsync (LD_PRELOAD) shows why — the IPC-imported source pointer is
reported as host/unregistered:

hipMemcpyPeerAsync(dst, dstDev=1, src, srcDev=5, n=9216)
  hipPointerGetAttributes(src) -> type=0 (UNREGISTERED/host), device=-2   <-- IPC-imported, offset sub-region
  hipPointerGetAttributes(dst) -> type=2 (DEVICE),            device=1
  hipDeviceCanAccessPeer(5,1)=1, (1,5)=1

Because HIP classifies the (mis-computed, out-of-range) IPC source pointer as
type=0/unregistered, hipMemcpyPeerAsync takes a CPU-copy path and dereferences
the GPU address on the host → SIGSEGV. The buffer registration itself succeeds;
the corruption/crash is entirely in the BatchReadWrite transfer path.

Likely root cause

XgmiBackendSession::BatchReadWrite appears to compute the remote/local transfer
addresses from the registered MemoryDesc without honoring the base offset
between the registered pointer and the underlying IPC allocation base. For
sub-region registrations:

• best case: it reads from the allocation base → silent wrong data (the deterministic repro above);
• production case: the computed source pointer is out of the mapped IPC range → hipPointerGetAttributes returns type=0 → hipMemcpyPeerAsync falls back to a host memcpy over a device address → SIGSEGV.

Registering the whole allocation + using the batch_read offset argument is
handled correctly, so the offset handling is missing specifically for the
registered region base on the XGMI path.

Suggested fixes

1. Honor the registered region's base offset (relative to the IPC allocation base)
   when computing transfer addresses in XgmiBackendSession::BatchReadWrite.
2. Avoid hipMemcpyPeerAsync for IPC-imported peer memory whose
   hipPointerGetAttributes is type != hipMemoryTypeDevice — it can silently
   degrade to a CPU memcpy over device addresses. Prefer the scatter/gather GPU
   kernel (already loaded via LoadScatterGatherModule) or an explicit
   hipMemcpyAsync(..., hipMemcpyDeviceToDevice, stream).
3. At minimum, validate pointer attributes / offset range and return a
   non-SUCCESS status instead of silently corrupting or segfaulting.