Yes, I already played with the batch settings and 32 seems to be a good batch size for the sequence length of 128. I did not play with the batch timeout yet, but latency is not my focus currently.
The Python version is very barebones and I don’t think that it performs any kind of batching at all. To batch I’d probably need to use something like Ray Serve: Scalable and Programmable Serving — Ray 2.56.0.
Ah that’s interesting and could very well explain what I’m seeing!
I tried this here (nx-sentence-transformer-bench/nx_serving_partitions.exs at main · SteffenDE/nx-sentence-transformer-bench · GitHub) and seeing lots of errors when benchmarking. Seems like EXLA does not like this:
17:49:30.868 [error] GenServer #PID<0.6906.0> terminating
** (stop) exited in: Nx.Serving.local_batched_run(MyServing, ["this is a test"])
** (EXIT) an exception was raised:
** (RuntimeError) Expected buffer to be placed on device 7
(exla 0.5.3) lib/exla/executable.ex:56: EXLA.Executable.unwrap!/1
(exla 0.5.3) lib/exla/executable.ex:19: EXLA.Executable.run/3
(exla 0.5.3) lib/exla/defn.ex:346: EXLA.Defn.maybe_outfeed/7
(stdlib 4.3.1.1) timer.erl:235: :timer.tc/1
(exla 0.5.3) lib/exla/defn.ex:283: anonymous fn/7 in EXLA.Defn.__compile__/4
(nx 0.5.3) lib/nx/defn.ex:313: anonymous fn/4 in Nx.Defn.compile/3
(nx 0.5.3) lib/nx/serving.ex:1107: anonymous fn/2 in Nx.Serving.Default.handle_batch/3
(nx 0.5.3) lib/nx/serving.ex:957: anonymous fn/3 in Nx.Serving.server_task_or_enqueue/3
(nx 0.5.3) lib/nx/serving.ex:620: Nx.Serving.local_batched_run!/3
nx_serving_partitions.exs:33: MyPlug.call/2
(bandit 1.0.0-pre.9) lib/bandit/pipeline.ex:110: Bandit.Pipeline.call_plug/2
(bandit 1.0.0-pre.9) lib/bandit/pipeline.ex:25: Bandit.Pipeline.run/6
(bandit 1.0.0-pre.9) lib/bandit/http1/handler.ex:27: Bandit.HTTP1.Handler.handle_data/3
(bandit 1.0.0-pre.9) lib/bandit/delegating_handler.ex:18: Bandit.DelegatingHandler.handle_data/3
(bandit 1.0.0-pre.9) /Users/steffen/Library/Caches/mix/installs/elixir-1.14.5-erts-13.2.2.1/9825c67976b12b773e3fa7c81710c74f/deps/thousand_island/lib/thousand_island/handler.ex:399: Bandit.DelegatingHandler.handle_continue/2
(stdlib 4.3.1.1) gen_server.erl:1123: :gen_server.try_dispatch/4
To reproduce:
$ XLA_FLAGS=--xla_force_host_platform_device_count=8 elixir nx_serving_partitions.exs
# then
$ wrk http://127.0.0.1:5001 -t 8 -c 32
