minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[kv-shard 2/4] Sharded pools (#37615)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Shunkangz
2026-09-16 19:22:23 +08:00
committed by GitHub
co-authored by Zhangheng
parent 76e06febab
commit e7f7447333
8 changed files with 1709 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+32
View File
@@ -22,6 +22,10 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.kv_shard_hooks import (
get_kv_shard_pool,
prepare_kv_shard_forward,
)
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_active
@@ -46,6 +50,7 @@ from sglang.srt.utils.common import get_device_capability
if TYPE_CHECKING:
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
from sglang.srt.model_executor.model_runner import ModelRunner
from sgl_kernel import merge_state_v2
@@ -145,6 +150,10 @@ class FlashAttentionBackend(AttentionBackend):
needs_cpu_seq_lens: bool = False
supports_ragged_verify_graph: bool = True
# Set from the pool type in __init__; the class default keeps the extend
# metadata guard readable on instances built with __new__ (test stubs).
_kv_shard_pool: Optional[PageInterleaveKVPoolMixin] = None
# Chunked-prefix attention reads the stable ForwardBatch cu-seqlens and
# KV-index buffers directly, so it needs no backend-private replay state.
supports_full_cuda_graph_chunked_prefix = True
@@ -205,6 +214,12 @@ class FlashAttentionBackend(AttentionBackend):
and model_runner.token_to_kv_pool.swa_layer_nums > 0
)
self._kv_shard_pool = get_kv_shard_pool(self.token_to_kv_pool)
# begin_shard_extend builds the owner-major gather plan from host-side
# prefix/final lengths. Normal FA3 metadata is device-only, so opt the
# sharded variant back into FutureMap's CPU mirror publication.
self.needs_cpu_seq_lens = self._kv_shard_pool is not None
self.topk = get_spec().speculative_eagle_topk or 0
self.speculative_num_steps = speculative_num_steps
self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
@@ -1154,6 +1169,23 @@ class FlashAttentionBackend(AttentionBackend):
)
)
# Logical-page KV sharding: capture the batch's gather plan and swap the
# page table to scratch rows. During a sharded extend, attention reads
# the assembled [prefix | chunk] scratch, never the striped pool rows;
# the plan capture also kicks the first layer's prefix gather.
#
# Runs after KVIndexTranslator and before the `// page_size` reduction.
# The unified-memory UnifiedKVPool and page-interleaved pools are
# alternatives, so at most one translation fires.
if self._kv_shard_pool is not None and prepare_kv_shard_forward(
self._kv_shard_pool,
self.req_to_token,
forward_batch,
):
metadata.page_table = self._kv_shard_pool.translate_loc_to_scratch(
metadata.page_table
).to(torch.int32)
# Convert the page table to a strided format which is needed by FA3 API
if self.page_size > 1 and not _unified_read:
self.strided_indices = torch.arange(
python/sglang/srt/layers/attention/kv_shard_hooks.py
+67
View File
@@ -0,0 +1,67 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Logical-page KV-shard helpers shared by attention backends.
Page-interleaved pools build one gather plan per extend batch. Keeping pool
detection and that begin/end lifecycle here prevents each compatible attention
backend from implementing a subtly different version of the contract.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def get_kv_shard_pool(token_to_kv_pool) -> Optional[PageInterleaveKVPoolMixin]:
"""Return the page-interleaved pool, if logical-page sharding is active."""
return (
token_to_kv_pool
if isinstance(token_to_kv_pool, PageInterleaveKVPoolMixin)
else None
)
def prepare_kv_shard_forward(
pool: PageInterleaveKVPoolMixin,
req_to_token: torch.Tensor,
forward_batch: ForwardBatch,
) -> bool:
"""Update the pool's gather plan and report whether this is an extend."""
if not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed():
pool.end_shard_extend()
return False
req_pool_indices = forward_batch.req_pool_indices
prefix_lens = forward_batch.extend_prefix_lens_cpu
seq_lens = forward_batch.seq_lens_cpu
if req_pool_indices is None or prefix_lens is None or seq_lens is None:
raise RuntimeError(
"KV-sharded attention requires request indices and CPU length metadata"
)
pool.begin_shard_extend(
req_to_token,
req_pool_indices,
prefix_lens,
seq_lens,
)
return True
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+13
View File
@@ -50,6 +50,10 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAAttnBackend,
FlashInferMLAMultiStepDraftBackend,
)
from sglang.srt.layers.attention.kv_shard_hooks import (
get_kv_shard_pool,
prepare_kv_shard_forward,
)
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.layers.dcp.layout import get_dcp_lens
from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
@@ -239,6 +243,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.q_data_type = model_runner.dtype
self.page_size = model_runner.page_size
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self._kv_shard_pool = get_kv_shard_pool(model_runner.token_to_kv_pool)
self.needs_cpu_seq_lens |= self._kv_shard_pool is not None
# Workspace allocation
self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
@@ -806,6 +812,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
if self._kv_shard_pool is not None:
prepare_kv_shard_forward(
self._kv_shard_pool,
self.req_to_token,
forward_batch,
)
self._decode_kernel_loc = None
# Delegate to parent for non-decode modes.
if (
python/sglang/srt/mem_cache/page_interleave.py
+29 -2
View File
@@ -23,8 +23,8 @@ sees only its own physical pages; the boundary is the pure bijection below.
The shard group is the group across which KV storage is replicated today and
therefore can be striped without extra compute-time communication:
- GQA/MHA models: the **attention CP group** — prefill CP already allgathers
the full chunk's K/V to every CP rank (``cp_allgather_and_save_kv_cache``).
- GQA/MHA models: the **attention CP group** — the prefill CP strategy
materializes the full chunk's K/V on every CP rank before the pool write.
- MLA models: the **attention TP group** — the latent KV projection is
``ReplicatedLinear``, so every attn-TP rank computes identical latent KV.
"""
@@ -32,10 +32,16 @@ therefore can be striped without extra compute-time communication:
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import msgspec
import torch
from sglang.srt.runtime_context import get_parallel
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state import GroupCoordinator
logger = logging.getLogger(__name__)
@@ -84,3 +90,24 @@ class PageInterleavePlacement:
def filter_local(self, loc: torch.Tensor, rank: int) -> torch.Tensor:
"""Logical slots -> this rank's physical pool rows, order-preserving."""
return self.local_index(loc[self.local_mask(loc, rank)])
def get_kv_shard_group(use_mla_backend: bool) -> GroupCoordinator:
"""The group KV pages are striped across — the axis that replicates KV
at rest, chosen by topology:
- An active attention-CP group takes precedence: prefill CP replicates
KV storage across CP ranks for every attention type (GQA via the
full-chunk allgather, MLA via the CP latent-KV rebuild).
- Without CP, MLA latent KV is still replicated across attention-TP
(ReplicatedLinear projection), so the attn-TP group is the shard axis.
- GQA without CP has no replicated axis (KV is head-sharded across TP);
the returned trivial CP group has world_size 1, which disables
sharding in get_kv_shard_group_info.
"""
cp_group = get_parallel().attn_cp_group
if cp_group.world_size > 1:
return cp_group
if use_mla_backend:
return get_parallel().attn_tp_group
return cp_group
python/sglang/srt/mem_cache/page_interleave_pool.py
+703
View File
@@ -0,0 +1,703 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""KV pools striped at logical-page granularity across a shard group.
Per-layer pool tensors keep the stock shape and size — identical rows on every
rank; only the written rows differ: each rank persists the rows it owns under
the pure placement bijection (``mem_cache/page_interleave.py``).
Central consequence for the forward pass: during a sharded prefill, extend
attention never reads the KV pool — a rank's pool holds only its stripe of
the prefix *and* of the current chunk. Attention reads a per-layer assembled
scratch slot laid out ``[prefix region | chunk region | trash page]``:
- The prefix is NCCL-allgathered from the shard group into the slot one layer
ahead, on a dedicated side stream with a dedicated ``PyNcclCommunicator``
ahead. Under rotated
owner-classed allocation the owners of a cached prefix's pages are exactly
cyclic, so per-rank owned counts differ by <= 1: each rank sends its own
prefix pages in local-page order, padded to ``ceil(prefix_pages / N)`` pages
with the (never referenced) trash page — a *regular* in-place allgather,
owner-major output, no reorder pass.
- The chunk region is staged locally on the compute stream where the write
path already holds the full chunk (GQA CP allgathers K/V before the pool
write; MLA latent is replicated across the shard group at write time).
- The trash page absorbs padded/dummy locations (the reserved logical pages
covering loc < N*ps, and any location outside the current batch's plan).
Scratch addressing goes through a per-batch ``logical page -> scratch
page`` lookup (``_page_pos``, one int32 per logical page, rebuilt in
``begin_shard_extend``): purely local, mirrored by construction, and valid
for any consumer index vector (attention page tables, MLA prefix
``kv_indices``), not just whole-prefix position arithmetic.
Layer-ahead pipelining needs no model changes: acquiring layer ``l``'s slot
for reading kicks layer ``l+1``'s gather. Reads happen in SPMD lockstep on
every rank, so the collective order is symmetric by construction.
"""
from __future__ import annotations
import logging
from contextlib import nullcontext
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
import torch
from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator
from sglang.srt.mem_cache.memory_pool import (
GPU_MEMORY_TYPE_KV_CACHE,
MHATokenToKVPool,
MLATokenToKVPool,
RadixAttention,
unwrap_write_loc,
)
from sglang.srt.mem_cache.page_interleave import (
PageInterleavePlacement,
PageShardSpec,
)
from sglang.srt.mem_cache.utils import (
get_mla_kv_buffer_triton,
set_mla_kv_buffer_triton,
)
from sglang.srt.utils import ceil_div, get_bool_env_var
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state import GroupCoordinator
logger = logging.getLogger(__name__)
class _ScratchSlot:
"""One half of the double buffer: scratch tensors + ready event.
``resident_key`` identifies what the slot currently holds — ``(layer_id,
epoch)`` — so repeated reads of the same layer skip the re-gather and a
new batch (epoch bump) retires stale residency without any invalidation
walk.
"""
def __init__(self, tensors: Dict[str, torch.Tensor], device_module):
self.tensors = tensors
self.ready = device_module.Event()
self.resident_key: Optional[Tuple[int, int]] = None
class PageInterleaveKVPoolMixin:
"""Shard-generic state and mechanics; mixed into concrete pools below.
Subclasses call ``_init_page_shard_state`` after the base pool has created
its buffers, and implement ``_scratch_tensor_specs`` (per-slot tensors) and
``_gather_pairs`` (pool buffer -> scratch tensor pairs of one layer).
"""
# ---- init ---------------------------------------------------------------
def _init_page_shard_state(
self, shard_spec: PageShardSpec, shard_group: GroupCoordinator
):
spec = shard_spec
assert spec.shard_size > 1, "page-interleave sharding needs shard_size > 1"
assert spec.shard_size == shard_group.world_size
assert spec.shard_rank == shard_group.rank_in_group
assert spec.page_size == self.page_size
# The prefix region must fit N * ceil(prefix_pages / N) pages for any
# prefix, i.e. be a multiple of the full-group span; the chunk region
# is per-page.
assert spec.max_prefix_tokens % spec.logical_page_size == 0
assert spec.chunk_tokens % spec.page_size == 0
self.shard_spec = spec
self.placement = PageInterleavePlacement(spec)
self.shard_rank = spec.shard_rank
self.shard_size = spec.shard_size
self.device_module = torch.get_device_module(self.device)
# Dedicated communicator + stream so the layer-ahead gathers never
# interleave with the group's main collectives.
self.kv_gather_comm: PyNcclCommunicator = PyNcclCommunicator(
group=shard_group.cpu_group, device=shard_group.device
)
self.kv_gather_stream = self.device_module.Stream()
# Scratch: [prefix | chunk | trash page], double-buffered.
scratch_rows = spec.max_prefix_tokens + spec.chunk_tokens + spec.page_size
self._chunk_base = spec.max_prefix_tokens
self._trash_base = spec.max_prefix_tokens + spec.chunk_tokens
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.custom_mem_pool
else nullcontext()
):
self._slots = [
_ScratchSlot(
self._scratch_tensor_specs(scratch_rows), self.device_module
)
for _ in range(2)
]
# Per-batch absolute scratch-page index of every logical page.
# Unplanned pages initially target the trash page. Logical page
# ids run [0, N * (size/ps + 1)) — N per physical page of one
# rank's pool incl. the reserved padded page.
self._page_pos = torch.full(
(spec.shard_size * (self.size // spec.page_size + 2),),
self._trash_base // spec.page_size,
dtype=torch.int32,
device=self.device,
)
self._epoch = 0
self._shard_extend_active = False
self._block_pages = 0
# Strictly larger than any local physical page id: the owner-major
# sort key of logical page l is (l % N) * stride + l // N.
self._local_page_stride = self.size // spec.page_size + 2
self._debug_plan_checks = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
self._send_rows: Optional[torch.Tensor] = None
self._write_plan_key = None
self._write_plan: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
self._translate_cache: Dict[Tuple[int, int], torch.Tensor] = {}
logger.info(
"Page-interleave KV sharding enabled: shard_rank=%d shard_size=%d "
"page_size=%d scratch_rows=%d x2",
self.shard_rank,
self.shard_size,
spec.page_size,
scratch_rows,
)
def set_kv_buffer_prefix_valid(self, *args, **kwargs):
raise NotImplementedError(
"prefix-valid commit is unsupported under logical-page KV sharding "
"(it writes pool rows directly, bypassing the ownership filter)"
)
# ---- subclass hooks -------------------------------------------------------
def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
raise NotImplementedError
def _gather_pairs(self, local_layer: int) -> List[Tuple[torch.Tensor, str]]:
"""(pool buffer of ``local_layer``, scratch tensor name) pairs."""
raise NotImplementedError
# ---- per-batch plan -------------------------------------------------------
def begin_shard_extend(
self,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
prefix_lens_cpu: List[int],
seq_lens_cpu: List[int],
) -> None:
"""Capture the batch's gather plan and kick the first layer's gather.
Called once per scheduled extend batch, at attention-metadata build
time (after ``req_to_token`` holds the batch's allocation). All inputs
are mirrored across ranks, so every rank derives the identical plan.
"""
spec = self.shard_spec
ps, shard_size = spec.page_size, spec.shard_size
# Per-request page samples. One logical page per position-page
# (whole-page draws + ps-aligned chunk starts), so stride-ps sampling
# enumerates each request's pages in position order — no run
# collapsing. block_pages (the padded allgather block) is the
# sync-free host bound sum_i ceil(K_i / N): each request's prefix is
# ONE cyclic rotation run (graft-declined inserts + per-request
# bases), so a rank owns at most ceil(K_i / N) of its pages, and the
# cross-request dedup below only shrinks per-rank counts.
empty = torch.empty((0,), dtype=torch.int64, device=self.device)
prefix_parts = []
chunk_parts = []
block_pages = 0
for i in range(len(prefix_lens_cpu)):
prefix_len, seq_len = int(prefix_lens_cpu[i]), int(seq_lens_cpu[i])
assert prefix_len % ps == 0, (
f"sharded prefill requires physical-page-aligned prefixes "
f"(the radix-tree match quantum), got prefix_len="
f"{prefix_len}, page_size={ps}"
)
row = req_to_token[req_pool_indices[i]]
if prefix_len:
pages_i = row[:prefix_len:ps].long() // ps
prefix_parts.append(pages_i)
block_pages += ceil_div(prefix_len // ps, shard_size)
if self._debug_plan_checks:
# Owner-congruence: each request's prefix owners must be
# exactly cyclic — the block bound above relies on it. A
# broken rotation would overflow a rank's block and the
# translation would silently address the wrong scratch
# rows. Fail loud instead (debug mode syncs).
owners = pages_i % shard_size
expected = (
int(owners[0])
+ torch.arange(
owners.numel(), dtype=torch.int64, device=self.device
)
) % shard_size
assert torch.equal(owners, expected), (
f"sharded prefix owners of request {i} are not "
f"cyclic — rotation base out of sync "
f"(owners[:16]={owners[:16].tolist()})"
)
if seq_len > prefix_len:
chunk_parts.append(row[prefix_len:seq_len:ps].long() // ps)
# A prefix page shared by several requests gathers into ONE scratch
# slot (torch.unique sorts — deterministic, mirrored); chunk pages
# are per-request fresh allocations, disjoint by construction.
prefix_pages = torch.unique(torch.cat(prefix_parts)) if prefix_parts else empty
chunk_pages = torch.cat(chunk_parts) if chunk_parts else empty
n_prefix = prefix_pages.numel()
n_chunk = chunk_pages.numel()
n_prefix_slots = shard_size * block_pages
assert n_prefix_slots * ps <= spec.max_prefix_tokens, (
f"prefix ({n_prefix} pages, padded gather span "
f"{n_prefix_slots * ps} tokens) exceeds the scratch prefix "
f"capacity ({spec.max_prefix_tokens}) — the PrefillAdder scratch "
f"reservation should have deferred this batch"
)
assert n_chunk * ps <= spec.chunk_tokens, (
f"chunk ({n_chunk * ps} tokens) exceeds the scratch chunk "
f"capacity ({spec.chunk_tokens})"
)
self._page_pos.fill_(self._trash_base // ps)
if n_prefix:
# Owner-major slot assignment: sort the batch's unique pages by
# (owner, local page) so rank r's pages are contiguous at
# r * block_pages in local-page order — the same order rank r
# packs its send block below, on every rank (mirrored).
owners = prefix_pages % shard_size
local_pages = prefix_pages // shard_size
order = torch.argsort(owners * self._local_page_stride + local_pages)
sorted_pages = prefix_pages[order]
sorted_owners = owners[order]
counts = torch.bincount(sorted_owners, minlength=shard_size)
starts = torch.cumsum(counts, 0) - counts
within = (
torch.arange(n_prefix, dtype=torch.int64, device=self.device)
- starts[sorted_owners]
)
self._page_pos[sorted_pages] = (sorted_owners * block_pages + within).to(
torch.int32
)
if n_chunk:
self._page_pos[chunk_pages] = torch.arange(
self._chunk_base // ps,
self._chunk_base // ps + n_chunk,
dtype=torch.int32,
device=self.device,
)
self._epoch += 1
self._block_pages = block_pages
self._shard_extend_active = True
self._translate_cache.clear()
if block_pages:
# This rank's owned pages in slot order (already local-sorted) —
# logical page l is its local physical page l // N — padded to
# the regular allgather block with the reserved trash page
# (local page 0), whose rows the plan never references.
own_local = sorted_pages[sorted_owners == self.shard_rank] // shard_size
n_pad = block_pages - own_local.numel()
if n_pad:
own_local = torch.cat(
[
own_local,
torch.zeros((n_pad,), dtype=torch.int64, device=self.device),
]
)
self._send_rows = (
own_local[:, None] * ps
+ torch.arange(ps, dtype=torch.int64, device=self.device)
).reshape(-1)
self._prefetch_layer(self.start_layer)
else:
self._send_rows = None
def end_shard_extend(self) -> None:
"""Mark no sharded extend in flight (non-extend forward modes)."""
self._shard_extend_active = False
# ---- translation ----------------------------------------------------------
def translate_loc_to_scratch(self, loc: torch.Tensor) -> torch.Tensor:
"""Logical token slots -> rows of the current batch's scratch slots.
``_page_pos`` holds every logical page's absolute scratch page slot:
prefix pages owner-major in ``[0, N * block_pages)`` (rank ``l % N``'s
pages contiguous at ``rank * block_pages``, local-page order), chunk
pages in batch-sequence order at ``_chunk_base``, and anything outside
the plan at the trash page.
"""
ps = self.shard_spec.page_size
loc64 = loc.long()
scratch_page = self._page_pos[loc64 // ps].long()
return scratch_page * ps + loc64 % ps
# ---- the layer-ahead gather -------------------------------------------------
def _prefetch_layer(self, layer_id: int) -> None:
"""Kick the allgather assembling ``layer_id``'s prefix into its slot,
on the gather stream. Idempotent per ``(layer_id, epoch)``."""
local_layer = layer_id - self.start_layer
if local_layer >= self.layer_num:
return
slot = self._slots[layer_id % 2]
key = (layer_id, self._epoch)
if slot.resident_key == key:
return
block = self._block_pages * self.shard_spec.page_size
# Order the gather after all prior compute-stream work: the
# previous tenant's reads (attention of layer_id - 2) and the pool
# writes that produced the prefix rows.
self.kv_gather_stream.wait_stream(self.device_module.current_stream())
with self.device_module.stream(self.kv_gather_stream):
for pool_buf, name in self._gather_pairs(local_layer):
scratch = slot.tensors[name]
send = scratch[self.shard_rank * block : (self.shard_rank + 1) * block]
torch.index_select(pool_buf, 0, self._send_rows, out=send)
# In-place regular allgather: send is exactly the rank's
# block of the output, every rank contributes `block` rows.
with self.kv_gather_comm.change_state(enable=True):
self.kv_gather_comm.all_gather(
scratch[: self.shard_size * block], send
)
slot.ready.record(self.kv_gather_stream)
slot.resident_key = key
def _acquire_slot_for_read(self, layer_id: int) -> _ScratchSlot:
"""Block the compute stream on the slot's ready event (a no-op when
the layer-ahead gather already landed) and kick the next layer's
gather. A residency miss is a caller bug — the batch plan captured by
``begin_shard_extend`` must have prefetched this layer."""
assert self._shard_extend_active, (
"sharded pool read outside an active sharded extend batch "
"(begin_shard_extend not called?)"
)
slot = self._slots[layer_id % 2]
if self._block_pages:
assert slot.resident_key == (layer_id, self._epoch), (
f"prefix scratch miss for layer {layer_id} "
f"(resident={slot.resident_key}, epoch={self._epoch})"
)
self.device_module.current_stream().wait_event(slot.ready)
self._prefetch_layer(layer_id + 1)
return slot
def _translate_loc_cached(self, loc: torch.Tensor) -> torch.Tensor:
"""Per-batch memoized ``translate_loc_to_scratch`` for the per-layer
callers: the same loc tensors (``out_cache_loc`` at every layer's
``set_kv_buffer``; the prefix ``kv_indices`` at every layer's
``get_mla_kv_buffer``) arrive at all layers of a forward, and the
plan is frozen per batch — so each distinct loc tensor is translated
once per batch instead of once per layer. ``begin_shard_extend`` clears
the cache when it installs a new plan."""
key = (loc.data_ptr(), loc.numel())
rows = self._translate_cache.get(key)
if rows is None:
rows = self.translate_loc_to_scratch(loc)
self._translate_cache[key] = rows
return rows
# ---- write plan (owner filter), cached per (loc tensor, epoch) --------------
def _get_write_plan(self, loc: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""(owned positions into ``loc``, their local physical rows).
The same ``out_cache_loc`` tensor is passed for every layer of a
forward, so the owner filter is computed once per batch, not per
layer.
"""
key = (loc.data_ptr(), loc.numel(), self._epoch)
if self._write_plan_key != key:
owned_idx = torch.nonzero(
self.placement.local_mask(loc, self.shard_rank)
).squeeze(1)
local_rows = self.placement.local_index(loc[owned_idx])
self._write_plan_key = key
self._write_plan = (owned_idx, local_rows)
return self._write_plan
class PageInterleaveMHATokenToKVPool(PageInterleaveKVPoolMixin, MHATokenToKVPool):
"""MHA/GQA pool striped across the attention CP group.
The write path relies on the prefill CP strategy materializing the full
chunk's K/V before ``set_kv_buffer`` receives the full (unsplit)
``out_cache_loc`` on every rank. Each rank persists only its stripe and
stages the full chunk into the current layer's scratch chunk region
(extend attention reads prefix *and* chunk through the scratch page table).
"""
def __init__(
self,
*args,
shard_spec: PageShardSpec,
shard_group: GroupCoordinator,
**kwargs,
):
super().__init__(*args, **kwargs)
assert not self.use_hnd, "KV sharding supports the NHD layout only"
assert self.kv_cache_layout == "nhd", (
f"KV sharding supports the NHD layout only, got {self.kv_cache_layout}"
)
assert not self.post_capture_active
self._init_page_shard_state(shard_spec, shard_group)
def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
return {
"k": torch.zeros(
(rows, self.head_num, self.head_dim),
dtype=self.store_dtype,
device=self.device,
),
"v": torch.zeros(
(rows, self.head_num, self.v_head_dim),
dtype=self.store_dtype,
device=self.device,
),
}
def _gather_pairs(self, local_layer: int):
return [
(self.k_buffer[local_layer], "k"),
(self.v_buffer[local_layer], "v"),
]
def set_kv_buffer(
self,
layer: RadixAttention,
loc_info,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
layer_id_override: Optional[int] = None,
dcp_kv_mask: Optional[torch.Tensor] = None,
):
assert dcp_kv_mask is None, "DCP is mutually exclusive with KV sharding"
loc, _, _ = unwrap_write_loc(loc_info)
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
if cache_k.dtype != self.dtype:
if k_scale is not None:
cache_k = cache_k / k_scale
if v_scale is not None:
cache_v = cache_v / v_scale
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
if self._shard_extend_active:
# Stage the full chunk into this layer's slot on the compute
# stream: ordered before this layer's attention read for free
# and disjoint from the gather stream, which only writes the
# prefix region. Padded locations translate to the trash page.
slot = self._slots[layer_id % 2]
rows = self._translate_loc_cached(loc)
slot.tensors["k"][rows] = cache_k
slot.tensors["v"][rows] = cache_v
owned_idx, local_rows = self._get_write_plan(loc)
if owned_idx.numel() > 0:
self._store_kv_layer(
layer_id - self.start_layer,
local_rows,
cache_k.index_select(0, owned_idx),
cache_v.index_select(0, owned_idx),
)
def get_key_buffer(self, layer_id: int):
if self._shard_extend_active:
slot = self._acquire_slot_for_read(layer_id)
k = slot.tensors["k"]
return k.view(self.dtype) if self.store_dtype != self.dtype else k
return super().get_key_buffer(layer_id)
def get_value_buffer(self, layer_id: int):
if self._shard_extend_active:
slot = self._acquire_slot_for_read(layer_id)
v = slot.tensors["v"]
return v.view(self.dtype) if self.store_dtype != self.dtype else v
return super().get_value_buffer(layer_id)
class PageInterleaveMLATokenToKVPool(PageInterleaveKVPoolMixin, MLATokenToKVPool):
"""MLA latent pool striped across its shard group (the attn-CP group when
prefill CP is active, the attn-TP group otherwise — see
``page_interleave.get_kv_shard_group``).
The latent KV reaching the write path is identical on every shard-group
rank (``ReplicatedLinear`` projection across attn-TP; the prefill CP
strategy's latent-KV rebuild across attn-CP), so the write filter needs no
compute-time communication. Read consumers, all served from the assembled
scratch: the chunked-prefix MHA path fetches prefix rows through
``get_mla_kv_buffer``; the absorbed-MLA and one-shot paths read
``get_key_buffer``/``get_value_buffer`` through the translated page
table. The current chunk is staged into the slot at write time so the
page-table consumers cover ``[prefix | chunk]`` uniformly.
"""
def __init__(
self,
*args,
shard_spec: PageShardSpec,
shard_group: GroupCoordinator,
**kwargs,
):
super().__init__(*args, **kwargs)
assert not self.use_dsa, "DSA models are not supported by KV sharding yet"
self._init_page_shard_state(shard_spec, shard_group)
def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
return {
"kv": torch.zeros(
(rows, 1, self.kv_cache_dim),
dtype=self.store_dtype,
device=self.device,
),
}
def _gather_pairs(self, local_layer: int):
return [(self.kv_buffer[local_layer], "kv")]
def _scratch_kv(self, slot: _ScratchSlot) -> torch.Tensor:
kv = slot.tensors["kv"]
if self.store_dtype != self.dtype:
return kv.view(self.dtype)
return kv
def set_kv_buffer(
self,
layer: RadixAttention,
loc_info,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
layer_id_override: Optional[int] = None,
):
loc, _, _ = unwrap_write_loc(loc_info)
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
if self._shard_extend_active:
# Stage the full chunk into this layer's slot (compute
# stream): the absorbed / one-shot readers cover the current
# chunk through the translated page table too.
slot = self._slots[layer_id % 2]
rows = self._translate_loc_cached(loc)
staged_k = cache_k
if staged_k.dtype != self.dtype:
staged_k = staged_k.to(self.dtype)
self._scratch_kv(slot)[rows] = staged_k
owned_idx, local_rows = self._get_write_plan(loc)
if owned_idx.numel() == 0:
return
super().set_kv_buffer(
layer,
local_rows,
cache_k.index_select(0, owned_idx),
cache_v,
layer_id_override=layer_id_override,
)
def set_mla_kv_buffer(
self,
layer: RadixAttention,
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
layer_id_override: Optional[int] = None,
):
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
if self._shard_extend_active:
slot = self._slots[layer_id % 2]
rows = self._translate_loc_cached(loc)
staged_nope, staged_rope = cache_k_nope, cache_k_rope
if staged_nope.dtype != self.dtype:
staged_nope = staged_nope.to(self.dtype)
staged_rope = staged_rope.to(self.dtype)
if self.store_dtype != self.dtype:
staged_nope = staged_nope.view(self.store_dtype)
staged_rope = staged_rope.view(self.store_dtype)
set_mla_kv_buffer_triton(slot.tensors["kv"], rows, staged_nope, staged_rope)
owned_idx, local_rows = self._get_write_plan(loc)
if owned_idx.numel() == 0:
return
super().set_mla_kv_buffer(
layer,
local_rows,
cache_k_nope.index_select(0, owned_idx),
cache_k_rope.index_select(0, owned_idx),
layer_id_override=layer_id_override,
)
def get_kv_buffer_shape(self):
# Shape probes (e.g. the eager runner's DCP-metadata prep) must not
# route through the attention getters below — they may run before
# this batch's metadata build while the previous batch's shard-extend
# flag is still set.
k = self.kv_buffer[0]
return k.shape, k[..., : self.kv_lora_rank].shape
def get_key_buffer(self, layer_id: int):
# During a sharded extend the pool holds only this rank's stripe;
# page-table readers (absorbed MLA, incl. the CP zigzag wrapper) get
# the assembled scratch — metadata.page_table is already translated
# to scratch rows.
if self._shard_extend_active:
return self._scratch_kv(self._acquire_slot_for_read(layer_id))
return super().get_key_buffer(layer_id)
def get_value_buffer(self, layer_id: int):
if self._shard_extend_active:
kv = self._scratch_kv(self._acquire_slot_for_read(layer_id))
return kv[..., : self.kv_lora_rank]
return super().get_value_buffer(layer_id)
def get_mla_kv_buffer(
self,
layer: RadixAttention,
loc: torch.Tensor,
dst_dtype: Optional[torch.dtype] = None,
):
slot = self._acquire_slot_for_read(layer.layer_id)
rows = self._translate_loc_cached(loc)
kv_buffer = slot.tensors["kv"]
if self.store_dtype != self.dtype:
kv_buffer = kv_buffer.view(self.dtype)
dst_dtype = dst_dtype or self.dtype
cache_k_nope = torch.empty(
(loc.shape[0], 1, self.kv_lora_rank),
dtype=dst_dtype,
device=kv_buffer.device,
)
cache_k_rope = torch.empty(
(loc.shape[0], 1, self.qk_rope_head_dim),
dtype=dst_dtype,
device=kv_buffer.device,
)
get_mla_kv_buffer_triton(kv_buffer, rows, cache_k_nope, cache_k_rope)
return cache_k_nope, cache_k_rope
test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py
+19 -2
View File
@@ -25,6 +25,7 @@ import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -32,7 +33,13 @@ from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64):
class _FakeShardPool(PageInterleaveKVPoolMixin):
pass
def _make_prefill_aware_swa_runner(
*, pool_size: int, max_context_len: int = 64, token_to_kv_pool=None
):
"""A minimal fake ModelRunner that reaches FlashAttentionBackend.__init__'s
is_prefill_aware_swa branch (mirrors how models like
python/sglang/srt/models/unlimited_ocr.py opt in)."""
@@ -63,7 +70,7 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
enable_prefill_cp=False,
enable_dp_attention=False,
)
token_to_kv_pool = object()
token_to_kv_pool = token_to_kv_pool if token_to_kv_pool is not None else object()
token_to_kv_pool_allocator = object()
return SimpleNamespace(
sliding_window_size=None,
@@ -94,6 +101,16 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestPrefillAwareSwaPrefillLensBound(CustomTestCase):
def test_sharded_pool_requests_cpu_sequence_lengths(self):
runner = _make_prefill_aware_swa_runner(
pool_size=8, token_to_kv_pool=_FakeShardPool()
)
with get_context().override_server_args():
backend = FlashAttentionBackend(runner)
self.assertTrue(backend.needs_cpu_seq_lens)
def test_buffer_covers_full_req_pool_idx_range(self):
pool_size = 8
runner = _make_prefill_aware_swa_runner(pool_size=pool_size)
test/registered/unit/layers/attention/test_kv_shard_hooks.py
+112
View File
@@ -0,0 +1,112 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import sys
from types import SimpleNamespace
import pytest
import torch
from sglang.srt.layers.attention.kv_shard_hooks import (
get_kv_shard_pool,
prepare_kv_shard_forward,
)
from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class _RecordingPool(PageInterleaveKVPoolMixin):
def __init__(self):
self.begin_args = None
self.begin_calls = 0
self.end_calls = 0
def begin_shard_extend(self, *args):
self.begin_args = args
self.begin_calls += 1
def end_shard_extend(self):
self.end_calls += 1
def _batch(mode: ForwardMode):
return SimpleNamespace(
forward_mode=mode,
req_pool_indices=torch.tensor([3], dtype=torch.int64),
extend_prefix_lens_cpu=[64],
seq_lens_cpu=torch.tensor([128], dtype=torch.int64),
)
def test_detects_only_page_interleaved_pools():
pool = _RecordingPool()
assert get_kv_shard_pool(pool) is pool
assert get_kv_shard_pool(object()) is None
@pytest.mark.parametrize(
"mode, active",
[
(ForwardMode.EXTEND, True),
(ForwardMode.MIXED, True),
(ForwardMode.SPLIT_PREFILL, True),
(ForwardMode.DECODE, False),
(ForwardMode.IDLE, False),
(ForwardMode.TARGET_VERIFY, False),
(ForwardMode.DRAFT_EXTEND_V2, False),
],
)
def test_prepare_updates_the_pool_lifecycle(mode, active):
pool = _RecordingPool()
req_to_token = torch.arange(8)
batch = _batch(mode)
assert prepare_kv_shard_forward(pool, req_to_token, batch) is active
assert pool.begin_calls == int(active)
assert pool.end_calls == int(not active)
if active:
assert all(
actual is expected
for actual, expected in zip(
pool.begin_args,
(
req_to_token,
batch.req_pool_indices,
batch.extend_prefix_lens_cpu,
batch.seq_lens_cpu,
),
)
)
else:
assert pool.begin_args is None
@pytest.mark.parametrize(
"missing_field",
["req_pool_indices", "extend_prefix_lens_cpu", "seq_lens_cpu"],
)
def test_extend_requires_host_metadata(missing_field):
batch = _batch(ForwardMode.EXTEND)
setattr(batch, missing_field, None)
with pytest.raises(RuntimeError, match="requires request indices and CPU"):
prepare_kv_shard_forward(_RecordingPool(), torch.empty(0), batch)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/mem_cache/test_page_interleave_shard.py
+734 -8
View File
@@ -11,9 +11,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Unit tests for logical-page KV cache sharding (CPU only).
"""Tests for logical-page KV cache sharding.
Pins the pure arithmetic that rotated owner-classed allocation hangs on:
Two sections. The first (CPU only, what the CPU CI job runs) pins the pure
arithmetic that rotated owner-classed allocation hangs on:
1. The placement bijection ``loc = Q*(N*ps) + r*ps + o`` — owner / local-row
round-trip, disjoint equal partition across ranks.
@@ -24,15 +25,33 @@ Pins the pure arithmetic that rotated owner-classed allocation hangs on:
3. The host rotation base on ``UnifiedTreeNode`` — stamped at insert, copied
on split, read through ``last_node``, and the pre-flight that declines an
insert whose pages carry a different base than the chain it would join.
4. ``translate_loc_to_scratch`` — the per-batch page->scratch-page lookup mapping
any consumer index vector onto the owner-major ``[prefix | chunk | trash]``
scratch, checked against a brute-force reference.
5. ``begin_shard_extend`` plan capture (page positions, padded send rows,
owner-congruence guard) with the gather stubbed out, following the
SimpleNamespace binding pattern of ``test_dsa_layer_shard_utils.py``.
The second section (``TestPageInterleaveGatherMultiGpu``, at the bottom) drives
real pools over a real 2-rank process group. It is the only check that the plan
the CPU stub validates actually addresses the bytes NCCL delivers, so it is
skipped rather than dropped when fewer than 2 CUDA devices are visible — which
is every run of the CPU suite this file is registered to.
"""
import os
import unittest
import unittest.mock
from array import array
from types import SimpleNamespace
import torch
import torch.multiprocessing as mp
from sglang.srt.distributed import (
init_distributed_environment,
initialize_model_parallel,
)
from sglang.srt.mem_cache.allocator.page_interleave import (
PageInterleavePoolAllocator,
page_interleave_shard_size,
@@ -40,19 +59,30 @@ from sglang.srt.mem_cache.allocator.page_interleave import (
from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictResult,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import _evict_until_allocatable
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
from sglang.srt.mem_cache.page_interleave import (
PageInterleavePlacement,
PageShardSpec,
get_kv_shard_group,
)
from sglang.srt.mem_cache.page_interleave_pool import (
PageInterleaveKVPoolMixin,
PageInterleaveMHATokenToKVPool,
PageInterleaveMLATokenToKVPool,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
from sglang.srt.runtime_context import get_parallel, publish
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import ceil_div
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -391,8 +421,6 @@ class TestEvictUntilAllocatable(CustomTestCase):
return alloc, outs
def _tree_stub(self, alloc, frees):
from sglang.srt.mem_cache.base_prefix_cache import EvictResult
stub = SimpleNamespace(calls=0)
def evict(params):
@@ -407,8 +435,6 @@ class TestEvictUntilAllocatable(CustomTestCase):
return stub
def test_iterates_until_min_class_covers(self):
from sglang.srt.mem_cache.common import _evict_until_allocatable
alloc, outs = self._allocator_with_tight_class()
# Each round frees ONE class-3 page (a whole 1-page chain): reaching
# a min-class floor of 2 pages takes 2 rounds.
@@ -418,8 +444,6 @@ class TestEvictUntilAllocatable(CustomTestCase):
self.assertEqual(tree.calls, 2)
def test_terminates_when_tree_dry(self):
from sglang.srt.mem_cache.common import _evict_until_allocatable
alloc, _ = self._allocator_with_tight_class()
tree = self._tree_stub(alloc, []) # nothing evictable
_evict_until_allocatable(tree, alloc, PS)
@@ -767,5 +791,707 @@ class TestRotationGraftDecline(CustomTestCase):
self.assertEqual(set(released.tolist()), set(own_locs[:8].tolist()))
def _chain_pages(base, n_pages, local_start=5):
"""Logical page ids of one chain: page P has owner (base + P) % N and an
arbitrary (here: increasing) local page on its owner."""
counter = {r: local_start for r in range(N)}
pages = []
for p in range(n_pages):
r = (base + p) % N
pages.append(counter[r] * N + r)
counter[r] += 1
return pages
def _chain_row(pages, seq_len):
row = torch.empty(seq_len, dtype=torch.int32)
for i in range(seq_len):
row[i] = pages[i // PS] * PS + i % PS
return row
def _make_pool_stub(spec, shard_rank=0, debug=True, table_pages=4096):
"""A SimpleNamespace carrying exactly the state begin_shard_extend /
translate_loc_to_scratch read."""
stub = SimpleNamespace()
stub.shard_spec = spec
stub.shard_rank = shard_rank
stub.device = "cpu"
stub.start_layer = 0
stub._chunk_base = spec.max_prefix_tokens
stub._trash_base = spec.max_prefix_tokens + spec.chunk_tokens
stub._page_pos = torch.full(
(table_pages,), stub._trash_base // PS, dtype=torch.int32
)
stub._local_page_stride = table_pages
stub._epoch = 0
stub._write_plan_key = stub._write_plan = None
stub._translate_cache = {}
stub._debug_plan_checks = debug
stub.translate_loc_to_scratch = lambda loc: (
PageInterleaveKVPoolMixin.translate_loc_to_scratch(stub, loc)
)
stub.prefetched = []
stub._prefetch_layer = lambda layer_id: stub.prefetched.append(layer_id)
return stub
def _run_begin(stub, prefix_lens, seq_lens, rows):
width = max(r.numel() for r in rows)
padded = [
torch.cat([r, torch.zeros(width - r.numel(), dtype=torch.int32)]) for r in rows
]
PageInterleaveKVPoolMixin.begin_shard_extend(
stub,
torch.stack(padded),
torch.arange(len(rows)),
prefix_lens,
seq_lens,
)
return stub
def _reference_prefix_slots(per_request_prefix_pages):
"""Brute-force reference of the owner-major slot assignment: the batch's
unique prefix pages sorted by (owner, local page), rank r's pages
contiguous at r * block; block = sum of per-request ceil(K_i / N)."""
block = sum(ceil_div(len(pages), N) for pages in per_request_prefix_pages)
uniq = sorted({p for pages in per_request_prefix_pages for p in pages})
slots = {}
counts = {r: 0 for r in range(N)}
for page in sorted(uniq, key=lambda p: (p % N, p // N)):
owner = page % N
slots[page] = owner * block + counts[owner]
counts[owner] += 1
return slots, block
class TestBeginShardExtendPlan(CustomTestCase):
def test_plan_with_rotated_prefix(self):
"""7 prefix pages of a base-2 chain + 9 chunk pages (last partial):
owner-major slots, send rows owner-filtered in the same order and
padded to the block bound ceil(7/4) = 2 pages."""
pages = _chain_pages(base=2, n_pages=16)
prefix_len, seq_len = 7 * PS, 16 * PS - 5
row = _chain_row(pages, seq_len)
slots, block = _reference_prefix_slots([pages[:7]])
for rank in range(N):
stub = _run_begin(
_make_pool_stub(_make_spec(), rank), [prefix_len], [seq_len], [row]
)
self.assertEqual(stub._block_pages, block)
self.assertTrue(stub._shard_extend_active)
self.assertEqual(stub._epoch, 1)
self.assertEqual(stub.prefetched, [0]) # first layer kicked
for page, slot in slots.items():
self.assertEqual(int(stub._page_pos[page]), slot)
for j, page in enumerate(pages[7:]):
self.assertEqual(int(stub._page_pos[page]), stub._chunk_base // PS + j)
own = sorted((p for p in pages[:7] if p % N == rank), key=lambda p: p // N)
expect = torch.cat(
[torch.arange((p // N) * PS, (p // N + 1) * PS) for p in own]
)
if len(own) < block: # padded with the trash page (local page 0)
expect = torch.cat([expect, torch.arange((block - len(own)) * PS)])
self.assertTrue(torch.equal(stub._send_rows, expect))
def test_multi_request_plan_shared_prefix_dedup(self):
"""bs > 1: request 0 and request 1 share a 3-page cached prefix
(request 1 extends it by 2 pages); request 2 is an unrelated base-2
chain. Shared pages must gather into ONE slot (no duplicate plan
entries), the block is the per-request ceil sum, and every request's
locs translate through the same table."""
chain_a = _chain_pages(base=0, n_pages=5)
chain_c = _chain_pages(base=2, n_pages=4, local_start=20)
# rows: request 0 = A[:3] prefix + 1 chunk page; request 1 = A[:5]
# prefix + 2 chunk pages; request 2 = C[:2] prefix + 2 chunk pages.
chunk0 = _chain_pages(base=3, n_pages=1, local_start=40)
chunk1 = _chain_pages(base=1, n_pages=2, local_start=50)
chunk2 = _chain_pages(base=0, n_pages=2, local_start=60)
rows = [
_chain_row(chain_a[:3] + chunk0, 4 * PS),
_chain_row(chain_a[:5] + chunk1, 7 * PS),
_chain_row(chain_c[:2] + chunk2, 4 * PS - 3),
]
stub = _run_begin(
_make_pool_stub(_make_spec()),
[3 * PS, 5 * PS, 2 * PS],
[4 * PS, 7 * PS, 4 * PS - 3],
rows,
)
slots, block = _reference_prefix_slots([chain_a[:3], chain_a[:5], chain_c[:2]])
self.assertEqual(block, 1 + 2 + 1)
self.assertEqual(stub._block_pages, block)
for page, slot in slots.items():
self.assertEqual(int(stub._page_pos[page]), slot)
# Chunk slots are absolute scratch pages in batch order.
for j, page in enumerate(chunk0 + chunk1 + chunk2):
self.assertEqual(int(stub._page_pos[page]), stub._chunk_base // PS + j)
# Shared pages: both requests' locs hit the SAME scratch rows.
shared_loc_r0 = rows[0][:PS].long()
shared_loc_r1 = rows[1][:PS].long()
t0 = PageInterleaveKVPoolMixin.translate_loc_to_scratch(stub, shared_loc_r0)
t1 = PageInterleaveKVPoolMixin.translate_loc_to_scratch(stub, shared_loc_r1)
self.assertTrue(torch.equal(t0, t1))
# Per-rank send lists fit the block and pad with the trash page.
all_prefix = sorted(set(chain_a[:5] + chain_c[:2]))
for rank in range(N):
stub_r = _run_begin(
_make_pool_stub(_make_spec(), rank),
[3 * PS, 5 * PS, 2 * PS],
[4 * PS, 7 * PS, 4 * PS - 3],
rows,
)
own = sorted((p for p in all_prefix if p % N == rank), key=lambda p: p // N)
self.assertLessEqual(len(own), block)
self.assertEqual(stub_r._send_rows.numel(), block * PS)
expect_head = torch.cat(
[torch.arange((p // N) * PS, (p // N + 1) * PS) for p in own]
)
self.assertTrue(
torch.equal(stub_r._send_rows[: len(own) * PS], expect_head)
)
def test_send_order_follows_local_page_not_position(self):
"""A freed-and-reused page can give a chain a LOWER local page id at
a later position. Slot assignment and send packing must both order
by local page id (they only need to agree — a mismatch reads the
wrong rank rows)."""
# Owner-0 pages appear at positions 0 and 4 with locals 9 then 3.
pages = [9 * N + 0, 5 * N + 1, 5 * N + 2, 5 * N + 3, 3 * N + 0]
row = _chain_row(pages, 5 * PS)
stub = _run_begin(
_make_pool_stub(_make_spec(), shard_rank=0),
[5 * PS],
[5 * PS + PS],
[torch.cat([row, _chain_row([7 * N + 1], PS)])],
)
slots, block = _reference_prefix_slots([pages])
self.assertEqual(block, 2)
# local 3 gets owner-0's first slot although it sits at position 4.
self.assertEqual(int(stub._page_pos[3 * N + 0]), 0)
self.assertEqual(int(stub._page_pos[9 * N + 0]), 1)
expect = torch.cat(
[torch.arange(3 * PS, 4 * PS), torch.arange(9 * PS, 10 * PS)]
)
self.assertTrue(torch.equal(stub._send_rows, expect))
def test_plan_without_prefix(self):
pages = _chain_pages(base=0, n_pages=2)
stub = _run_begin(
_make_pool_stub(_make_spec()), [0], [PS + 5], [_chain_row(pages, PS + 5)]
)
self.assertEqual(stub._block_pages, 0)
self.assertTrue(stub._shard_extend_active)
self.assertEqual(stub.prefetched, []) # nothing to gather
self.assertIsNone(stub._send_rows)
self.assertEqual(int(stub._page_pos[pages[0]]), stub._chunk_base // PS)
self.assertEqual(int(stub._page_pos[pages[1]]), stub._chunk_base // PS + 1)
def test_unaligned_prefix_rejected(self):
# The tree quantum is the PHYSICAL page: a prefix that is not a
# ps-multiple can never come out of match_prefix.
pages = _chain_pages(base=0, n_pages=4)
with self.assertRaises(AssertionError):
_run_begin(
_make_pool_stub(_make_spec()),
[PS + 3],
[4 * PS],
[_chain_row(pages, 4 * PS)],
)
def test_owner_congruence_guard(self):
"""A rotation-base bug that breaks a request's prefix-owner
cyclicity invalidates the sync-free block bound (a rank can own more
than ceil(K/N) pages); the debug guard must catch it at plan time."""
pages = _chain_pages(base=1, n_pages=8)
pages[2], pages[5] = pages[5], pages[2] # same multiset, not cyclic
with self.assertRaises(AssertionError) as ctx:
_run_begin(
_make_pool_stub(_make_spec()),
[6 * PS],
[8 * PS],
[_chain_row(pages, 8 * PS)],
)
self.assertIn("cyclic", str(ctx.exception))
class TestScratchTranslation(CustomTestCase):
def _plan(self, base=2, n_prefix=7, n_chunk=9, rank=1):
pages = _chain_pages(base=base, n_pages=n_prefix + n_chunk)
seq_len = (n_prefix + n_chunk) * PS
stub = _run_begin(
_make_pool_stub(_make_spec(), rank),
[n_prefix * PS],
[seq_len],
[_chain_row(pages, seq_len)],
)
return stub, pages[:n_prefix], pages[n_prefix:]
def _reference_row(self, stub, prefix_pages, chunk_pages, loc):
"""Brute-force reference: owner-major (owner, local-page)-sorted
prefix slots, sequence-order chunk."""
spec = stub.shard_spec
page, off = loc // PS, loc % PS
if page in prefix_pages:
slots, _ = _reference_prefix_slots([prefix_pages])
return slots[page] * PS + off
if page in chunk_pages:
k = chunk_pages.index(page)
return spec.max_prefix_tokens + k * PS + off
return stub._trash_base + off
def test_translation_matches_reference(self):
stub, prefix_pages, chunk_pages = self._plan()
locs = (
[p * PS + o for p in prefix_pages + chunk_pages for o in (0, 3, PS - 1)]
+ list(range(0, N)) # reserved pages -> trash
+ [3000, 3001] # off-plan -> trash
)
got = PageInterleaveKVPoolMixin.translate_loc_to_scratch(
stub, torch.tensor(locs, dtype=torch.int64)
)
expect = torch.tensor(
[self._reference_row(stub, prefix_pages, chunk_pages, l) for l in locs],
dtype=torch.int64,
)
self.assertTrue(torch.equal(got, expect))
def test_translation_is_injective_over_the_plan(self):
stub, prefix_pages, chunk_pages = self._plan(base=3, n_prefix=5, n_chunk=4)
locs = [p * PS + o for p in prefix_pages + chunk_pages for o in range(PS)]
rows = PageInterleaveKVPoolMixin.translate_loc_to_scratch(
stub, torch.tensor(locs, dtype=torch.int64)
)
self.assertEqual(len(torch.unique(rows)), len(locs))
# Prefix rows stay inside the (padded) gather span, chunk rows inside
# the chunk region.
n_prefix_tokens = len(prefix_pages) * PS
self.assertTrue(
bool((rows[:n_prefix_tokens] < N * stub._block_pages * PS).all())
)
self.assertTrue(
bool(
(rows[n_prefix_tokens:] >= stub.shard_spec.max_prefix_tokens).all()
and (rows[n_prefix_tokens:] < stub._trash_base).all()
)
)
def test_int32_page_table_input(self):
stub, prefix_pages, chunk_pages = self._plan(base=0, n_prefix=4, n_chunk=1)
table = torch.tensor(
[prefix_pages[0] * PS, prefix_pages[1] * PS, chunk_pages[0] * PS, 0],
dtype=torch.int32,
)
rows = PageInterleaveKVPoolMixin.translate_loc_to_scratch(stub, table)
self.assertEqual(rows.dtype, torch.int64)
# Page-aligned inputs land on page-aligned scratch rows (the FA3
# stride-divide contract).
self.assertTrue(bool((rows[:3] % PS == 0).all()))
self.assertEqual(int(rows[3]), stub._trash_base)
def test_translation_cache_cleared_with_new_plan(self):
pages = _chain_pages(base=0, n_pages=2)
stub = _make_pool_stub(_make_spec())
# The first batch treats page 0 as part of the current chunk.
_run_begin(stub, [0], [PS], [_chain_row(pages[:1], PS)])
loc = _chain_row(pages[:1], PS).long()
first = PageInterleaveKVPoolMixin._translate_loc_cached(stub, loc)
again = PageInterleaveKVPoolMixin._translate_loc_cached(stub, loc)
self.assertIs(again, first)
# The next batch reuses the same loc tensor after page 0 becomes a
# cached prefix. Installing the new plan must discard the old mapping.
_run_begin(stub, [PS], [2 * PS], [_chain_row(pages, 2 * PS)])
fresh = PageInterleaveKVPoolMixin._translate_loc_cached(stub, loc)
self.assertIsNot(fresh, first)
self.assertFalse(torch.equal(fresh, first))
class TestWritePlan(CustomTestCase):
def test_owner_filter_cached_per_loc_tensor(self):
spec = _make_spec(shard_rank=2)
stub = SimpleNamespace()
stub.placement = PageInterleavePlacement(spec)
stub.shard_rank = 2
stub._epoch = 1
stub._write_plan_key = stub._write_plan = None
loc = torch.arange(5 * GS, 7 * GS) # two whole groups
owned_idx, local_rows = PageInterleaveKVPoolMixin._get_write_plan(stub, loc)
self.assertEqual(owned_idx.numel(), 2 * PS)
# Owned rows are ps-contiguous runs at [Q*ps, (Q+1)*ps).
self.assertTrue(
torch.equal(
local_rows,
torch.cat([torch.arange(5 * PS, 6 * PS), torch.arange(6 * PS, 7 * PS)]),
)
)
# Same tensor + same epoch -> cached (identity).
again = PageInterleaveKVPoolMixin._get_write_plan(stub, loc)
self.assertIs(again[0], owned_idx)
# Epoch bump invalidates.
stub._epoch = 2
fresh = PageInterleaveKVPoolMixin._get_write_plan(stub, loc)
self.assertIsNot(fresh[0], owned_idx)
def test_partial_tail_page_may_own_nothing(self):
spec = _make_spec(shard_rank=3)
stub = SimpleNamespace()
stub.placement = PageInterleavePlacement(spec)
stub.shard_rank = 3
stub._epoch = 1
stub._write_plan_key = stub._write_plan = None
# 10 tokens: all inside owner-0's page of the group.
loc = torch.arange(8 * GS, 8 * GS + 10)
owned_idx, local_rows = PageInterleaveKVPoolMixin._get_write_plan(stub, loc)
self.assertEqual(owned_idx.numel(), 0)
self.assertEqual(local_rows.numel(), 0)
# =============================================================================
# Multi-GPU: the real NCCL layer-ahead gather (2 GPUs).
#
# Everything above is pure arithmetic on a CPU stub. This section drives real
# pools over a real process group, which is the only check that the plan the
# stub validates actually addresses the bytes the collective delivers:
#
# 1. MLA pool sharded across the attention-TP group: replicated writes are
# owner-filtered into disjoint pool stripes; a later batch's chunked-prefix
# read (get_mla_kv_buffer) assembles the full prefix from all ranks via the
# layer-ahead NCCL allgather and must return the canonical bytes.
# 2. MHA pool sharded across the attention-CP group: the post-allgather full
# chunk is staged into the scratch chunk region and owner-persisted; a later
# batch reads prefix+chunk through the translated page table (the scratch),
# and the assembled rows must match the canonical bytes.
#
# Skipped unless 2 CUDA devices are visible, so it is inert on the CPU runner
# this file is registered to. Run it explicitly with:
# CUDA_VISIBLE_DEVICES=0,1 python3 test/registered/unit/mem_cache/\
# test_page_interleave_shard.py TestPageInterleaveGatherMultiGpu
# =============================================================================
_GATHER_WORLD = 2
_GATHER_LAYER_NUM = 4
_GATHER_PAGE_SIZE = 16
_GATHER_GRANULE = _GATHER_WORLD * _GATHER_PAGE_SIZE
_GATHER_SIZE = _GATHER_PAGE_SIZE * 64 # physical token slots per rank
_GATHER_KV_LORA_RANK = 128
_GATHER_QK_ROPE = 32
_GATHER_HEAD_NUM = 2
_GATHER_HEAD_DIM = 32
_GATHER_DTYPE = torch.bfloat16
def _mla_value(loc, dim):
"""Deterministic canonical latent value for logical slot ``loc``."""
loc = loc.to(torch.float32)
return (loc.unsqueeze(-1) + torch.arange(dim, device=loc.device) * 0.001).to(
_GATHER_DTYPE
)
def _dist_init(rank, world, port, attn_cp_size):
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(world)
os.environ.setdefault("no_proxy", "127.0.0.1,localhost")
torch.cuda.set_device(rank)
init_distributed_environment(
world_size=world,
rank=rank,
local_rank=rank,
distributed_init_method=f"tcp://127.0.0.1:{port}",
backend="nccl",
)
# Publish the widths the groups below are about to be built at. The derived
# quotients (attn_tp_size, attn_dcp_size, ...) are projected from these
# leaves at publish; initialize_model_parallel no longer supplies them, and
# MLATokenToKVPool.set_mla_kv_buffer reads attn_dcp_size on the write path.
publish(
ServerArgs(model_path="dummy", tp_size=world, attn_cp_size=attn_cp_size),
role="scheduler",
)
initialize_model_parallel(
tensor_model_parallel_size=world,
attention_context_model_parallel_size=attn_cp_size,
)
def _gather_make_spec(shard_rank, max_prefix_groups=16, chunk_groups=4):
return PageShardSpec(
shard_rank=shard_rank,
shard_size=_GATHER_WORLD,
page_size=_GATHER_PAGE_SIZE,
max_prefix_tokens=max_prefix_groups * _GATHER_GRANULE,
chunk_tokens=chunk_groups * _GATHER_GRANULE,
)
def _fake_req_to_token(groups, seq_len, device):
"""req_to_token row where sequence group j is allocator group groups[j]."""
row = torch.zeros(
(1, len(groups) * _GATHER_GRANULE), dtype=torch.int32, device=device
)
for j, q in enumerate(groups):
row[0, j * _GATHER_GRANULE : (j + 1) * _GATHER_GRANULE] = torch.arange(
q * _GATHER_GRANULE,
(q + 1) * _GATHER_GRANULE,
dtype=torch.int32,
device=device,
)
return row[:, :seq_len] if seq_len < row.shape[1] else row
def _check(rank, name, got, expect, atol=0.0):
ok = torch.allclose(got.float(), expect.float(), atol=atol, rtol=0)
max_err = (got.float() - expect.float()).abs().max().item()
print(f"[rank {rank}] {name}: max_err={max_err:.6f} {'OK' if ok else 'FAIL'}")
assert ok, f"[rank {rank}] {name} mismatch (max_err={max_err})"
def _run_mla(rank, world, port):
_dist_init(rank, world, port, attn_cp_size=1)
group = get_parallel().attn_tp_group
assert group.world_size == world
# Topology-first shard-group selection: no CP here, so MLA falls back to
# the attn-TP axis, while GQA has no replicated axis (world_size 1).
assert get_kv_shard_group(use_mla_backend=True) is group
assert get_kv_shard_group(use_mla_backend=False).world_size == 1
spec = _gather_make_spec(shard_rank=group.rank_in_group)
pool = PageInterleaveMLATokenToKVPool(
_GATHER_SIZE,
page_size=_GATHER_PAGE_SIZE,
dtype=_GATHER_DTYPE,
kv_lora_rank=_GATHER_KV_LORA_RANK,
qk_rope_head_dim=_GATHER_QK_ROPE,
layer_num=_GATHER_LAYER_NUM,
device=f"cuda:{rank}",
enable_memory_saver=False,
start_layer=0,
end_layer=_GATHER_LAYER_NUM - 1,
shard_spec=spec,
shard_group=group,
)
device = pool.kv_buffer[0].device
# ---- chunk 1: replicated write, owner-filtered persist -----------------
# "Allocator" hands out fragmented groups (identical on every rank).
chunk1_groups = [5, 2, 9]
chunk1_locs = _fake_req_to_token(chunk1_groups, 3 * _GATHER_GRANULE, device)[
0
].long()
for layer_id in range(_GATHER_LAYER_NUM):
layer = SimpleNamespace(layer_id=layer_id)
vals = _mla_value(
chunk1_locs + layer_id * 1000, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE
)
pool.set_mla_kv_buffer(
layer,
chunk1_locs,
vals[:, :_GATHER_KV_LORA_RANK].unsqueeze(1),
vals[:, _GATHER_KV_LORA_RANK:].unsqueeze(1),
)
torch.cuda.synchronize()
torch.distributed.barrier()
# Pool holds only the owned stripe: group Q sits at local rows [Q*ps,(Q+1)*ps)
# on every rank, holding that rank's page of the group.
for q in chunk1_groups:
local_rows = torch.arange(
q * _GATHER_PAGE_SIZE, (q + 1) * _GATHER_PAGE_SIZE, device=device
)
owned_locs = (
q * _GATHER_GRANULE
+ group.rank_in_group * _GATHER_PAGE_SIZE
+ torch.arange(_GATHER_PAGE_SIZE, device=device)
)
got = pool.kv_buffer[0][local_rows, 0, :].view(_GATHER_DTYPE)
expect = _mla_value(owned_locs, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE)
_check(rank, f"mla owned stripe g{q}", got, expect)
# ---- chunk 2: prefix gather + staged chunk, both read styles -----------
seq_groups = chunk1_groups + [12] # one new chunk group
prefix_len = 3 * _GATHER_GRANULE
seq_len = prefix_len + _GATHER_GRANULE
req_to_token = _fake_req_to_token(seq_groups, seq_len, device)
chunk2_locs = req_to_token[0, prefix_len:seq_len].long()
pool.begin_shard_extend(req_to_token, torch.tensor([0]), [prefix_len], [seq_len])
for layer_id in range(_GATHER_LAYER_NUM):
layer = SimpleNamespace(layer_id=layer_id)
# Write the current chunk (stages it into the slot + persists the
# owned stripe), like the extend forward does before attention.
chunk_vals = _mla_value(
chunk2_locs + layer_id * 1000, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE
)
pool.set_mla_kv_buffer(
layer,
chunk2_locs,
chunk_vals[:, :_GATHER_KV_LORA_RANK].unsqueeze(1),
chunk_vals[:, _GATHER_KV_LORA_RANK:].unsqueeze(1),
)
# Chunked-prefix MHA style: fetch an arbitrary sub-range of the
# prefix through get_mla_kv_buffer.
sub = chunk1_locs[_GATHER_PAGE_SIZE // 2 : prefix_len - 3]
k_nope, k_rope = pool.get_mla_kv_buffer(layer, sub, _GATHER_DTYPE)
expect = _mla_value(
sub + layer_id * 1000, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE
)
_check(
rank,
f"mla prefix read l{layer_id}",
k_nope[:, 0, :],
expect[:, :_GATHER_KV_LORA_RANK],
)
_check(
rank,
f"mla prefix rope l{layer_id}",
k_rope[:, 0, :],
expect[:, _GATHER_KV_LORA_RANK:],
)
# Absorbed-MLA style (what MLA-under-CP uses): read [prefix | chunk]
# from get_key_buffer through the translated page table.
all_locs = req_to_token[0, :seq_len].long()
rows = pool.translate_loc_to_scratch(all_locs)
kv_scratch = pool.get_key_buffer(layer_id)
_check(
rank,
f"mla absorbed read l{layer_id}",
kv_scratch[rows, 0, :],
_mla_value(
all_locs + layer_id * 1000, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE
),
)
torch.distributed.barrier()
if rank == 0:
print("PASS: MLA page-interleave shard (attn-TP axis)")
def _run_mha(rank, world, port):
_dist_init(rank, world, port, attn_cp_size=world)
group = get_parallel().attn_cp_group
assert group.world_size == world
# Topology-first shard-group selection: with an active CP group, both
# GQA and MLA shard across CP (CP replicates KV for every attention
# type; the TP axis is only the no-CP MLA fallback).
assert get_kv_shard_group(use_mla_backend=False) is group
assert get_kv_shard_group(use_mla_backend=True) is group
spec = _gather_make_spec(shard_rank=group.rank_in_group)
pool = PageInterleaveMHATokenToKVPool(
_GATHER_SIZE,
page_size=_GATHER_PAGE_SIZE,
dtype=_GATHER_DTYPE,
head_num=_GATHER_HEAD_NUM,
head_dim=_GATHER_HEAD_DIM,
layer_num=_GATHER_LAYER_NUM,
device=f"cuda:{rank}",
enable_memory_saver=False,
start_layer=0,
end_layer=_GATHER_LAYER_NUM - 1,
enable_alt_stream=False,
shard_spec=spec,
shard_group=group,
)
device = pool.k_buffer[0].device
def kv_value(locs, layer_id, is_v):
base = locs.to(torch.float32) + layer_id * 1000 + (500000 if is_v else 0)
return (
base.view(-1, 1, 1)
+ torch.arange(_GATHER_HEAD_NUM, device=device).view(1, -1, 1) * 0.01
+ torch.arange(_GATHER_HEAD_DIM, device=device).view(1, 1, -1) * 0.0001
).to(_GATHER_DTYPE)
# ---- chunk 1 (prefix-less batch): stage + owner-persist ----------------
chunk1_groups = [7, 3]
chunk1_locs = _fake_req_to_token(chunk1_groups, 2 * _GATHER_GRANULE, device)[
0
].long()
req_to_token = _fake_req_to_token(chunk1_groups, 2 * _GATHER_GRANULE, device)
pool.begin_shard_extend(req_to_token, torch.tensor([0]), [0], [2 * _GATHER_GRANULE])
for layer_id in range(_GATHER_LAYER_NUM):
layer = SimpleNamespace(layer_id=layer_id)
pool.set_kv_buffer(
layer,
chunk1_locs,
kv_value(chunk1_locs, layer_id, False),
kv_value(chunk1_locs, layer_id, True),
)
# The current chunk must be readable through the scratch right away.
k_scratch = pool.get_key_buffer(layer_id)
rows = pool.translate_loc_to_scratch(chunk1_locs)
_check(
rank,
f"mha chunk stage l{layer_id}",
k_scratch[rows],
kv_value(chunk1_locs, layer_id, False),
)
torch.cuda.synchronize()
torch.distributed.barrier()
# ---- chunk 2: prefix gathered from peers via translated page table -----
seq_groups = chunk1_groups + [11]
prefix_len = 2 * _GATHER_GRANULE
seq_len = prefix_len + _GATHER_GRANULE
req_to_token = _fake_req_to_token(seq_groups, seq_len, device)
chunk2_locs = req_to_token[0, prefix_len:seq_len].long()
pool.begin_shard_extend(req_to_token, torch.tensor([0]), [prefix_len], [seq_len])
for layer_id in range(_GATHER_LAYER_NUM):
layer = SimpleNamespace(layer_id=layer_id)
pool.set_kv_buffer(
layer,
chunk2_locs,
kv_value(chunk2_locs, layer_id, False),
kv_value(chunk2_locs, layer_id, True),
)
all_locs = req_to_token[0, :seq_len].long()
rows = pool.translate_loc_to_scratch(all_locs)
k_scratch = pool.get_key_buffer(layer_id)
v_scratch = pool.get_value_buffer(layer_id)
_check(
rank,
f"mha seq read k l{layer_id}",
k_scratch[rows],
kv_value(all_locs, layer_id, False),
)
_check(
rank,
f"mha seq read v l{layer_id}",
v_scratch[rows],
kv_value(all_locs, layer_id, True),
)
torch.distributed.barrier()
if rank == 0:
print("PASS: MHA page-interleave shard (attn-CP axis)")
@unittest.skipIf(
torch.cuda.device_count() < 2, "page-interleave gather needs 2 CUDA devices"
)
class TestPageInterleaveGatherMultiGpu(CustomTestCase):
"""Real pools, real NCCL, 2 ranks — one mp.spawn per phase.
Separate spawns (and separate ports) because each phase builds its own
process group with a different attention-CP width.
"""
def test_mla_shard_over_attention_tp(self):
mp.spawn(_run_mla, args=(_GATHER_WORLD, 29811), nprocs=_GATHER_WORLD, join=True)
def test_mha_shard_over_attention_cp(self):
mp.spawn(_run_mha, args=(_GATHER_WORLD, 29812), nprocs=_GATHER_WORLD, join=True)
if __name__ == "__main__":
unittest.main()