Overlap CP shared KV prefix materialization for cached MLA prefill
Shared CP KV materialization remained on the critical path for cached NSA/MLA prefill batches. This change introduces a one-layer-ahead prefetcher that materializes the cached prefix for the next layer on a separate CUDA stream and consumes it when that layer reaches attention. The prefetch path keeps the existing dense page-table semantics, defers waiting until the prefetched buffer is actually consumed, and uses the TAI optimized materialize/remap helpers when enabled before falling back to the torch implementation. The implementation is intentionally gated by environment variables and keeps layer-2-only probe logging for functional confirmation without making normal profiling noisy. Constraint: Prefill CP shared KV must preserve existing page-table and dense KV semantics for NSA paged topk attention Constraint: The production performance path requires SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 and logging disabled Rejected: Wait immediately after the producer layer attention | this truncated the overlap window and hid less work Rejected: Torch-only prefetch materialize | it bypassed the optimized TAI materialize/remap path and could erase the expected win Confidence: medium Scope-risk: moderate Directive: Do not evaluate Phase8 throughput with SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=1; use it only to confirm create/start/consume_hit behavior Tested: Local AST parse for modified Python files Tested: Local git diff --check Tested: Remote g0034 container AST parse for modified files under /sgl-workspace/sglang-tai Tested: Remote g0034 container pytest target covering Phase8 log env, TAI range materialize, optimized slot inverse/remap, and existing token TAI path Not-tested: Full prefill/decode/router throughput after the TAI prefetch-path fix
This commit is contained in:
laoyao0822
@@ -205,6 +205,9 @@ class Envs:
|
||||
SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)
|
||||
SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
|
||||
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
|
||||
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
|
||||
SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION = EnvBool(False)
|
||||
SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
|
||||
SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
|
||||
SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)
|
||||
|
||||
431
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
Normal file
431
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
Normal file
@@ -0,0 +1,431 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
_all_reduce_materialized_buffer_async,
|
||||
_all_reduce_materialized_buffer_range,
|
||||
build_shared_token_kv_slot_remap,
|
||||
cp_shared_kv_debug_enabled,
|
||||
cp_shared_kv_mla_prefetch_enabled,
|
||||
cp_shared_kv_mla_prefetch_log,
|
||||
cp_shared_kv_mla_prefetch_should_log_layer,
|
||||
cp_shared_kv_mla_prefetch_wait_after_attention_enabled,
|
||||
filter_locs_mappable_to_physical_pool,
|
||||
materialize_local_token_kv_page_slots_into,
|
||||
remap_logical_locs_to_slot_dense_locs_optimized,
|
||||
slot_range_to_token_slice,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
|
||||
from sglang.srt.layers.dp_attention import get_attention_cp_group
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prefetch_log(message: str, *args) -> None:
|
||||
cp_shared_kv_mla_prefetch_log(message, *args)
|
||||
|
||||
|
||||
def _is_cuda_stream_capturing() -> bool:
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
try:
|
||||
return torch.cuda.is_current_stream_capturing()
|
||||
except RuntimeError:
|
||||
return False
|
||||
|
||||
|
||||
@dataclass
|
||||
class CpSharedKVMlaPrefetchHandle:
|
||||
layer_id: int
|
||||
dense_kv_cache: torch.Tensor
|
||||
event: torch.cuda.Event
|
||||
|
||||
|
||||
class CpSharedKVMlaPrefetcher:
|
||||
"""One-layer-ahead MLA prefix materialize prefetch for CP shared KV.
|
||||
|
||||
This object is per-forward-batch. It only materializes historical prefix
|
||||
pages, because current/suffix pages for layer L+1 are not written until that
|
||||
layer's MLA prepare has run.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
prefix_pages: int,
|
||||
slot_logical_pages: torch.Tensor,
|
||||
page_inverse: torch.Tensor,
|
||||
dense_num_pages: int,
|
||||
) -> None:
|
||||
self.layout = layout
|
||||
self.page_size = page_size
|
||||
self.prefix_pages = prefix_pages
|
||||
self.slot_logical_pages = slot_logical_pages
|
||||
self.page_inverse = page_inverse
|
||||
self.dense_num_pages = dense_num_pages
|
||||
self.total_slots = int(slot_logical_pages.numel())
|
||||
self.stream = torch.cuda.Stream()
|
||||
self.handles: dict[int, CpSharedKVMlaPrefetchHandle] = {}
|
||||
self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None
|
||||
self.disabled = False
|
||||
|
||||
@classmethod
|
||||
def maybe_create(
|
||||
cls,
|
||||
*,
|
||||
forward_batch: Any,
|
||||
metadata: Any,
|
||||
topk_transform_is_paged: bool,
|
||||
) -> Optional["CpSharedKVMlaPrefetcher"]:
|
||||
if not cp_shared_kv_mla_prefetch_enabled():
|
||||
return None
|
||||
if cp_shared_kv_debug_enabled():
|
||||
_prefetch_log("create_skip reason=debug_enabled")
|
||||
return None
|
||||
if not torch.cuda.is_available() or _is_cuda_stream_capturing():
|
||||
_prefetch_log(
|
||||
"create_skip reason=cuda_unavailable_or_stream_capturing cuda_available=%s",
|
||||
torch.cuda.is_available(),
|
||||
)
|
||||
return None
|
||||
if not getattr(forward_batch, "uses_cp_shared_kv", False):
|
||||
_prefetch_log("create_skip reason=not_cp_shared_kv")
|
||||
return None
|
||||
if getattr(forward_batch, "hisparse_coordinator", None) is not None:
|
||||
_prefetch_log("create_skip reason=hisparse")
|
||||
return None
|
||||
forward_mode = getattr(forward_batch, "forward_mode", None)
|
||||
if forward_mode is None or not forward_mode.is_context_parallel_extend():
|
||||
_prefetch_log("create_skip reason=not_context_parallel_extend")
|
||||
return None
|
||||
if not is_nsa_prefill_cp_in_seq_split():
|
||||
_prefetch_log("create_skip reason=not_in_seq_split")
|
||||
return None
|
||||
if not topk_transform_is_paged:
|
||||
_prefetch_log("create_skip reason=not_paged_topk")
|
||||
return None
|
||||
if int(getattr(forward_batch, "batch_size", 0)) != 1:
|
||||
_prefetch_log(
|
||||
"create_skip reason=batch_size batch_size=%s",
|
||||
getattr(forward_batch, "batch_size", None),
|
||||
)
|
||||
return None
|
||||
|
||||
token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
|
||||
if token_to_kv_pool is None:
|
||||
_prefetch_log("create_skip reason=missing_token_to_kv_pool")
|
||||
return None
|
||||
if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None:
|
||||
_prefetch_log("create_skip reason=layer_transfer_active")
|
||||
return None
|
||||
|
||||
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
|
||||
if layout is None:
|
||||
_prefetch_log("create_skip reason=missing_layout")
|
||||
return None
|
||||
|
||||
extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
|
||||
if extend_prefix_lens_cpu is None or len(extend_prefix_lens_cpu) != 1:
|
||||
_prefetch_log("create_skip reason=bad_prefix_lens_metadata")
|
||||
return None
|
||||
page_size = int(getattr(token_to_kv_pool, "page_size", 1))
|
||||
if page_size <= 1:
|
||||
_prefetch_log("create_skip reason=bad_page_size page_size=%s", page_size)
|
||||
return None
|
||||
extend_prefix_len = int(extend_prefix_lens_cpu[0])
|
||||
if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0:
|
||||
_prefetch_log(
|
||||
"create_skip reason=prefix_not_page_aligned prefix_len=%s page_size=%s",
|
||||
extend_prefix_len,
|
||||
page_size,
|
||||
)
|
||||
return None
|
||||
prefix_pages = extend_prefix_len // page_size
|
||||
|
||||
real_page_table = getattr(metadata, "real_page_table", None)
|
||||
page_table_1 = getattr(metadata, "page_table_1", None)
|
||||
if real_page_table is None or page_table_1 is None:
|
||||
_prefetch_log("create_skip reason=missing_page_tables")
|
||||
return None
|
||||
if prefix_pages <= 0 or prefix_pages > int(real_page_table.numel()):
|
||||
_prefetch_log(
|
||||
"create_skip reason=prefix_pages_out_of_range prefix_pages=%s real_pages=%s",
|
||||
prefix_pages,
|
||||
int(real_page_table.numel()),
|
||||
)
|
||||
return None
|
||||
|
||||
cp_group = get_attention_cp_group()
|
||||
if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
|
||||
_prefetch_log(
|
||||
"create_skip reason=missing_pynccl cp_rank=%s cp_size=%s",
|
||||
layout.cp_rank,
|
||||
layout.cp_size,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0))
|
||||
kv_cache = token_to_kv_pool.get_key_buffer(first_layer_id)
|
||||
remap = build_shared_token_kv_slot_remap(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=None,
|
||||
remap_logical_pages=real_page_table,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to initialize CP shared KV MLA prefetcher.")
|
||||
return None
|
||||
|
||||
_prefetch_log(
|
||||
"create cp_rank=%s cp_size=%s prefix_pages=%s total_slots=%s dense_pages=%s page_size=%s",
|
||||
layout.cp_rank,
|
||||
layout.cp_size,
|
||||
prefix_pages,
|
||||
int(remap.slot_logical_pages.numel()),
|
||||
remap.dense_num_pages,
|
||||
page_size,
|
||||
)
|
||||
|
||||
return cls(
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
prefix_pages=prefix_pages,
|
||||
slot_logical_pages=remap.slot_logical_pages,
|
||||
page_inverse=remap.page_inverse,
|
||||
dense_num_pages=remap.dense_num_pages,
|
||||
)
|
||||
|
||||
def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
|
||||
start_layer = int(getattr(token_to_kv_pool, "start_layer", 0))
|
||||
kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
|
||||
if kv_buffer is None:
|
||||
return layer_id >= start_layer
|
||||
return start_layer <= layer_id < start_layer + len(kv_buffer)
|
||||
|
||||
def consume(
|
||||
self,
|
||||
*,
|
||||
layer_id: int,
|
||||
kv_cache: torch.Tensor,
|
||||
logical_locs: torch.Tensor,
|
||||
) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
|
||||
if self.disabled:
|
||||
self._log_layer(
|
||||
layer_id,
|
||||
"consume_skip reason=disabled layer=%s",
|
||||
layer_id,
|
||||
)
|
||||
return None
|
||||
|
||||
handle = self.handles.pop(layer_id, None)
|
||||
if handle is None:
|
||||
self._log_layer(layer_id, "consume_miss layer=%s", layer_id)
|
||||
return None
|
||||
if self.pending_attention_handle is handle:
|
||||
self.pending_attention_handle = None
|
||||
if handle.layer_id != layer_id:
|
||||
self.disabled = True
|
||||
self._log_layer(
|
||||
layer_id,
|
||||
"consume_skip reason=layer_mismatch expected=%s actual=%s",
|
||||
layer_id,
|
||||
handle.layer_id,
|
||||
)
|
||||
return None
|
||||
|
||||
torch.cuda.current_stream().wait_event(handle.event)
|
||||
dense_kv_cache = handle.dense_kv_cache
|
||||
suffix_slots = self.total_slots - self.prefix_pages
|
||||
|
||||
if self.prefix_pages < self.total_slots:
|
||||
materialize_local_token_kv_page_slots_into(
|
||||
kv_cache=kv_cache,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
slot_logical_pages=self.slot_logical_pages,
|
||||
layout=self.layout,
|
||||
page_size=self.page_size,
|
||||
start_slot=self.prefix_pages,
|
||||
end_slot=self.total_slots,
|
||||
)
|
||||
suffix_rows = slot_range_to_token_slice(
|
||||
self.page_size,
|
||||
self.prefix_pages,
|
||||
self.total_slots,
|
||||
)
|
||||
_all_reduce_materialized_buffer_range(
|
||||
dense_kv_cache,
|
||||
self.layout.cp_size,
|
||||
suffix_rows.start,
|
||||
suffix_rows.stop,
|
||||
)
|
||||
|
||||
self._log_layer(
|
||||
layer_id,
|
||||
"consume_hit layer=%s prefix_pages=%s suffix_slots=%s dense_rows=%s",
|
||||
layer_id,
|
||||
self.prefix_pages,
|
||||
suffix_slots,
|
||||
int(dense_kv_cache.shape[0]),
|
||||
)
|
||||
|
||||
logical_locs = filter_locs_mappable_to_physical_pool(
|
||||
logical_locs=logical_locs,
|
||||
layout=self.layout,
|
||||
physical_token_capacity=kv_cache.shape[0],
|
||||
)
|
||||
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
|
||||
logical_locs,
|
||||
page_inverse=self.page_inverse,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
return dense_kv_cache, dense_locs
|
||||
|
||||
def start_next_layer_prefix(
|
||||
self,
|
||||
*,
|
||||
next_layer_id: int,
|
||||
token_to_kv_pool: Any,
|
||||
) -> None:
|
||||
if self.disabled:
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_skip reason=disabled next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
if next_layer_id in self.handles:
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_skip reason=already_started next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
if not self._layer_in_pool(token_to_kv_pool, next_layer_id):
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_skip reason=layer_out_of_pool next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
kv_cache = token_to_kv_pool.get_key_buffer(next_layer_id)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to get next-layer KV cache for CP shared KV MLA prefetch."
|
||||
)
|
||||
self.disabled = True
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_disable reason=get_kv_failed next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
|
||||
current_stream = torch.cuda.current_stream()
|
||||
self.stream.wait_stream(current_stream)
|
||||
try:
|
||||
with torch.cuda.stream(self.stream):
|
||||
dense_kv_cache = kv_cache.new_zeros(
|
||||
(self.dense_num_pages * self.page_size, *kv_cache.shape[1:])
|
||||
)
|
||||
materialize_local_token_kv_page_slots_into(
|
||||
kv_cache=kv_cache,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
slot_logical_pages=self.slot_logical_pages,
|
||||
layout=self.layout,
|
||||
page_size=self.page_size,
|
||||
start_slot=0,
|
||||
end_slot=self.prefix_pages,
|
||||
)
|
||||
prefix_rows = slot_range_to_token_slice(
|
||||
self.page_size,
|
||||
0,
|
||||
self.prefix_pages,
|
||||
)
|
||||
event = _all_reduce_materialized_buffer_async(
|
||||
dense_kv_cache[prefix_rows],
|
||||
cp_size=self.layout.cp_size,
|
||||
stream=self.stream,
|
||||
)
|
||||
if event is None:
|
||||
self.disabled = True
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_disable reason=async_reduce_unavailable next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("Failed to start CP shared KV MLA prefix prefetch.")
|
||||
self.disabled = True
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start_disable reason=start_exception next_layer=%s",
|
||||
next_layer_id,
|
||||
)
|
||||
return
|
||||
|
||||
handle = CpSharedKVMlaPrefetchHandle(
|
||||
layer_id=next_layer_id,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
event=event,
|
||||
)
|
||||
self.handles[next_layer_id] = handle
|
||||
self.pending_attention_handle = handle
|
||||
self._log_next_layer(
|
||||
next_layer_id,
|
||||
"start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s",
|
||||
next_layer_id,
|
||||
self.prefix_pages,
|
||||
prefix_rows.stop - prefix_rows.start,
|
||||
int(dense_kv_cache.shape[0]),
|
||||
)
|
||||
|
||||
def wait_attention_window(self) -> None:
|
||||
if not cp_shared_kv_mla_prefetch_wait_after_attention_enabled():
|
||||
handle = self.pending_attention_handle
|
||||
if handle is not None:
|
||||
self._log_next_layer(
|
||||
handle.layer_id,
|
||||
"attention_wait_deferred next_layer=%s",
|
||||
handle.layer_id,
|
||||
)
|
||||
return
|
||||
|
||||
handle = self.pending_attention_handle
|
||||
self.pending_attention_handle = None
|
||||
if handle is None:
|
||||
return
|
||||
torch.cuda.current_stream().wait_event(handle.event)
|
||||
self._log_next_layer(
|
||||
handle.layer_id,
|
||||
"attention_wait next_layer=%s",
|
||||
handle.layer_id,
|
||||
)
|
||||
|
||||
def _log_layer(self, layer_id: int, message: str, *args) -> None:
|
||||
if cp_shared_kv_mla_prefetch_should_log_layer(layer_id):
|
||||
self._log(message, *args)
|
||||
|
||||
def _log_next_layer(self, next_layer_id: int, message: str, *args) -> None:
|
||||
if cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id):
|
||||
self._log(message, *args)
|
||||
|
||||
def _log(self, message: str, *args) -> None:
|
||||
_prefetch_log(
|
||||
"cp_rank=%s cp_size=%s " + message,
|
||||
self.layout.cp_rank,
|
||||
self.layout.cp_size,
|
||||
*args,
|
||||
)
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
@@ -13,6 +14,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_DEBUG_LOG_COUNTS: dict[str, int] = {}
|
||||
_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
|
||||
_MLA_PREFETCH_LOG_PROBE_LAYER = 2
|
||||
|
||||
|
||||
def cp_shared_kv_debug_enabled() -> bool:
|
||||
@@ -27,6 +29,36 @@ def cp_shared_kv_tai_materialize_enabled() -> bool:
|
||||
return envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get()
|
||||
|
||||
|
||||
def cp_shared_kv_mla_prefetch_enabled() -> bool:
|
||||
return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
|
||||
|
||||
|
||||
def cp_shared_kv_mla_prefetch_log_enabled() -> bool:
|
||||
return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get()
|
||||
|
||||
|
||||
def cp_shared_kv_mla_prefetch_wait_after_attention_enabled() -> bool:
|
||||
return envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.get()
|
||||
|
||||
|
||||
def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None:
|
||||
if cp_shared_kv_mla_prefetch_log_enabled():
|
||||
logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args)
|
||||
|
||||
|
||||
def cp_shared_kv_mla_prefetch_should_log_layer(layer_id: int) -> bool:
|
||||
return int(layer_id) == _MLA_PREFETCH_LOG_PROBE_LAYER
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SharedTokenKVSlotRemap:
|
||||
slot_logical_pages: torch.Tensor
|
||||
page_inverse: torch.Tensor
|
||||
dense_locs: torch.Tensor | None
|
||||
logical_page_capacity: int
|
||||
dense_num_pages: int
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_materialize_kernels():
|
||||
try:
|
||||
@@ -142,6 +174,157 @@ def _try_tai_materialize_token_kv_pages_and_locs(
|
||||
return None
|
||||
|
||||
|
||||
def _try_tai_build_slot_page_inverse(
|
||||
slot_logical_pages: torch.Tensor,
|
||||
logical_page_capacity: int,
|
||||
) -> torch.Tensor | None:
|
||||
if not _tai_materialize_runtime_enabled():
|
||||
return None
|
||||
|
||||
kernels = _load_tai_materialize_kernels()
|
||||
if kernels is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return kernels.build_slot_page_inverse(
|
||||
_contiguous_for_tai(slot_logical_pages.reshape(-1)),
|
||||
logical_page_capacity,
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"page_inverse_failed",
|
||||
"CP shared KV tai page inverse build failed; falling back to torch "
|
||||
"remap. error=%s",
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def build_slot_page_inverse_optimized(
|
||||
slot_logical_pages: torch.Tensor,
|
||||
logical_page_capacity: int,
|
||||
) -> torch.Tensor:
|
||||
tai_result = _try_tai_build_slot_page_inverse(
|
||||
slot_logical_pages,
|
||||
logical_page_capacity,
|
||||
)
|
||||
if tai_result is not None:
|
||||
return tai_result
|
||||
return build_slot_page_inverse(
|
||||
slot_logical_pages,
|
||||
logical_page_capacity=logical_page_capacity,
|
||||
)
|
||||
|
||||
|
||||
def remap_logical_locs_to_slot_dense_locs_optimized(
|
||||
logical_locs: torch.Tensor,
|
||||
page_inverse: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> torch.Tensor:
|
||||
if _tai_materialize_runtime_enabled():
|
||||
kernels = _load_tai_materialize_kernels()
|
||||
if kernels is not None:
|
||||
try:
|
||||
return kernels.remap_logical_locs_to_slot_dense_locs(
|
||||
_contiguous_for_tai(logical_locs),
|
||||
_contiguous_for_tai(page_inverse),
|
||||
page_size=page_size,
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"loc_remap_failed",
|
||||
"CP shared KV tai loc remap failed; falling back to torch "
|
||||
"remap. error=%s",
|
||||
exc,
|
||||
)
|
||||
return remap_logical_locs_to_slot_dense_locs(
|
||||
logical_locs,
|
||||
page_inverse=page_inverse,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
|
||||
def _copy_tai_dense_slot_range_body(
|
||||
*,
|
||||
tai_dense_kv_cache: torch.Tensor,
|
||||
dense_kv_cache: torch.Tensor,
|
||||
page_size: int,
|
||||
start_slot: int,
|
||||
end_slot: int,
|
||||
) -> None:
|
||||
if start_slot == end_slot:
|
||||
return
|
||||
dst_rows = slot_range_to_token_slice(page_size, start_slot, end_slot)
|
||||
src_rows = slot_range_to_token_slice(page_size, 0, end_slot - start_slot)
|
||||
dense_kv_cache[dst_rows].copy_(tai_dense_kv_cache[src_rows])
|
||||
|
||||
|
||||
def _try_tai_materialize_token_kv_page_slots_into(
|
||||
*,
|
||||
kv_cache: torch.Tensor,
|
||||
dense_kv_cache: torch.Tensor,
|
||||
slot_logical_pages: torch.Tensor,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
start_slot: int,
|
||||
end_slot: int,
|
||||
) -> bool:
|
||||
if not _tai_materialize_runtime_enabled():
|
||||
return False
|
||||
|
||||
kernels = _load_tai_materialize_kernels()
|
||||
if kernels is None:
|
||||
return False
|
||||
|
||||
flat_slot_logical_pages = slot_logical_pages.reshape(-1)
|
||||
slot_logical_pages_range = _contiguous_for_tai(
|
||||
flat_slot_logical_pages[start_slot:end_slot]
|
||||
)
|
||||
if slot_logical_pages_range.numel() == 0:
|
||||
return True
|
||||
|
||||
try:
|
||||
materialize_into = getattr(
|
||||
kernels,
|
||||
"materialize_shared_token_kv_pages_into",
|
||||
None,
|
||||
)
|
||||
if materialize_into is not None:
|
||||
materialize_into(
|
||||
kv_cache,
|
||||
slot_logical_pages_range,
|
||||
dense_kv_cache,
|
||||
page_size=page_size,
|
||||
start_slot=start_slot,
|
||||
cp_rank=layout.cp_rank,
|
||||
cp_size=layout.cp_size,
|
||||
)
|
||||
else:
|
||||
tai_dense_kv_cache = kernels.materialize_shared_token_kv_pages(
|
||||
kv_cache,
|
||||
slot_logical_pages_range,
|
||||
page_size=page_size,
|
||||
cp_rank=layout.cp_rank,
|
||||
cp_size=layout.cp_size,
|
||||
)
|
||||
_copy_tai_dense_slot_range_body(
|
||||
tai_dense_kv_cache=tai_dense_kv_cache,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
page_size=page_size,
|
||||
start_slot=start_slot,
|
||||
end_slot=end_slot,
|
||||
)
|
||||
return True
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"token_range_failed",
|
||||
"CP shared KV tai token range materialize failed; falling back to "
|
||||
"torch materialize. error=%s",
|
||||
exc,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def is_current_only_extend_batch(forward_batch) -> bool:
|
||||
"""Return whether an extend batch has no cached/history tokens.
|
||||
|
||||
@@ -472,6 +655,59 @@ def remap_logical_locs_to_slot_dense_locs(
|
||||
return torch.where(mapped, dense_values, dense_locs)
|
||||
|
||||
|
||||
def build_shared_token_kv_slot_remap(
|
||||
kv_cache: torch.Tensor,
|
||||
logical_locs: torch.Tensor | None,
|
||||
remap_logical_pages: torch.Tensor,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
) -> SharedTokenKVSlotRemap:
|
||||
"""Build the fixed slot-layout remap used by shared token KV materialize.
|
||||
|
||||
The slot layout is intentionally the same as `build_slot_page_remap`: dense
|
||||
page 0 is the dummy page and dense page `slot + 1` corresponds to
|
||||
`remap_logical_pages.reshape(-1)[slot]`. Phase 8 uses the same remap to
|
||||
materialize prefix/suffix ranges into one dense buffer without changing
|
||||
attention page-table semantics.
|
||||
"""
|
||||
|
||||
_debug_assert_no_negative_tensor_values(
|
||||
remap_logical_pages,
|
||||
context="CP shared KV token materialize page remap",
|
||||
tensor_name="remap_logical_pages",
|
||||
)
|
||||
remap_logical_pages = filter_pages_mappable_to_physical_pool(
|
||||
logical_pages=remap_logical_pages,
|
||||
layout=layout,
|
||||
physical_page_capacity=kv_cache.shape[0] // page_size,
|
||||
)
|
||||
logical_page_capacity = _logical_page_capacity_from_physical_page_capacity(
|
||||
kv_cache.shape[0] // page_size,
|
||||
layout,
|
||||
)
|
||||
slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
|
||||
page_inverse = build_slot_page_inverse_optimized(
|
||||
slot_logical_pages,
|
||||
logical_page_capacity=logical_page_capacity,
|
||||
)
|
||||
dense_locs = (
|
||||
remap_logical_locs_to_slot_dense_locs_optimized(
|
||||
logical_locs,
|
||||
page_inverse=page_inverse,
|
||||
page_size=page_size,
|
||||
)
|
||||
if logical_locs is not None
|
||||
else None
|
||||
)
|
||||
return SharedTokenKVSlotRemap(
|
||||
slot_logical_pages=slot_logical_pages,
|
||||
page_inverse=page_inverse,
|
||||
dense_locs=dense_locs,
|
||||
logical_page_capacity=logical_page_capacity,
|
||||
dense_num_pages=int(slot_logical_pages.numel()) + 1,
|
||||
)
|
||||
|
||||
|
||||
def remap_logical_locs_to_dense_locs(
|
||||
logical_locs: torch.Tensor,
|
||||
unique_logical_pages: torch.Tensor,
|
||||
@@ -675,7 +911,65 @@ def materialize_local_token_kv_page_slots(
|
||||
if slot_logical_pages.numel() == 0:
|
||||
return dense_kv_cache
|
||||
|
||||
logical_pages = slot_logical_pages.reshape(-1).to(torch.long)
|
||||
materialize_local_token_kv_page_slots_into(
|
||||
kv_cache=kv_cache,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
slot_logical_pages=slot_logical_pages,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
start_slot=0,
|
||||
end_slot=int(slot_logical_pages.numel()),
|
||||
)
|
||||
return dense_kv_cache
|
||||
|
||||
|
||||
def materialize_local_token_kv_page_slots_into(
|
||||
kv_cache: torch.Tensor,
|
||||
dense_kv_cache: torch.Tensor,
|
||||
slot_logical_pages: torch.Tensor,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
start_slot: int,
|
||||
end_slot: int | None = None,
|
||||
) -> None:
|
||||
"""Materialize a slot range into an existing dense token KV buffer.
|
||||
|
||||
`start_slot`/`end_slot` are page-table slots, not dense page ids. Dense
|
||||
page 0 is the dummy page, so slot `i` writes dense token rows for page
|
||||
`i + 1`.
|
||||
"""
|
||||
|
||||
flat_slot_logical_pages = slot_logical_pages.reshape(-1)
|
||||
total_slots = int(flat_slot_logical_pages.numel())
|
||||
if end_slot is None:
|
||||
end_slot = total_slots
|
||||
if start_slot < 0 or end_slot < start_slot or end_slot > total_slots:
|
||||
raise ValueError(
|
||||
"Invalid CP shared KV slot materialize range: "
|
||||
f"start_slot={start_slot} end_slot={end_slot} total_slots={total_slots}"
|
||||
)
|
||||
if start_slot == end_slot:
|
||||
return
|
||||
|
||||
expected_rows = (total_slots + 1) * page_size
|
||||
if dense_kv_cache.shape[0] < expected_rows:
|
||||
raise ValueError(
|
||||
"CP shared KV dense token buffer is too small for slot materialize: "
|
||||
f"dense_rows={dense_kv_cache.shape[0]} expected_at_least={expected_rows}"
|
||||
)
|
||||
|
||||
if _try_tai_materialize_token_kv_page_slots_into(
|
||||
kv_cache=kv_cache,
|
||||
dense_kv_cache=dense_kv_cache,
|
||||
slot_logical_pages=flat_slot_logical_pages,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
start_slot=start_slot,
|
||||
end_slot=end_slot,
|
||||
):
|
||||
return
|
||||
|
||||
logical_pages = flat_slot_logical_pages[start_slot:end_slot].to(torch.long)
|
||||
owned_mask = layout.owned_pages_mask(logical_pages)
|
||||
physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long)
|
||||
safe_physical_pages = torch.where(
|
||||
@@ -686,16 +980,29 @@ def materialize_local_token_kv_page_slots(
|
||||
page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long)
|
||||
src_tokens = (safe_physical_pages[:, None] * page_size + page_offsets).reshape(-1)
|
||||
|
||||
dense_body = dense_kv_cache[page_size:].view(
|
||||
dense_num_pages - 1,
|
||||
dense_body = dense_kv_cache[page_size:expected_rows].view(
|
||||
total_slots,
|
||||
page_size,
|
||||
*kv_cache.shape[1:],
|
||||
)
|
||||
gathered = kv_cache[src_tokens].view_as(dense_body)
|
||||
dense_range = dense_body[start_slot:end_slot]
|
||||
gathered = kv_cache[src_tokens].view_as(dense_range)
|
||||
owned_view = owned_mask.view(-1, *([1] * (dense_body.ndim - 1)))
|
||||
zero = torch.zeros((), dtype=kv_cache.dtype, device=kv_cache.device)
|
||||
dense_body.copy_(torch.where(owned_view, gathered, zero))
|
||||
return dense_kv_cache
|
||||
dense_range.copy_(torch.where(owned_view, gathered, zero))
|
||||
|
||||
|
||||
def slot_range_to_token_slice(
|
||||
page_size: int,
|
||||
start_slot: int,
|
||||
end_slot: int,
|
||||
) -> slice:
|
||||
if start_slot < 0 or end_slot < start_slot:
|
||||
raise ValueError(
|
||||
"Invalid CP shared KV slot token slice range: "
|
||||
f"start_slot={start_slot} end_slot={end_slot}"
|
||||
)
|
||||
return slice((start_slot + 1) * page_size, (end_slot + 1) * page_size)
|
||||
|
||||
|
||||
def token_page_copy_debug_checksum(
|
||||
@@ -824,6 +1131,65 @@ def _all_reduce_materialized_buffer(buffer: torch.Tensor, cp_size: int) -> torch
|
||||
return buffer
|
||||
|
||||
|
||||
def _all_reduce_materialized_buffer_range(
|
||||
buffer: torch.Tensor,
|
||||
cp_size: int,
|
||||
start_row: int,
|
||||
end_row: int,
|
||||
) -> torch.Tensor:
|
||||
if start_row < 0 or end_row < start_row or end_row > buffer.shape[0]:
|
||||
raise ValueError(
|
||||
"Invalid CP shared KV materialize reduce row range: "
|
||||
f"start_row={start_row} end_row={end_row} rows={buffer.shape[0]}"
|
||||
)
|
||||
if start_row == end_row:
|
||||
return buffer
|
||||
_all_reduce_materialized_buffer(buffer[start_row:end_row], cp_size)
|
||||
return buffer
|
||||
|
||||
|
||||
def _all_reduce_materialized_buffer_async(
|
||||
buffer: torch.Tensor,
|
||||
cp_size: int,
|
||||
stream: torch.cuda.Stream,
|
||||
) -> torch.cuda.Event | None:
|
||||
"""Enqueue an in-place CP all-reduce on `stream`.
|
||||
|
||||
Returns a CUDA event recorded after the collective, or `None` when the
|
||||
async pynccl path is unavailable. Callers must fallback before launching
|
||||
rank-divergent collectives if this returns `None`.
|
||||
"""
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return None
|
||||
event = torch.cuda.Event()
|
||||
if cp_size <= 1 or buffer.numel() == 0:
|
||||
with torch.cuda.stream(stream):
|
||||
event.record(stream)
|
||||
return event
|
||||
|
||||
cp_group = get_attention_cp_group()
|
||||
pynccl_comm = getattr(cp_group, "pynccl_comm", None)
|
||||
if pynccl_comm is None:
|
||||
return None
|
||||
|
||||
comm_buffer = _comm_view(buffer)
|
||||
try:
|
||||
with pynccl_comm.change_state(enable=True, stream=stream):
|
||||
pynccl_comm.all_reduce(comm_buffer, stream=stream)
|
||||
event.record(stream)
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"prefetch_async_allreduce_failed",
|
||||
"CP shared KV MLA prefetch async all-reduce is unavailable; "
|
||||
"falling back to sync materialize. error=%s",
|
||||
exc,
|
||||
limit=4,
|
||||
)
|
||||
return None
|
||||
return event
|
||||
|
||||
|
||||
def materialize_shared_token_kv_buffer(
|
||||
kv_cache: torch.Tensor,
|
||||
logical_locs: torch.Tensor,
|
||||
@@ -896,11 +1262,11 @@ def materialize_shared_token_kv_buffer(
|
||||
)
|
||||
if tai_result is None:
|
||||
materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
|
||||
page_inverse = build_slot_page_inverse(
|
||||
page_inverse = build_slot_page_inverse_optimized(
|
||||
materialized_logical_pages,
|
||||
logical_page_capacity=logical_page_capacity,
|
||||
)
|
||||
dense_locs = remap_logical_locs_to_slot_dense_locs(
|
||||
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
|
||||
logical_locs,
|
||||
page_inverse=page_inverse,
|
||||
page_size=page_size,
|
||||
|
||||
@@ -9,11 +9,17 @@ import torch
|
||||
from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
|
||||
CpSharedKVMlaPrefetcher,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
build_current_loc_remap,
|
||||
cp_shared_kv_debug_enabled,
|
||||
cp_shared_kv_debug_log,
|
||||
cp_shared_kv_current_reuse_enabled,
|
||||
cp_shared_kv_mla_prefetch_log,
|
||||
cp_shared_kv_mla_prefetch_log_enabled,
|
||||
cp_shared_kv_mla_prefetch_should_log_layer,
|
||||
filter_owned_logical_locs,
|
||||
is_current_only_extend_batch,
|
||||
materialize_shared_token_kv_buffer,
|
||||
@@ -585,6 +591,7 @@ class NativeSparseAttnBackend(
|
||||
"""Init the metadata for a forward pass."""
|
||||
batch_size = forward_batch.batch_size
|
||||
device = forward_batch.seq_lens.device
|
||||
forward_batch.cp_shared_kv_mla_prefetcher = None
|
||||
|
||||
if forward_batch.forward_mode.is_target_verify():
|
||||
draft_token_num = self.speculative_num_draft_tokens
|
||||
@@ -857,6 +864,15 @@ class NativeSparseAttnBackend(
|
||||
token_to_batch_idx=token_to_batch_idx,
|
||||
)
|
||||
self.forward_metadata = metadata
|
||||
forward_batch.cp_shared_kv_mla_prefetcher = (
|
||||
CpSharedKVMlaPrefetcher.maybe_create(
|
||||
forward_batch=forward_batch,
|
||||
metadata=metadata,
|
||||
topk_transform_is_paged=(
|
||||
topk_transform_method == TopkTransformMethod.PAGED
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def _cal_indexer_k_start_end(
|
||||
self,
|
||||
@@ -1594,6 +1610,9 @@ class NativeSparseAttnBackend(
|
||||
and topk_transform_method == TopkTransformMethod.PAGED
|
||||
):
|
||||
assert forward_batch.cp_shared_kv_layout is not None
|
||||
mla_prefetcher = getattr(
|
||||
forward_batch, "cp_shared_kv_mla_prefetcher", None
|
||||
)
|
||||
can_reuse_current_kv = (
|
||||
cp_shared_kv_current_reuse_enabled()
|
||||
and is_current_only_extend_batch(forward_batch)
|
||||
@@ -1602,6 +1621,39 @@ class NativeSparseAttnBackend(
|
||||
and k.shape[0] == forward_batch.out_cache_loc.numel()
|
||||
and k_rope.shape[0] == forward_batch.out_cache_loc.numel()
|
||||
)
|
||||
if cp_shared_kv_mla_prefetch_log_enabled():
|
||||
if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id):
|
||||
prefix_lens_cpu = getattr(
|
||||
forward_batch, "extend_prefix_lens_cpu", None
|
||||
)
|
||||
extend_lens_cpu = getattr(
|
||||
forward_batch, "extend_seq_lens_cpu", None
|
||||
)
|
||||
prefix_lens = (
|
||||
[int(x) for x in prefix_lens_cpu]
|
||||
if prefix_lens_cpu is not None
|
||||
else None
|
||||
)
|
||||
extend_lens = (
|
||||
[int(x) for x in extend_lens_cpu]
|
||||
if extend_lens_cpu is not None
|
||||
else None
|
||||
)
|
||||
cp_shared_kv_mla_prefetch_log(
|
||||
"forward_layer cp_rank=%s layer=%s cache_hit=%s "
|
||||
"has_prefetcher=%s can_current_reuse=%s prefix_lens=%s "
|
||||
"extend_lens=%s page_table_shape=%s",
|
||||
forward_batch.cp_shared_kv_layout.cp_rank,
|
||||
layer.layer_id,
|
||||
any(prefix_len > 0 for prefix_len in prefix_lens or []),
|
||||
mla_prefetcher is not None,
|
||||
can_reuse_current_kv,
|
||||
prefix_lens,
|
||||
extend_lens,
|
||||
tuple(page_table_1.shape)
|
||||
if page_table_1 is not None
|
||||
else None,
|
||||
)
|
||||
if can_reuse_current_kv:
|
||||
logical_page_table_1 = page_table_1
|
||||
current_mask, page_table_1 = build_current_loc_remap(
|
||||
@@ -1630,100 +1682,124 @@ class NativeSparseAttnBackend(
|
||||
)
|
||||
kv_cache = _cat([k, k_rope], dim=-1)
|
||||
else:
|
||||
kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1,
|
||||
remap_logical_locs=metadata.page_table_1,
|
||||
remap_logical_pages=metadata.real_page_table,
|
||||
layout=forward_batch.cp_shared_kv_layout,
|
||||
page_size=forward_batch.token_to_kv_pool.page_size,
|
||||
)
|
||||
|
||||
if nsa_impl == "tilelang":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
return self._forward_tilelang(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
)
|
||||
elif nsa_impl == "flashmla_sparse":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
|
||||
if topk_transform_method == TopkTransformMethod.RAGGED:
|
||||
if any(forward_batch.extend_prefix_lens_cpu):
|
||||
page_table_1_flattened = (
|
||||
self.forward_metadata.page_table_1_flattened
|
||||
)
|
||||
assert page_table_1_flattened is not None
|
||||
if forward_batch.uses_cp_shared_kv:
|
||||
assert forward_batch.cp_shared_kv_layout is not None
|
||||
kv_cache, page_table_1_flattened = (
|
||||
materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1_flattened,
|
||||
layout=forward_batch.cp_shared_kv_layout,
|
||||
page_size=forward_batch.token_to_kv_pool.page_size,
|
||||
)
|
||||
)
|
||||
kv_cache = dequantize_k_cache_paged(
|
||||
kv_cache, page_table_1_flattened
|
||||
prefetched_kv = None
|
||||
if mla_prefetcher is not None:
|
||||
prefetched_kv = mla_prefetcher.consume(
|
||||
layer_id=layer.layer_id,
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1,
|
||||
)
|
||||
if prefetched_kv is not None:
|
||||
kv_cache, page_table_1 = prefetched_kv
|
||||
else:
|
||||
kv_cache = _cat([k, k_rope], dim=-1)
|
||||
page_table_1 = topk_indices
|
||||
kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1,
|
||||
remap_logical_locs=metadata.page_table_1,
|
||||
remap_logical_pages=metadata.real_page_table,
|
||||
layout=forward_batch.cp_shared_kv_layout,
|
||||
page_size=forward_batch.token_to_kv_pool.page_size,
|
||||
)
|
||||
|
||||
return self._forward_flashmla_sparse(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
)
|
||||
elif nsa_impl == "flashmla_kv":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
return self._forward_flashmla_kv(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
# TODO optimize args
|
||||
layer=layer,
|
||||
metadata=metadata,
|
||||
page_table_1=page_table_1,
|
||||
)
|
||||
elif nsa_impl == "fa3":
|
||||
return self._forward_fa3(
|
||||
q_rope=q_rope,
|
||||
kv_cache=kv_cache,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
q_nope=q_nope,
|
||||
page_table=page_table_1,
|
||||
cache_seqlens=metadata.nsa_cache_seqlens_int32,
|
||||
cu_seqlens_q=metadata.nsa_cu_seqlens_q,
|
||||
cu_seqlens_k=metadata.nsa_cu_seqlens_k,
|
||||
max_seqlen_q=metadata.nsa_max_seqlen_q,
|
||||
sm_scale=layer.scaling,
|
||||
logit_cap=layer.logit_cap,
|
||||
page_size=1,
|
||||
)
|
||||
elif nsa_impl == "aiter":
|
||||
if q_rope is not None:
|
||||
q_all = torch.cat([q_nope, q_rope], dim=-1)
|
||||
return self._forward_aiter_extend(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
layer=layer,
|
||||
)
|
||||
if mla_prefetcher is not None:
|
||||
mla_prefetcher.start_next_layer_prefix(
|
||||
next_layer_id=layer.layer_id + 1,
|
||||
token_to_kv_pool=forward_batch.token_to_kv_pool,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend."
|
||||
)
|
||||
mla_prefetcher = None
|
||||
|
||||
try:
|
||||
if nsa_impl == "tilelang":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
attn_output = self._forward_tilelang(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
)
|
||||
elif nsa_impl == "flashmla_sparse":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
|
||||
if topk_transform_method == TopkTransformMethod.RAGGED:
|
||||
if any(forward_batch.extend_prefix_lens_cpu):
|
||||
page_table_1_flattened = (
|
||||
self.forward_metadata.page_table_1_flattened
|
||||
)
|
||||
assert page_table_1_flattened is not None
|
||||
if forward_batch.uses_cp_shared_kv:
|
||||
assert forward_batch.cp_shared_kv_layout is not None
|
||||
kv_cache, page_table_1_flattened = (
|
||||
materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1_flattened,
|
||||
layout=forward_batch.cp_shared_kv_layout,
|
||||
page_size=forward_batch.token_to_kv_pool.page_size,
|
||||
)
|
||||
)
|
||||
kv_cache = dequantize_k_cache_paged(
|
||||
kv_cache, page_table_1_flattened
|
||||
)
|
||||
else:
|
||||
kv_cache = _cat([k, k_rope], dim=-1)
|
||||
page_table_1 = topk_indices
|
||||
|
||||
attn_output = self._forward_flashmla_sparse(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
)
|
||||
elif nsa_impl == "flashmla_kv":
|
||||
if q_rope is not None:
|
||||
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
|
||||
attn_output = self._forward_flashmla_kv(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
sm_scale=layer.scaling,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
# TODO optimize args
|
||||
layer=layer,
|
||||
metadata=metadata,
|
||||
page_table_1=page_table_1,
|
||||
)
|
||||
elif nsa_impl == "fa3":
|
||||
attn_output = self._forward_fa3(
|
||||
q_rope=q_rope,
|
||||
kv_cache=kv_cache,
|
||||
v_head_dim=layer.v_head_dim,
|
||||
q_nope=q_nope,
|
||||
page_table=page_table_1,
|
||||
cache_seqlens=metadata.nsa_cache_seqlens_int32,
|
||||
cu_seqlens_q=metadata.nsa_cu_seqlens_q,
|
||||
cu_seqlens_k=metadata.nsa_cu_seqlens_k,
|
||||
max_seqlen_q=metadata.nsa_max_seqlen_q,
|
||||
sm_scale=layer.scaling,
|
||||
logit_cap=layer.logit_cap,
|
||||
page_size=1,
|
||||
)
|
||||
elif nsa_impl == "aiter":
|
||||
if q_rope is not None:
|
||||
q_all = torch.cat([q_nope, q_rope], dim=-1)
|
||||
attn_output = self._forward_aiter_extend(
|
||||
q_all=q_all,
|
||||
kv_cache=kv_cache,
|
||||
page_table_1=page_table_1,
|
||||
layer=layer,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend."
|
||||
)
|
||||
finally:
|
||||
if mla_prefetcher is not None:
|
||||
mla_prefetcher.wait_attention_window()
|
||||
|
||||
return attn_output
|
||||
|
||||
def forward_decode(
|
||||
self,
|
||||
|
||||
@@ -32,7 +32,7 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum, auto
|
||||
from functools import total_ordering
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import triton
|
||||
@@ -424,6 +424,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
cp_shared_kv_layout: Optional[CpSharedKVLayout] = None
|
||||
cp_local_out_cache_loc: Optional[torch.Tensor] = None
|
||||
cp_shared_mla_direct_write_done: bool = False
|
||||
cp_shared_kv_mla_prefetcher: Optional[Any] = None
|
||||
|
||||
# For hidden states before normal
|
||||
return_hidden_states_before_norm: bool = False
|
||||
|
||||
Reference in New Issue
Block a user