Skip to content

Run parallel Engram disk workers on Linux with io_uring lanes - #216

Closed
jayinsine wants to merge 1 commit into
NVIDIA:mainfrom
jayinsine:linux-engram-workers
Closed

jayinsine wants to merge 1 commit into
NVIDIA:mainfrom
jayinsine:linux-engram-workers

Conversation

@jayinsine

Copy link
Copy Markdown

With disk-backed embedding tables, cold long prefill waits on Engram
gathers: on a fresh 106k-token prompt the GPUs spent 5.0-5.9 s of
10.1-10.7 s waiting for table rows. The Linux reader submits every
gather through one io_uring on the delivery thread, and
--embedding-table-workers above one was rejected on Linux. fio with the
same pattern (4 KiB direct random reads, 128 per batch, drained per
batch) tops out at about 264k reads/s on one ring, independent of depth
or refill policy, while several rings sharing the same 128-slot budget
reach 386k (2), 499k (4) and 590k (8) reads/s on the same files.

Linux now uses the Windows worker design: each of 1, 2, 4, or 8 lanes
owns an io_uring and registered staging, the lanes share the 128 slots
and 8 MiB of staging per table, large gathers split their planned reads
across scoped helper threads, and small gathers stay on the delivery
thread. A lane refills its window as reads complete and, after an
error or panic, drains reads the kernel still owns before the next
gather can reuse the ring. Row planning, coalescing, output ownership
and counters are unchanged; worker phase counters sum overlapping
intervals as on Windows. The Linux default stays at one worker.

The two portable worker tests move from the Windows module to the
shared reader tests, so Linux also checks that 1, 2, 4 and 8 workers
return identical rows within guard bands with identical read counts and
recover after a truncated file. A Linux test checks the shared slot and
staging budget.

Validation:

  • Helios library tests with runtime-package,server and unlimited
    memlock, on main 972e28e: 1284 passed and 4 failed with this change;
    1281 passed and the same 4 failed without it. The 4 failures are
    template tests that need FLARE_TEST_TEMPLATE_DIR, which is not set
    up on this host. The 3 additional passes are the moved worker tests
    and the Linux budget test.
  • GPU, measured with this change on Flare cc696a6 (not repeated after
    rebasing onto 972e28e; the Engram reader files did not conflict):
    TP4 on four RTX PRO 6000 Blackwell, a DeepSeek-V4.1-Flash-1017
    derivative package, 64 GiB row cache, resident scales. A fresh server
    per arm (empty row cache) served the same fresh prompts.
    Engram on Btrfs RAID0 over two PCIe 5.0 NVMe drives, workers 1, 2, 4,
    8, then 1 again, in two rounds:
    random 106k tokens: 10.1-10.7 s (1) -> 9.1-10.0 s (2),
    9.2-9.7 s (4), 9.1-9.9 s (8)
    Python stdlib 77k tokens: 6.9-7.0 s -> 6.2 s (2), 6.2-6.4 s (4),
    6.3 s (8)
    GPU wait on Engram, random 106k: 5.0-5.9 s -> 4.1-4.8 s
    The repeated one-worker arm reproduced the first. Decode throughput
    and a fixed arithmetic answer were unchanged.
    Engram on XFS on one Crucial T705 (PCIe 4.0 x4 link), same package:
    random 106k tokens: 7.7-7.9 s (1), 7.5-7.7 s (2), 7.7-8.4 s (4)
    Python stdlib 77k tokens: 5.3 s (1), 5.3 s (2), 5.4 s (4)

Limitations: on these drives the end-to-end gain is at most about 10%
on cold prefill from Btrfs and within noise on XFS, and it does not
grow beyond two workers. Gathers average about 40k reads, and summed
lane wait divided by lanes stays near 7 s on Btrfs for 2, 4 and 8
workers. Drives whose random-read throughput scales with the number of
submitting threads at a fixed total depth are the case this change
targets; none was tested here. The Linux default stays at one worker.
Windows was not built or run here: its reader code is unchanged, and
only its two portable tests moved.

This is the change from closed #184, rebased onto main 972e28e4.

With disk-backed embedding tables, cold long prefill waits on Engram
gathers: on a fresh 106k-token prompt the GPUs spent 5.0-5.9 s of
10.1-10.7 s waiting for table rows. The Linux reader submits every
gather through one io_uring on the delivery thread, and
--embedding-table-workers above one was rejected on Linux. fio with the
same pattern (4 KiB direct random reads, 128 per batch, drained per
batch) tops out at about 264k reads/s on one ring, independent of depth
or refill policy, while several rings sharing the same 128-slot budget
reach 386k (2), 499k (4) and 590k (8) reads/s on the same files.

