chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fold the symm compose into one barrier-gated mega gather

Browse Source 
The publish-variant staging exchange measured ~equal to the default
compact-current AR (90.4 vs 86.5 ms/batch on the traced scenario): the
publish copy and the barrier serialized behind the 0.65 ms prefix
gather ate the transport win the isolated current exchange showed
(0.196 vs 0.354 ms).  Fix the structure instead of the copy: current
rows are now written straight INTO the staging — the fill kernels take
their write destinations solely from page_inverse, so a per-batch
staging-remapped page inverse on the plan retargets them with zero
kernel changes — then cp_symm_barrier, then ONE slot-dense gather
covers prefix pages (pool pointers) and ALL current pages (staging
pointers, including this rank's own) through a concatenated 2*cp
pointer table where current slots carry owner = cp_size + writer and
src = staging slot.  No publish copy, no prefix pre-gather, no second
gather.

The fused fill's loc outputs are dense-geometry-bound, so the token-KV
path computes mixed_locs/staging row indices once per batch (they are
layer-invariant) and the per-layer fill collapses to a single
index_copy_ into the zeroed staging span.

Benchmark (g0033 8xH200, byte-exact, idle-checked): 62.8 ms/batch vs
84.4 default Step A (-26%) and 60.8 ideal; publish variant was 88.0.
151 unit tests; 8-rank GPU byte-exactness vs v2 across 8 layers,
arena on and off.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-11 22:09:29 +00:00
parent e8e5edb230
commit d63fbd4d79
4 changed files with 362 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py
View File

@@ -52,8 +52,9 @@ def cp_shared_kv_compose_arena_enabled() -> bool:
def cp_shared_kv_compose_symm_enabled() -> bool:
"""Step B: exchange current pages via the compact symm staging (zero
NCCL): publish my current pages to staging, barrier, gather peers'.
"""Step B: zero-NCCL compose via the compact symm staging — current rows
are written straight into the staging, then barrier + ONE slot-dense
gather covers prefix (pool) and current (staging) pages together.
Independent of the arena (dense buffers stay rank-local)."""
return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.get())
@@ -74,16 +75,19 @@ class ComposePlan:
dummy-page convention) of all current slots in merged-span order.
When the caller provides per-current-page writer (compute-owner) ranks,
the plan splits the current pages by writer for the symm staging exchange.
The staging slot of current page ``i`` is ``i`` (its index in merged-span
order) — identical on every rank, so no per-batch offset handshake:
the plan additionally carries the symm staging descriptors. The staging
slot of current page ``i`` is ``i`` (its index in merged-span order) —
identical on every rank, so no per-batch offset handshake:
- ``local_current_*``: pages THIS rank wrote — published from the dense
buffer into staging slots before the barrier. The writer-ranks tensor
is carried (all entries == cp_rank) so the publish kernel can index a
uniform pointer table without a per-layer allocation.
- ``remote_current_*``: pages OTHER ranks wrote — gathered after the
barrier from ``staging[writer][slot]`` into ``dense[page]``.
- ``symm_all_owner_ranks`` / ``symm_all_src_pages``: ONE slot-dense
gather descriptor pair covering the whole dense buffer, indexed against
a concatenated ``[pool peer ptrs | staging peer ptrs]`` table — prefix
slots keep the pool owner/page, current slots carry
``cp_size + writer`` / staging slot ``i``, everything else stays -1
(sentinel zero-fill).
- ``staging_page_inverse``: ``page_inverse`` with current pages remapped
to their staging slot (others -1), so the unchanged fill kernels write
current rows straight INTO the staging instead of the dense buffer.
"""
prefix_owner_ranks: torch.Tensor
@@ -92,12 +96,9 @@ class ComposePlan:
total_slots: int
num_prefix_slots: int
num_current_pages: int
local_current_writer_ranks: torch.Tensor | None = None
local_current_slot_indices: torch.Tensor | None = None
local_current_dense_pages: torch.Tensor | None = None
remote_current_writer_ranks: torch.Tensor | None = None
remote_current_slot_indices: torch.Tensor | None = None
remote_current_dense_pages: torch.Tensor | None = None
symm_all_owner_ranks: torch.Tensor | None = None
symm_all_src_pages: torch.Tensor | None = None
staging_page_inverse: torch.Tensor | None = None
def _spans_key(spans: list[tuple[int, int]] | None) -> tuple[tuple[int, int], ...]:
@@ -106,6 +107,34 @@ def _spans_key(spans: list[tuple[int, int]] | None) -> tuple[tuple[int, int], ..
return tuple((int(s), int(e)) for s, e in spans)
def _build_staging_page_inverse(
page_inverse: torch.Tensor,
current_dense_pages: torch.Tensor,
) -> torch.Tensor:
"""Remap ``page_inverse`` so current pages point at their staging slot.
``page_inverse`` maps (request row, logical page) -> dense page id; the
fill kernels derive every current-row write destination from it. In the
staging copy a current dense page becomes its index in the (ascending)
merged-span order — the staging slot — and every other entry becomes -1,
so the unchanged fill kernels write current rows into the compact staging
viewed as a page-major buffer (no dummy page). Current locs only ever
map to current pages, so the -1 entries are never exercised by the fill.
"""
if page_inverse.dim() == 1:
page_inverse = page_inverse.unsqueeze(0)
num_current = int(current_dense_pages.numel())
invalid = torch.full_like(page_inverse, -1)
if num_current == 0:
return invalid
pos = torch.searchsorted(
current_dense_pages, page_inverse.clamp(min=0).contiguous()
).clamp(max=num_current - 1)
is_current = (page_inverse > 0) & (current_dense_pages[pos] == page_inverse)
return torch.where(is_current, pos, invalid)
def get_or_build_compose_plan(
*,
slot_remap,
@@ -182,12 +211,9 @@ def get_or_build_compose_plan(
else:
current_dense_pages = torch.empty(0, device=device, dtype=torch.long)
local_writer_ranks = None
local_slot_indices = None
local_dense_pages = None
remote_writer_ranks = None
remote_slot_indices = None
remote_dense_pages = None
symm_all_owner_ranks = None
symm_all_src_pages = None
staging_page_inverse = None
if current_page_writer_ranks is not None:
num_current = int(current_dense_pages.numel())
if len(current_page_writer_ranks) != num_current:
@@ -196,18 +222,22 @@ def get_or_build_compose_plan(
f"count mismatch: writers={len(current_page_writer_ranks)} "
f"current_pages={num_current} (page-aligned split required)"
)
writers = torch.tensor(
current_page_writer_ranks, dtype=torch.long, device=device
)
slot_indices = torch.arange(num_current, dtype=torch.long, device=device)
remote = writers != int(layout.cp_rank)
local = ~remote
local_writer_ranks = writers[local].contiguous()
local_slot_indices = slot_indices[local].contiguous()
local_dense_pages = current_dense_pages[local].contiguous()
remote_writer_ranks = writers[remote].contiguous()
remote_slot_indices = slot_indices[remote].contiguous()
remote_dense_pages = current_dense_pages[remote].contiguous()
symm_all_owner_ranks = prefix_owner_ranks.clone()
symm_all_src_pages = prefix_src_pages.clone()
if num_current > 0:
writers = torch.tensor(
current_page_writer_ranks, dtype=torch.long, device=device
)
current_slots = current_dense_pages - 1
symm_all_owner_ranks[current_slots] = writers + int(layout.cp_size)
symm_all_src_pages[current_slots] = torch.arange(
num_current, dtype=torch.long, device=device
)
page_inverse = getattr(slot_remap, "page_inverse", None)
if page_inverse is not None:
staging_page_inverse = _build_staging_page_inverse(
page_inverse, current_dense_pages
)
plan = ComposePlan(
prefix_owner_ranks=prefix_owner_ranks,
@@ -216,12 +246,9 @@ def get_or_build_compose_plan(
total_slots=total_slots,
num_prefix_slots=num_prefix_slots,
num_current_pages=int(current_dense_pages.numel()),
local_current_writer_ranks=local_writer_ranks,
local_current_slot_indices=local_slot_indices,
local_current_dense_pages=local_dense_pages,
remote_current_writer_ranks=remote_writer_ranks,
remote_current_slot_indices=remote_slot_indices,
remote_current_dense_pages=remote_dense_pages,
symm_all_owner_ranks=symm_all_owner_ranks,
symm_all_src_pages=symm_all_src_pages,
staging_page_inverse=staging_page_inverse,
)
plans[key] = plan
return plan
@@ -335,7 +362,7 @@ class CpComposeStaging(_RoundParity):
same on every rank), so ``peer_base + slot * page_nbytes`` addresses any
peer's copy of page ``i`` directly.
Reuse safety is the double-buffer parity argument: my publish into half
Reuse safety is the double-buffer parity argument: my fill into half
H at round R+2 starts only after my round-R+1 barrier completed, which
proves every peer's R+1 barrier kernel ran, which (stream order) proves
every peer's round-R gather of half H finished.

363
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -4556,41 +4556,90 @@ def _symm_staging_ready_or_register(
return staging.registered
def _symm_exchange_current_pages(
dense_buffer: torch.Tensor,
plan,
layout: CpSharedKVLayout,
class _SymmTokenFillState:
"""Layer-invariant pieces of the symm token-KV compose, built once per
batch and anchored on the (batch-lifetime) compose plan.
``staging_current_rows`` maps each current KV row to its row in the
compact staging (via the staging page inverse), so the per-layer fill is
a single ``index_copy_``. ``mixed_locs`` is the page-slack-masked loc
remap the fused fill kernel used to recompute every layer — it depends
only on the batch's locs, never on buffer contents.
"""
__slots__ = ("staging_current_rows", "mixed_locs")
def __init__(self, staging_current_rows, mixed_locs):
self.staging_current_rows = staging_current_rows
self.mixed_locs = mixed_locs
def _get_or_build_symm_token_fill_state(
*,
page_nbytes: int,
plan,
slot_remap,
layout: CpSharedKVLayout,
kv_cache: torch.Tensor,
logical_locs: torch.Tensor,
current_locs: torch.Tensor,
loc_req_id: torch.Tensor,
current_req_id: torch.Tensor,
page_size: int,
) -> "_SymmTokenFillState":
state = getattr(plan, "_symm_token_fill_state", None)
if state is not None:
return state
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
layout=layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
slot_remap=slot_remap,
page_size=page_size,
loc_req_id=loc_req_id,
)
# Same masking as the fused fill kernel / torch reference: page-slack
# rows inside current pages become -1 (invisible to attention).
current_mask, _ = build_current_loc_remap(logical_locs, current_locs)
current_page_mask = build_current_page_mask(
logical_locs, current_locs, page_size=page_size
)
mixed_locs = torch.where(
current_page_mask & (~current_mask),
torch.full_like(dense_locs, -1),
dense_locs,
)
staging_current_rows = remap_logical_locs_to_slot_dense_locs_optimized(
current_locs.reshape(-1),
page_inverse=plan.staging_page_inverse,
page_size=page_size,
loc_req_id=current_req_id.reshape(-1),
).to(torch.long)
state = _SymmTokenFillState(staging_current_rows, mixed_locs)
object.__setattr__(plan, "_symm_token_fill_state", state)
return state
def _symm_begin_current_staging(
*,
staging,
plan,
kind: str,
layer_id: int,
) -> None:
"""Step B current-page exchange through the compact symm staging:
page_nbytes: int,
) -> tuple[int, torch.Tensor]:
"""Open this round's staging span for the current-row fill.
publish my written current pages: dense[page] -> staging[slot]
barrier cp_symm_barrier (counting; release/acquire at system scope)
gather peers' pages: peer_staging[writer][slot] -> dense[page]
Validates the batch against the registered staging layout (both checks
are batch-logical, hence rank-uniform), advances the round parity, and
returns the zeroed span — zeroed so that tail-slack rows inside current
pages stay zero in every peer's gathered copy, exactly like the sentinel
zero-fill did on the dense buffer.
"""
The staging slot of current page ``i`` is ``i`` (merged-span order),
identical on every rank, so peers address each other's staging without a
per-batch handshake. The barrier runs even with zero local/remote pages
— barrier COUNTS must match across ranks.
INVARIANT: every gate on the path to this call (compose env flags,
page_aligned metadata, staging.registered, the agreed IPC capability,
prefetch absence) MUST be rank-uniform. A per-rank divergence here
desyncs the barrier counting and hangs the CP group — never add a
per-rank condition without routing it through a group agreement first
(see _agreed_tai_ipc_peer_ptrs)."""
from tai_kernel.nsa_prefill.ipc import (
cp_symm_barrier,
gather_cuda_ipc_peer_pages,
)
staging = get_compose_staging(dense_buffer.device)
if int(plan.num_current_pages) > staging.capacity_pages:
# Batch-logical quantity -> raises uniformly on every rank.
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_symm] batch current pages "
f"exceed the symm staging capacity: pages={plan.num_current_pages} "
@@ -4604,37 +4653,64 @@ def _symm_exchange_current_pages(
f"the registered staging layout: kind={kind} page_nbytes="
f"{page_nbytes} registered={staging.page_nbytes(kind)}"
)
if plan.staging_page_inverse is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_symm] compose plan has no "
"staging page inverse (slot_remap.page_inverse missing?)"
)
parity = staging.begin_round(int(layer_id), kind)
if (
plan.local_current_slot_indices is not None
and plan.local_current_slot_indices.numel() > 0
):
# Local copy via the gather kernel: a uniform pointer table aimed at
# my own dense buffer turns it into dense[page] -> staging[slot].
self_ptrs = torch.full(
(staging.cp_size,), dense_buffer.data_ptr(), dtype=torch.int64
)
gather_cuda_ipc_peer_pages(
self_ptrs,
staging.buffer(kind, parity),
plan.local_current_writer_ranks,
plan.local_current_dense_pages,
plan.local_current_slot_indices,
page_nbytes=page_nbytes,
)
span = staging.buffer(kind, parity)[
: int(plan.num_current_pages) * page_nbytes
]
span.zero_()
return parity, span
def _symm_barrier_and_gather_all(
*,
kernels,
pool_peer_ptrs: torch.Tensor,
dense_buffer: torch.Tensor,
plan,
layout: CpSharedKVLayout,
staging,
kind: str,
parity: int,
page_nbytes: int,
) -> None:
"""Step B compose: barrier, then ONE slot-dense gather for everything.
The current rows were already written into this rank's staging span (the
fill kernels were pointed there via ``plan.staging_page_inverse``), so
after the barrier a single gather against the concatenated
``[pool peer ptrs | staging peer ptrs]`` table materializes prefix pages
(owner < cp_size, from the KV pools), current pages (owner = cp_size +
writer, from the stagings — including this rank's own), and zero-fills
the remaining sentinel slots. No publish copy, no second gather.
The barrier runs even with zero current pages — barrier COUNTS must
match across ranks.
INVARIANT: every gate on the path to this call (compose env flags,
page_aligned metadata, staging.registered, the agreed IPC capability,
prefetch absence) MUST be rank-uniform. A per-rank divergence here
desyncs the barrier counting and hangs the CP group — never add a
per-rank condition without routing it through a group agreement first
(see _agreed_tai_ipc_peer_ptrs)."""
from tai_kernel.nsa_prefill.ipc import cp_symm_barrier
cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank))
if (
plan.remote_current_slot_indices is not None
and plan.remote_current_slot_indices.numel() > 0
):
gather_cuda_ipc_peer_pages(
staging.peer_region_ptrs(kind, parity),
dense_buffer,
plan.remote_current_writer_ranks,
plan.remote_current_slot_indices,
plan.remote_current_dense_pages,
page_nbytes=page_nbytes,
)
combined_ptrs = torch.cat(
[pool_peer_ptrs, staging.peer_region_ptrs(kind, parity)]
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
combined_ptrs,
dense_buffer,
plan.symm_all_owner_ranks,
plan.symm_all_src_pages,
page_nbytes=page_nbytes,
)
def _resolve_partial_current_spans(
@@ -4759,6 +4835,7 @@ def _compose_token_kv_partial_current_v2(
page_size=page_size,
)
gathered = ipc_state is not None
kv_page_nbytes = _token_kv_page_nbytes(kv_cache, page_size)
if gathered:
kernels, peer_ptrs = ipc_state
dense_kv_cache = acquire_dense_buffer(
@@ -4768,13 +4845,14 @@ def _compose_token_kv_partial_current_v2(
layer_id=layer_id,
kind="token_kv",
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_kv_cache,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
)
if not use_symm:
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_kv_cache,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=kv_page_nbytes,
)
else:
dense_kv_cache = kv_cache.new_zeros((dense_rows, *kv_cache.shape[1:]))
for prefix_start_slot, prefix_end_slot in prefix_spans:
@@ -4788,6 +4866,70 @@ def _compose_token_kv_partial_current_v2(
end_slot=prefix_end_slot,
)
if use_symm:
# Step B: write current rows straight into this round's staging span
# (one index_copy; the loc remap and slack masking are layer-invariant
# and cached on the plan), then barrier + ONE slot-dense gather for
# prefix AND current pages. The prefix-only gather above was skipped
# — the mega gather covers it.
staging = get_compose_staging(kv_cache.device)
parity, staging_span = _symm_begin_current_staging(
staging=staging,
plan=plan,
kind="token_kv",
layer_id=layer_id,
page_nbytes=kv_page_nbytes,
)
fill_state = _get_or_build_symm_token_fill_state(
plan=plan,
slot_remap=slot_remap,
layout=layout,
kv_cache=kv_cache,
logical_locs=logical_locs,
current_locs=current_locs,
loc_req_id=loc_req_id,
current_req_id=current_req_id,
page_size=page_size,
)
num_rows = int(fill_state.staging_current_rows.numel())
if num_rows > 0:
staging_rows = staging_span.view(kv_cache.dtype).reshape(
int(plan.num_current_pages) * page_size, *kv_cache.shape[1:]
)
staging_rows.index_copy_(
0,
fill_state.staging_current_rows,
current_kv_cache[:num_rows],
)
_symm_barrier_and_gather_all(
kernels=kernels,
pool_peer_ptrs=peer_ptrs,
dense_buffer=dense_kv_cache,
plan=plan,
layout=layout,
staging=staging,
kind="token_kv",
parity=parity,
page_nbytes=kv_page_nbytes,
)
mixed_kv_cache = dense_kv_cache
mixed_locs = fill_state.mixed_locs
log_cp_shared_kv_bs_gt1_timing(
"mla_partial_current_compose_v2",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_rows=%s "
"prefix_span_pages=%s current_pages=%s symm=1 kv_dtype=%s",
layout.cp_rank,
layer_id,
layout.cp_size,
plan.total_slots,
int(mixed_kv_cache.shape[0]),
plan.num_prefix_slots,
plan.num_current_pages,
kv_cache.dtype,
)
return mixed_kv_cache, mixed_locs
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
layout=layout,
@@ -4812,16 +4954,7 @@ def _compose_token_kv_partial_current_v2(
)
if gathered:
if use_symm:
_symm_exchange_current_pages(
mixed_kv_cache,
plan,
layout,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
kind="token_kv",
layer_id=layer_id,
)
elif plan.num_current_pages > 0:
if plan.num_current_pages > 0:
_reduce_current_pages_compact(
mixed_kv_cache,
int(slot_remap.dense_num_pages),
@@ -5161,13 +5294,14 @@ def _compose_index_partial_current_v2(
layer_id=layer_id,
kind="index",
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_page_buffer,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
if not use_symm:
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_page_buffer,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
else:
dense_page_buffer = page_buffer.new_zeros(
(dense_num_pages, *page_buffer.shape[1:])
@@ -5182,6 +5316,60 @@ def _compose_index_partial_current_v2(
end_slot=prefix_end_slot,
)
if use_symm:
# Step B: fill current index rows straight into this round's staging
# span (the unchanged fill kernel — its write destinations all come
# from the page inverse, which the plan remapped to staging slots),
# then barrier + ONE slot-dense gather for prefix AND current pages.
index_page_nbytes = _page_nbytes_from_page_tensor(page_buffer)
staging = get_compose_staging(page_buffer.device)
parity, staging_span = _symm_begin_current_staging(
staging=staging,
plan=plan,
kind="index",
layer_id=layer_id,
page_nbytes=index_page_nbytes,
)
if plan.num_current_pages > 0:
staging_pages = staging_span.view(page_buffer.dtype).reshape(
int(plan.num_current_pages), *page_buffer.shape[1:]
)
fill_current_index_page_slots(
dense_page_buffer=staging_pages,
current_index_k=current_index_k,
current_index_scale=current_index_scale,
current_locs=current_locs,
page_inverse=plan.staging_page_inverse,
page_size=page_size,
index_head_dim=index_head_dim,
current_req_id=current_req_id,
)
_symm_barrier_and_gather_all(
kernels=kernels,
pool_peer_ptrs=peer_ptrs,
dense_buffer=dense_page_buffer,
plan=plan,
layout=layout,
staging=staging,
kind="index",
parity=parity,
page_nbytes=index_page_nbytes,
)
log_cp_shared_kv_bs_gt1_timing(
"index_partial_current_compose_v2",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_pages=%s "
"prefix_span_pages=%s current_pages=%s symm=1",
layout.cp_rank,
layer_id,
layout.cp_size,
plan.total_slots,
dense_num_pages,
plan.num_prefix_slots,
plan.num_current_pages,
)
return dense_page_buffer, slot_remap.dense_pages
dense_page_buffer = fill_current_index_page_slots(
dense_page_buffer=dense_page_buffer,
current_index_k=current_index_k,
@@ -5194,16 +5382,7 @@ def _compose_index_partial_current_v2(
)
if gathered:
if use_symm:
_symm_exchange_current_pages(
dense_page_buffer,
plan,
layout,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
kind="index",
layer_id=layer_id,
)
elif plan.num_current_pages > 0:
if plan.num_current_pages > 0:
_reduce_current_pages_compact(
dense_page_buffer,
dense_num_pages,