perf(disagg): coarse-grained per-layer transfer + skip re-registration (lever A)

For large bs (production target: bs~10 x ~100k tokens), per-layer granularity does
10x79 = 790 submitTransfer calls + CUDA events + enqueues per forward on the forward
thread. Two overhead cuts:
- Group SGLANG_CP_SHARED_KV_PER_LAYER_GROUP (default 8) consecutive layers into ONE
  RDMA submit: ~num_layers/K submits + events + enqueues instead of per-layer; same
  bytes (page index lists are identical across layers). on_layer_end is O(1) at
  non-boundary layers. The last partial group enqueues via the num_layers boundary;
  any misses fall back to one batched sync submit.
- Scheduler hook skips reqs already registered (bs>1 batch-forming re-iterates the
  same reqs ~9x -> was rebuilding the CP filter + context every time).

27 unit tests pass incl. grouping-boundary + batched-fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-07 09:51:05 +00:00
co-authored by Claude Opus 4.8
parent 5fdf439c7e
commit e12afe8ced
4 changed files with 119 additions and 59 deletions
+4
View File
@@ -215,6 +215,10 @@ class Envs:
# asynchronously (pipelined batch_transfer_async) instead of one monolithic
# all-layers batch_transfer_sync. Default off until validated.
SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER = EnvBool(False)
# Lever A granularity: group this many consecutive layers into one RDMA submit.
# Larger = fewer submitTransfer calls / events / enqueues (less CPU overhead at
# large bs) but coarser overlap. 1 = per-layer.
SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)