Linux now uses the Windows worker design: each of 1, 2, 4, or 8 lanes
owns an io_uring and registered staging, the lanes share the 128 slots
and 8 MiB of staging per table, large gathers split their planned reads
across scoped helper threads, and small gathers stay on the delivery
thread. A lane refills its window as reads complete and, after an
error or panic, drains reads the kernel still owns before the next
gather can reuse the ring. Row planning, coalescing, output ownership
and counters are unchanged; worker phase counters sum overlapping
intervals as on Windows. The Linux default stays at one worker.

The two portable worker tests move from the Windows module to the
shared reader tests, so Linux also checks that 1, 2, 4 and 8 workers
return identical rows within guard bands with identical read counts and
recover after a truncated file. A Linux test checks the shared slot and
staging budget.

Validation:
- Helios library tests with runtime-package,server and unlimited
  memlock, on main 972e28e: 1284 passed and 4 failed with this change;
  1281 passed and the same 4 failed without it. The 4 failures are
  template tests that need FLARE_TEST_TEMPLATE_DIR, which is not set
  up on this host. The 3 additional passes are the moved worker tests
  and the Linux budget test.
- GPU, measured with this change on Flare cc696a6 (not repeated after
  rebasing onto 972e28e; the Engram reader files did not conflict):
  TP4 on four RTX PRO 6000 Blackwell, a DeepSeek-V4.1-Flash-1017
  derivative package, 64 GiB row cache, resident scales. A fresh server
  per arm (empty row cache) served the same fresh prompts.
  Engram on Btrfs RAID0 over two PCIe 5.0 NVMe drives, workers 1, 2, 4,
  8, then 1 again, in two rounds:
    random 106k tokens: 10.1-10.7 s (1) -> 9.1-10.0 s (2),
      9.2-9.7 s (4), 9.1-9.9 s (8)
    Python stdlib 77k tokens: 6.9-7.0 s -> 6.2 s (2), 6.2-6.4 s (4),
      6.3 s (8)
    GPU wait on Engram, random 106k: 5.0-5.9 s -> 4.1-4.8 s
  The repeated one-worker arm reproduced the first. Decode throughput
  and a fixed arithmetic answer were unchanged.
  Engram on XFS on one Crucial T705 (PCIe 4.0 x4 link), same package:
    random 106k tokens: 7.7-7.9 s (1), 7.5-7.7 s (2), 7.7-8.4 s (4)
    Python stdlib 77k tokens: 5.3 s (1), 5.3 s (2), 5.4 s (4)

Limitations: on these drives the end-to-end gain is at most about 10%
on cold prefill from Btrfs and within noise on XFS, and it does not
grow beyond two workers. Gathers average about 40k reads, and summed
lane wait divided by lanes stays near 7 s on Btrfs for 2, 4 and 8
workers. Drives whose random-read throughput scales with the number of
submitting threads at a fixed total depth are the case this change
targets; none was tested here. The Linux default stays at one worker.
Windows was not built or run here: its reader code is unchanged, and
only its two portable tests moved.

Design-Signed-Off-By: Claude Opus 5.5 (agent for Josh Cartu); design-review=4d3a
Signed-off-by: Josh Cartu <j@mail99.me>
@jayinsine

Copy link
Copy Markdown
Author

@original-el8 this removes the Linux restriction on --embedding-table-workers that was discussed in the Discord channel, with one io_uring per worker lane as on Windows.

Drives where this should matter are those whose 4 KiB random-read throughput depends on the number of submitting threads at a fixed total queue depth. One such report, from a user on a WD SN8100 (Silicon Motion SM2508): fio at a total depth of 512 reached 440k IOPS at 1,150 µs with 4 threads and 2,312k IOPS at 220 µs with 16 threads. That drive was not tested here. On the two configurations measured here (Btrfs RAID0 over two PCIe 5.0 drives, and XFS on one Crucial T705 at PCIe 4.0 x4), the commit message reports the results: about 10% faster cold prefill on Btrfs and no measurable change on XFS. The default stays at one worker.

@tastyrice90

Copy link
Copy Markdown

I'm the SN8100 report you cited. Happy to be the missing test - same drive,
same total depth, and my box is bound on exactly this.

