fix: various bugs to eliminate crashes when SGLang overlap scheduling is enabled

[ivanium]

This PR aims to address issue #5

1. Key fix: synchronize CUDA operations before unmapping pages. This avoids potential illegal memory accesses caused when the worker still runs kernels accessing the KV cache page while the scheduler thread tries to unmap it.
2. Key fix: reserve the 0-th token slot for SGLang as the padded slot.
3. Key fix: allocate 0-th token slot and launch pre-alloc thread after ensuring all workers have created the KV tensors
4. Release Python GIL in C++ code to avoid kvcached vmm operations blocking other threads.
5. Protect FTensorAllocator with a mutex to make it thread safe when sglang overlap scheduling is enabled.
6. Fix a type error in the sglang patch to always pass in free_index as List[int]
7. Protect KVCacheManager with a lock for multi-threaded scheduler (overlap-schedule)
8. Update interfaces for overlapped scheduling.

Close #5

We can optimize the performance and reduce the overhead of CUDA synchronization, but that will be done in another separate PR.

[jiarong0907]

Hi @ivanium, as I tested, the same illegal memory access error still happened. I added --sharegpt-output-len 300 parameter for the client of SGLang. It could make it easier to trigger the problem.

[2025-07-12 20:27:19] Scheduler hit an exception: Traceback (most recent call last):
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/python/sglang/srt/managers/scheduler.py", line 2229, in run_scheduler_process
    scheduler.event_loop_overlap()
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/python/sglang/srt/managers/scheduler.py", line 662, in event_loop_overlap
    batch = self.get_next_batch_to_run()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/python/sglang/srt/managers/scheduler.py", line 1309, in get_next_batch_to_run
    new_batch = self.get_new_batch_prefill()
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/python/sglang/srt/managers/scheduler.py", line 1458, in get_new_batch_prefill
    new_batch.prepare_for_extend()
  File "/home/jrxing/kvcached-yifan/engine_integration/sglang-v0.4.6.post2/python/sglang/srt/managers/schedule_batch.py", line 1045, in prepare_for_extend
    prefix_lens_tensor = torch.tensor(
                         ^^^^^^^^^^^^^
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


[2025-07-12 20:27:19] Received sigquit from a child process. It usually means the child failed.
./start_server.sh: line 29:  9944 Killed                  python -m sglang.launch_server --model "$MODEL" --disable-radix-cache --trust-remote-code --port "$SGL_PORT"

[ivanium]

Thanks let me take a deeper look

[ivanium]

Benchmark results

./start_server.sh sglang and ./start_client.sh sglang

With this PR:

============ Serving Benchmark Result ============
Backend:                                 sglang-oai
Traffic request rate:                    10.0
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  105.30
Total input tokens:                      296523
Total generated tokens:                  186737
Total generated tokens (retokenized):    183313
Request throughput (req/s):              9.50
Input token throughput (tok/s):          2815.93
Output token throughput (tok/s):         1773.35
Total token throughput (tok/s):          4589.27
Concurrency:                             16.05
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   1690.62
Median E2E Latency (ms):                 1033.31
---------------Time to First Token----------------
Mean TTFT (ms):                          37.44
Median TTFT (ms):                        34.30
P99 TTFT (ms):                           102.59
---------------Inter-Token Latency----------------
Mean ITL (ms):                           9.06
Median ITL (ms):                         7.70
P95 ITL (ms):                            18.58
P99 ITL (ms):                            29.55
Max ITL (ms):                            162.24
==================================================

This PR without torch.cuda.synchronize

============ Serving Benchmark Result ============
Backend:                                 sglang-oai
Traffic request rate:                    10.0
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  105.36
Total input tokens:                      296523
Total generated tokens:                  186737
Total generated tokens (retokenized):    183334
Request throughput (req/s):              9.49
Input token throughput (tok/s):          2814.31
Output token throughput (tok/s):         1772.32
Total token throughput (tok/s):          4586.63
Concurrency:                             16.02
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   1687.82
Median E2E Latency (ms):                 1038.57
---------------Time to First Token----------------
Mean TTFT (ms):                          38.14
Median TTFT (ms):                        33.94
P99 TTFT (ms):                           105.96
---------------Inter-Token Latency----------------
Mean ITL (ms):                           9.04
Median ITL (ms):                         7.69
P95 ITL (ms):                            18.45
P99 ITL (ms):                            29.20
Max ITL (ms):                            141.47
==================================================

Previously (without this PR)

============ Serving Benchmark Result ============
Backend:                                 sglang-oai
Traffic request rate:                    10.0
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  109.08
Total input tokens:                      296523
Total generated tokens:                  186737
Total generated tokens (retokenized):    183303
Request throughput (req/s):              9.17
Input token throughput (tok/s):          2718.43
Output token throughput (tok/s):         1711.95
Total token throughput (tok/s):          4430.38
Concurrency:                             22.84
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   2491.31
Median E2E Latency (ms):                 1478.15
---------------Time to First Token----------------
Mean TTFT (ms):                          40.02
Median TTFT (ms):                        32.20
P99 TTFT (ms):                           157.36
---------------Inter-Token Latency----------------
Mean ITL (ms):                           13.43
Median ITL (ms):                         11.28
P95 ITL (ms):                            30.60
P99 ITL (ms):                            53.61
Max ITL (ms):                            341.02
==================================================

Somehow it looks not slower but actually a bit faster :)

[jiarong0907]

While I don't know why it gets faster, it is always good to be fast. -:)

[jiarong0907]

I will test this on GCP later this afternoon.

[jiarong0907]

I tested several times on GCP A100, and I didn't get the error anymore. This seems to work well for now.

[Copilot]

Pull Request Overview

This PR addresses thread safety issues and implements overlap scheduling support for the KV cache management system. The changes enable concurrent operations between worker threads handling CUDA kernels and scheduler threads managing memory allocation/deallocation, while ensuring data consistency and preventing illegal memory accesses.

• Synchronizes CUDA operations before unmapping pages to prevent illegal memory access during concurrent operations
• Implements thread-safe FTensorAllocator and KVCacheManager with proper mutex protection
• Adds support for asynchronous scheduling with null block reservation and background initialization

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
kvcached/tp_ipc_util.py	Adds IPC support for checking KV tensor creation status across workers
kvcached/kv_cache_manager.py	Major refactoring to add thread safety, async scheduling, and background initialization
kvcached/integration/vllm/interfaces.py	Updates interface to support async scheduling parameter
kvcached/integration/sglang/interfaces.py	Updates interface with async scheduling and null block reservation support
engine_integration/scripts/kvcached-sglang-v0.4.6.post2.patch	Fixes type conversion and enables async scheduling in SGLang integration
engine_integration/benchmark/start_server.sh	Removes overlap schedule disable flag to enable the feature
csrc/torch_bindings.cpp	Adds GIL release for thread safety and new KV tensor status checking
csrc/inc/ftensor.hpp	Updates comment to reflect thread safety status
csrc/inc/constants.hpp	Adds start address constant for virtual memory allocation
csrc/inc/allocator.hpp	Adds thread safety with mutex and reorganizes method visibility
csrc/ftensor.cpp	Implements deterministic virtual memory allocation using atomic offset
csrc/allocator.cpp	Implements comprehensive thread safety with mutex protection