Batch CP shared-KV index work for bs>1 fast paths
The bs>1 path needs index top-k, shared-index prepare, current-index compact, and current-slot compose to consume flattened batch descriptors instead of falling back to per-request or per-segment Python/Torch work. This change wires SGLang to the new TAI batch prepare kernels, keeps fallback explicit, and records the remaining HiCache/load-backup gaps in the bs>1 workstream docs. Constraint: CP shared-KV bs>1 must reuse fast paths rather than adding slow batch-only fallbacks Constraint: No new collective operations were introduced Rejected: Leave current-only cp_index as Python slice/cat | it keeps per-segment overhead in the short-extend bs>1 case Rejected: Infer max segment lengths from CUDA descriptor tensors | .item() would add CPU synchronization on the hot path Confidence: medium Scope-risk: moderate Directive: Do not remove the explicit fallback warnings without verifying the corresponding TAI symbols are present in production Tested: local py_compile for touched SGLang files Tested: remote g0034 test_nsa_cp_utils.py passed, 53 tests Tested: remote g0034 test_fill_current_index_page_slots_uses_tai_kernel_when_available passed Not-tested: full ETE bs>1 traffic with HiCache load/backup and draft/EAGLE enabled
This commit is contained in:
laoyao0822
@@ -577,6 +577,26 @@ def _load_tai_index_mqa_prepare_kernel():
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_index_mqa_prepare_batch_kernel():
|
||||
try:
|
||||
from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range_batch
|
||||
|
||||
return prepare_cp_mqa_kv_and_range_batch
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_index_mqa_current_prepare_batch_kernel():
|
||||
try:
|
||||
from tai_kernel.nsa_prefill import prepare_cp_mqa_current_kv_and_range_batch
|
||||
|
||||
return prepare_cp_mqa_current_kv_and_range_batch
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_index_mqa_range_kernel():
|
||||
try:
|
||||
@@ -800,6 +820,214 @@ def try_tai_prepare_cp_mqa_index(
|
||||
return None
|
||||
|
||||
|
||||
def try_tai_prepare_cp_mqa_current_index_batch(
|
||||
*,
|
||||
current_index_k: torch.Tensor,
|
||||
current_index_scale: torch.Tensor,
|
||||
current_bases: torch.Tensor,
|
||||
kv_lens: torch.Tensor,
|
||||
q_starts: torch.Tensor,
|
||||
q_lens: torch.Tensor,
|
||||
k_bases: torch.Tensor,
|
||||
q_bases: torch.Tensor,
|
||||
total_kv_len: int,
|
||||
total_q_count: int,
|
||||
max_kv_len: int,
|
||||
max_q_len: int,
|
||||
index_head_dim: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
|
||||
"""Try the TAI batched current-index compact + MQA range prepare path."""
|
||||
|
||||
if not cp_shared_kv_tai_index_mqa_prepare_enabled():
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"current_batch_env_disabled",
|
||||
"CP shared KV TAI batched current-index prepare fast path is disabled; "
|
||||
"falling back to Python slice/cat. segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s index_head_dim=%s",
|
||||
int(current_bases.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
index_head_dim,
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
kernel = _load_tai_index_mqa_current_prepare_batch_kernel()
|
||||
if kernel is None:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"current_batch_kernel_missing",
|
||||
"CP shared KV TAI batched current-index prepare kernel is unavailable; "
|
||||
"falling back to Python slice/cat. segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s index_head_dim=%s",
|
||||
int(current_bases.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
index_head_dim,
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
if index_head_dim != 128:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"current_batch_unsupported_layout",
|
||||
"CP shared KV TAI batched current-index prepare supports "
|
||||
"index_head_dim=128 only; falling back to Python slice/cat. "
|
||||
"index_head_dim=%s",
|
||||
index_head_dim,
|
||||
limit=4,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
return kernel(
|
||||
_contiguous_for_tai(current_index_k),
|
||||
_contiguous_for_tai(current_index_scale),
|
||||
_contiguous_for_tai(current_bases),
|
||||
_contiguous_for_tai(kv_lens),
|
||||
_contiguous_for_tai(q_starts),
|
||||
_contiguous_for_tai(q_lens),
|
||||
_contiguous_for_tai(k_bases),
|
||||
_contiguous_for_tai(q_bases),
|
||||
total_kv_len=int(total_kv_len),
|
||||
total_q_count=int(total_q_count),
|
||||
max_kv_len=int(max_kv_len),
|
||||
max_q_len=int(max_q_len),
|
||||
index_head_dim=int(index_head_dim),
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"current_batch_kernel_failed",
|
||||
"CP shared KV TAI batched current-index prepare failed; falling back "
|
||||
"to Python slice/cat. error=%s segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s",
|
||||
exc,
|
||||
int(current_bases.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def try_tai_prepare_cp_mqa_index_batch(
|
||||
*,
|
||||
index_buffer: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
batch_indices: torch.Tensor,
|
||||
kv_lens: torch.Tensor,
|
||||
q_starts: torch.Tensor,
|
||||
q_lens: torch.Tensor,
|
||||
k_bases: torch.Tensor,
|
||||
q_bases: torch.Tensor,
|
||||
total_kv_len: int,
|
||||
total_q_count: int,
|
||||
max_kv_len: int,
|
||||
max_q_len: int,
|
||||
page_size: int,
|
||||
index_head_dim: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
|
||||
"""Try the TAI batched GetK/GetS + MQA range prepare path.
|
||||
|
||||
This is the bs>1 / multi-CP-segment counterpart of
|
||||
:func:`try_tai_prepare_cp_mqa_index`: all segment descriptors are flattened
|
||||
and one kernel family prepares the concatenated K/S buffer and ks/ke ranges.
|
||||
"""
|
||||
|
||||
if not cp_shared_kv_tai_index_mqa_prepare_enabled():
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_env_disabled",
|
||||
"CP shared KV TAI batched index MQA prepare fast path is disabled; "
|
||||
"falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s page_size=%s index_head_dim=%s",
|
||||
int(batch_indices.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
page_size,
|
||||
index_head_dim,
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
kernel = _load_tai_index_mqa_prepare_batch_kernel()
|
||||
if kernel is None:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_kernel_missing",
|
||||
"CP shared KV TAI batched index MQA prepare kernel is unavailable; "
|
||||
"falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s page_size=%s index_head_dim=%s",
|
||||
int(batch_indices.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
page_size,
|
||||
index_head_dim,
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
if index_buffer.dtype != torch.uint8:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_unsupported_dtype",
|
||||
"CP shared KV TAI batched index MQA prepare requires uint8 page "
|
||||
"buffer; falling back to per-segment GetK/GetS. dtype=%s",
|
||||
index_buffer.dtype,
|
||||
limit=4,
|
||||
)
|
||||
return None
|
||||
if not index_buffer.is_contiguous() or not block_tables.is_contiguous():
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_non_contiguous",
|
||||
"CP shared KV TAI batched index MQA prepare requires contiguous "
|
||||
"index buffer and block tables; falling back to per-segment GetK/GetS. "
|
||||
"index_shape=%s index_stride=%s table_shape=%s table_stride=%s",
|
||||
tuple(index_buffer.shape),
|
||||
tuple(index_buffer.stride()),
|
||||
tuple(block_tables.shape),
|
||||
tuple(block_tables.stride()),
|
||||
limit=4,
|
||||
)
|
||||
return None
|
||||
if index_head_dim != 128 or page_size != 64:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_unsupported_layout",
|
||||
"CP shared KV TAI batched index MQA prepare supports page_size=64 "
|
||||
"and index_head_dim=128 only; falling back to per-segment GetK/GetS. "
|
||||
"page_size=%s index_head_dim=%s",
|
||||
page_size,
|
||||
index_head_dim,
|
||||
limit=4,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
return kernel(
|
||||
index_buffer,
|
||||
_contiguous_for_tai(block_tables),
|
||||
_contiguous_for_tai(batch_indices),
|
||||
_contiguous_for_tai(kv_lens),
|
||||
_contiguous_for_tai(q_starts),
|
||||
_contiguous_for_tai(q_lens),
|
||||
_contiguous_for_tai(k_bases),
|
||||
_contiguous_for_tai(q_bases),
|
||||
total_kv_len=int(total_kv_len),
|
||||
total_q_count=int(total_q_count),
|
||||
max_kv_len=int(max_kv_len),
|
||||
max_q_len=int(max_q_len),
|
||||
page_size=int(page_size),
|
||||
index_head_dim=int(index_head_dim),
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_index_mqa_prepare_fallback(
|
||||
"batch_kernel_failed",
|
||||
"CP shared KV TAI batched index MQA prepare failed; falling back to "
|
||||
"per-segment GetK/GetS. error=%s segments=%s total_kv_len=%s "
|
||||
"total_q_count=%s",
|
||||
exc,
|
||||
int(batch_indices.numel()),
|
||||
total_kv_len,
|
||||
total_q_count,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def try_tai_prepare_cp_mqa_range(
|
||||
*,
|
||||
valid_q_count: int,
|
||||
@@ -1222,6 +1450,64 @@ def fill_current_kv_page_slots_and_remap_locs(
|
||||
return dense_kv_cache, mixed_locs, current_mask
|
||||
|
||||
|
||||
def _try_tai_fill_current_index_page_slots(
|
||||
*,
|
||||
dense_page_buffer: torch.Tensor,
|
||||
current_index_k: torch.Tensor,
|
||||
current_index_scale: torch.Tensor,
|
||||
current_locs: torch.Tensor,
|
||||
page_inverse: torch.Tensor,
|
||||
page_size: int,
|
||||
index_head_dim: int,
|
||||
) -> torch.Tensor | None:
|
||||
if not _tai_materialize_runtime_enabled():
|
||||
_log_tai_materialize_runtime_disabled("fill_current_index_page_slots")
|
||||
return None
|
||||
|
||||
kernels = _load_tai_materialize_kernels()
|
||||
if kernels is None:
|
||||
return None
|
||||
fill_kernel = getattr(kernels, "fill_current_index_page_slots", None)
|
||||
if fill_kernel is None:
|
||||
_log_tai_materialize_fallback(
|
||||
"index_fill_current_missing",
|
||||
"CP shared KV tai index current-slot fill kernel is unavailable; "
|
||||
"falling back to torch reference. Upgrade tai-kernel to keep this "
|
||||
"hot path off PyTorch. page_size=%s index_head_dim=%s current_rows=%s "
|
||||
"dense_pages=%s",
|
||||
page_size,
|
||||
index_head_dim,
|
||||
int(current_locs.numel()),
|
||||
int(dense_page_buffer.shape[0]),
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
return fill_kernel(
|
||||
_contiguous_for_tai(dense_page_buffer),
|
||||
_contiguous_for_tai(current_index_k),
|
||||
_contiguous_for_tai(current_index_scale),
|
||||
_contiguous_for_tai(current_locs.reshape(-1)),
|
||||
_contiguous_for_tai(page_inverse),
|
||||
page_size=int(page_size),
|
||||
index_head_dim=int(index_head_dim),
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"index_fill_current_failed",
|
||||
"CP shared KV tai index current-slot fill failed; falling back to "
|
||||
"torch reference. error=%s page_size=%s index_head_dim=%s "
|
||||
"current_rows=%s dense_pages=%s",
|
||||
exc,
|
||||
page_size,
|
||||
index_head_dim,
|
||||
int(current_locs.numel()),
|
||||
int(dense_page_buffer.shape[0]),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def fill_current_index_page_slots(
|
||||
*,
|
||||
dense_page_buffer: torch.Tensor,
|
||||
@@ -1244,6 +1530,17 @@ def fill_current_index_page_slots(
|
||||
current_rows = int(current_locs.numel())
|
||||
if current_rows == 0:
|
||||
return dense_page_buffer
|
||||
tai_result = _try_tai_fill_current_index_page_slots(
|
||||
dense_page_buffer=dense_page_buffer,
|
||||
current_index_k=current_index_k,
|
||||
current_index_scale=current_index_scale,
|
||||
current_locs=current_locs,
|
||||
page_inverse=page_inverse,
|
||||
page_size=page_size,
|
||||
index_head_dim=index_head_dim,
|
||||
)
|
||||
if tai_result is not None:
|
||||
return tai_result
|
||||
if dense_page_buffer.is_cuda:
|
||||
_log_tai_materialize_fallback(
|
||||
"index_current_fill_torch_reference_cuda",
|
||||
|
||||
@@ -32,7 +32,9 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
should_reuse_current_extend_kv,
|
||||
tensor_debug_checksum,
|
||||
tensor_debug_summary,
|
||||
try_tai_prepare_cp_mqa_current_index_batch,
|
||||
try_tai_prepare_cp_mqa_index,
|
||||
try_tai_prepare_cp_mqa_index_batch,
|
||||
try_tai_prepare_cp_mqa_range,
|
||||
)
|
||||
from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
|
||||
@@ -1139,8 +1141,15 @@ class Indexer(MultiPlatformOp):
|
||||
actual_seq_q_list = []
|
||||
batch_idx_list = []
|
||||
|
||||
if current_index_kv is not None and cp_index is not None:
|
||||
current_index_kv = None
|
||||
if (
|
||||
current_index_kv is not None
|
||||
and cp_index is not None
|
||||
and (shared_index_buffer is not None or shared_block_tables is not None)
|
||||
):
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_cp_index_ambiguous_current_and_shared_buffer"
|
||||
)
|
||||
if current_index_kv is None:
|
||||
if shared_index_buffer is not None or shared_block_tables is not None:
|
||||
if shared_index_buffer is None or shared_block_tables is None:
|
||||
@@ -1177,32 +1186,234 @@ class Indexer(MultiPlatformOp):
|
||||
and forward_batch.extend_seq_lens_cpu is not None
|
||||
)
|
||||
if cp_index is not None:
|
||||
# TODO Multi-batch support has accuracy issues
|
||||
for batch_idx, start_seq_position, end_seq_position in cp_index:
|
||||
current_req_offsets: Optional[List[int]] = None
|
||||
if current_index_kv is not None:
|
||||
current_req_offsets = []
|
||||
current_cursor = 0
|
||||
for extend_len in forward_batch.extend_seq_lens_cpu:
|
||||
current_req_offsets.append(current_cursor)
|
||||
current_cursor += int(extend_len)
|
||||
|
||||
segment_records: List[Tuple[int, int, int, int, int, int, int]] = []
|
||||
batch_idx_list = []
|
||||
kv_lens_list = []
|
||||
q_starts_list = []
|
||||
q_lens_list = []
|
||||
k_bases_list = []
|
||||
q_bases_list = []
|
||||
k_cursor = 0
|
||||
q_cursor = 0
|
||||
for raw_batch_idx, start_seq_position, end_seq_position in cp_index:
|
||||
batch_idx = int(raw_batch_idx)
|
||||
pre_chunk_offset = (
|
||||
forward_batch.seq_lens_cpu[batch_idx].item()
|
||||
- forward_batch.extend_seq_lens_cpu[batch_idx]
|
||||
)
|
||||
start_seq_position += pre_chunk_offset
|
||||
end_seq_position += pre_chunk_offset
|
||||
if offset == 0 and batch_idx != 0:
|
||||
offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1]
|
||||
k_fp8 = index_buf_accessor.GetK.execute(
|
||||
forward_batch.token_to_kv_pool,
|
||||
index_buffer,
|
||||
seq_len=end_seq_position,
|
||||
page_indices=block_tables[batch_idx],
|
||||
)
|
||||
k_scale = index_buf_accessor.GetS.execute(
|
||||
forward_batch.token_to_kv_pool,
|
||||
index_buffer,
|
||||
seq_len=end_seq_position,
|
||||
page_indices=block_tables[batch_idx],
|
||||
if end_seq_position < start_seq_position:
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_cp_index_bad_segment "
|
||||
f"batch_idx={batch_idx} start={start_seq_position} "
|
||||
f"end={end_seq_position}"
|
||||
)
|
||||
extend_seq_len = int(end_seq_position - start_seq_position)
|
||||
kv_len_i = int(end_seq_position)
|
||||
segment_records.append(
|
||||
(
|
||||
batch_idx,
|
||||
int(start_seq_position),
|
||||
int(end_seq_position),
|
||||
extend_seq_len,
|
||||
kv_len_i,
|
||||
k_cursor,
|
||||
q_cursor,
|
||||
)
|
||||
)
|
||||
batch_idx_list.append(batch_idx)
|
||||
kv_lens_list.append(kv_len_i)
|
||||
q_starts_list.append(int(start_seq_position))
|
||||
q_lens_list.append(extend_seq_len)
|
||||
k_bases_list.append(k_cursor)
|
||||
q_bases_list.append(q_cursor)
|
||||
k_cursor += kv_len_i
|
||||
q_cursor += extend_seq_len
|
||||
|
||||
if current_index_kv is None:
|
||||
assert index_buffer is not None
|
||||
assert block_tables is not None
|
||||
descriptor_device = q_fp8.device
|
||||
tai_batch_prepared = try_tai_prepare_cp_mqa_index_batch(
|
||||
index_buffer=index_buffer,
|
||||
block_tables=block_tables,
|
||||
batch_indices=torch.tensor(
|
||||
batch_idx_list, dtype=torch.int64, device=descriptor_device
|
||||
),
|
||||
kv_lens=torch.tensor(
|
||||
kv_lens_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_starts=torch.tensor(
|
||||
q_starts_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_lens=torch.tensor(
|
||||
q_lens_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
k_bases=torch.tensor(
|
||||
k_bases_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_bases=torch.tensor(
|
||||
q_bases_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
total_kv_len=k_cursor,
|
||||
total_q_count=q_cursor,
|
||||
max_kv_len=max(kv_lens_list, default=0),
|
||||
max_q_len=max(q_lens_list, default=0),
|
||||
page_size=page_size,
|
||||
index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
|
||||
)
|
||||
if tai_batch_prepared is not None:
|
||||
k_fp8_u8, k_scale, ks, ke_offset = tai_batch_prepared
|
||||
k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
|
||||
kv_fp8 = (k_fp8, k_scale)
|
||||
actual_seq_q = torch.tensor(
|
||||
q_lens_list, dtype=torch.int32, device=q_fp8.device
|
||||
)
|
||||
ke = ks + ke_offset
|
||||
with self._with_real_sm_count():
|
||||
logits = deep_gemm.fp8_mqa_logits(
|
||||
q_fp8,
|
||||
kv_fp8,
|
||||
weights,
|
||||
ks,
|
||||
ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
topk_result = metadata.topk_transform(
|
||||
logits,
|
||||
self.index_topk,
|
||||
ks=ks,
|
||||
cu_seqlens_q=actual_seq_q,
|
||||
ke_offset=ke_offset,
|
||||
batch_idx_list=batch_idx_list,
|
||||
)
|
||||
return topk_result
|
||||
else:
|
||||
assert current_req_offsets is not None
|
||||
descriptor_device = q_fp8.device
|
||||
current_bases_list = [
|
||||
int(current_req_offsets[batch_idx]) for batch_idx in batch_idx_list
|
||||
]
|
||||
current_index_head_dim = getattr(
|
||||
forward_batch.token_to_kv_pool,
|
||||
"index_head_dim",
|
||||
int(current_index_kv[0].reshape(current_index_kv[0].shape[0], -1).shape[1]),
|
||||
)
|
||||
tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch(
|
||||
current_index_k=current_index_kv[0],
|
||||
current_index_scale=current_index_kv[1],
|
||||
current_bases=torch.tensor(
|
||||
current_bases_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
kv_lens=torch.tensor(
|
||||
kv_lens_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_starts=torch.tensor(
|
||||
q_starts_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_lens=torch.tensor(
|
||||
q_lens_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
k_bases=torch.tensor(
|
||||
k_bases_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
q_bases=torch.tensor(
|
||||
q_bases_list, dtype=torch.int32, device=descriptor_device
|
||||
),
|
||||
total_kv_len=k_cursor,
|
||||
total_q_count=q_cursor,
|
||||
max_kv_len=max(kv_lens_list, default=0),
|
||||
max_q_len=max(q_lens_list, default=0),
|
||||
index_head_dim=current_index_head_dim,
|
||||
)
|
||||
if tai_current_prepared is not None:
|
||||
k_fp8_u8, k_scale, ks, ke_offset = tai_current_prepared
|
||||
k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
|
||||
kv_fp8 = (k_fp8, k_scale)
|
||||
actual_seq_q = torch.tensor(
|
||||
q_lens_list, dtype=torch.int32, device=q_fp8.device
|
||||
)
|
||||
ke = ks + ke_offset
|
||||
with self._with_real_sm_count():
|
||||
logits = deep_gemm.fp8_mqa_logits(
|
||||
q_fp8,
|
||||
kv_fp8,
|
||||
weights,
|
||||
ks,
|
||||
ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
topk_result = metadata.topk_transform(
|
||||
logits,
|
||||
self.index_topk,
|
||||
ks=ks,
|
||||
cu_seqlens_q=actual_seq_q,
|
||||
ke_offset=ke_offset,
|
||||
batch_idx_list=batch_idx_list,
|
||||
)
|
||||
return topk_result
|
||||
|
||||
for (
|
||||
batch_idx,
|
||||
start_seq_position,
|
||||
end_seq_position,
|
||||
extend_seq_len,
|
||||
_kv_len_i,
|
||||
segment_k_base,
|
||||
_segment_q_base,
|
||||
) in segment_records:
|
||||
if current_index_kv is None:
|
||||
k_fp8 = index_buf_accessor.GetK.execute(
|
||||
forward_batch.token_to_kv_pool,
|
||||
index_buffer,
|
||||
seq_len=end_seq_position,
|
||||
page_indices=block_tables[batch_idx],
|
||||
)
|
||||
k_scale = index_buf_accessor.GetS.execute(
|
||||
forward_batch.token_to_kv_pool,
|
||||
index_buffer,
|
||||
seq_len=end_seq_position,
|
||||
page_indices=block_tables[batch_idx],
|
||||
)
|
||||
else:
|
||||
if pre_chunk_offset != 0:
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_cp_index_current_with_prefix "
|
||||
f"batch_idx={batch_idx} prefix={pre_chunk_offset}"
|
||||
)
|
||||
assert current_req_offsets is not None
|
||||
req_base = current_req_offsets[batch_idx]
|
||||
req_end = req_base + int(end_seq_position)
|
||||
if (
|
||||
req_end > int(current_index_kv[0].shape[0])
|
||||
or req_end > int(current_index_kv[1].shape[0])
|
||||
):
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_cp_index_current_rows_short "
|
||||
f"batch_idx={batch_idx} req_end={req_end} "
|
||||
f"k_rows={int(current_index_kv[0].shape[0])} "
|
||||
f"scale_rows={int(current_index_kv[1].shape[0])}"
|
||||
)
|
||||
k_fp8 = current_index_kv[0][req_base:req_end].contiguous()
|
||||
k_scale = current_index_kv[1][req_base:req_end].contiguous()
|
||||
|
||||
extend_seq_len = end_seq_position - start_seq_position
|
||||
ks = torch.full(
|
||||
(extend_seq_len,), offset, dtype=torch.int32, device="cuda"
|
||||
(extend_seq_len,),
|
||||
segment_k_base,
|
||||
dtype=torch.int32,
|
||||
device=q_fp8.device,
|
||||
)
|
||||
k_fp8_list.append(k_fp8)
|
||||
k_scale_list.append(k_scale)
|
||||
@@ -1211,14 +1422,13 @@ class Indexer(MultiPlatformOp):
|
||||
start_seq_position + 1,
|
||||
end_seq_position + 1,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
device=q_fp8.device,
|
||||
)
|
||||
ke_offset_list.append(ke_offset)
|
||||
actual_seq_q = torch.tensor(
|
||||
[extend_seq_len], dtype=torch.int32, device="cuda"
|
||||
[extend_seq_len], dtype=torch.int32, device=q_fp8.device
|
||||
)
|
||||
actual_seq_q_list.append(actual_seq_q)
|
||||
batch_idx_list.append(batch_idx)
|
||||
|
||||
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
|
||||
k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
|
||||
@@ -1553,68 +1763,83 @@ class Indexer(MultiPlatformOp):
|
||||
)
|
||||
)
|
||||
|
||||
outputs = []
|
||||
cursor = 0
|
||||
output_cursor = 0
|
||||
cp_index: List[Tuple[int, int, int]] = []
|
||||
compact_q_chunks = []
|
||||
compact_weight_chunks = []
|
||||
compact_output_spans: List[Tuple[int, int]] = []
|
||||
current_only = (
|
||||
current_index_kv_for_topk is not None
|
||||
and is_current_only_extend_batch(forward_batch)
|
||||
)
|
||||
page_table_1 = None if current_only else metadata.get_page_table_1()
|
||||
|
||||
def call_segment(
|
||||
def collect_segment(
|
||||
*,
|
||||
req_id: int,
|
||||
segment_len: int,
|
||||
kv_len: int,
|
||||
) -> torch.Tensor:
|
||||
nonlocal cursor
|
||||
) -> None:
|
||||
nonlocal cursor, output_cursor
|
||||
segment_len = int(segment_len)
|
||||
kv_len = int(kv_len)
|
||||
q_segment = q_fp8[cursor : cursor + segment_len]
|
||||
weights_segment = weights[cursor : cursor + segment_len]
|
||||
cursor += segment_len
|
||||
if segment_len == 0:
|
||||
return torch.empty(
|
||||
(0, self.index_topk),
|
||||
dtype=torch.int32,
|
||||
device=q_fp8.device,
|
||||
return
|
||||
|
||||
seq_len = int(forward_batch.seq_lens_cpu[req_id].item())
|
||||
extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id])
|
||||
cp_kv_end = seq_len - extend_seq_len + kv_len
|
||||
if current_only:
|
||||
logical_kv_limit = seq_len
|
||||
else:
|
||||
if page_table_1 is None:
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_gt1_missing_page_table_1"
|
||||
)
|
||||
if page_table_1.dim() == 1:
|
||||
logical_kv_limit = min(seq_len, int(page_table_1.numel()))
|
||||
else:
|
||||
logical_kv_limit = min(seq_len, int(page_table_1.shape[1]))
|
||||
valid_q_count = _compute_contiguous_valid_cp_query_count(
|
||||
cp_kv_end=cp_kv_end,
|
||||
actual_seq_q=segment_len,
|
||||
logical_kv_limit=logical_kv_limit,
|
||||
)
|
||||
if valid_q_count <= 0:
|
||||
return
|
||||
|
||||
start_abs = cp_kv_end - segment_len
|
||||
end_abs = start_abs + valid_q_count
|
||||
pre_chunk_offset = seq_len - extend_seq_len
|
||||
cp_index.append(
|
||||
(
|
||||
req_id,
|
||||
start_abs - pre_chunk_offset,
|
||||
end_abs - pre_chunk_offset,
|
||||
)
|
||||
actual_seq_q_tensor = torch.tensor(
|
||||
[segment_len],
|
||||
dtype=torch.int32,
|
||||
device=q_fp8.device,
|
||||
)
|
||||
actual_seq_q_cu_tensor = torch.tensor(
|
||||
[0, segment_len],
|
||||
dtype=torch.int32,
|
||||
device=q_fp8.device,
|
||||
)
|
||||
return self._get_topk_ragged_with_cp(
|
||||
forward_batch,
|
||||
layer_id,
|
||||
q_segment,
|
||||
weights_segment,
|
||||
metadata,
|
||||
kv_len,
|
||||
segment_len,
|
||||
current_index_kv=current_index_kv_for_topk,
|
||||
shared_index_buffer=shared_index_buffer,
|
||||
shared_block_tables=shared_block_tables,
|
||||
actual_seq_q_tensor=actual_seq_q_tensor,
|
||||
actual_seq_q_cu_tensor=actual_seq_q_cu_tensor,
|
||||
batch_idx=req_id,
|
||||
)
|
||||
compact_q_chunks.append(q_segment[:valid_q_count])
|
||||
compact_weight_chunks.append(weights_segment[:valid_q_count])
|
||||
compact_output_spans.append((output_cursor, valid_q_count))
|
||||
|
||||
for req_id in range(batch_size):
|
||||
outputs.append(
|
||||
call_segment(
|
||||
req_id=req_id,
|
||||
segment_len=request_actual_seq_q_prev[req_id],
|
||||
kv_len=request_kv_len_prev[req_id],
|
||||
)
|
||||
collect_segment(
|
||||
req_id=req_id,
|
||||
segment_len=request_actual_seq_q_prev[req_id],
|
||||
kv_len=request_kv_len_prev[req_id],
|
||||
)
|
||||
outputs.append(
|
||||
call_segment(
|
||||
req_id=req_id,
|
||||
segment_len=request_actual_seq_q_next[req_id],
|
||||
kv_len=request_kv_len_next[req_id],
|
||||
)
|
||||
output_cursor += int(request_actual_seq_q_prev[req_id])
|
||||
collect_segment(
|
||||
req_id=req_id,
|
||||
segment_len=request_actual_seq_q_next[req_id],
|
||||
kv_len=request_kv_len_next[req_id],
|
||||
)
|
||||
output_cursor += int(request_actual_seq_q_next[req_id])
|
||||
|
||||
if cursor != int(q_fp8.shape[0]) or cursor != int(weights.shape[0]):
|
||||
raise RuntimeError(
|
||||
@@ -1624,9 +1849,46 @@ class Indexer(MultiPlatformOp):
|
||||
f"q_tokens={int(q_fp8.shape[0])} weights_tokens={int(weights.shape[0])}"
|
||||
)
|
||||
|
||||
if not outputs:
|
||||
return torch.empty((0, self.index_topk), dtype=torch.int32, device=q_fp8.device)
|
||||
return torch.cat(outputs, dim=0)
|
||||
result = torch.full(
|
||||
(output_cursor, self.index_topk),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
device=q_fp8.device,
|
||||
)
|
||||
if not compact_q_chunks:
|
||||
return result
|
||||
|
||||
compact_q = torch.cat(compact_q_chunks, dim=0)
|
||||
compact_weights = torch.cat(compact_weight_chunks, dim=0)
|
||||
compact_rows = int(compact_q.shape[0])
|
||||
compact_topk = self._get_topk_ragged_with_cp(
|
||||
forward_batch,
|
||||
layer_id,
|
||||
compact_q,
|
||||
compact_weights,
|
||||
metadata,
|
||||
0,
|
||||
compact_rows,
|
||||
cp_index=cp_index,
|
||||
current_index_kv=current_index_kv_for_topk,
|
||||
shared_index_buffer=shared_index_buffer,
|
||||
shared_block_tables=shared_block_tables,
|
||||
actual_seq_q_tensor=None,
|
||||
actual_seq_q_cu_tensor=None,
|
||||
)
|
||||
if int(compact_topk.shape[0]) != compact_rows:
|
||||
raise RuntimeError(
|
||||
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
|
||||
"reason=batch_gt1_compact_topk_rows_mismatch "
|
||||
f"expected={compact_rows} got={int(compact_topk.shape[0])}"
|
||||
)
|
||||
compact_cursor = 0
|
||||
for output_start, valid_q_count in compact_output_spans:
|
||||
result[output_start : output_start + valid_q_count] = compact_topk[
|
||||
compact_cursor : compact_cursor + valid_q_count
|
||||
]
|
||||
compact_cursor += valid_q_count
|
||||
return result
|
||||
|
||||
def forward_indexer(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user