Skip to content

Add Ascend Cache IO aggregation copy path - #1038

Open
NaganooMei wants to merge 1 commit into
ModelEngine-Group:developfrom
NaganooMei:io-aggregation-upstream-pr-v2
Open

Add Ascend Cache IO aggregation copy path#1038
NaganooMei wants to merge 1 commit into
ModelEngine-Group:developfrom
NaganooMei:io-aggregation-upstream-pr-v2

Conversation

@NaganooMei

@NaganooMei NaganooMei commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Purpose

This PR adds an Ascend CacheStore IO aggregation path for small H2D/D2H shard copies.

The original CacheStore transfer path submits one async copy per tensor fragment. For workloads where one CacheStore shard contains many tensor fragments, this creates many small H2D/D2H submissions and increases CPU/runtime scheduling overhead.

This PR introduces an Ascend-only aggregation stream. It treats one CacheStore shard as one IO object, uses a staging buffer on device, and launches FFTS SDMA descriptors to scatter/gather between the staging buffer and device tensor addresses.

For PLATFORM=ascend / RUNTIME_ENVIRONMENT=ascend builds, CacheStore enables IO aggregation by default. When use_layerwise=true, the vLLM connector disables CacheStore IO aggregation to preserve the existing layerwise path behavior. Non-Ascend runtimes do not compile or allow this feature.

Modifications

  • Add an explicit trans-layer factory for CacheStore IO aggregation:
    • IoAggregationStreamConfig
    • MakeIoAggregationStream(config)
  • Keep CopyStream as the CacheStore-side stream owner, but let IO aggregation use one outer stream object with internal lanes.
  • Add shard-level multi-fragment H2D/D2H stream APIs so LoadQueue and DumpQueue can submit one CacheStore shard as a scatter/gather object instead of looping over tensor fragments.
  • Add AscendIoAggregationStream and AscendShardIOAggregator.
    • H2D load: host shard buffer -> device staging buffer -> device tensor fragments.
    • D2H dump: device tensor fragments -> device staging buffer -> host shard buffer.
    • Internal lane count comes from cache_stream_number.
    • Pipeline depth and max ready lanes keep default values and remain configurable for advanced tuning.
  • Add FftsSdmaDispatcher to build and launch FFTS SDMA copy descriptors.
  • Add Ascend runtime/build isolation:
    • Ascend builds enable UCM_RUNTIME_ASCEND_IO_AGGREGATION.
    • Other runtimes reject cache_io_aggregation=true.
  • Add config validation for unsupported combinations:
    • IO aggregation requires Ascend runtime support.
    • IO aggregation is incompatible with use_gdr.
    • Invalid pipeline depth / ready-lane values fail early.
  • Keep existing GDR, CUDA, simulation, and normal Ascend copy paths unchanged when IO aggregation is disabled.

Config Behavior

  • cache_io_aggregation defaults to enabled in Ascend builds.
  • cache_io_aggregation is forced to disabled by the vLLM connector when use_layerwise=true.
  • cache_io_aggregation_pipeline_depth defaults to 2.
  • cache_io_aggregation_max_ready_lanes defaults to 8.
  • The example YAML does not expose IO aggregation tuning knobs; normal users only need to build with Ascend runtime support.

Test

  • Verified Ascend build path with IO aggregation enabled.
  • Verified default CacheStore load/dump behavior with IO aggregation disabled.
  • Verified CacheStore load H2D path with IO aggregation enabled.
  • Verified CacheStore dump D2H path with IO aggregation enabled.
  • Verified unsupported runtime/config combinations fail during config validation.

Comment thread ucm/shared/trans/ascend/ffts/ffts_sdma_dispatcher.cc
Comment thread ucm/shared/trans/ascend/ascend_io_aggregation_stream.cc Outdated
@NaganooMei
NaganooMei force-pushed the io-aggregation-upstream-pr-v2 branch from 5aafeb1 to 7382b4e Compare June 20, 2026 08:28
@NaganooMei
NaganooMei requested a review from qyh111 as a code owner June 22, 2026 11:07
@NaganooMei
NaganooMei force-pushed the io-aggregation-upstream-pr-v2 branch 13 times, most recently from ca7a54c to d184a64 Compare June 27, 2026 01:28
@NaganooMei
NaganooMei force-pushed the io-aggregation-upstream-pr-v2 branch 2 times, most recently from c59591c to 5982518 Compare July 8, 2026 02:59
@NaganooMei
NaganooMei force-pushed the io-aggregation-upstream-pr-v2 branch from 5982518 to a577e5c Compare July 8, 2026 03:40
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