will run workers 1/2/4/8 cold 250k prefill, fresh server per arm, and post
engram counters + gather_total for each.

@tastyrice90

Copy link
Copy Markdown

full curve on the target drive (WD SN8100 / SM2508, XFS), matched runs -
hit rate and reads/token identical across all four arms:

build workers hit reads/tok gather tok/s
main 1 26.5% 34.7 42.7s 6,915
lanes 2 26.5% 34.7 46.9s 6,525 -10%
lanes 4 26.3% 34.7 49.1s 6,328 -15%
lanes 8 26.1% 34.8 51.3s 6,143 -20%

monotonic across four points, so it's the lane model not run variance. the
common factor in your design is that lanes share the 128 slots and 8 MiB of
staging per table - two lanes get 64 slots each, eight get 16, and total
outstanding never rises above 128.

fio on this drive wants 512 outstanding to show its advantage (16 threads at
512 = 2,312k IOPS vs 4 threads at 512 = 440k). so more submitters on a queue
that stays 128 deep looks like pure overhead here.

question: does the Windows path raise the slot budget with worker count? if
128 is fixed there too, the fix may be workers x QUEUE_DEPTH rather than a
shared budget.

caveat: workers=1 is the release build; the PR keeps the default at 1 and the
reader should be identical, but say the word and I'll run this PR at 1 to close it.

@tastyrice90

Copy link
Copy Markdown

follow-up: I patched uring.rs:63 from let depth = DEPTH / workers; to
let depth = DEPTH; and rebuilt, to test whether the shared 128-slot budget
was the cause. It was 2.7% of a 20% regression:

build workers depth/lane gather tok/s
main 1 128 42.7s 6,915
lanes 2 64 46.9s 6,525 -10%
lanes 4 32 49.1s 6,328 -15%
lanes 8 16 51.3s 6,143 -20%
deep 8 128 49.9s 6,270 -17% (my patch)

matched throughout - hit rate 26.1-26.5%, 34.7 reads/token, 4,108B reads.

so it isn't the shared budget. eight lanes at full depth is still 17% worse
than one lane, which suggests the per-gather helper-thread split costs more
than the extra submission parallelism buys on this drive. your gathers average
~40k reads, so splitting each one across lanes is a lot of handoff.

not arguing against the PR - it clearly helps on Btrfs RAID0 (+10%) and is
neutral on XFS/T705. just that SM2508/XFS is the case it was written for and
it doesn't help there. happy to re-test any variant.

@tastyrice90

Copy link
Copy Markdown

re-ran the worker matrix on the new base — c2606aa "Keep disk embedding reads
ahead of prefill":
randomvariable@c2606aae

same box (4x RTX PRO 6000, WD SN8100 / SM2508 on XFS, cold 250k random prompt,
32G row cache, prefetch on):

workers hit gather tok/s
1 82.4% 11.3s 11,385
4 72.3% 19.6s 9,759

workers still lose, and the reason is in this PR's design: the lanes share the
256-slot depth (start_depth(requested / workers)), so four workers means 64
slots each — more submitters, the same reads in flight. each lane also reads its
slice of the read-ahead more slowly, so fewer rows land before they're needed:
prefetch_rows_loaded fell 3.11M -> 2.79M, and the lost hit rate follows
directly from that.

why depth is the binding term on this drive: fio gives 425k IOPS at 4 threads /
depth 128, 440k at 4 threads / depth 512, and 2,312k at 16 threads / 512. the
lanes deliver the threads without the depth.

if depth scaled with worker count instead of being divided, that's the
configuration worth re-testing — many submitters and a deep queue at once, which
nothing in flare currently offers.

happy to test any variant.

@jayinsine

Copy link
Copy Markdown
Author

Closing. @tastyrice90's measurements on the drive this change targeted (WD SN8100 / SM2508 on XFS, cold 250k-token random prefill) show lanes reduce prefill throughput: 6,915 tok/s with one worker on main, against 6,525, 6,328 and 6,143 tok/s with 2, 4 and 8 lanes. Giving each of 8 lanes the full 128-slot depth still measured 17% below one lane. On the read-ahead branch randomvariable/spark-local-inference@c2606aae, one worker reached 11,385 tok/s and four workers 9,759 tok/s. On the configurations measured here, the gain was about 10% on Btrfs RAID0 and none on XFS. Splitting each gather across lanes does not improve throughput on any tested XFS configuration, so this change does not justify its additional reader code. Thanks to @tastyrice90 for the full matrix.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants