diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index e36841e44..67ba07afd 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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( diff --git a/python/sglang/srt/layers/attention/kv_shard_hooks.py b/python/sglang/srt/layers/attention/kv_shard_hooks.py new file mode 100644 index 000000000..d22724a37 --- /dev/null +++ b/python/sglang/srt/layers/attention/kv_shard_hooks.py @@ -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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 4a6cc4a81..2b5a4c0e1 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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 ( diff --git a/python/sglang/srt/mem_cache/page_interleave.py b/python/sglang/srt/mem_cache/page_interleave.py index 9092b4ed0..2ef3c8fb9 100644 --- a/python/sglang/srt/mem_cache/page_interleave.py +++ b/python/sglang/srt/mem_cache/page_interleave.py @@ -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 diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py new file mode 100644 index 000000000..f4aab3973 --- /dev/null +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -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 diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py index f00bda40e..c3789e787 100644 --- a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -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) diff --git a/test/registered/unit/layers/attention/test_kv_shard_hooks.py b/test/registered/unit/layers/attention/test_kv_shard_hooks.py new file mode 100644 index 000000000..f1ba77775 --- /dev/null +++ b/test/registered/unit/layers/attention/test_kv_shard_hooks.py @@ -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"])) diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index f9b909db9..508204574 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -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()