Still too many threads

[sharpobject]

You probably already know this, but the moonchat still spawns way too many threads. I ran strace ./zig-out/bin/moonchat 2> blah.txt to collect a list of syscalls in a text file. I opened an image and submitted two text prompts and then quit. The program called clone3 22622 times.

[sharpobject]

here's the count of each syscall I observed:

$ grep -o "^[^(]\+(" blah.txt  | sort | uniq -c | sort -n
      1 access(
      1 execve(
      1 getrandom(
      1 rseq(
      1 set_robust_list(
      1 set_tid_address(
      2 arch_prctl(
      2 rt_sigaction(
      3 prlimit64(
      4 pread64(
      6 newfstatat(
      7 brk(
      9 close(
      9 mprotect(
      9 openat(
     38 lseek(
    257 read(
   4512 munmap(
  11348 sched_getaffinity(
  21106 futex(
  22622 clone3(
  30682 write(
  45245 rt_sigprocmask(
 271006 mmap(

[snowclipsed]

Goodness, I didn't know there were so many futex calls. Sync bottlenecks?

Maybe I need to do a big refactor with more global thread management.

[snowclipsed]

I am sure a lot of these are the multithreaded attention mechanism with its own thread pool calling the multithreaded GEMM, causing some big sync issues. I tried eliminating them as much as I could but a lot of it was poor design choice in hindsight.

[sharpobject]

You probably want to make a single persistent thread pool and define every kind of job somewhere, then the threads can run a loop that accepts a job and does different stuff based on what it is. Here's an example from some other application that offloads DNS (because async DNS is annoying) and various setSockOpt calls (because at the time these were not actually possible to make over io_uring. They could still not be possible to make over io_uring, idk). It also has a facility to offload pwrite. That's only used for debugging the io_uring-based runtime on the main thread, but it's good that it does two things, because an example that does only one thing wouldn't be very informative.

This example uses a pair of SPSC queues, but for a thread pool you may want a fancier kind of queue than that. I don't know a lot about what kind of queues are good for this sort of thing.

This example also won't work with recent versions of Zig (even if you implement all the missing files) because a lot of the setSockOpt stuff has been ripped out of the stdlib and replaced by nothing.

const std = @import("std");
const spsc = @import("common/spsc.zig");
const runtime = @import("runtime.zig");
const sync = @import("sync.zig");
const http = @import("http");
const BoundedQueue = @import("common/spsc.zig").BoundedQueue;

const mem = std.mem;
const math = std.math;
const builtin = std.builtin;
const testing = std.testing;
const os = std.os;
const ip = std.x.net.ip;
const tcp = std.x.net.tcp;
const nanosleep = std.c.nanosleep;

const assert = std.debug.assert;

const Socket = std.x.os.Socket;
const Context = runtime.Context;

const ConnectResult = struct {
    client: tcp.Client,
    host: ip.Address,
};
const Result = union(enum) {
    connect_result: ConnectResult,
    write_result: usize,
};
const ResultType = anyerror!Result;

const PWriteRequest = struct {
    fd: os.fd_t,
    buf: []const u8,
    offset: usize,
};

const Request = union(enum) {
    shutdown: void,
    connect: http.Uri,
    pwrite: PWriteRequest,
};

const SQE = struct {
    parker_idx: usize = undefined,
    request: Request = undefined,
};

const CQE = struct {
    parker_idx: usize,
    result: ResultType,
};

pub const OffloadThread = struct {
    const Self = @This();

    submission_queue: BoundedQueue(SQE, 1024),
    completion_queue: BoundedQueue(CQE, 1024),
    parkers: [2100]sync.Parker(ResultType) = [_]sync.Parker(ResultType){.{}} ** 2100,
    in_use: [2100]bool = [_]bool{false} ** 2100,
    connect_buf: [2 * 1024 * 1024]u8 = undefined,
    spawned: bool = false,
    thread: std.Thread = undefined,

    pub fn init(gpa: mem.Allocator) !OffloadThread {
        const submission_queue = try BoundedQueue(SQE, 1024).init(gpa);
        const completion_queue = try BoundedQueue(CQE, 1024).init(gpa);
        return Self{ .submission_queue = submission_queue, .completion_queue = completion_queue };
    }

    fn acquireParker(self: *Self) usize {
        for (self.in_use, 0..) |*used, i| {
            if (!used.*) {
                used.* = true;
                return i;
            }
        }
        @panic("no free parking");
    }

    fn releaseParker(self: *Self, idx: usize) void {
        self.in_use[idx] = false;
    }

    pub fn resolveAndConnect(self: *Self, ctx: *Context, uri: http.Uri) !ConnectResult {
        const parker_idx = self.acquireParker();
        defer self.releaseParker(parker_idx);
        const sqe: SQE = .{ .parker_idx = parker_idx, .request = .{ .connect = uri }};
        while (!self.submission_queue.push(sqe)) {
            try runtime.timeout(ctx, .{ .nanoseconds = 100_000_000 });
        }
        return (try try self.parkers[parker_idx].park(ctx)).connect_result;
    }

    pub fn pwrite(self: *Self, ctx: *Context, fd: os.fd_t, buf: []const u8, offset: usize) !usize {
        const parker_idx = self.acquireParker();
        defer self.releaseParker(parker_idx);
        const sqe: SQE = .{ 
            .parker_idx = parker_idx,
            .request = .{ .pwrite = .{ .fd = fd, .buf = buf, .offset = offset }},
        };
        while (!self.submission_queue.push(sqe)) {
            try runtime.timeout(ctx, .{ .nanoseconds = 100_000_000 });
        }
        return (try try self.parkers[parker_idx].park(ctx)).write_result;
    }

    fn doResolveAndConnect(self: *Self, gpa: mem.Allocator, uri: http.Uri) ResultType {
        _ = self;
        const host = switch (uri.host) {
            .ip => |address| xd: {
                const ptr: *align(4) *@TypeOf(address.any) = &address.any;
                break :xd ip.Address.from(Socket.Address.fromNative(ptr));
            },
            .name => |name| resolve: {
                const result = try std.net.getAddressList(gpa, name, uri.port.?);
                defer result.deinit();

                if (result.addrs.len == 0) {
                    return error.UnknownHost;
                }

                const ptr: *align(4) *@TypeOf(result.addrs[0].any) = &result.addrs[0].any;
                break :resolve ip.Address.from(Socket.Address.fromNative(ptr));
            },
        };

        var client = try tcp.Client.init(switch (host) {
            .ipv4 => .ip,
            .ipv6 => .ipv6,
        }, .{ .close_on_exec = true });
        errdefer client.deinit();

        try client.setNoDelay(true);
        try client.setKeepAlive(true);
        try client.setQuickACK(true);
        try client.setLinger(0);
        try client.connect(host);

        return Result{ .connect_result = .{ .client = client, .host = host }};
    }

    fn doPwrite(self: *Self, fd: os.fd_t, buf: []const u8, offset: usize) ResultType {
        _ = self;
        const written = try os.pwrite(fd, buf, offset);
        return Result{ .write_result = written };
    }

    fn run(self: *Self) void {
        while (true) {
            var fba = std.heap.FixedBufferAllocator.init(self.connect_buf[0..]);
            const gpa = fba.allocator();
            if (self.submission_queue.pop()) |sqe| {
                var cqe: ?CQE = null;
                switch(sqe.request) {
                    .connect => {
                        const result = self.doResolveAndConnect(gpa, sqe.request.connect);
                        cqe = CQE{ .parker_idx = sqe.parker_idx, .result = result };
                    },
                    .pwrite => {
                        const result = self.doPwrite(sqe.request.pwrite.fd, sqe.request.pwrite.buf, sqe.request.pwrite.offset);
                        cqe = CQE{ .parker_idx = sqe.parker_idx, .result = result };
                    },
                    .shutdown => {
                        return;
                    },
                }
                if (cqe) |cqe_really| {
                    while (!self.completion_queue.push(cqe_really)) {
                        nanosleep(0, 100_000_000);
                    }
                }
            } else {
                nanosleep(0, 100_000_000);
            }
        }
    }

    pub fn handleCompletions(self: *Self) void {
        while (self.completion_queue.pop()) |cqe| {
            self.parkers[cqe.parker_idx].broadcast(cqe.result);
        }
    }

    pub fn spawn(self: *Self) !void {
        assert(!self.spawned);
        self.thread = try std.Thread.spawn(.{}, Self.run, .{self});
        self.spawned = true;
    }

    fn shutdown(self: *Self) void {
        const sqe: SQE = .{ .parker_idx = 0, .request = .{ .shutdown = {} }};
        while (!self.submission_queue.push(sqe)) {
            nanosleep(0, 100_000_000);
        }
    }

    fn join(self: *Self) void {
        if (self.spawned) {
            self.thread.join();
            self.spawned = false;
        }
    }

    pub fn deinit(self: *Self, gpa: mem.Allocator) void {
        self.shutdown();
        self.join();
        self.submission_queue.deinit(gpa);
        self.completion_queue.deinit(gpa);
    }
};

[sharpobject]

According to a chatbot, calling clone3 will cause you to also see some mmaps in your strace output. here's the number of occurrences of each mmap size for a different run of the moonchat which called clone3 22929 times.

$ grep -o "^[^(]\+([^,]\+,[^,]\+," strace.txt | grep mmap | awk '{print $2}' | sort | uniq -c | sort -n
      1 106496,
      1 1118208,
      1 13111296,
      1 134217728,
      1 143360,
      1 1642496,
      1 1658880,
      1 192512,
      1 20480,
      1 208896,
      1 2232320,
      1 2260560,
      1 24576,
      1 274432,
      1 282624,
      1 28672,
      1 3280896,
      1 360448,
      1 36864,
      1 372736,
      1 413696,
      1 4460544,
      1 507904,
      1 52816,
      1 53248,
      1 561152,
      1 56615,
      1 6557696,
      1 67108864,
      1 73728,
      1 823296,
      1 8916992,
      1 942344,
      2 12288,
      3 1048576,
      3 1073741824,
      4 32768,
      6 131072,
      6 268435456,
      8 2097152,
      8 65536,
     10 524288,
     11 262144,
     33 8192,
     41 4194304,
     57 16777216,
     73 33554432,
     80 16384,
    151 8388608,
  11417 16781312,
 267640 4096,

This was collected after increasing the deque size in SlabReusingAllocator to 1000, so we have 0 munmaps.

Almost all of them are 16781312 and 4096! As a chatbot has told me, 16781312=4096*4097 is the size of the mmap used by clone3 to set up storage for a stack and guard pages, and clone3 also calls mmap(4096) some number of times to set up space for thread-local storage and to configure the guard pages.

So any further effort about memory allocation is mostly not useful compared to effort about using a persistent thread pool.

I have confirmed the above claim by making a branch in which shapes are stored in arrays on Tensors and TensorViews rather than dynamically allocated, which has too-small-to-measure impact on performance https://github.com/sharpobject/moondream-zig/tree/shape_arr

[snowclipsed]

did a simple persistent thread implementation and clone3 is down from ~120k to ~12k on my machine

[snowclipsed]

update, clone3 down to just 16 calls

[sharpobject]

very nice, any speedup overall?

[snowclipsed]

We're now up to 16 tok/s on my machine from 11.5-12 tok/s