[Bug]: Large context DMA errors

[joerunde]

Describe the bug

Symptom: at boot, when the on-device KV cache is sized via num_gpu_blocks_override = max_num_seqs * ceil(max_model_len / block_size) in platform.py, the worker fails inside TorchSpyreModelRunner.initialize_kv_cache_tensors with RAS::VFIO::MapDMAFailed once the user-facing max_model_len × max_num_seqs are above a host-dependent threshold.

Current workaround in code: spyre_inference/platform.py hard-caps MAX_MODEL_LEN_CAP=128, MAX_NUM_SEQS_CAP=8. These are far below anything useful.

How to reproduce

run any model with max_model_len=16384, max_num_seqs=256

Error origin (where RAS::VFIO::MapDMAFailed is thrown)

The symbol lives in libsenlib.so / libVF.so / libsenlib-dd2.so at /opt/ibm/spyre/senlib/lib/ (it is NOT in torch-spyre or in spyre-inference). Its fields are dma_iova, dma_vaddr, dma_size, len_remaining. The class is thrown when the kernel's VFIO_IOMMU_MAP_DMA ioctl returns nonzero.

Reproduced JSON blob from a failing run (max_model_len=16384, max_num_seqs=256):

{
  "action":"information",
  "category":"linuxconfig",
  "code":"0xad15",
  "description":"The Linux VFIO kernel module was unable to map a memory region to the device's address space for DMA.  This is an unexpected condition.",
  "dma_iova":70386725949440,
  "dma_size":262144,
  "dma_vaddr":94799475351552,
  "errno":"No space left on device",
  "len_remaining":262144,
  "message":"Unable to map DMA to the device address",
  "name":"RAS::VFIO::MapDMAFailed",
  "severity":"ERROR",
  "step":"Open ticket and consider restarting Linux",
  "type":"runtime_error"
}

The decisive field is errno: "No space left on device": that's ENOSPC from the kernel ioctl, which the Linux VFIO type1 IOMMU returns when the per-container DMA-entry table is full. On this host:

$ cat /sys/module/vfio_iommu_type1/parameters/dma_entry_limit
65535

The dma_size=262144 field exactly matches one Granite-1b KV page: num_kv_heads(8) × block_size(128) × head_size(128) × 2 bytes = 262144. So the failing call is one of the per-page allocations.

Why each KV-page allocation consumes an IOMMU entry

The worker traceback at failure points at spyre_model_runner.py:379 (one of the per-page torch.zeros(..., device="spyre") calls) → torch_spyre/ops/eager.py:211 spyre__copy_from → torch_spyre._C.copy_tensor(self, dst, non_blocking). So the allocation IS a host-to-device copy under the hood.

The reason is torch-spyre's implementation of aten::zero_ and aten::fill_.Scalar at torch_spyre/ops/eager.py:127-156:

@torch.library.register_kernel("aten::zero_", ["spyre"])
def spyre__zero_(self: torch.Tensor) -> torch.Tensor:
    """Zero out the tensor in-place."""
    # Create zeros on CPU
    tmp = torch.zeros(self.size(), dtype=self.dtype, device="cpu")
    # Copy to device
    self.copy_(tmp)
    # TODO: Can we zero out tensors in-place without copy
    return self

@torch.library.register_kernel("aten::fill_.Scalar", ["spyre"])
def spyre__fill_scalar(self, other):
    tmp = torch.ones(self.size(), dtype=self.dtype) * other
    self.copy_(tmp)
    return self

So every way to zero or fill a Spyre tensor — both torch.zeros(shape, device="spyre") (decomposes to aten::empty + aten::zero_) and the explicit tensor.zero_() on a Spyre tensor — goes through a host buffer + H2D copy. The torch-spyre author's own TODO: Can we zero out tensors in-place without copy on the same function is the existing acknowledgement.

Why H2D consumes an IOMMU entry per call

Inside the H2D path, flex calls senlib's VfioWrapper::allocate_pin_and_map(size) or the modern senlib::v2::DmaMemoryHandle::allocate(..., manager) path. Both pin host pages and call VFIO_IOMMU_MAP_DMA on them. The mapping is held by the returned DmaAllocationPtr (std::shared_ptr); it's only released when refcount drops to zero. For in-flight async copies the handle is held until the copy completes, so concurrent copies hold concurrent mappings.

Per-page allocation count: num_blocks × 2 (K+V) × num_layers. For Granite-1b (40 layers) at the failing config: 32768 × 2 × 40 = 2,621,440 calls. Even with handles released after each copy, the in-flight peak grows with the loop's allocation rate and quickly crosses 65535.

Additional context

No response

Checklist

• I have searched for similar issues.