chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
f50e2b1e00a27b2fd2ae0b86b56fd5a8ef03a832
sglang/python/sglang/srt/layers/attention/nsa/utils.py
laoyao0822 f50e2b1e00 Prevent stale CP shared-KV contracts from corrupting prefill 
CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata.

Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order.

Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads.

Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local.

Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout.

Tested: Remote container py_compile for touched runtime/test files.

Tested: Remote PYTHONPATH=python pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed).

Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-06-04 20:22:29 +08:00

2798 lines
107 KiB
Python
Raw Blame History

# temp NSA debugging environ
import logging
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
attn_cp_all_gather_into_tensor,
get_attention_cp_group,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
is_allocation_symmetric,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils.common import ceil_align, ceil_div
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
_CP_DRAFT_SHARED_KV_DEBUG_COUNTS = {}
def log_cp_draft_shared_kv_debug(
key: str,
message: str,
*args,
limit: int = 128,
) -> None:
if not envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
return
count = _CP_DRAFT_SHARED_KV_DEBUG_COUNTS.get(key, 0)
if count >= limit:
return
_CP_DRAFT_SHARED_KV_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
def log_cp_shared_kv_direct_write_fallback(
reason: str,
message: str,
*args,
) -> None:
"""Log every direct-write fallback event.
Warmup can hit the same fallback reason as a later real request, so
de-duplicating by reason hides correctness/performance issues after startup.
"""
logger.warning(
"[CP_SHARED_KV_FALLBACK][direct_write] reason=%s " + message,
reason,
*args,
)
def raise_cp_shared_kv_direct_write_error(
reason: str,
message: str,
*args,
) -> None:
"""Fail fast when CP shared-KV direct-write invariants are violated."""
formatted = message % args if args else message
error_msg = (
f"[CP_SHARED_KV_FAIL_FAST][direct_write] reason={reason} {formatted}"
)
logger.error(error_msg)
raise RuntimeError(error_msg)
def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
return original_seq_lens.clamp(max=nsa_index_topk)
def is_nsa_enable_prefill_cp():
return get_global_server_args().enable_nsa_prefill_context_parallel
def is_nsa_prefill_cp_in_seq_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "in-seq-split"
)
def is_nsa_prefill_cp_round_robin_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "round-robin-split"
)
def _is_cp_shared_kv_forward_batch(forward_batch: "ForwardBatch") -> bool:
return bool(getattr(forward_batch, "uses_cp_shared_kv", False))
def _is_cp_shared_kv_draft_extend(forward_batch: "ForwardBatch") -> bool:
"""Return whether EAGLE/NextN draft extend should keep CP-local semantics."""
if not _is_cp_shared_kv_forward_batch(forward_batch):
return False
if not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return False
forward_mode = getattr(forward_batch, "forward_mode", None)
is_draft_extend = getattr(forward_mode, "is_draft_extend", None)
if not callable(is_draft_extend):
return False
try:
return bool(is_draft_extend(include_v2=True))
except TypeError:
return bool(is_draft_extend())
def _fail_if_cp_shared_kv_round_robin(
forward_batch: "ForwardBatch",
*,
op: str,
) -> None:
if forward_batch is None or not _is_cp_shared_kv_forward_batch(forward_batch):
return
try:
round_robin_split = is_nsa_prefill_cp_round_robin_split()
except ValueError:
round_robin_split = False
if not round_robin_split:
return
error_msg = (
"[CP_SHARED_KV_FAIL_FAST][round_robin_unsupported] "
"CP shared KV only supports in-seq zigzag split. "
f"op={op} nsa_prefill_cp_mode=round-robin-split"
)
logger.error(error_msg)
raise RuntimeError(error_msg)
def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
_fail_if_cp_shared_kv_round_robin(
forward_batch,
op="can_nsa_prefill_cp_round_robin_split",
)
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_nsa_prefill_cp_round_robin_split()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
)
def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
# for torch device tensor
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
if dp_padding_mode.is_max_len():
tokens = max(global_num_tokens)
elif len(global_num_tokens) > 1:
tokens = global_num_tokens[get_attention_dp_rank()]
else:
tokens = global_num_tokens[0]
if can_nsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
attn_cp_size = get_attention_cp_size()
needs_cp_pad = attn_cp_size > 1 and can_nsa_prefill_cp_round_robin_split(
forward_batch
)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return nsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - nsa_cache_seqlens.shape[0]
if pad_len > 0:
nsa_cache_seqlens = torch.cat(
[
nsa_cache_seqlens,
nsa_cache_seqlens.new_zeros(pad_len, *nsa_cache_seqlens.shape[1:]),
]
)
return nsa_cache_seqlens
@dataclass(frozen=True)
class PageAlignedCacheExtent:
valid_tokens: int
padded_pages: int
padded_tokens: int
padding_tokens: int
def build_page_aligned_cache_extent(
*, valid_tokens: int, page_size: int
) -> PageAlignedCacheExtent:
if valid_tokens < 0:
raise ValueError(f"valid_tokens must be non-negative, got {valid_tokens}")
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
padded_pages = ceil_div(valid_tokens, page_size) if valid_tokens > 0 else 0
padded_tokens = padded_pages * page_size
return PageAlignedCacheExtent(
valid_tokens=valid_tokens,
padded_pages=padded_pages,
padded_tokens=padded_tokens,
padding_tokens=padded_tokens - valid_tokens,
)
@dataclass
class PageAlignedInSeqSplitInfo:
page_aligned: bool = False
page_size: int = 1
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
extend_valid_tokens: int = 0
extend_padded_pages: int = 0
extend_padded_tokens: int = 0
extend_padding_tokens: int = 0
@dataclass
class NSAContextParallelMetadata:
split_list: List[int] = None
split_list_tensor: torch.Tensor = None
split_prefix_tensor: torch.Tensor = None
max_rank_len: List[int] = None
zigzag_index: List[int] = None
per_rank_actual_token: List[int] = None
reverse_split_len: List[int] = None
cp_reverse_index: List[int] = None
kv_len_prev: int = -1
kv_len_next: int = -1
actual_seq_q_prev: int = -1
actual_seq_q_next: int = -1
kv_len_prev_tensor: torch.Tensor = None
kv_len_next_tensor: torch.Tensor = None
actual_seq_q_prev_tensor: torch.Tensor = None
actual_seq_q_next_tensor: torch.Tensor = None
actual_seq_q_prev_cu_tensor: torch.Tensor = None
actual_seq_q_next_cu_tensor: torch.Tensor = None
total_seq_lens: torch.Tensor = None
page_aligned: bool = False
page_size: int = 1
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
extend_valid_tokens: int = 0
extend_padded_pages: int = 0
extend_padded_tokens: int = 0
extend_padding_tokens: int = 0
batch_size: int = 1
request_token_offsets: List[int] = None
request_padded_token_offsets: List[int] = None
request_page_offsets: List[int] = None
request_extend_lens: List[int] = None
request_prefix_lens: List[int] = None
request_padded_pages: List[int] = None
request_padded_tokens: List[int] = None
request_padding_tokens: List[int] = None
request_split_lists: List[List[int]] = None
request_zigzag_indices: List[List[int]] = None
request_segment_page_starts: List[List[int]] = None
request_segment_page_ends: List[List[int]] = None
request_rank_local_tokens: List[int] = None
request_rank_local_offsets: List[int] = None
request_kv_len_prev: List[int] = None
request_kv_len_next: List[int] = None
request_actual_seq_q_prev: List[int] = None
request_actual_seq_q_next: List[int] = None
request_kv_len_prev_tensor: torch.Tensor = None
request_kv_len_next_tensor: torch.Tensor = None
request_actual_seq_q_prev_tensor: torch.Tensor = None
request_actual_seq_q_next_tensor: torch.Tensor = None
request_actual_seq_q_prev_cu_tensor: torch.Tensor = None
request_actual_seq_q_next_cu_tensor: torch.Tensor = None
request_last_token_owner: List[int] = None
request_last_token_local_offset: List[int] = None
output_collect_mode: str = None
flat_split_list: List[int] = None
flat_zigzag_index: List[int] = None
flat_segment_request_ids: List[int] = None
flat_segment_offsets: List[int] = None
compute_padding_enabled: bool = False
request_valid_split_lists: List[List[int]] = None
request_valid_segment_page_starts: List[List[int]] = None
request_valid_segment_page_ends: List[List[int]] = None
request_valid_padded_pages: List[int] = None
request_valid_padded_tokens: List[int] = None
request_valid_padding_tokens: List[int] = None
request_valid_rank_local_tokens: List[int] = None
request_valid_rank_local_offsets: List[int] = None
request_valid_actual_seq_q_prev: List[int] = None
request_valid_actual_seq_q_next: List[int] = None
request_valid_seq_q_prev: List[int] = None
request_valid_seq_q_next: List[int] = None
request_valid_query_row_spans: List[List[Tuple[int, int]]] = None
request_compute_split_lists: List[List[int]] = None
request_compute_segment_page_starts: List[List[int]] = None
request_compute_segment_page_ends: List[List[int]] = None
request_compute_padded_pages: List[int] = None
request_compute_padded_tokens: List[int] = None
request_compute_padding_tokens: List[int] = None
request_compute_padded_token_offsets: List[int] = None
request_compute_rank_local_tokens: List[int] = None
request_compute_rank_local_offsets: List[int] = None
request_compute_actual_seq_q_prev: List[int] = None
request_compute_actual_seq_q_next: List[int] = None
request_compute_seq_q_prev: List[int] = None
request_compute_seq_q_next: List[int] = None
batch_plan: object = None
@dataclass(frozen=True)
class CPSharedKVBatchPlan:
batch_size: int
page_size: int
cp_size: int
cp_rank: int
request_token_offsets: List[int]
request_padded_token_offsets: List[int]
request_page_offsets: List[int]
request_extend_lens: List[int]
request_prefix_lens: List[int]
request_padded_pages: List[int]
request_padded_tokens: List[int]
request_padding_tokens: List[int]
request_split_infos: List[PageAlignedInSeqSplitInfo]
request_split_lists: List[List[int]]
request_zigzag_indices: List[List[int]]
request_segment_page_starts: List[List[int]]
request_segment_page_ends: List[List[int]]
request_rank_local_tokens: List[int]
request_rank_local_offsets: List[int]
request_kv_len_prev: List[int]
request_kv_len_next: List[int]
request_actual_seq_q_prev: List[int]
request_actual_seq_q_next: List[int]
request_last_token_owner: List[int]
request_last_token_local_offset: List[int]
output_collect_mode: str
flat_split_list: List[int]
flat_zigzag_index: List[int]
flat_segment_request_ids: List[int]
flat_segment_offsets: List[int]
compute_padding_enabled: bool
request_valid_split_lists: List[List[int]]
request_valid_segment_page_starts: List[List[int]]
request_valid_segment_page_ends: List[List[int]]
request_valid_padded_pages: List[int]
request_valid_padded_tokens: List[int]
request_valid_padding_tokens: List[int]
request_valid_rank_local_tokens: List[int]
request_valid_rank_local_offsets: List[int]
request_valid_actual_seq_q_prev: List[int]
request_valid_actual_seq_q_next: List[int]
request_valid_seq_q_prev: List[int]
request_valid_seq_q_next: List[int]
request_valid_query_row_spans: List[List[Tuple[int, int]]]
request_compute_split_lists: List[List[int]]
request_compute_segment_page_starts: List[List[int]]
request_compute_segment_page_ends: List[List[int]]
request_compute_padded_pages: List[int]
request_compute_padded_tokens: List[int]
request_compute_padding_tokens: List[int]
request_compute_padded_token_offsets: List[int]
request_compute_rank_local_tokens: List[int]
request_compute_rank_local_offsets: List[int]
request_compute_actual_seq_q_prev: List[int]
request_compute_actual_seq_q_next: List[int]
request_compute_seq_q_prev: List[int]
request_compute_seq_q_next: List[int]
def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
if total_len < 0:
raise ValueError(f"total_len must be non-negative, got {total_len}")
cp_segment_num = cp_size * 2
base = total_len // cp_segment_num
remainder = total_len % cp_segment_num
return [base + (1 if i < remainder else 0) for i in range(cp_segment_num)]
def _prefix_offsets(lengths: List[int]) -> List[int]:
offsets: List[int] = []
cursor = 0
for length in lengths:
offsets.append(cursor)
cursor += int(length)
return offsets
def _build_full_page_unit_split(
*,
page_units: int,
extend_prefix_len: int,
page_size: int,
cp_size: int,
) -> Tuple[List[int], List[int], List[int]]:
"""Split full physical pages across in-seq CP segments.
This is the compute-side counterpart of
`build_page_aligned_in_seq_split_list`: every assigned page contributes a
full `page_size` rows because dummy compute rows fill any valid-token tail.
"""
if page_units < 0:
raise ValueError(f"page_units must be non-negative, got {page_units}")
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
if extend_prefix_len < 0 or extend_prefix_len % page_size != 0:
raise ValueError(
"extend_prefix_len must be non-negative and page-aligned, "
f"got extend_prefix_len={extend_prefix_len} page_size={page_size}"
)
cp_segment_num = cp_size * 2
base_units = page_units // cp_segment_num
remainder_units = page_units % cp_segment_num
unit_counts = [
base_units + (1 if i < remainder_units else 0)
for i in range(cp_segment_num)
]
split_list: List[int] = []
segment_page_starts: List[int] = []
segment_page_ends: List[int] = []
unit_cursor = 0
base_page = extend_prefix_len // page_size
for unit_count in unit_counts:
segment_page_starts.append(base_page + unit_cursor)
unit_cursor += unit_count
segment_page_ends.append(base_page + unit_cursor)
split_list.append(unit_count * page_size)
return split_list, segment_page_starts, segment_page_ends
def build_batch_page_aligned_in_seq_split_plan(
*,
extend_lens: List[int],
prefix_lens: List[int],
page_size: int,
cp_size: int,
cp_rank: int,
) -> CPSharedKVBatchPlan:
"""Build per-request page-aligned in-seq CP metadata for a real batch.
The contract is intentionally request-first: each request is split and
page-rounded independently, then flattened for runtime consumers. This
preserves phase1 narrow-output collection for bs>1 without treating the
batch as one long sequence.
"""
if len(extend_lens) != len(prefix_lens):
raise ValueError(
"extend_lens and prefix_lens must have the same length, "
f"got {len(extend_lens)} and {len(prefix_lens)}"
)
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
if cp_rank < 0 or cp_rank >= cp_size:
raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}")
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
request_extend_lens = [int(x) for x in extend_lens]
request_prefix_lens = [int(x) for x in prefix_lens]
for req_id, (extend_len, prefix_len) in enumerate(
zip(request_extend_lens, request_prefix_lens)
):
if extend_len <= 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_empty_extend] "
f"CP shared-KV batch planning requires positive extend lens. "
f"req_id={req_id} extend_len={extend_len}"
)
if prefix_len < 0:
raise ValueError(f"prefix_len must be non-negative, got {prefix_len}")
if prefix_len % page_size != 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_non_page_aligned_prefix] "
"CP shared-KV batch planning requires page-aligned prefixes. "
"The radix/HiCache match path should floor cache hits to the "
"previous page boundary before planning. "
f"req_id={req_id} prefix_len={prefix_len} "
f"extend_len={extend_len} page_size={page_size}"
)
cp_segment_num = cp_size * 2
request_split_infos: List[PageAlignedInSeqSplitInfo] = []
request_split_lists: List[List[int]] = []
request_zigzag_indices: List[List[int]] = []
request_segment_page_starts: List[List[int]] = []
request_segment_page_ends: List[List[int]] = []
request_padded_pages: List[int] = []
request_padded_tokens: List[int] = []
request_padding_tokens: List[int] = []
request_rank_local_tokens: List[int] = []
request_kv_len_prev: List[int] = []
request_kv_len_next: List[int] = []
request_actual_seq_q_prev: List[int] = []
request_actual_seq_q_next: List[int] = []
request_last_token_owner: List[int] = []
request_last_token_local_offset: List[int] = []
request_valid_split_lists: List[List[int]] = []
request_valid_segment_page_starts: List[List[int]] = []
request_valid_segment_page_ends: List[List[int]] = []
request_valid_padded_pages: List[int] = []
request_valid_padded_tokens: List[int] = []
request_valid_padding_tokens: List[int] = []
request_valid_rank_local_tokens: List[int] = []
request_valid_actual_seq_q_prev: List[int] = []
request_valid_actual_seq_q_next: List[int] = []
request_valid_seq_q_prev: List[int] = []
request_valid_seq_q_next: List[int] = []
request_valid_query_row_spans: List[List[Tuple[int, int]]] = []
request_compute_split_lists: List[List[int]] = []
request_compute_segment_page_starts: List[List[int]] = []
request_compute_segment_page_ends: List[List[int]] = []
request_compute_padded_pages: List[int] = []
request_compute_padded_tokens: List[int] = []
request_compute_padding_tokens: List[int] = []
request_compute_rank_local_tokens: List[int] = []
request_compute_actual_seq_q_prev: List[int] = []
request_compute_actual_seq_q_next: List[int] = []
request_compute_seq_q_prev: List[int] = []
request_compute_seq_q_next: List[int] = []
flat_split_list: List[int] = []
flat_zigzag_index: List[int] = []
flat_segment_request_ids: List[int] = []
flat_segment_offsets: List[int] = []
for req_id, (extend_len, prefix_len) in enumerate(
zip(request_extend_lens, request_prefix_lens)
):
split_list, split_info = build_page_aligned_in_seq_split_list(
total_len=extend_len,
extend_len=extend_len,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
)
if not split_info.page_aligned:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_page_split_fallback] "
"CP shared-KV batch planning must not use token-balanced "
f"fallback. req_id={req_id} prefix_len={prefix_len} "
f"extend_len={extend_len} page_size={page_size}"
)
compute_pages = max(split_info.extend_padded_pages, cp_size)
compute_split_list, compute_page_starts, compute_page_ends = (
_build_full_page_unit_split(
page_units=compute_pages,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
)
)
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
split_list=compute_split_list,
cp_size=cp_size,
actual_token_count=extend_len,
)
zigzag_index = [cp_rank, cp_segment_num - cp_rank - 1]
prefix_sum_list = list(accumulate(split_list))
mirror_idx = cp_segment_num - cp_rank - 1
rank_local_tokens = (
split_list[cp_rank] + split_list[mirror_idx]
)
compute_rank_local_tokens = (
compute_split_list[cp_rank] + compute_split_list[mirror_idx]
)
split_prefix_list = [0] + prefix_sum_list[:-1]
request_split_infos.append(split_info)
request_split_lists.append(split_list)
request_zigzag_indices.append(zigzag_index)
request_segment_page_starts.append(split_info.segment_page_starts)
request_segment_page_ends.append(split_info.segment_page_ends)
request_padded_pages.append(split_info.extend_padded_pages)
request_padded_tokens.append(split_info.extend_padded_tokens)
request_padding_tokens.append(split_info.extend_padding_tokens)
request_rank_local_tokens.append(rank_local_tokens)
request_kv_len_prev.append(prefix_sum_list[cp_rank])
request_kv_len_next.append(prefix_sum_list[mirror_idx])
request_actual_seq_q_prev.append(compute_split_list[cp_rank])
request_actual_seq_q_next.append(compute_split_list[mirror_idx])
request_last_token_owner.append(owner)
request_last_token_local_offset.append(local_offset)
request_valid_split_lists.append(split_list)
request_valid_segment_page_starts.append(split_info.segment_page_starts)
request_valid_segment_page_ends.append(split_info.segment_page_ends)
request_valid_padded_pages.append(split_info.extend_padded_pages)
request_valid_padded_tokens.append(split_info.extend_padded_tokens)
request_valid_padding_tokens.append(split_info.extend_padding_tokens)
request_valid_rank_local_tokens.append(rank_local_tokens)
request_valid_actual_seq_q_prev.append(split_list[cp_rank])
request_valid_actual_seq_q_next.append(split_list[mirror_idx])
request_valid_seq_q_prev.append(split_list[cp_rank])
request_valid_seq_q_next.append(split_list[mirror_idx])
request_valid_query_row_spans.append(
[
(0, split_list[cp_rank]),
(compute_split_list[cp_rank], split_list[mirror_idx]),
]
)
request_compute_split_lists.append(compute_split_list)
request_compute_segment_page_starts.append(compute_page_starts)
request_compute_segment_page_ends.append(compute_page_ends)
request_compute_padded_pages.append(compute_pages)
request_compute_padded_tokens.append(compute_pages * page_size)
request_compute_padding_tokens.append(compute_pages * page_size - extend_len)
request_compute_rank_local_tokens.append(compute_rank_local_tokens)
request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank])
request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx])
request_compute_seq_q_prev.append(compute_split_list[cp_rank])
request_compute_seq_q_next.append(compute_split_list[mirror_idx])
flat_split_list.extend(split_list)
segment_base = req_id * cp_segment_num
flat_zigzag_index.extend(segment_base + idx for idx in zigzag_index)
flat_segment_request_ids.extend([req_id] * cp_segment_num)
flat_segment_offsets.extend(split_prefix_list)
return CPSharedKVBatchPlan(
batch_size=len(request_extend_lens),
page_size=page_size,
cp_size=cp_size,
cp_rank=cp_rank,
request_token_offsets=_prefix_offsets(request_extend_lens),
request_padded_token_offsets=_prefix_offsets(request_padded_tokens),
request_page_offsets=_prefix_offsets(request_padded_pages),
request_extend_lens=request_extend_lens,
request_prefix_lens=request_prefix_lens,
request_padded_pages=request_padded_pages,
request_padded_tokens=request_padded_tokens,
request_padding_tokens=request_padding_tokens,
request_split_infos=request_split_infos,
request_split_lists=request_split_lists,
request_zigzag_indices=request_zigzag_indices,
request_segment_page_starts=request_segment_page_starts,
request_segment_page_ends=request_segment_page_ends,
request_rank_local_tokens=request_rank_local_tokens,
request_rank_local_offsets=_prefix_offsets(request_rank_local_tokens),
request_kv_len_prev=request_kv_len_prev,
request_kv_len_next=request_kv_len_next,
request_actual_seq_q_prev=request_actual_seq_q_prev,
request_actual_seq_q_next=request_actual_seq_q_next,
request_last_token_owner=request_last_token_owner,
request_last_token_local_offset=request_last_token_local_offset,
output_collect_mode="narrow_last_token",
flat_split_list=flat_split_list,
flat_zigzag_index=flat_zigzag_index,
flat_segment_request_ids=flat_segment_request_ids,
flat_segment_offsets=flat_segment_offsets,
compute_padding_enabled=any(
compute_tokens != valid_tokens
for compute_tokens, valid_tokens in zip(
request_compute_padded_tokens, request_valid_padded_tokens
)
),
request_valid_split_lists=request_valid_split_lists,
request_valid_segment_page_starts=request_valid_segment_page_starts,
request_valid_segment_page_ends=request_valid_segment_page_ends,
request_valid_padded_pages=request_valid_padded_pages,
request_valid_padded_tokens=request_valid_padded_tokens,
request_valid_padding_tokens=request_valid_padding_tokens,
request_valid_rank_local_tokens=request_valid_rank_local_tokens,
request_valid_rank_local_offsets=_prefix_offsets(request_valid_rank_local_tokens),
request_valid_actual_seq_q_prev=request_valid_actual_seq_q_prev,
request_valid_actual_seq_q_next=request_valid_actual_seq_q_next,
request_valid_seq_q_prev=request_valid_seq_q_prev,
request_valid_seq_q_next=request_valid_seq_q_next,
request_valid_query_row_spans=request_valid_query_row_spans,
request_compute_split_lists=request_compute_split_lists,
request_compute_segment_page_starts=request_compute_segment_page_starts,
request_compute_segment_page_ends=request_compute_segment_page_ends,
request_compute_padded_pages=request_compute_padded_pages,
request_compute_padded_tokens=request_compute_padded_tokens,
request_compute_padding_tokens=request_compute_padding_tokens,
request_compute_padded_token_offsets=_prefix_offsets(
request_compute_padded_tokens
),
request_compute_rank_local_tokens=request_compute_rank_local_tokens,
request_compute_rank_local_offsets=_prefix_offsets(
request_compute_rank_local_tokens
),
request_compute_actual_seq_q_prev=request_compute_actual_seq_q_prev,
request_compute_actual_seq_q_next=request_compute_actual_seq_q_next,
request_compute_seq_q_prev=request_compute_seq_q_prev,
request_compute_seq_q_next=request_compute_seq_q_next,
)
def get_cp_shared_kv_batch_plan(forward_batch: "ForwardBatch"):
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
if metadata is None:
return None
plan = getattr(metadata, "batch_plan", None)
if plan is not None:
return plan
if isinstance(metadata, CPSharedKVBatchPlan):
return metadata
if getattr(metadata, "batch_size", 1) > 1:
return metadata
return None
def _get_forward_batch_static_padded_tokens(
forward_batch: "ForwardBatch",
) -> Optional[int]:
static_padded_tokens = getattr(forward_batch, "extend_num_tokens", None)
if static_padded_tokens is None:
return None
try:
return int(static_padded_tokens)
except (TypeError, ValueError):
return None
def split_tensor_by_cp_batch_plan(
tensor: torch.Tensor,
plan,
*,
mode: str = "data",
split_kind: str = "compute",
static_padded_tokens: Optional[int] = None,
) -> torch.Tensor:
"""Split a flattened batch tensor by per-request in-seq CP plan.
`mode` is kept explicit for future shape-specific kernels. The current
CPU/Python planner path splits along dim0 for 1d, position, and data views.
`split_kind="compute"` materializes padded compute rows. Cache writes must
use `split_kind="valid"` so dummy compute rows never receive cache locs.
"""
if mode not in ("1d", "data", "position"):
raise ValueError(f"unsupported CP batch split mode={mode!r}")
if split_kind not in ("compute", "valid"):
raise ValueError(f"unsupported CP batch split_kind={split_kind!r}")
request_extend_lens = getattr(plan, "request_extend_lens", None)
compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False))
if split_kind == "valid":
request_split_lists = getattr(
plan, "request_valid_split_lists", None
) or getattr(plan, "request_split_lists", None)
request_target_lens = request_extend_lens
elif compute_padding_enabled:
request_split_lists = getattr(plan, "request_compute_split_lists", None)
request_target_lens = getattr(plan, "request_compute_padded_tokens", None)
else:
request_split_lists = getattr(plan, "request_split_lists", None)
request_target_lens = request_extend_lens
request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
batch_size = int(getattr(plan, "batch_size", 1) or 1)
if (
request_extend_lens is None
or request_target_lens is None
or request_split_lists is None
or request_zigzag_indices is None
or len(request_extend_lens) != batch_size
or len(request_target_lens) != batch_size
or len(request_split_lists) != batch_size
or len(request_zigzag_indices) != batch_size
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_split_metadata] "
"CP shared-KV batch split requires per-request split metadata."
)
expected_tokens = sum(int(x) for x in request_extend_lens)
input_tokens = int(tensor.shape[0])
if input_tokens != expected_tokens:
max_compute_tokens = (
sum(int(x) for x in request_target_lens)
if split_kind == "compute" and compute_padding_enabled
else expected_tokens
)
max_static_tokens = (
int(static_padded_tokens)
if static_padded_tokens is not None
else expected_tokens
)
max_allowed_tokens = max(max_compute_tokens, max_static_tokens)
if (
split_kind == "compute"
and input_tokens > expected_tokens
and input_tokens <= max_allowed_tokens
and (
compute_padding_enabled
or (
static_padded_tokens is not None
and int(static_padded_tokens) > expected_tokens
)
)
):
tensor = tensor[:expected_tokens]
else:
expected_detail = f"expected={expected_tokens}"
if split_kind == "compute" and compute_padding_enabled:
expected_detail += f" max_compute={max_compute_tokens}"
if split_kind == "compute" and static_padded_tokens is not None:
expected_detail += f" static_padded={int(static_padded_tokens)}"
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] "
f"input tokens={input_tokens} {expected_detail}"
)
local_chunks = []
request_tensors = torch.split(tensor, [int(x) for x in request_extend_lens], dim=0)
for req_id, (req_tensor, target_len, split_list, zigzag_index) in enumerate(
zip(
request_tensors,
request_target_lens,
request_split_lists,
request_zigzag_indices,
)
):
req_tensor = _pad_cp_request_tensor_for_split(
req_tensor,
target_len=int(target_len),
mode=mode,
req_id=req_id,
)
split_total = sum(int(x) for x in split_list)
if split_total != int(req_tensor.shape[0]):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch] "
"request split rows must equal padded request rows. "
f"req_id={req_id} split_total={split_total} "
f"request_rows={int(req_tensor.shape[0])} mode={mode}"
)
req_segments = list(torch.split(req_tensor, [int(x) for x in split_list], dim=0))
local_chunks.extend(req_segments[int(index)] for index in zigzag_index)
if not local_chunks:
return tensor.new_empty((0, *tensor.shape[1:]))
return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:])
def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device):
cached = getattr(forward_batch, "cp_local_valid_row_indices_for_cache_write", None)
cached_expected_rows = getattr(
forward_batch, "cp_local_valid_compute_rows_for_cache_write", None
)
if cached is not None and cached_expected_rows is not None:
if cached.device == device:
return cached, int(cached_expected_rows)
batch_size = int(getattr(plan, "batch_size", 1) or 1)
request_compute_split_lists = getattr(plan, "request_compute_split_lists", None)
request_valid_split_lists = getattr(
plan, "request_valid_split_lists", None
) or getattr(plan, "request_split_lists", None)
request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
if (
request_compute_split_lists is None
or request_valid_split_lists is None
or request_zigzag_indices is None
or len(request_compute_split_lists) != batch_size
or len(request_valid_split_lists) != batch_size
or len(request_zigzag_indices) != batch_size
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
"CP shared-KV cache writes require valid and compute split metadata "
"when compute padding is enabled."
)
chunks = []
local_cursor = 0
for req_id in range(batch_size):
compute_split = request_compute_split_lists[req_id]
valid_split = request_valid_split_lists[req_id]
if len(compute_split) != len(valid_split):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
"valid and compute split metadata disagree. "
f"req_id={req_id} compute_segments={len(compute_split)} "
f"valid_segments={len(valid_split)}"
)
for segment_id in request_zigzag_indices[req_id]:
segment_id = int(segment_id)
compute_len = int(compute_split[segment_id])
valid_len = int(valid_split[segment_id])
if valid_len > compute_len:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
"valid segment is longer than compute segment. "
f"req_id={req_id} segment_id={segment_id} "
f"valid_len={valid_len} compute_len={compute_len}"
)
if valid_len > 0:
chunks.append(
torch.arange(
local_cursor,
local_cursor + valid_len,
device=device,
dtype=torch.long,
)
)
local_cursor += compute_len
if chunks:
indices = torch.cat(chunks, dim=0)
else:
indices = torch.empty((0,), device=device, dtype=torch.long)
forward_batch.cp_local_valid_row_indices_for_cache_write = indices
forward_batch.cp_local_valid_compute_rows_for_cache_write = local_cursor
return indices, local_cursor
def select_cp_local_valid_rows_for_cache_write(
forward_batch,
local_tensor: torch.Tensor,
) -> torch.Tensor:
"""Drop compute-padding rows before writing CP shared KV into persistent cache."""
plan = get_cp_shared_kv_batch_plan(forward_batch)
if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)):
return local_tensor
indices, expected_compute_rows = _get_cp_local_valid_row_indices_cache(
forward_batch,
plan,
local_tensor.device,
)
local_rows = int(local_tensor.shape[0])
if local_rows != expected_compute_rows:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_shape] "
"CP shared-KV cache write tensor must contain local compute rows "
"before valid-row selection. "
f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} "
f"valid_rows={int(indices.numel())}"
)
if indices.numel() == local_rows:
return local_tensor
if indices.numel() == 0:
return local_tensor.new_empty((0, *local_tensor.shape[1:]))
return local_tensor.index_select(0, indices)
def _pad_cp_request_tensor_for_split(
tensor: torch.Tensor,
*,
target_len: int,
mode: str,
req_id: int,
) -> torch.Tensor:
current_len = int(tensor.shape[0])
if target_len < current_len:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_len_mismatch] "
"target split length is shorter than request valid rows. "
f"req_id={req_id} target_len={target_len} current_len={current_len}"
)
pad_len = target_len - current_len
if pad_len == 0:
return tensor
if current_len <= 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_empty_request] "
f"cannot compute-pad an empty request. req_id={req_id}"
)
if mode == "position":
if tensor.dim() != 1:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_position_padding_rank] "
"position compute padding expects a 1D position tensor. "
f"req_id={req_id} shape={tuple(tensor.shape)}"
)
start = tensor[-1] + 1
padding = torch.arange(
pad_len,
device=tensor.device,
dtype=tensor.dtype,
) + start
else:
padding = tensor.new_zeros((pad_len, *tensor.shape[1:]))
return torch.cat([tensor, padding], dim=0)
def build_flat_page_owner_plan(plan) -> List[int]:
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
)
owners: List[int] = []
for req_id, (extend_len, prefix_len) in enumerate(
zip(plan.request_extend_lens, plan.request_prefix_lens)
):
request_owners = build_in_seq_page_compute_owners(
extend_len=int(extend_len),
extend_prefix_len=int(prefix_len),
page_size=int(plan.page_size),
cp_size=int(plan.cp_size),
)
if request_owners is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_page_owner_plan_unavailable] "
f"req_id={req_id} extend_len={extend_len} prefix_len={prefix_len} "
f"page_size={plan.page_size} cp_size={plan.cp_size}"
)
owners.extend(request_owners)
return owners
def _fallback_page_aligned_split_info(
*,
page_size: int,
extend_prefix_len: int,
) -> PageAlignedInSeqSplitInfo:
return PageAlignedInSeqSplitInfo(
page_aligned=False,
page_size=page_size,
extend_prefix_len=extend_prefix_len,
segment_page_starts=[],
segment_page_ends=[],
extend_valid_tokens=0,
extend_padded_pages=0,
extend_padded_tokens=0,
extend_padding_tokens=0,
)
def build_page_aligned_in_seq_split_list(
*,
total_len: int,
extend_len: int,
extend_prefix_len: int,
page_size: int,
cp_size: int,
) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
"""Build an in-seq split list whose real-token boundaries do not cut pages.
The split remains valid-token based, but page metadata covers the physical
page units. Short chunks keep the page-aligned contract by assigning zero
valid tokens to surplus zigzag segments instead of falling back to
token-balanced splits that would poison later shared-KV/HiCache reuse.
"""
if extend_len < 0:
raise ValueError(f"extend_len must be non-negative, got {extend_len}")
if total_len < extend_len:
raise ValueError(
f"total_len must be >= extend_len, got total_len={total_len} "
f"extend_len={extend_len}"
)
fallback_split = build_token_balanced_in_seq_split_list(total_len, cp_size)
fallback_info = _fallback_page_aligned_split_info(
page_size=page_size,
extend_prefix_len=extend_prefix_len,
)
if page_size <= 1 or extend_len <= 0 or extend_prefix_len % page_size != 0:
return fallback_split, fallback_info
extent = build_page_aligned_cache_extent(
valid_tokens=extend_len,
page_size=page_size,
)
full_pages = extend_len // page_size
tail_tokens = extend_len % page_size
num_page_units = extent.padded_pages
cp_segment_num = cp_size * 2
base_units = num_page_units // cp_segment_num
remainder_units = num_page_units % cp_segment_num
unit_counts = [
base_units + (1 if i < remainder_units else 0)
for i in range(cp_segment_num)
]
split_list: List[int] = []
segment_page_starts: List[int] = []
segment_page_ends: List[int] = []
unit_cursor = 0
base_page = extend_prefix_len // page_size
for unit_count in unit_counts:
segment_page_starts.append(base_page + unit_cursor)
token_count = 0
for _ in range(unit_count):
if unit_cursor < full_pages:
token_count += page_size
else:
token_count += tail_tokens
unit_cursor += 1
segment_page_ends.append(base_page + unit_cursor)
split_list.append(token_count)
padding_tokens = total_len - extend_len
if padding_tokens > 0:
split_list[-1] += padding_tokens
return split_list, PageAlignedInSeqSplitInfo(
page_aligned=True,
page_size=page_size,
extend_prefix_len=extend_prefix_len,
segment_page_starts=segment_page_starts,
segment_page_ends=segment_page_ends,
extend_valid_tokens=extent.valid_tokens,
extend_padded_pages=extent.padded_pages,
extend_padded_tokens=extent.padded_tokens,
extend_padding_tokens=extent.padding_tokens,
)
def _build_in_seq_split_for_forward_batch(
*,
total_len: int,
cp_size: int,
forward_batch: "ForwardBatch" = None,
page_size: int = None,
) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
use_page_aligned_split = (
forward_batch is not None
and getattr(forward_batch, "uses_cp_shared_kv", False)
and page_size is not None
and getattr(forward_batch, "extend_seq_lens_cpu", None) is not None
and getattr(forward_batch, "extend_prefix_lens_cpu", None) is not None
and len(forward_batch.extend_seq_lens_cpu) == 1
and len(forward_batch.extend_prefix_lens_cpu) == 1
)
if use_page_aligned_split:
return build_page_aligned_in_seq_split_list(
total_len=total_len,
extend_len=int(forward_batch.extend_seq_lens_cpu[0]),
extend_prefix_len=int(forward_batch.extend_prefix_lens_cpu[0]),
page_size=int(page_size),
cp_size=cp_size,
)
split_list = build_token_balanced_in_seq_split_list(total_len, cp_size)
return split_list, _fallback_page_aligned_split_info(
page_size=int(page_size or 1),
extend_prefix_len=0,
)
def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
communication_split_lists = (
plan.request_compute_split_lists
if plan.compute_padding_enabled
else plan.request_split_lists
)
per_rank_actual_token = []
for rank in range(plan.cp_size):
rank_tokens = 0
mirror = plan.cp_size * 2 - rank - 1
for split_list in communication_split_lists:
rank_tokens += split_list[rank] + split_list[mirror]
per_rank_actual_token.append(rank_tokens)
max_rank_token = max(per_rank_actual_token) if per_rank_actual_token else 0
max_rank_len = [max_rank_token for _ in range(plan.cp_size)]
# Scalar fields remain populated for compatibility, but scalar-only
# consumers must not use them when batch_size > 1.
first_split = communication_split_lists[0] if communication_split_lists else []
first_valid_split = (
plan.request_valid_split_lists[0]
if plan.request_valid_split_lists
else plan.request_split_lists[0]
if plan.request_split_lists
else []
)
first_info = plan.request_split_infos[0] if plan.request_split_infos else None
first_zigzag = plan.request_zigzag_indices[0] if plan.request_zigzag_indices else []
first_valid_prefix_sum = list(accumulate(first_valid_split))
first_kv_len_prev = (
first_valid_prefix_sum[plan.cp_rank] if first_valid_prefix_sum else 0
)
first_mirror = plan.cp_size * 2 - plan.cp_rank - 1
first_kv_len_next = (
first_valid_prefix_sum[first_mirror] if first_valid_prefix_sum else 0
)
first_actual_seq_q_prev = first_split[plan.cp_rank] if first_split else 0
first_actual_seq_q_next = first_split[first_mirror] if first_split else 0
flat_communication_split_list = [
token_count
for split_list in communication_split_lists
for token_count in split_list
]
first_reverse_split_len = [
element
for i in range(plan.cp_size)
for element in (first_split[i], first_split[plan.cp_size * 2 - i - 1])
]
first_cp_reverse_index = (
list(range(0, plan.cp_size * 2, 2))
+ list(range(plan.cp_size * 2 - 1, 0, -2))
if first_split
else []
)
return NSAContextParallelMetadata(
split_list=first_split,
split_list_tensor=torch.tensor(
flat_communication_split_list, device="cuda", dtype=torch.int32
),
split_prefix_tensor=torch.tensor(
[0] + list(accumulate(flat_communication_split_list))[:-1],
device="cuda",
dtype=torch.int32,
),
max_rank_len=max_rank_len,
zigzag_index=first_zigzag,
per_rank_actual_token=per_rank_actual_token,
reverse_split_len=first_reverse_split_len,
cp_reverse_index=first_cp_reverse_index,
kv_len_prev=first_kv_len_prev,
kv_len_next=first_kv_len_next,
actual_seq_q_prev=first_actual_seq_q_prev,
actual_seq_q_next=first_actual_seq_q_next,
kv_len_prev_tensor=torch.tensor(
first_kv_len_prev, device="cuda", dtype=torch.int32
),
kv_len_next_tensor=torch.tensor(
first_kv_len_next, device="cuda", dtype=torch.int32
),
actual_seq_q_prev_tensor=torch.tensor(
first_actual_seq_q_prev, device="cuda", dtype=torch.int32
),
actual_seq_q_next_tensor=torch.tensor(
first_actual_seq_q_next, device="cuda", dtype=torch.int32
),
actual_seq_q_prev_cu_tensor=torch.tensor(
[0, first_actual_seq_q_prev], device="cuda", dtype=torch.int32
),
actual_seq_q_next_cu_tensor=torch.tensor(
[0, first_actual_seq_q_next], device="cuda", dtype=torch.int32
),
total_seq_lens=torch.tensor(max_rank_token * plan.cp_size),
page_aligned=True,
page_size=plan.page_size,
extend_prefix_len=(
first_info.extend_prefix_len if first_info is not None else 0
),
segment_page_starts=(
first_info.segment_page_starts if first_info is not None else []
),
segment_page_ends=(
first_info.segment_page_ends if first_info is not None else []
),
extend_valid_tokens=(
first_info.extend_valid_tokens if first_info is not None else 0
),
extend_padded_pages=(
first_info.extend_padded_pages if first_info is not None else 0
),
extend_padded_tokens=(
first_info.extend_padded_tokens if first_info is not None else 0
),
extend_padding_tokens=(
first_info.extend_padding_tokens if first_info is not None else 0
),
batch_size=plan.batch_size,
request_token_offsets=plan.request_token_offsets,
request_padded_token_offsets=plan.request_padded_token_offsets,
request_page_offsets=plan.request_page_offsets,
request_extend_lens=plan.request_extend_lens,
request_prefix_lens=plan.request_prefix_lens,
request_padded_pages=plan.request_padded_pages,
request_padded_tokens=plan.request_padded_tokens,
request_padding_tokens=plan.request_padding_tokens,
request_split_lists=plan.request_split_lists,
request_zigzag_indices=plan.request_zigzag_indices,
request_segment_page_starts=plan.request_segment_page_starts,
request_segment_page_ends=plan.request_segment_page_ends,
request_rank_local_tokens=plan.request_rank_local_tokens,
request_rank_local_offsets=plan.request_rank_local_offsets,
request_kv_len_prev=plan.request_kv_len_prev,
request_kv_len_next=plan.request_kv_len_next,
request_actual_seq_q_prev=plan.request_actual_seq_q_prev,
request_actual_seq_q_next=plan.request_actual_seq_q_next,
request_kv_len_prev_tensor=torch.tensor(
plan.request_kv_len_prev, device="cuda", dtype=torch.int32
),
request_kv_len_next_tensor=torch.tensor(
plan.request_kv_len_next, device="cuda", dtype=torch.int32
),
request_actual_seq_q_prev_tensor=torch.tensor(
plan.request_actual_seq_q_prev, device="cuda", dtype=torch.int32
),
request_actual_seq_q_next_tensor=torch.tensor(
plan.request_actual_seq_q_next, device="cuda", dtype=torch.int32
),
request_actual_seq_q_prev_cu_tensor=torch.tensor(
[0] + list(accumulate(plan.request_actual_seq_q_prev)),
device="cuda",
dtype=torch.int32,
),
request_actual_seq_q_next_cu_tensor=torch.tensor(
[0] + list(accumulate(plan.request_actual_seq_q_next)),
device="cuda",
dtype=torch.int32,
),
request_last_token_owner=plan.request_last_token_owner,
request_last_token_local_offset=plan.request_last_token_local_offset,
output_collect_mode=plan.output_collect_mode,
flat_split_list=plan.flat_split_list,
flat_zigzag_index=plan.flat_zigzag_index,
flat_segment_request_ids=plan.flat_segment_request_ids,
flat_segment_offsets=plan.flat_segment_offsets,
compute_padding_enabled=plan.compute_padding_enabled,
request_valid_split_lists=plan.request_valid_split_lists,
request_valid_segment_page_starts=plan.request_valid_segment_page_starts,
request_valid_segment_page_ends=plan.request_valid_segment_page_ends,
request_valid_padded_pages=plan.request_valid_padded_pages,
request_valid_padded_tokens=plan.request_valid_padded_tokens,
request_valid_padding_tokens=plan.request_valid_padding_tokens,
request_valid_rank_local_tokens=plan.request_valid_rank_local_tokens,
request_valid_rank_local_offsets=plan.request_valid_rank_local_offsets,
request_valid_actual_seq_q_prev=plan.request_valid_actual_seq_q_prev,
request_valid_actual_seq_q_next=plan.request_valid_actual_seq_q_next,
request_valid_seq_q_prev=plan.request_valid_seq_q_prev,
request_valid_seq_q_next=plan.request_valid_seq_q_next,
request_valid_query_row_spans=plan.request_valid_query_row_spans,
request_compute_split_lists=plan.request_compute_split_lists,
request_compute_segment_page_starts=plan.request_compute_segment_page_starts,
request_compute_segment_page_ends=plan.request_compute_segment_page_ends,
request_compute_padded_pages=plan.request_compute_padded_pages,
request_compute_padded_tokens=plan.request_compute_padded_tokens,
request_compute_padding_tokens=plan.request_compute_padding_tokens,
request_compute_padded_token_offsets=plan.request_compute_padded_token_offsets,
request_compute_rank_local_tokens=plan.request_compute_rank_local_tokens,
request_compute_rank_local_offsets=plan.request_compute_rank_local_offsets,
request_compute_actual_seq_q_prev=plan.request_compute_actual_seq_q_prev,
request_compute_actual_seq_q_next=plan.request_compute_actual_seq_q_next,
request_compute_seq_q_prev=plan.request_compute_seq_q_prev,
request_compute_seq_q_next=plan.request_compute_seq_q_next,
batch_plan=plan,
)
def should_use_replicated_compute_for_short_radix_hit(
forward_batch: "ForwardBatch",
cp_size: int,
) -> bool:
"""Return whether a short radix-hit suffix should avoid CP splitting.
Kept as a compatibility hook for older callers. The page-aligned cache
contract no longer uses replicated compute for short suffixes: zero-length
CP segments are preferred over breaking the physical page-owner pattern.
"""
return False
def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
forward_batch: "ForwardBatch",
cp_size: int,
) -> bool:
"""Compatibility hook for the old tiny-suffix skip gate.
Compute padding now handles suffixes with fewer physical pages than CP
lanes, so this function validates the page-aligned shared-KV contract and
no longer requests a skip.
"""
if not _is_cp_shared_kv_forward_batch(forward_batch) or cp_size <= 1:
return False
_validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size)
return False
def _validate_cp_shared_kv_cp_split_plan_inputs(
forward_batch: "ForwardBatch",
cp_size: int,
) -> bool:
if not _is_cp_shared_kv_forward_batch(forward_batch):
return False
if cp_size <= 1:
return False
extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
page_size = int(getattr(token_to_kv_pool, "page_size", 0) or 0)
if (
extend_seq_lens_cpu is None
or extend_prefix_lens_cpu is None
or page_size <= 0
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_split_missing_page_plan_inputs] "
"CP shared KV NSA in-seq split requires extend lengths, prefix "
"lengths, and token_to_kv_pool.page_size before planning. "
f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
f"page_size={page_size} cp_size={cp_size}"
)
if len(extend_seq_lens_cpu) != len(extend_prefix_lens_cpu):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_split_length_mismatch] "
"CP shared KV NSA in-seq split requires one prefix length per "
"extend length. "
f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
f"cp_size={cp_size}"
)
if len(extend_seq_lens_cpu) == 0:
return False
has_extend = False
for req_id, (extend_len_raw, prefix_len_raw) in enumerate(
zip(extend_seq_lens_cpu, extend_prefix_lens_cpu)
):
extend_len = int(extend_len_raw)
prefix_len = int(prefix_len_raw)
if extend_len < 0 or prefix_len < 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_split_negative_len] "
"CP shared KV NSA in-seq split received a negative length. "
f"req_id={req_id} prefix_len={prefix_len} "
f"extend_len={extend_len} page_size={page_size} "
f"cp_size={cp_size}"
)
if prefix_len % page_size != 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
"CP shared KV NSA in-seq split requires a page-aligned prefix. "
"The radix/HiCache match path should floor cache hits to the "
"previous page boundary before CP split planning. "
f"req_id={req_id} prefix_len={prefix_len} "
f"extend_len={extend_len} page_size={page_size} "
f"cp_size={cp_size}"
)
has_extend = has_extend or extend_len > 0
return has_extend
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
_fail_if_cp_shared_kv_round_robin(forward_batch, op="can_cp_split")
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert seq_len % cp_size == 0, (
f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
)
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
if should_use_replicated_compute_for_short_radix_hit(forward_batch, cp_size):
return False
if _is_cp_shared_kv_forward_batch(forward_batch):
cur_cp_seq_len = (
1
if _validate_cp_shared_kv_cp_split_plan_inputs(
forward_batch, cp_size
)
else 0
)
else:
cur_cp_seq_len = seq_len // (cp_size * 2)
extend_token_count = sum(forward_batch.extend_seq_lens_cpu)
if _is_cp_shared_kv_forward_batch(forward_batch):
min_extend_token_count = 1
else:
min_extend_token_count = cp_size
is_context_parallel_extend = (
forward_batch.forward_mode.is_context_parallel_extend()
or _is_cp_shared_kv_draft_extend(forward_batch)
)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_nsa
and is_context_parallel_extend
and is_nsa_enable_prefill_cp()
and extend_token_count >= min_extend_token_count
):
return True
else:
return False
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
_fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_data")
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert input_.shape[0] % cp_size == 0, (
f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(input_)
metadata = forward_batch.nsa_cp_metadata
if get_cp_shared_kv_batch_plan(forward_batch) is not None:
return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
input_list = list(
torch.split(input_, metadata.split_list, dim=0)
)
result = torch.cat(
[input_list[i] for i in metadata.zigzag_index], dim=0
).view(-1, *input_.shape[1:])
return result
def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
_fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_1d")
try:
round_robin_split = is_nsa_prefill_cp_round_robin_split()
except ValueError:
round_robin_split = False
if round_robin_split:
return nsa_cp_round_robin_split_data(input_)
metadata = forward_batch.nsa_cp_metadata
if get_cp_shared_kv_batch_plan(forward_batch) is not None:
return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
input_list = list(
torch.split(input_, metadata.split_list, dim=0)
)
return torch.cat(
[input_list[i] for i in metadata.zigzag_index], dim=0
).view(-1)
def _cp_split_and_rebuild_batch_in_seq(forward_batch, input_: torch.Tensor):
return split_tensor_by_cp_batch_plan(
input_,
get_cp_shared_kv_batch_plan(forward_batch),
mode="1d" if input_.dim() == 1 else "data",
static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch),
)
def get_cp_local_embedding_padded_token_count(forward_batch, local_num_tokens: int):
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
max_rank_len = getattr(metadata, "max_rank_len", None)
if not max_rank_len:
return None
try:
padded_token_count = int(max_rank_len[0])
except (TypeError, ValueError, IndexError):
return None
if padded_token_count < int(local_num_tokens):
return None
return padded_token_count
def pad_cp_local_input_ids_for_embedding(
forward_batch,
local_input_ids: torch.Tensor,
*,
pad_token_id: int = 0,
):
local_num_tokens = local_input_ids.shape[0]
padded_token_count = get_cp_local_embedding_padded_token_count(
forward_batch, local_num_tokens
)
if padded_token_count is None:
return None
if padded_token_count == local_num_tokens:
return local_input_ids
pad_input_ids = local_input_ids.new_full(
(padded_token_count - local_num_tokens,), pad_token_id
)
return torch.cat((local_input_ids, pad_input_ids), dim=0)
def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
"""Return this CP rank's local logical out_cache_loc for direct writes.
`None` means the batch should keep using the compatibility path. This path
is intentionally conservative: it only enables direct writes after Phase 4
page-aligned split and after the logical page ids prove they are owned by
this CP rank.
"""
cached = getattr(forward_batch, "cp_local_out_cache_loc", None)
if cached is not None:
return cached
if not getattr(forward_batch, "uses_cp_shared_kv", False):
return None
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
if metadata is None:
raise_cp_shared_kv_direct_write_error(
"missing_metadata",
"nsa_cp_metadata is missing",
)
if layout is None:
raise_cp_shared_kv_direct_write_error(
"missing_layout",
"cp_shared_kv_layout is missing",
)
if out_cache_loc is None:
raise_cp_shared_kv_direct_write_error(
"missing_out_cache_loc",
"out_cache_loc is missing",
)
if not getattr(metadata, "page_aligned", False):
raise_cp_shared_kv_direct_write_error(
"not_page_aligned",
"metadata is not page-aligned: page_size=%s extend_prefix_len=%s",
getattr(metadata, "page_size", None),
getattr(metadata, "extend_prefix_len", None),
)
try:
in_seq_split = is_nsa_prefill_cp_in_seq_split()
except ValueError:
in_seq_split = True
if not in_seq_split:
raise_cp_shared_kv_direct_write_error(
"not_in_seq_split",
"nsa_prefill_cp_mode is not in-seq-split",
)
batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
if batch_plan is not None:
split_tokens = sum(int(x) for x in getattr(batch_plan, "request_extend_lens", []))
mismatch_reason = "batch_split_out_cache_len_mismatch"
else:
split_tokens = sum(int(x) for x in metadata.split_list)
mismatch_reason = "split_out_cache_len_mismatch"
out_cache_tokens = int(out_cache_loc.numel())
if split_tokens != out_cache_tokens:
static_padded_tokens = _get_forward_batch_static_padded_tokens(forward_batch)
if (
static_padded_tokens is not None
and out_cache_tokens > split_tokens
and out_cache_tokens <= static_padded_tokens
):
# Model-runner/speculative warmup can append one global block of
# static padding rows after the valid flattened batch. These rows
# may carry dummy cache locs, but CP shared-KV direct-write is a
# valid-token operation: never split/write dummy compute rows.
#
# Do not use CP compute-padding metadata as an implicit allowance
# here. Page-tail/owner-lane compute padding is produced by
# split_tensor_by_cp_batch_plan() after valid input rows are split;
# only forward_batch.extend_num_tokens proves that out_cache_loc
# already contains global trailing static padding rows.
out_cache_loc = out_cache_loc[:split_tokens]
else:
raise_cp_shared_kv_direct_write_error(
mismatch_reason,
"split_list tokens=%s out_cache_loc tokens=%s static_padded=%s",
split_tokens,
out_cache_tokens,
static_padded_tokens,
)
if batch_plan is not None:
local_out_cache_loc = split_tensor_by_cp_batch_plan(
out_cache_loc.contiguous(),
batch_plan,
mode="1d",
split_kind="valid",
)
else:
local_out_cache_loc = cp_split_and_rebuild_1d(
forward_batch,
out_cache_loc.contiguous(),
)
if local_out_cache_loc.numel() == 0:
forward_batch.cp_local_out_cache_loc = local_out_cache_loc
return local_out_cache_loc
valid_locs = local_out_cache_loc[local_out_cache_loc > 0]
if valid_locs.numel() > 0 and not torch.all(layout.owned_by_this_rank(valid_locs)):
raise_cp_shared_kv_direct_write_error(
"local_loc_owner_mismatch",
"local out_cache_loc contains pages not owned by this rank: cp_rank=%s cp_size=%s page_size=%s",
layout.cp_rank,
layout.cp_size,
layout.page_size,
)
forward_batch.cp_local_out_cache_loc = local_out_cache_loc
return local_out_cache_loc
def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch"):
"""Return cached physical rows for this CP rank's shared-KV direct writes.
`get_cp_shared_kv_local_out_cache_loc` returns logical shared-KV locs because
radix/scheduler/PD metadata are expressed in the global logical address
space. Persistent writes into this rank's compact physical pool need the
same locs remapped through `CpSharedKVLayout`. The logical loc tensor is
batch-scoped, not layer-scoped, so cache the physical remap on ForwardBatch
and reuse it for both MLA KV and NSA index K/scale writes across layers.
"""
cached = getattr(forward_batch, "cp_local_physical_out_cache_loc", None)
if cached is not None:
return cached
local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
if local_out_cache_loc is None:
return None
if local_out_cache_loc.numel() == 0:
forward_batch.cp_local_physical_out_cache_loc = local_out_cache_loc
return local_out_cache_loc
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
if layout is None:
log_cp_shared_kv_direct_write_fallback(
"missing_layout",
"cp_shared_kv_layout is missing while building physical out_cache_loc",
)
return None
physical_out_cache_loc = layout.logical_locs_to_physical(
local_out_cache_loc
).contiguous()
log_cp_draft_shared_kv_debug(
"physical_out_loc",
"physical_out_loc cp_rank=%s cp_size=%s page_size=%s tokens=%s "
"physical_tokens=%s pool=%s",
layout.cp_rank,
layout.cp_size,
layout.page_size,
local_out_cache_loc.numel(),
physical_out_cache_loc.numel(),
getattr(forward_batch, "token_to_kv_pool", None).__class__.__name__,
)
forward_batch.cp_local_physical_out_cache_loc = physical_out_cache_loc
return physical_out_cache_loc
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
_fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_position")
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(positions)
if get_cp_shared_kv_batch_plan(forward_batch) is not None:
return split_tensor_by_cp_batch_plan(
positions,
get_cp_shared_kv_batch_plan(forward_batch),
mode="position",
static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch),
)
position_id_list = list(
torch.split(positions, forward_batch.nsa_cp_metadata.split_list, dim=-1)
)
positions = torch.cat(
[position_id_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index],
dim=-1,
)
return positions
def split_in_seq_cp_local_pair(
input_: torch.Tensor,
prev_len: int,
next_len: int,
*,
name: str = "input",
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Split a local in-seq CP tensor by its two logical segment lengths.
In-seq CP assigns each rank two logical segments: `rank` and
`2 * cp_size - rank - 1`. The legacy token-balanced split made these
segments almost equal, so splitting the local tensor in half happened to
work. Page-aligned split can intentionally make the two segment lengths
different, so consumers must split by metadata lengths instead of by half.
"""
prev_len = int(prev_len)
next_len = int(next_len)
expected = prev_len + next_len
actual = int(input_.shape[0])
if actual != expected:
raise RuntimeError(
f"{name} local in-seq CP length mismatch: actual={actual}, "
f"expected={expected}, prev_len={prev_len}, next_len={next_len}"
)
return torch.split(input_, (prev_len, next_len), dim=0)
@triton.jit
def nsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
tokens: tl.constexpr,
cp_size: tl.constexpr,
cp_rank: tl.constexpr,
):
extra_seq = 0
bs_idx = 0
for bs in range(tokens):
cur_len = tl.load(in_seqs_ptr + bs)
cur_len += extra_seq
cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
if cur_seq > 0:
tl.store(bs_idx_ptr + bs_idx, bs)
tl.store(out_seqs_ptr + bs_idx, cur_seq)
bs_idx += 1
extra_seq = cur_len - cur_seq * cp_size
def nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def nsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
grid = (1,)
nsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
if nsa_enable_prefill_cp is None:
nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
forward_mode = getattr(forward_batch, "forward_mode", None)
is_context_parallel_extend = (
forward_mode is not None
and forward_mode.is_context_parallel_extend()
) or _is_cp_shared_kv_draft_extend(forward_batch)
if (
forward_batch.nsa_cp_metadata is not None
and nsa_enable_prefill_cp
and is_context_parallel_extend
):
return True
else:
return False
def _cp_attn_tp_all_gather_padded_tensor(
input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op
):
"""
Allgather communication for context_parallel(kv_cache, index_k, hidden_states).
This implementation mainly consists of three parts:
Step 1, padding the input shape to unify the shape for allgather communication (the shape must be the same).
Step 2, allgather communication(async).
Step 3, removing the padding and reassembling the data according to the actual tokens.
"""
# step1
max_len = (total_len + attn_tp_size - 1) // attn_tp_size
pad_size = max_len - input_.shape[0]
if pad_size > 0:
input_ = torch.cat(
(
input_,
input_.new_zeros((pad_size, *input_.shape[1:])),
),
dim=0,
)
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
input_tensor_all = torch.empty(
max_len * attn_tp_size,
*input_.shape[1:],
device=input_.device,
dtype=input_.dtype,
)
# step2
get_attention_cp_group().cp_all_gather_into_tensor_async(
input_tensor_all, input_, stream_op
)
return input_tensor_all
def _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch):
outputs_list_max = list(
torch.split(input_tensor_all, forward_batch.nsa_cp_metadata.max_rank_len, dim=0)
)
return torch.cat(
[
outputs_list_max[index][:per_rank_len]
for index, per_rank_len in enumerate(
forward_batch.nsa_cp_metadata.per_rank_actual_token
)
],
dim=0,
)
def cp_attn_tp_all_gather_reorganazied_into_tensor(
input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op
):
input_tensor_all = _cp_attn_tp_all_gather_padded_tensor(
input_, total_len, attn_tp_size, forward_batch, stream_op
)
# step3
outputs = _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch)
return outputs
def _log_tai_in_seq_rerange_fallback(reason: str, message: str, *args) -> None:
logger.warning(
"[CP_SHARED_KV_FALLBACK][tai_in_seq_rerange] reason=%s " + message,
reason,
*args,
)
def _try_tai_in_seq_all_gather_rerange(
input_tensor_all: torch.Tensor,
forward_batch: "ForwardBatch",
*,
hidden_size: int,
cp_size: int,
) -> torch.Tensor | None:
metadata = forward_batch.nsa_cp_metadata
split_lens = getattr(metadata, "split_list_tensor", None)
split_prefix = getattr(metadata, "split_prefix_tensor", None)
split_list = getattr(metadata, "split_list", None)
max_rank_len = getattr(metadata, "max_rank_len", None)
if split_lens is None or split_prefix is None:
_log_tai_in_seq_rerange_fallback(
"missing_metadata",
"split_list_tensor or split_prefix_tensor is missing",
)
return None
if split_list is None or max_rank_len is None:
_log_tai_in_seq_rerange_fallback(
"missing_cpu_metadata",
"split_list or max_rank_len is missing",
)
return None
if not input_tensor_all.is_cuda or input_tensor_all.dim() != 2:
return None
total_tokens = int(sum(split_list))
if total_tokens == 0:
return input_tensor_all.new_empty((0, hidden_size))
max_segment_len = int(max(split_list))
max_rank_token = int(max_rank_len[0]) if len(max_rank_len) > 0 else 0
if max_segment_len <= 0 or max_rank_token <= 0:
return None
if input_tensor_all.shape[0] < max_rank_token * cp_size:
return None
try:
from tai_kernel.nsa_prefill import in_seq_all_gather_rerange
except Exception as exc:
_log_tai_in_seq_rerange_fallback(
"import_failed",
"failed to import tai_kernel.nsa_prefill.in_seq_all_gather_rerange: %s",
exc,
)
return None
try:
return in_seq_all_gather_rerange(
input_tensor_all,
split_lens,
split_prefix,
total_tokens=total_tokens,
hidden_size=hidden_size,
max_segment_len=max_segment_len,
max_rank_token=max_rank_token,
cp_size=cp_size,
)
except Exception as exc:
_log_tai_in_seq_rerange_fallback(
"kernel_failed",
"tai-kernel in-seq rerange failed: %s",
exc,
)
return None
def _torch_in_seq_all_gather_rerange(
input_tensor_all: torch.Tensor,
forward_batch: "ForwardBatch",
hidden_size: int,
) -> torch.Tensor:
output_tensor = _trim_cp_rank_padding_after_all_gather(
input_tensor_all, forward_batch
)
outputs_list = list(
torch.split(
output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0
)
)
output_tensor = torch.cat(
[outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0
)
return output_tensor.view(-1, hidden_size)
def _raise_batch_rerange_error(reason: str, message: str, *args) -> None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_full_rerange_unsupported] "
+ f"reason={reason} "
+ (message % args if args else message)
)
def _torch_batch_in_seq_all_gather_rerange(
input_tensor_all: torch.Tensor,
forward_batch: "ForwardBatch",
*,
cp_size: int,
) -> torch.Tensor:
"""Restore rank-major batched in-seq CP all-gather rows to request order.
The row payload is intentionally opaque. MLA bf16 latent rows, packed fp8
KV rows, and NSA index byte rows all share the same first-dimension token
order contract, so this helper only slices/copies rows and never interprets
the dtype or payload shape.
"""
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
batch_size = int(getattr(metadata, "batch_size", 1) or 1)
request_split_lists = getattr(metadata, "request_split_lists", None)
max_rank_len = getattr(metadata, "max_rank_len", None)
if metadata is None:
_raise_batch_rerange_error("missing_metadata", "nsa_cp_metadata is missing")
if batch_size <= 1:
_raise_batch_rerange_error(
"not_batch",
"batch-aware rerange requires batch_size > 1, got %s",
batch_size,
)
if request_split_lists is None or len(request_split_lists) != batch_size:
_raise_batch_rerange_error(
"missing_request_split_lists",
"request_split_lists is missing or incomplete. batch_size=%s value=%s",
batch_size,
request_split_lists,
)
if max_rank_len is None or len(max_rank_len) < cp_size:
_raise_batch_rerange_error(
"missing_max_rank_len",
"max_rank_len is missing or shorter than cp_size. cp_size=%s value=%s",
cp_size,
max_rank_len,
)
if cp_size <= 0:
_raise_batch_rerange_error("bad_cp_size", "cp_size must be positive: %s", cp_size)
split_lists: List[List[int]] = []
for req_id, split_list in enumerate(request_split_lists):
if split_list is None or len(split_list) != cp_size * 2:
_raise_batch_rerange_error(
"bad_request_split",
"request split must have 2 * cp_size entries. req_id=%s "
"cp_size=%s split=%s",
req_id,
cp_size,
split_list,
)
split_lists.append([int(x) for x in split_list])
max_rank_token = int(max_rank_len[0])
if max_rank_token < 0:
_raise_batch_rerange_error(
"bad_max_rank_token",
"max_rank_len[0] must be non-negative, got %s",
max_rank_token,
)
required_rows = max_rank_token * cp_size
if int(input_tensor_all.shape[0]) < required_rows:
_raise_batch_rerange_error(
"input_rows_short",
"input rows=%s required=%s max_rank_token=%s cp_size=%s",
int(input_tensor_all.shape[0]),
required_rows,
max_rank_token,
cp_size,
)
rank_request_offsets: List[List[int]] = []
for source_rank in range(cp_size):
mirror = cp_size * 2 - source_rank - 1
offsets: List[int] = []
cursor = 0
for split_list in split_lists:
offsets.append(cursor)
cursor += split_list[source_rank] + split_list[mirror]
if cursor > max_rank_token:
_raise_batch_rerange_error(
"rank_payload_exceeds_max",
"rank local rows exceed max_rank_token. rank=%s rows=%s "
"max_rank_token=%s",
source_rank,
cursor,
max_rank_token,
)
rank_request_offsets.append(offsets)
total_tokens = sum(sum(split_list) for split_list in split_lists)
output_tensor = input_tensor_all.new_empty(
(total_tokens, *input_tensor_all.shape[1:])
)
if total_tokens == 0:
return output_tensor
output_request_base = 0
for req_id, split_list in enumerate(split_lists):
segment_prefix = [0] + list(accumulate(split_list))[:-1]
for segment_id, segment_len in enumerate(split_list):
if segment_len <= 0:
continue
if segment_id < cp_size:
source_rank = segment_id
source_segment_offset = 0
else:
source_rank = cp_size * 2 - segment_id - 1
source_segment_offset = split_list[source_rank]
source_start = (
source_rank * max_rank_token
+ rank_request_offsets[source_rank][req_id]
+ source_segment_offset
)
output_start = output_request_base + segment_prefix[segment_id]
output_tensor[output_start : output_start + segment_len].copy_(
input_tensor_all[source_start : source_start + segment_len]
)
output_request_base += sum(split_list)
return output_tensor
def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
"""
# for in-seq-split
| +-----------before allgather------------+|
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
|
| +----------before rerange---------------+|
| block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
|
| +--------------result-------------------+
| block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+
# for round-robin-split
| +-----------before allgather------------+|
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
|
| +--------------result-------------------+
| token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+
"""
_fail_if_cp_shared_kv_round_robin(
forward_batch,
op="cp_all_gather_rerange_output",
)
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
batch_size = int(getattr(metadata, "batch_size", 1) or 1)
if batch_size > 1:
if input_tensor.dim() < 2:
_raise_batch_rerange_error(
"bad_input_dim",
"input tensor must have token rows plus payload dims, got shape=%s",
tuple(input_tensor.shape),
)
if getattr(metadata, "request_split_lists", None) is None:
_raise_batch_rerange_error(
"missing_request_split_lists",
"request_split_lists is required for batch-aware current rerange. "
"batch_size=%s",
batch_size,
)
total_seq_lens = getattr(metadata, "total_seq_lens", None)
if total_seq_lens is None:
_raise_batch_rerange_error(
"missing_total_seq_lens",
"total_seq_lens is required for batch-aware current rerange. "
"batch_size=%s",
batch_size,
)
input_tensor_all = _cp_attn_tp_all_gather_padded_tensor(
input_tensor,
total_seq_lens,
cp_size,
forward_batch,
stream,
)
return _torch_batch_in_seq_all_gather_rerange(
input_tensor_all,
forward_batch,
cp_size=cp_size,
)
if is_nsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric()
):
output_tensor = input_tensor.new_empty(
(input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
)
attn_cp_all_gather_into_tensor(
output_tensor,
input_tensor,
)
out_shape = output_tensor.shape
output_tensor = (
output_tensor.view(cp_size, -1, *out_shape[1:])
.transpose(0, 1)
.reshape(out_shape)
)
return output_tensor
bs_seq_len, hidden_size = input_tensor.shape
input_tensor_all = _cp_attn_tp_all_gather_padded_tensor(
input_tensor,
forward_batch.nsa_cp_metadata.total_seq_lens,
cp_size,
forward_batch,
stream,
)
output_tensor = _try_tai_in_seq_all_gather_rerange(
input_tensor_all,
forward_batch,
hidden_size=hidden_size,
cp_size=cp_size,
)
if output_tensor is not None:
return output_tensor
return _torch_in_seq_all_gather_rerange(
input_tensor_all,
forward_batch,
hidden_size,
)
def calculate_cp_seq_idx(cp_chunks_len, seqs_len):
"""Used to obtain the index of the seq corresponding
to each cp block in the forwardbatch, and the starting
and ending positions of the corresponding seq in the cp block"""
j = 0
tuple_len = [] # Only keep this result list
cumulative = {} # Used to track cumulative values for each index
for i in range(len(cp_chunks_len)):
current_dict = {}
current_tuples = []
c_val = cp_chunks_len[i]
while j < len(seqs_len):
s_val = seqs_len[j]
if s_val == c_val:
idx = j
current_dict[idx] = s_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + s_val
j += 1
break
elif s_val > c_val:
idx = j
current_dict[idx] = c_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + c_val
seqs_len[j] = s_val - c_val
break
else: # s_val < c_val
idx = j
current_dict[idx] = s_val
# Update cumulative value for this index
cumulative[idx] = cumulative.get(idx, 0) + s_val
c_val -= s_val
j += 1
# Build tuple: (index, historical cumulative, historical+current)
for idx, val in current_dict.items():
# Subtract current value to get historical cumulative
prev_cum = cumulative.get(idx, 0) - val
current_cum = prev_cum + val
current_tuples.append((idx, prev_cum, current_cum))
tuple_len.append(current_tuples)
return tuple_len
def prepare_input_dp_with_cp_dsa(
kv_len,
cp_rank,
cp_size,
seqs_len,
*,
forward_batch: "ForwardBatch" = None,
page_size: int = None,
):
_fail_if_cp_shared_kv_round_robin(
forward_batch,
op="prepare_input_dp_with_cp_dsa",
)
if is_nsa_prefill_cp_round_robin_split():
return True
"""prepare_input_dp_with_cp_dsa-zigzag index
Example (DP_ATTENT_TP == CP_SIZE == 4):
Description:
1. Start with a full-length request.
2. Split the request into multiple blocks (block0 to block7).
3. Rearrange these blocks to balance computational
load across different DP ranks.
4. Assign the rearranged blocks to different DP attention
time points (dp_atten_tp0 to dp_atten_tp3).
+---------------------------------+
| cp_split_tokens |
+---------------------------------+
| |
| request_with_full_length |
| | split (cp_size * 2) |
| +-------------------------+ |
| | block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+ |
| | rerange |
| +---------------------------------+
| | block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
| +---------------------------------+
| |
| +-------------------------+
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
| +-------------------------+
Why zigzag rearrange?
- Attention calculations must follow causal attention principles.
- Simply slicing by rank order can lead to computational load imbalance:
* First rank may focus on fewer historical key-value tokens (less computation)
* Last rank may focus on more tokens (more computation)
- To mitigate uneven load, the input hissenstate needs to be sliced by cp_size*2 and rearranged.
"""
if (
forward_batch is not None
and getattr(forward_batch, "uses_cp_shared_kv", False)
and getattr(forward_batch, "extend_seq_lens_cpu", None) is not None
and getattr(forward_batch, "extend_prefix_lens_cpu", None) is not None
):
if page_size is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_shared_missing_page_size] "
"CP shared-KV batch planning requires token_to_kv_pool.page_size"
)
batch_plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[int(x) for x in forward_batch.extend_seq_lens_cpu],
prefix_lens=[int(x) for x in forward_batch.extend_prefix_lens_cpu],
page_size=int(page_size),
cp_size=cp_size,
cp_rank=cp_rank,
)
return _build_batch_metadata_from_plan(batch_plan)
# scalar compatibility path
kv_len_int = int(kv_len)
kv_len = torch.tensor(kv_len_int)
bs_per_cp_group = 1
# get zigzag index
cp_segment_num = cp_size * 2
split_list, page_split_info = _build_in_seq_split_for_forward_batch(
total_len=kv_len_int,
cp_size=cp_size,
forward_batch=forward_batch,
page_size=page_size,
)
zigzag_index = list(
range(cp_rank, cp_rank + bs_per_cp_group * cp_segment_num, cp_segment_num)
) + list(
range(
cp_segment_num - cp_rank - 1,
bs_per_cp_group * cp_segment_num,
cp_segment_num,
)
)
per_rank_actual_token = list(
split_list[i] + split_list[cp_size * 2 - i - 1] for i in range(cp_size)
)
max_rank_token = max(per_rank_actual_token) if per_rank_actual_token else 0
max_rank_len = [max_rank_token for _ in range(cp_size)]
comm_total_seq_lens = torch.tensor(max_rank_token * cp_size)
reverse_split_len = [
element
for i in range(cp_size)
for element in (split_list[i], split_list[cp_size * 2 - i - 1])
]
# get zigzag reverse index
cp_reverse_index = []
for batch_id in range(bs_per_cp_group):
cp_reverse_index.extend(
list(range(batch_id, cp_segment_num * bs_per_cp_group, 2 * bs_per_cp_group))
+ list(
range(
(cp_segment_num - 1) * bs_per_cp_group + batch_id,
0,
-2 * bs_per_cp_group,
)
)
)
prefix_sum_list = list(accumulate(split_list))
split_prefix_list = [0] + prefix_sum_list[:-1]
# TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
# cp_seq_index = calculate_cp_seq_idx(split_list[:], seqs_len[:])
kv_len_prev = prefix_sum_list[cp_rank]
kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1]
actual_seq_q_prev = split_list[cp_rank]
actual_seq_q_next = split_list[cp_size * 2 - cp_rank - 1]
kv_len_prev_tensor = torch.tensor(kv_len_prev).to(device="cuda", dtype=torch.int32)
kv_len_next_tensor = torch.tensor(kv_len_next).to(device="cuda", dtype=torch.int32)
actual_seq_q_prev_tensor = torch.tensor(actual_seq_q_prev).to(
device="cuda", dtype=torch.int32
)
actual_seq_q_next_tensor = torch.tensor(actual_seq_q_next).to(
device="cuda", dtype=torch.int32
)
actual_seq_q_prev_cu_tensor = torch.tensor(
[0, actual_seq_q_prev], device="cuda", dtype=torch.int32
)
actual_seq_q_next_cu_tensor = torch.tensor(
[0, actual_seq_q_next], device="cuda", dtype=torch.int32
)
nsa_cp_metadata = NSAContextParallelMetadata(
split_list=split_list,
split_list_tensor=torch.tensor(split_list, device="cuda", dtype=torch.int32),
split_prefix_tensor=torch.tensor(
split_prefix_list, device="cuda", dtype=torch.int32
),
max_rank_len=max_rank_len,
zigzag_index=zigzag_index,
per_rank_actual_token=per_rank_actual_token,
reverse_split_len=reverse_split_len,
cp_reverse_index=cp_reverse_index,
kv_len_prev=kv_len_prev,
kv_len_next=kv_len_next,
actual_seq_q_prev=actual_seq_q_prev,
actual_seq_q_next=actual_seq_q_next,
kv_len_prev_tensor=kv_len_prev_tensor,
kv_len_next_tensor=kv_len_next_tensor,
actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
actual_seq_q_next_tensor=actual_seq_q_next_tensor,
actual_seq_q_prev_cu_tensor=actual_seq_q_prev_cu_tensor,
actual_seq_q_next_cu_tensor=actual_seq_q_next_cu_tensor,
total_seq_lens=comm_total_seq_lens,
page_aligned=page_split_info.page_aligned,
page_size=page_split_info.page_size,
extend_prefix_len=page_split_info.extend_prefix_len,
segment_page_starts=page_split_info.segment_page_starts,
segment_page_ends=page_split_info.segment_page_ends,
extend_valid_tokens=page_split_info.extend_valid_tokens,
extend_padded_pages=page_split_info.extend_padded_pages,
extend_padded_tokens=page_split_info.extend_padded_tokens,
extend_padding_tokens=page_split_info.extend_padding_tokens,
)
return nsa_cp_metadata
def cp_collect_last_token_hidden(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
_fail_if_cp_shared_kv_round_robin(
forward_batch,
op="cp_collect_last_token_hidden",
)
if is_nsa_prefill_cp_round_robin_split():
return _round_robin_collect_last_token(hidden_states, forward_batch, cp_size)
return _in_seq_collect_last_token(hidden_states, forward_batch, cp_size)
def _round_robin_collect_last_token(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
cp_rank = get_attention_cp_rank()
bs = len(forward_batch.extend_seq_lens_cpu)
request_offsets = _prefix_offsets([int(x) for x in forward_batch.extend_seq_lens_cpu])
owners = []
local_offsets = []
for req_offset, req_len in zip(request_offsets, forward_batch.extend_seq_lens_cpu):
req_len = int(req_len)
if req_len <= 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_empty_extend] "
f"round-robin last-token collect got non-positive req_len={req_len}"
)
global_last = int(req_offset) + req_len - 1
owners.append(global_last % cp_size)
local_offsets.append(global_last // cp_size)
local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
for req_id, (owner, local_offset) in enumerate(zip(owners, local_offsets)):
if cp_rank != owner:
continue
if local_offset >= int(hidden_states.shape[0]):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_last_token_oob] "
"round-robin last-token owner metadata points past local hidden. "
f"req_id={req_id} local_offset={local_offset} "
f"local_tokens={int(hidden_states.shape[0])}"
)
local_last[req_id] = hidden_states[local_offset]
gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
attn_cp_all_gather_into_tensor(gathered, local_last)
gather_indices = torch.tensor(
[owner * bs + req_id for req_id, owner in enumerate(owners)],
device=hidden_states.device,
dtype=torch.long,
)
return gathered.index_select(0, gather_indices)
def _get_in_seq_last_token_owner_and_offset(
split_list: List[int],
cp_size: int,
actual_token_count: int,
) -> Tuple[int, int]:
"""Return the CP rank and local offset for the real last token.
In in-seq split, each rank owns two logical segments:
rank r gets segment r followed by segment (2 * cp_size - r - 1).
`split_list` is built from the padded model input length, while
`actual_token_count` is the non-padded extend length. The real last token
can therefore sit in a middle segment instead of rank 0's trailing segment.
"""
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
cp_segment_num = cp_size * 2
if len(split_list) != cp_segment_num:
raise ValueError(
"in-seq split_list length must equal 2 * cp_size, "
f"got len={len(split_list)} cp_size={cp_size}"
)
if actual_token_count <= 0:
raise ValueError(
f"actual_token_count must be positive, got {actual_token_count}"
)
total_split_tokens = sum(split_list)
if actual_token_count > total_split_tokens:
raise ValueError(
"actual_token_count exceeds split capacity: "
f"actual={actual_token_count} split_total={total_split_tokens}"
)
token_idx = actual_token_count - 1
segment_idx = 0
offset_in_segment = token_idx
for idx, segment_len in enumerate(split_list):
if offset_in_segment < segment_len:
segment_idx = idx
break
offset_in_segment -= segment_len
else:
raise ValueError(
"failed to locate last token in split_list: "
f"actual={actual_token_count} split_list={split_list}"
)
if segment_idx < cp_size:
return segment_idx, offset_in_segment
owner = cp_segment_num - segment_idx - 1
local_offset = split_list[owner] + offset_in_segment
return owner, local_offset
def _in_seq_collect_last_token(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
cp_rank = get_attention_cp_rank()
bs = len(forward_batch.extend_seq_lens_cpu)
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
if get_cp_shared_kv_batch_plan(forward_batch) is not None or bs > 1:
return _in_seq_collect_last_token_batch(hidden_states, metadata, cp_size, cp_rank, bs)
owner = 0
local_offset = hidden_states.shape[0] - bs
if bs == 1 and metadata is not None:
actual_token_count = sum(int(x) for x in forward_batch.extend_seq_lens_cpu)
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
split_list=metadata.split_list,
cp_size=cp_size,
actual_token_count=actual_token_count,
)
if cp_rank == owner and hidden_states.shape[0] >= local_offset + bs:
local_last = hidden_states[local_offset : local_offset + bs].contiguous()
else:
local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
attn_cp_all_gather_into_tensor(gathered, local_last)
return gathered[owner * bs : owner * bs + bs]
def _in_seq_collect_last_token_batch(
hidden_states: torch.Tensor,
metadata: NSAContextParallelMetadata,
cp_size: int,
cp_rank: int,
bs: int,
) -> torch.Tensor:
if metadata is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_metadata] "
"CP in-seq bs>1 last-token collection requires batch metadata."
)
plan = getattr(metadata, "batch_plan", None)
owners = _get_cp_last_token_metadata_list(
metadata, plan, "request_last_token_owner"
)
local_offsets = _get_cp_last_token_metadata_list(
metadata, plan, "request_last_token_local_offset"
)
compute_padding_enabled = bool(
getattr(metadata, "compute_padding_enabled", False)
or bool(getattr(plan, "compute_padding_enabled", False))
)
if compute_padding_enabled:
rank_offsets = _get_cp_last_token_metadata_list(
metadata, plan, "request_compute_rank_local_offsets"
)
else:
rank_offsets = _get_cp_last_token_metadata_list(
metadata, plan, "request_rank_local_offsets"
)
if (
owners is None
or local_offsets is None
or rank_offsets is None
or len(owners) != bs
or len(local_offsets) != bs
or len(rank_offsets) != bs
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_last_token_metadata] "
"CP in-seq bs>1 narrow output requires per-request owner, "
"local offset, and rank-local offset metadata."
)
local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
for req_id, (owner, local_offset, rank_offset) in enumerate(
zip(owners, local_offsets, rank_offsets)
):
owner = int(owner)
if owner < 0 or owner >= cp_size:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_invalid_last_token_owner] "
f"req_id={req_id} owner={owner} cp_size={cp_size}"
)
if cp_rank != owner:
continue
global_local_offset = int(rank_offset) + int(local_offset)
if global_local_offset >= int(hidden_states.shape[0]):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_last_token_oob] "
"in-seq last-token metadata points past local hidden. "
f"req_id={req_id} owner={owner} rank_offset={rank_offset} "
f"local_offset={local_offset} local_tokens={int(hidden_states.shape[0])}"
)
local_last[req_id] = hidden_states[global_local_offset]
gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
attn_cp_all_gather_into_tensor(gathered, local_last)
gather_indices = torch.tensor(
[int(owner) * bs + req_id for req_id, owner in enumerate(owners)],
device=hidden_states.device,
dtype=torch.long,
)
return gathered.index_select(0, gather_indices)
def _get_cp_last_token_metadata_list(
metadata: NSAContextParallelMetadata,
plan,
field_name: str,
):
value = getattr(metadata, field_name, None)
if value is not None:
return value
if plan is not None:
return getattr(plan, field_name, None)
return None
Reference in New Issue View Git Blame Copy Permalink