fork: make StateGet chunk size runtime-configurable, default 2MB

[ddvnguyen]

What changed

Two commits:

1. 8189debdc — initial fix: hardcoded CHUNK from 256 KB to 8 MB.
2. 5f9263508 — superseding redesign: replaced the hardcoded constant with a runtime-configurable value, since a compile-time constant doesn't fit the engine control-plane pattern this codebase already uses for other tunables (SET_EXPERT_MODE, SWAP_QUANT — Hydra tells the engine what to do via RPC, not via rebuild). Default is now 2 MB, not 8 MB.

Concretely (commit 2):

• llama_cparams::hydra_state_chunk_size (new field, default 2 * 1024 * 1024), mirroring the existing hydra_expert_mode field.
• llama_context::hydra_set_state_chunk_size(size_t) setter, clamped to [64 KiB, 64 MiB] since it's reachable from the network.
• llama_hydra_set_state_chunk_size / llama_hydra_get_state_chunk_size C API (include/llama-hydra.h).
• llama_io_write_socket constructor now takes chunk_size instead of using a static constexpr CHUNK.
• CONFIGURE (0x40) — previously received config_json and did nothing with it (pure stub returning {"success":true} unconditionally). Now actually parses it for a "state_chunk_size" field and applies it via the setter above. This is also the first production use of CONFIGURE; it was wired but unused before this PR.

Why

Closes #6 (tracks ddvnguyen/hydra_vortex#334).

Production StateGet traffic (the Hydra RPC wire protocol's engine→core save path) routes through llama_io_write_socket::write_tensor() — confirmed by reading the call path (hydra_handle_state_get() → state_seq_get_data_to_fd() → this class), not assumed. For every chunk of every K/V tensor, this loop does, strictly serially with no overlap:

1. ggml_backend_tensor_get() → on CUDA, ggml_backend_cuda_buffer_get_tensor() issues cudaMemcpyAsync immediately followed by cudaStreamSynchronize — effectively synchronous, fixed per-call driver overhead.
2. A blocking send() syscall.

At the original 256 KB, a 1.21 GB KV state meant ~4,800 of these serialized round trips on a single thread — measured production throughput was a flat ~200 MB/s regardless of state size, consistent with fixed per-call overhead dominating over actual data movement (256 KB over PCIe/loopback should take microseconds, not ~1.3ms). This also explains the "100% user CPU" observation in the Hydra issue: the thread is busy through syscall/driver overhead, not idle on I/O.

Note: the original Hydra-side issue suspected realloc-per-layer buffer growth as the root cause. That doesn't apply here — this is the zero-copy M2 streaming path (no buffer growth, just a fixed-size reused staging buffer), and is the path all real RPC traffic uses.

What this does NOT do (scoped out, see #6 for the rest)

• Double-buffer copy/send pipelining — only worth adding if the chunk-size change alone doesn't reach the ~600-800 MB/s target. Needs a live before/after measurement first; not implemented in this PR.
• Parallelizing the layer loop across threads — the original Hydra issue suggested this, but concurrent cudaMemcpy calls from multiple host threads against the same CUDA context add synchronization risk for unclear gain once fixed-call overhead is amortized, and the wire format has no per-layer framing for out-of-order reassembly. Not planned.

How to test

Build: verified clean against both sm_120 (RTX) and sm_60 (P100) locally.

cmake --build build_sm120 --config Release -j$(nproc) --target llama-engine
cmake --build build_sm60  --config Release -j$(nproc) --target llama-engine

Correctness (round-trip) — not yet done, needed before merge:
Round-trip a large (500 MB–1.2 GB) sequence's KV state through STATE_GET → STATE_PUT and confirm identical generation output / cache_n after restore. No existing automated test covers /slots/{id}/state in tools/server/tests/, so this needs either a manual check against a running llama-engine instance, or a small ad-hoc script issuing the Hydra RPC StateGet/StatePut opcodes directly. Also worth a basic CONFIGURE round-trip check: send {"state_chunk_size":N} and confirm llama_hydra_get_state_chunk_size reflects it (e.g. via a log line — SRV_INF("hydra: CONFIGURE state_chunk_size=...")).

Throughput — not yet done, needed before merge:
On the Hydra side (paired with the corresponding ddvnguyen/hydra_vortex PR, which sends CONFIGURE {"state_chunk_size": ...} once per worker connection — see ConfigureStateChunkSizeAsync in WorkerSchedulerService.cs), send a 100K+ token request through the coordinator in P/D split (cold_concurrency) mode and watch the hydra_save_kv_rpc_seconds Prometheus metric / save_kv_rpc_ms phase on the Grafana dashboard. Target: ~6.2s/1.21GB (199 MB/s) baseline → ~2s/1.21GB (600-800 MB/s) or better.

(This wasn't measured yet on my end because the RTX GPU was fully loaded — 14.4/16.3 GB — by an already-running hydra-system stack at the time of this change, and swapping the binary would have required restarting that live stack. Deferring the live measurement + merge decision to whoever picks this up next.)

Links

• Fork issue: [fork/hydra-334] StateGet: 256KB chunk size + serial copy/send caps throughput at ~200 MB/s #6
• Hydra issue: [M-Perf] C++ StateGet: 200 MB/s ceiling limits save_kv_rpc throughput for large KV states hydra_vortex#334

[ddvnguyen]

Pushed a follow-up commit (358955300) addressing review feedback:

• Extracted the clamp into a pure llama_hydra_clamp_state_chunk_size(), exposed via llama-hydra.h, with a dedicated unit test (tests/test-hydra-state-chunk-size.cpp, registered via llama_build_and_test — no model/context needed, runs on every platform). Confirmed green locally.
• CONFIGURE now echoes state_chunk_size_applied (post-clamp) in both the wire response (hydra_handle_configure) and to_json(), so a silent clamp is no longer invisible to the caller.
• Documented the benign cparams.hydra_state_chunk_size data race inline (same pattern/rationale as the pre-existing hydra_expert_mode field).
• Confirmed (no code change needed): all slots share one llama_context per engine process (server-context.cpp:1127), so CONFIGURE on slot "0" is correctly engine-wide.

Builds clean on sm_120 + sm_60. Still open: live throughput measurement + STATE_GET/PUT round-trip correctness (same as before — blocked on GPU availability).