[KV-offload][FS] : Batch store/load_block in C

[varun-sundar-rabindranath]

Purpose

We have pools of python threads to read and write KV files from/to disk. These python threads compete for the GIL and submit bursty read/write commands to the disk. This prevents the disk from reaching 100% utilization.

Changes:

1. Add C implementations for store_block and load_block that does the heavy-lifting inside a GIL-free region. Note that we still use the python thread pool, the optimization is just that the work done by the threads is now GIL-free.
2. Batching : This PR introduces a semantics change to thread work assignment.

• On main : 1 request == 1 job == many keys == many reads/writes == each thread consumes a single read/write at a time.
• This PR : 1 request == 1 job == many keys == many reads/writes == 1 thread does all the reads/writes.
  This 1 request to 1 thread mapping is likely detrimental at low concurrency. We can batch at the "keys" level and this I believe would be better to introduce in a followup PR and have this PR focus on the C implementation.

Performance

MODEL="openai/gpt-oss-120b"
TP_SIZE=2
CPU_BYTES=25769803776 # 25GB

    KV_TRANSFER_CONFIG=$(cat <<EOF
{
  "kv_connector": "OffloadingConnector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "spec_name": "TieringOffloadingSpec",
    "cpu_bytes_to_use": ${CPU_BYTES},
    "eviction_policy": "lru",
    "secondary_tiers": [{
       "type": "fs",
       "root_dir": "/mnt/nvme-storage/",
       "n_read_threads": ${THREADS},
       "n_write_threads": ${THREADS}
    }]
  }
}
EOF
)
    vllm serve "${MODEL}" \
      --tensor-parallel-size="${TP_SIZE}" \
      --kv-transfer-config "${KV_TRANSFER_CONFIG}" \
      --gpu-memory-utilization 0.7 \
      --enable-prefix-caching \
      --no-disable-hybrid-kv-cache-manager \
      --disable-uvicorn-access-log \
      --port 8000 

guidellm bench command

BENCH_RATE="64"
BENCH_RATE_TYPE="concurrent"
BENCH_MAX_SECONDS="700"
BENCH_RANDOM_SEED="889"
BENCH_TURNS=5
BENCH_PROMPT_TOKENS="4096"
BENCH_OUTPUT_TOKENS="512"
BENCH_PREFIX_TOKENS="10000"
PREFIX_COUNT=$((4 * BENCH_RATE))
#PREFIX_COUNT=$BENCH_RATE

DATA="{\"kind\":\"synthetic_text\",\"prompt_tokens\":${BENCH_PROMPT_TOKENS},\"output_tokens\":${BENCH_OUTPUT_TOKENS},\"turns\":${BENCH_TURNS},\"prefix_buckets\": [{\"bucket_weight\": 100, \"prefix_count\": ${PREFIX_COUNT}, \"prefix_tokens\": ${BENCH_PREFIX_TOKENS}}]}"
    guidellm run \
      --backend "kind=openai_http,target=${TARGET_URL},request_format=/v1/completions" \
      --profile "kind=concurrent,streams=${BENCH_RATE}" \
      --constraint "kind=max_duration,seconds=${BENCH_MAX_SECONDS}" \
      --seed "kind=static,value=${BENCH_RANDOM_SEED}" \
      --data "$DATA" 

The benchmark and vllm serve is setup to load KV from disk as much as possible.

main

n_read_threads=4, n_write_threads=4

ℹ Server Throughput Statistics (All Requests)                                                                                      
|============|=======|======|=========|==============|===============|==============|                                              
| Benchmark  | Requests             ||| Input Tokens | Output Tokens | Total Tokens |                                              
| Strategy   | Concurrency || Per Sec | Per Sec      | Per Sec       | Per Sec      |                                              
|            | Mdn   | Mean | Mean                                               ||||                                              
|------------|-------|------|---------|--------------|---------------|--------------|                                              
| concurrent | 64.0  | 63.7 | 1.7     | 41330.8      | 903.2         | 42234.0      |                                              
|============|=======|======|=========|==============|===============|==============|  

n_read_threads=16, n_write_threads=16

ℹ Server Throughput Statistics (All Requests)
|============|=======|======|=========|==============|===============|==============|
| Benchmark  | Requests             ||| Input Tokens | Output Tokens | Total Tokens |
| Strategy   | Concurrency || Per Sec | Per Sec      | Per Sec       | Per Sec      |
|            | Mdn   | Mean | Mean                                               ||||
|------------|-------|------|---------|--------------|---------------|--------------|
| concurrent | 64.0  | 63.7 | 1.8     | 42917.0      | 949.8         | 43866.8      |
|============|=======|======|=========|==============|===============|==============|

PR

n_read_threads=4, n_write_threads=4

ℹ Server Throughput Statistics (All Requests)                                                                                      
|============|=======|======|=========|==============|===============|==============|                                              
| Benchmark  | Requests             ||| Input Tokens | Output Tokens | Total Tokens |                                              
| Strategy   | Concurrency || Per Sec | Per Sec      | Per Sec       | Per Sec      |                                              
|            | Mdn   | Mean | Mean                                               ||||                                              
|------------|-------|------|---------|--------------|---------------|--------------|                                              
| concurrent | 64.0  | 63.7 | 2.3     | 54333.4      | 1209.0        | 55542.3      |                                              
|============|=======|======|=========|==============|===============|==============|  

