Skip to content

Support Cache|Posix H2H pipeline for SGLang MLA models - #1280

Open
lc1314555 wants to merge 2 commits into
ModelEngine-Group:developfrom
lc1314555:develop_lc_cache_h2h_clean
Open

Support Cache|Posix H2H pipeline for SGLang MLA models#1280
lc1314555 wants to merge 2 commits into
ModelEngine-Group:developfrom
lc1314555:develop_lc_cache_h2h_clean

Conversation

@lc1314555

Copy link
Copy Markdown
Contributor

Purpose

Enable SGLang MLA models to use the UCM Cache|Posix pipeline through
Host-to-Host transfers.

SGLang HiCache passes Host KV addresses to external storage, while CacheStore
previously treated all addresses as device pointers and always performed
D2H/H2D transfers.

This PR adds a Host-buffer mode to CacheStore and automatically selects
Cache|Posix for MLA models.

Supported scenarios:

  • Ascend MLA with page_first_kv_split
  • CUDA MLA with page_first
  • Primary MLA KV cache only; Indexer and auxiliary KV pools are out of scope

Non-MLA models continue to use the Posix-only pipeline.

Modifications

CacheStore H2H support

  • Add the cache_use_host_buffer configuration.
  • Gather one or multiple Host tensors into a contiguous CacheStore shard on dump.
  • Scatter a CacheStore shard back to Host tensors on load.
  • Support tensors with different sizes through tensor_size_list.
  • Skip device streams and D2H/H2D synchronization in Host-buffer mode.
  • Reject incompatible options:
    • cache_sdma_direct
    • use_gdr
    • gpu_kv_buffer_addrs
  • Preserve the existing device transfer path by default.

Data path:

Dump:
SGLang Host KV -> CacheStore H2H gather -> PosixStore

Load:
PosixStore -> CacheStore H2H scatter -> SGLang Host KV

### SGLang integration

- Automatically select the pipeline using is_mla_model:

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

- Support page_first, page_first_direct, and page_first_kv_split.
- Calculate per-page tensor sizes from the actual Host tensors.
- Support separate, variable-sized K/V buffers for page_first_kv_split.
- Use a single Host tensor per page for regular CUDA MLA layouts.
- Align shard_size to 4 KiB when io_direct=true.
- Generate layout- and tensor-size-aware cache identifiers.
- Disable SDMA/GDR paths when passing Host pointers.

Non-MLA behavior remains unchanged. On Ascend, the direct Posix path should use:

io_direct: false
posix_io_engine: psync

Direct I/O for Ascend non-MLA Host pools is outside the scope of this PR.

## Test

Added CacheStore unit tests covering:

- H2H gather and scatter
- Multiple tensors with different sizes
- Tensor ordering and data integrity

if (shard.addrs[i] == nullptr) {
return Status::InvalidParam("invalid null host source({})", i);
}
std::memcpy(dst + offset, shard.addrs[i], tensorSizes_[i]);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

CacheStore employs a single-threaded, asynchronous I/O scheduling model; synchronous memcpy operations would cause all access operations to be serialized.

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