[Perf] engram: Read host-table rows in address order on GPUs without host page tables

python/sglang/kernels/ops/embeddings/engram_gather.py

@@ -19,15 +19,18 @@ def _engram_gather_kernel(
     w_ptr,
     s_ptr,
     ids_ptr,
+    dst_ptr,
     out_ptr,
     row_lo,
     row_hi,
     DIM: tl.constexpr,
     BLK: tl.constexpr,
     E8M0_ZERO: tl.constexpr,
+    HAS_DST: tl.constexpr,
 ):
     row = tl.program_id(0).to(tl.int64)
     idx = tl.load(ids_ptr + row).to(tl.int64)
+    dst = tl.load(dst_ptr + row).to(tl.int64) if HAS_DST else row
     # The table holds rows [row_lo, row_hi); an id outside it is not read and
     # comes out as zeros, which is what the sharded all-reduce sums.
     owned = (idx >= row_lo) & (idx < row_hi)
@@ -43,7 +46,7 @@ def _engram_gather_kernel(
     scale = (exps << 23).to(tl.float32, bitcast=True)
     scale = tl.where(exps == 0, E8M0_ZERO, scale)
     out = tl.where(owned, vals * scale, 0.0)
-    tl.store(out_ptr + row * DIM + offs, out.to(tl.bfloat16))
+    tl.store(out_ptr + dst * DIM + offs, out.to(tl.bfloat16))
 
 
 def engram_gather(
@@ -55,12 +58,17 @@ def engram_gather(
     block_size: int,
     row_lo: int = 0,
     row_hi: int = 2**62,
+    sort_ids: bool = False,
 ) -> torch.Tensor:
     """Gather rows ``ids`` ([N] int) into ``out`` ([N, dim] bf16, contiguous).
 
     ``weight_ptr`` addresses [rows, dim] fp8 e4m3 bytes and ``scale_ptr``
     [rows, dim // block_size] e8m0 bytes for global rows [row_lo, row_hi); both
     may live in device or host memory. Ids outside the range produce zero rows.
+
+    ``sort_ids`` reads the rows in address order and writes each to its original
+    position. Random rows of a host table larger than the GPU's translation reach
+    miss on nearly every row; sorted, neighbouring programs share a page.
     """
     assert dim > 0 and dim & (dim - 1) == 0 and block_size > 0, (dim, block_size)
     assert dim % block_size == 0, (dim, block_size)
@@ -71,15 +79,20 @@ def engram_gather(
     assert out.device == ids.device and out.shape == (ids.numel(), dim), out.shape
     n = ids.numel()
     if n:
+        dst = None
+        if sort_ids:
+            ids, dst = torch.sort(ids)
         _engram_gather_kernel[(n,)](
             weight_ptr,
             scale_ptr,
             ids,
+            dst if dst is not None else ids,
             out,
             row_lo,
             row_hi,
             DIM=dim,
             BLK=block_size,
             E8M0_ZERO=_E8M0_ZERO,
+            HAS_DST=dst is not None,
         )
     return out

python/sglang/srt/layers/engram.py

@@ -627,6 +627,21 @@ def _hip_runtime() -> ctypes.CDLL:
     return lib
 
 
+@functools.cache
+def _sorted_host_reads(device: int) -> bool:
+    """Whether host-table rows are read in address order. Without ATS (no
+    NVLink-C2C) the GPU maps pinned host memory in its own page tables, and
+    random rows of a table far larger than its translation reach miss on nearly
+    every row; with ATS the host page tables serve the reads and order costs."""
+    if _is_hip:
+        return False
+    from cuda.bindings import driver
+
+    attr = driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES
+    err, uses_host_page_tables = driver.cuDeviceGetAttribute(attr, device)
+    return err == driver.CUresult.CUDA_SUCCESS and uses_host_page_tables == 0
+
+
 def _registered_device_ptr(host_ptr: int) -> int:
     """
     Address kernels must use for registered host memory.
@@ -813,6 +828,7 @@ def __init__(
         )
         self.rows = row_end - self.row_start
         self.host_table: Optional[_HostTable] = None
+        self._sort_host_reads = False
         if envs.SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE.get():
             self._init_host_table(num_embeddings, dim, layer_id)
         else:
@@ -847,6 +863,7 @@ def _init_host_table(self, num_embeddings: int, dim: int, layer_id: int):
         self.scale = nn.Parameter(scale, requires_grad=False)
         device_ptr = self.host_table.device_ptr
         self._host_table_ptrs = (device_ptr, device_ptr + w_bytes)
+        self._sort_host_reads = _sorted_host_reads(torch.cuda.current_device())
 
     @property
     def _shared(self) -> bool:
@@ -888,6 +905,7 @@ def forward(
                 out.view(-1, self.dim),
                 self.dim,
                 FP8_BLOCK_SIZE,
+                sort_ids=self._sort_host_reads,
             )
             return out
         if cp_all_tokens and self.tp_size > 1:
@@ -949,6 +967,7 @@ def _owned_rows(self, indices: torch.Tensor) -> torch.Tensor:
             FP8_BLOCK_SIZE,
             row_lo=self.row_start,
             row_hi=self.row_start + self.rows,
+            sort_ids=self._sort_host_reads,
         )
         return out