n_read_threads=16, n_write_threads=16

ℹ Server Throughput Statistics (All Requests)                                                                                      
|============|=======|======|=========|==============|===============|==============|
| Benchmark  | Requests             ||| Input Tokens | Output Tokens | Total Tokens |
| Strategy   | Concurrency || Per Sec | Per Sec      | Per Sec       | Per Sec      |
|            | Mdn   | Mean | Mean                                               ||||
|------------|-------|------|---------|--------------|---------------|--------------|
| concurrent | 64.0  | 63.7 | 2.3     | 53658.9      | 1181.1        | 54840.0      |
|============|=======|======|=========|==============|===============|==============|     

Things to note

• 1. increased throughput
• 2. Higher external tokens count
• 3. Disk :
     • main - 4 threads throughput : 3 GB/s ; utilization : mostly <60%
     • main - 16 threads throughput : 3.8 GB/s; utilization : mostly <60%
     • PR - 4 threads throughput : 4.5 GB/s; utilization : around 80%
     • PR - 16 threads throughput : 4.5 GB/s; utilization : around 70%
       Note that the PR doesn't rely that much on the number of threads.

Test Plan

Compare lm-evals without and with offloading.

without offloading:

  vllm serve openai/gpt-oss-120b \
    --tensor-parallel-size=2 \
    --gpu-memory-utilization 0.7 \
    --enable-prefix-caching \
    --no-disable-hybrid-kv-cache-manager \
    --disable-uvicorn-access-log \
    --port 8000 

with offloading:

build_kv_transfer_config() {
  cat <<EOF
{
  "kv_connector": "OffloadingConnector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "spec_name": "TieringOffloadingSpec",
    "cpu_bytes_to_use": 68719476736,
    "eviction_policy": "lru",
    "secondary_tiers": [{
       "type": "fs",
       "root_dir": "/mnt/nvme-storage/"
    }]
  }
}
EOF
}

  vllm serve openai/gpt-oss-120b \
    --tensor-parallel-size=2 \
    --gpu-memory-utilization 0.7 \
    --enable-prefix-caching \
    --no-disable-hybrid-kv-cache-manager \
    --disable-uvicorn-access-log \
    --port 8000 \
    --kv-transfer-config "$(build_kv_transfer_config)" 

lm-eval command:

  lm-eval \
    --model local-completions \
    --model_args "model=openai/gpt-oss-120b,base_url=http://127.0.0.1:8000/v1/completions,tokenized_requests=False,num_concurrent=1000,trust_remote_code=True,max_retries=10" \
    --tasks gsm8k \
    --seed 42 \
    --num_fewshot 25 \
    --gen_kwargs temperature=0.0 

Test Result

baseline: without offloading

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.4473|±  |0.0137|
|     |       |strict-match    |    25|exact_match|↑  |0.2767|±  |0.0123|

with offloading : run 1 - cold cache

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|

|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.4579|±  |0.0137|
|     |       |strict-match    |    25|exact_match|↑  |0.2889|±  |0.0125|

with offloading: run 2 - warm cache

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.4526|±  |0.0137|
|     |       |strict-match    |    25|exact_match|↑  |0.2835|±  |0.0124|

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[orozery]

Thanks @varun-sundar-rabindranath !
I see fs_io.cpp is growing significantly with this.
Can we somehow reduce/simplify?

This is what Claude suggested:

Remove the IOResult struct and safe_* wrappers (~70 lines)

These are an over-engineered error-reporting mechanism. The C functions could just return the index of the first failing operation (or -1 for success) and let Python reconstruct the exception from errno + the path at that index:

// Instead of IOResult + safe_open + safe_write + safe_close + safe_rename + safe_unlink...
static Py_ssize_t _store_block(const char* tmp_path, const char* dest_path,
                                const char* src, size_t size) {
    if (access(dest_path, F_OK) == 0) return 0;  // already exists

    int fd = open(tmp_path, O_CREAT | O_EXCL | O_WRONLY | kODirectFlag, 0644);
    if (fd < 0) return -1;

    ssize_t w = write(fd, src, size);
    close(fd);

    if (w < 0 || (size_t)w != size) { unlink(tmp_path); return -1; }
    if (rename(tmp_path, dest_path) != 0) { unlink(tmp_path); return -1; }
    return 0;
}

Then in the Python entry point: return the failing index, and Python raises using the path from its list + errno.

[varun-sundar-rabindranath]

Thanks for the review @orozery .

Not having the safe* functions and IOResult, made the store_block and load_block harder to read, as I had to track a lot of errno for the operations between actual work and cleanup. I agree that the file is big. I'll take another shot at it.

[varun-sundar-rabindranath]

Hi @orozery

I could shave off ~70 lines from fs_io.cpp. I believe the PR is ready for another round of review. PTAL ! Thanks 🙌

[orozery]

LGTM! Thanks @varun-sundar-rabindranath !

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @varun-sundar-rabindranath.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @varun-sundar-rabindranath, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @varun-sundar-rabindranath, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.