Skip to content

feat:add shard-level thread pool for parallel H2H transfers && sglang integration - #1292

Open
lc1314555 wants to merge 3 commits into
ModelEngine-Group:developfrom
lc1314555:develop_lc_h2h_thread_pool
Open

feat:add shard-level thread pool for parallel H2H transfers && sglang integration#1292
lc1314555 wants to merge 3 commits into
ModelEngine-Group:developfrom
lc1314555:develop_lc_h2h_thread_pool

Conversation

@lc1314555

Copy link
Copy Markdown
Contributor

Summary

This PR adds shard-level Host-to-Host (H2H) thread pools to CacheStore and integrates the H2H path with SGLang HiCache.

In the SGLang integration, one page corresponds to one CacheStore shard. Multiple pages can therefore be copied concurrently by different H2H workers.

CacheStore H2H thread pools

Added the following configuration options:

cache_use_host_buffer: true
cache_h2h_worker_number: 4
cache_h2h_queue_depth: 1024

- cache_use_host_buffer enables Host-to-Host Cache transfers.
- cache_h2h_worker_number configures the number of Load and Dump copy workers.
- cache_h2h_queue_depth limits the number of queued and running H2H shard jobs.

Host-buffer mode is incompatible with Device transfer paths, including:

- cache_io_aggregation
- cache_sdma_direct
- use_gdr
- gpu_kv_buffer_addrs

## Dump path

Each non-ready Cache shard is submitted as an independent H2H job.

Dump workers gather data from the SGLang Host tensors into the CacheStore shared buffer in parallel. After all shards of a request are ready, a dedicated completion worker:

1. Marks the Cache handles as ready.
2. Aggregates the shards into one backend task.
3. Submits the task to the downstream PosixStore.
4. Completes the request waiter.

For a split K/V layout, K and V tensors within one page are copied sequentially by one worker, while different pages are copied concurrently.

## Load path

Each Cache shard is submitted as an independent Load job.

A Load worker:

1. Waits for the corresponding backend Load task when required.
2. Marks the Cache buffer as ready.
3. Scatters the Cache payload into the target SGLang Host tensors.
4. Updates the shared request completion counter.

The request waiter is completed after all shards have reached a terminal state.

Cache-hit shards can be scattered directly without another backend Load.

## Queue backpressure

Load and Dump maintain separate atomic outstanding-shard counters.

Queue capacity includes both queued and running jobs. Capacity is reserved for the entire request:

- All shards are submitted when sufficient capacity is available.
- The complete request fails when capacity is insufficient.
- Partial request submission is not allowed.

The queue depth is measured in shards/pages rather than requests.

## Failure handling

The H2H path handles:

- H2H queue exhaustion.
- Invalid Host pointer counts.
- Null source or destination addresses.
- Gather and scatter failures.
- Backend Load/Dump submission failures.
- Backend Load wait failures.

Incomplete Cache handles are marked as failed when a request cannot finish successfully. This prevents Cache entries from remaining permanently in the LOADING state.

All submitted shard jobs participate in request completion, including failure paths.

## SGLang integration

SGLang selects the storage pipeline according to the model type:

 Model type                             Pipeline
━━━━━━━━━━━━━━━━━━━━━━━━━
 Non-MLA models                         Posix
───────────────────────────────────
 MLA models                        Cache|Posix
───────────────────────────────────

For MLA models, the connector automatically configures:

store_pipeline: Cache|Posix
cache_use_host_buffer: true
cache_h2h_worker_number: 4
cache_h2h_queue_depth: 1024
cache_io_aggregation: false
cache_sdma_direct: false
use_gdr: false

User-provided worker and queue-depth values take precedence.

## page_first_kv_split support

For Ascend MLA, the connector obtains one K address and one V address for each page and configures:

tensor_size_list = [K_page_bytes, V_page_bytes]

The Dump payload layout is:

Cache shard = K page | V page | optional alignment padding

Load performs the reverse scatter operation.

When Direct I/O is enabled, shard_size and block_size are aligned to 4096 bytes.

The MLA cache identity includes the model, Host layout, page size, TP size, and tensor-size fingerprint to prevent incompatible layouts from sharing Cache data.

## Scope

This PR supports the main MLA KV Cache.

The following are outside the current scope:

- Indexer Cache
- V2 HostPoolGroup aggregation
- Auxiliary KV pools
- Parallel copies within a single shard
- DMA-based asynchronous H2H

Each H2H worker performs synchronous CPU memcpy, while different pages/shards run concurrently.

## Tests

Added unit tests covering:

- Parallel execution of multiple Load shards.
- Multi-tensor K/V scatter.
- Multi-shard Dump gather.
- Variable tensor sizes.
- Backend payload validation.

@lc1314555 lc1314555 changed the title add shard-level thread pool for parallel H2H transfers && sglang integration feat:add shard-level thread pool for parallel H2H transfers && sglang integration Aug 27, 2026
@qyh111
qyh111 self-requested a review August 31, 2026 02:52
Comment thread ucm/store/cache/cc/dump_queue.cc Outdated
h2hCompletionPool_
.SetNWorker(1)
.SetWorkerFn([this](H2HDumpContextPtr& context, void* const&) {
CompleteH2HDump(std::move(context));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

建议考虑将本 PR 新增的 H2H 复制能力抽象为独立的 HostCopyStream,而不是由 LoadQueue 和 DumpQueue 分别直接持有 ThreadPool 并实现 scatter/gather。
当前 h2hCopyPool_、worker 初始化、任务排队以及 HostToHostScatter/Gather 分散在 Load/Dump 两条路径中,使业务队列同时负责请求状态编排和底层复制执行。后续若需要增加 NUMA affinity、替换 memcpy 实现、统一流控或增加同步接口,需要在两处重复修改。
建议让 HostCopy 抽象负责 worker 生命周期、H2H copy/scatter/gather 和 completion;LoadQueue/DumpQueue 只保留 backend 协调、buffer 状态以及请求级 pending/failure 聚合。另外需要明确该抽象是单 worker FIFO “stream”,还是多 worker 无序 “executor”,避免调用方错误依赖 stream ordering。某种程度上可以参考simustream的写法

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants