laoyao0822
c3fc3ff752
Cache-hit EAGLE/NextN draft extends can enter the draft DeepEP MoE immediately after CP shared-KV attention. The partial current-reuse path is kept for target layers, but draft cache-hit suffixes now use full materialization until draft has an explicit same-layer reuse contract. Next-layer MLA/index prefetch is also gated by the actual model depth, so the single-layer draft model does not enqueue unused next-layer async work. The temporary stage traces used to isolate the hang are removed. The retained draft current-reuse fallback is a bounded warning because it changes the runtime path intentionally. Constraint: EAGLE/NextN has one executable draft layer and mirrors target KV state. Rejected: Keep partial current reuse for draft cache-hit suffixes | reproduced hangs at draft layer0 before DeepEP MoE completion. Rejected: Keep temporary stage traces | useful for diagnosis but too noisy for normal runs. Confidence: medium Scope-risk: moderate Directive: Do not re-enable draft cache-hit partial current reuse without an explicit draft same-layer reuse contract and ETE validation with CP shared KV + HiCache + EAGLE. Tested: py_compile on edited Python files; git diff --check; temp trace grep returned no matches. Not-tested: Local targeted pytest is blocked by missing pybase64 in this environment; full ETE after log cleanup not run.
2060 lines
78 KiB
Python
2060 lines
78 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
|
|
from sglang.jit_kernel.fused_store_index_cache import (
|
|
can_use_nsa_fused_store,
|
|
fused_store_index_k_cache,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.attention.nsa import index_buf_accessor
|
|
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
|
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,
|
|
cp_shared_kv_should_prefetch_next_layer,
|
|
filter_owned_logical_locs,
|
|
get_or_build_shared_paged_buffer_slot_remap,
|
|
is_current_only_extend_batch,
|
|
log_cp_draft_shared_kv_debug,
|
|
materialize_shared_paged_buffer,
|
|
tensor_debug_checksum,
|
|
tensor_debug_summary,
|
|
try_tai_prepare_cp_mqa_index,
|
|
try_tai_prepare_cp_mqa_range,
|
|
)
|
|
from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
|
|
from sglang.srt.layers.layernorm import LayerNorm
|
|
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
|
|
from sglang.srt.layers.utils import MultiPlatformOp
|
|
from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
|
|
|
|
global _use_multi_stream
|
|
_is_cuda = is_cuda()
|
|
_is_hip = is_hip()
|
|
_is_npu = is_npu()
|
|
_is_fp8_fnuz = is_fp8_fnuz()
|
|
logger = logging.getLogger(__name__)
|
|
if _is_cuda:
|
|
try:
|
|
import deep_gemm
|
|
except ImportError as e:
|
|
deep_gemm = e
|
|
|
|
if is_npu():
|
|
import torch_npu
|
|
from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream
|
|
|
|
from sglang.srt.distributed import (
|
|
get_attn_context_model_parallel_rank,
|
|
get_attn_context_model_parallel_world_size,
|
|
)
|
|
from sglang.srt.distributed.parallel_state import get_pp_group
|
|
from sglang.srt.layers import deep_gemm_wrapper
|
|
from sglang.srt.layers.attention.nsa.utils import (
|
|
cp_all_gather_rerange_output,
|
|
get_cp_shared_kv_local_out_cache_loc,
|
|
get_cp_shared_kv_local_physical_out_cache_loc,
|
|
is_nsa_enable_prefill_cp,
|
|
is_nsa_prefill_cp_in_seq_split,
|
|
log_cp_shared_kv_direct_write_fallback,
|
|
nsa_use_prefill_cp,
|
|
split_in_seq_cp_local_pair,
|
|
)
|
|
from sglang.srt.layers.communicator import ScatterMode
|
|
from sglang.srt.layers.linear import ReplicatedLinear
|
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
|
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
|
|
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|
from sglang.srt.server_args import get_global_server_args
|
|
|
|
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
|
|
|
|
|
DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0
|
|
|
|
|
|
def _compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end: int,
|
|
actual_seq_q: int,
|
|
logical_kv_limit: int,
|
|
) -> int:
|
|
"""Return the visible prefix length for in-seq CP query rows.
|
|
|
|
`_get_topk_ragged_with_cp` constructs monotonically increasing `ke_offset`
|
|
values for one CP segment:
|
|
|
|
[cp_kv_end - actual_seq_q + 1, ..., cp_kv_end]
|
|
|
|
The validity predicate is `ke_offset <= logical_kv_limit`, so valid rows are
|
|
always a contiguous prefix. Compute the prefix length with CPU metadata that
|
|
already exists instead of building a CUDA mask and synchronizing it back via
|
|
`.item()`.
|
|
"""
|
|
|
|
actual_seq_q = int(actual_seq_q)
|
|
if actual_seq_q <= 0:
|
|
return 0
|
|
|
|
ke_start = int(cp_kv_end) - actual_seq_q + 1
|
|
valid_count = int(logical_kv_limit) - ke_start + 1
|
|
return max(0, min(actual_seq_q, valid_count))
|
|
|
|
|
|
def _log_cp_shared_kv_index_prefetch_fallback(
|
|
reason: str,
|
|
message: str,
|
|
*args,
|
|
) -> None:
|
|
"""Log every index prefetch fallback event.
|
|
|
|
Warmup and real requests can hit the same reason independently; do not
|
|
dedupe here or later real fallbacks become invisible during profiling.
|
|
"""
|
|
|
|
logger.warning(
|
|
"[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message,
|
|
reason,
|
|
*args,
|
|
)
|
|
|
|
|
|
class BaseIndexerMetadata(ABC):
|
|
@abstractmethod
|
|
def get_seqlens_int32(self) -> torch.Tensor:
|
|
"""
|
|
Return: (batch_size,) int32 tensor
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_page_table_64(self) -> torch.Tensor:
|
|
"""
|
|
Return: (batch_size, num_blocks) int32, page table.
|
|
The page size of the table is 64.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_page_table_1(self) -> torch.Tensor:
|
|
"""
|
|
Return: (batch_size, num_blocks) int32, page table.
|
|
The page size of the table is 1.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_seqlens_expanded(self) -> torch.Tensor:
|
|
"""
|
|
Return: (sum_extend_seq_len,) int32 tensor
|
|
"""
|
|
|
|
def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Return: (tokens, ), (tokens, ) int32, k_start and k_end in kv cache(token,xxx) for each token.
|
|
"""
|
|
|
|
def get_indexer_seq_len_cpu(self) -> torch.Tensor:
|
|
"""
|
|
Return: seq lens for each batch.
|
|
"""
|
|
|
|
def get_indexer_seq_len(self) -> torch.Tensor:
|
|
"""
|
|
Return: seq lens for each batch.
|
|
"""
|
|
|
|
def get_nsa_extend_len_cpu(self) -> List[int]:
|
|
"""
|
|
Return: extend seq lens for each batch.
|
|
"""
|
|
|
|
def get_token_to_batch_idx(self) -> torch.Tensor:
|
|
"""
|
|
Return: batch idx for each token.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def topk_transform(
|
|
self,
|
|
logits: torch.Tensor,
|
|
topk: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Perform topk selection on the logits and possibly transform the result.
|
|
|
|
NOTE that attention backend may override this function to do some
|
|
transformation, which means the result of this topk_transform may not
|
|
be the topk indices of the input logits.
|
|
|
|
Return: Anything, since it will be passed to the attention backend
|
|
for further processing on sparse attention computation.
|
|
Don't assume it is the topk indices of the input logits.
|
|
"""
|
|
|
|
|
|
def rotate_activation(x: torch.Tensor) -> torch.Tensor:
|
|
assert x.dtype == torch.bfloat16
|
|
# from sgl_kernel import hadamard_transform
|
|
if _is_hip:
|
|
from fast_hadamard_transform import hadamard_transform
|
|
else:
|
|
from sglang.jit_kernel.hadamard import hadamard_transform
|
|
|
|
hidden_size = x.size(-1)
|
|
assert (
|
|
hidden_size & (hidden_size - 1)
|
|
) == 0, "Hidden size must be a power of 2 for Hadamard transform."
|
|
return hadamard_transform(x, scale=hidden_size**-0.5)
|
|
|
|
|
|
class Indexer(MultiPlatformOp):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
index_n_heads: int,
|
|
index_head_dim: int,
|
|
rope_head_dim: int,
|
|
index_topk: int,
|
|
q_lora_rank: int,
|
|
max_position_embeddings: int,
|
|
rope_theta: float,
|
|
layer_id: int,
|
|
scale_fmt: Optional[str],
|
|
block_size: int = 128,
|
|
rope_scaling: Optional[Dict[str, Any]] = None,
|
|
is_neox_style: bool = True,
|
|
prefix: str = "",
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
alt_stream: Optional[torch.cuda.Stream] = None,
|
|
):
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.n_heads = index_n_heads
|
|
self.head_dim = index_head_dim
|
|
self.rope_head_dim = rope_head_dim
|
|
self.index_topk = index_topk
|
|
self.q_lora_rank = q_lora_rank
|
|
self.layer_id = layer_id
|
|
self.alt_stream = alt_stream
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
if self.nsa_enable_prefill_cp:
|
|
self.cp_size = get_attn_context_model_parallel_world_size()
|
|
self.cp_rank = get_attn_context_model_parallel_rank()
|
|
else:
|
|
self.cp_size = None
|
|
self.cp_rank = None
|
|
if _is_cuda:
|
|
self.sm_count = deep_gemm.get_num_sms()
|
|
self.half_device_sm_count = ceil_align(self.sm_count // 2, 8)
|
|
pp_size = get_global_server_args().pp_size
|
|
self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank
|
|
else:
|
|
self.logits_with_pp_recv = False
|
|
|
|
self.wq_b = ReplicatedLinear(
|
|
self.q_lora_rank,
|
|
self.n_heads * self.head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("wq_b", prefix),
|
|
)
|
|
|
|
self.wk = ReplicatedLinear(
|
|
self.hidden_size,
|
|
self.head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("wk", prefix),
|
|
)
|
|
self.weights_proj = ReplicatedLinear(
|
|
self.hidden_size,
|
|
self.n_heads,
|
|
bias=False,
|
|
params_dtype=torch.bfloat16 if _is_cuda else torch.float32,
|
|
prefix=add_prefix("weights_proj", prefix),
|
|
)
|
|
self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
|
|
self.rotary_emb = get_rope_wrapper(
|
|
rope_head_dim,
|
|
rotary_dim=rope_head_dim,
|
|
max_position=max_position_embeddings,
|
|
base=rope_theta, # type: ignore
|
|
rope_scaling=rope_scaling,
|
|
is_neox_style=is_neox_style,
|
|
device=get_global_server_args().device,
|
|
)
|
|
self.block_size = block_size
|
|
self.scale_fmt = scale_fmt
|
|
self.softmax_scale = self.head_dim**-0.5
|
|
|
|
def _maybe_materialize_shared_index_buffer(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
logical_page_table: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
index_buffer = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
|
|
layer_id=layer_id
|
|
)
|
|
if not forward_batch.uses_cp_shared_kv:
|
|
return index_buffer, logical_page_table
|
|
|
|
assert forward_batch.cp_shared_kv_layout is not None
|
|
layout = forward_batch.cp_shared_kv_layout
|
|
index_prefetcher = getattr(
|
|
forward_batch, "cp_shared_kv_index_prefetcher", None
|
|
)
|
|
if index_prefetcher is not None:
|
|
prefetched = index_prefetcher.consume(
|
|
layer_id=layer_id,
|
|
page_buffer=index_buffer,
|
|
logical_pages=logical_page_table,
|
|
)
|
|
if prefetched is not None:
|
|
if cp_shared_kv_debug_enabled():
|
|
cp_shared_kv_debug_log(
|
|
"index_materialize_prefetch_hit",
|
|
"NSA index materialize prefetch hit cp_rank=%s layer=%s dense_pages=%s buffer_ck=%s",
|
|
layout.cp_rank,
|
|
layer_id,
|
|
tensor_debug_summary(prefetched[1]),
|
|
tensor_debug_checksum(prefetched[0]),
|
|
)
|
|
return prefetched
|
|
start_layer = int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0))
|
|
if int(layer_id) > start_layer:
|
|
_log_cp_shared_kv_index_prefetch_fallback(
|
|
"consume_miss",
|
|
"prefetcher did not provide layer buffer; falling back to "
|
|
"sync paged materialize. layer=%s start_layer=%s cp_rank=%s "
|
|
"logical_page_table_shape=%s",
|
|
layer_id,
|
|
start_layer,
|
|
layout.cp_rank,
|
|
tuple(logical_page_table.shape),
|
|
)
|
|
|
|
if cp_shared_kv_debug_enabled():
|
|
cp_shared_kv_debug_log(
|
|
"index_materialize_call",
|
|
"NSA index materialize call cp_rank=%s layer=%s logical_page_table=%s",
|
|
layout.cp_rank,
|
|
layer_id,
|
|
tensor_debug_summary(logical_page_table),
|
|
)
|
|
slot_remap = get_or_build_shared_paged_buffer_slot_remap(
|
|
forward_batch,
|
|
page_buffer=index_buffer,
|
|
logical_pages=logical_page_table,
|
|
layout=layout,
|
|
)
|
|
materialized, dense_pages = materialize_shared_paged_buffer(
|
|
page_buffer=index_buffer,
|
|
logical_pages=logical_page_table,
|
|
layout=layout,
|
|
slot_remap=slot_remap,
|
|
nvtx_source="index.full_materialize",
|
|
nvtx_layer_id=layer_id,
|
|
)
|
|
if cp_shared_kv_debug_enabled():
|
|
cp_shared_kv_debug_log(
|
|
"index_materialize_result",
|
|
"NSA index materialize result cp_rank=%s layer=%s dense_pages=%s buffer_ck=%s",
|
|
layout.cp_rank,
|
|
layer_id,
|
|
tensor_debug_summary(dense_pages),
|
|
tensor_debug_checksum(materialized),
|
|
)
|
|
return materialized, dense_pages
|
|
|
|
def _maybe_start_next_layer_index_prefetch(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
) -> None:
|
|
next_layer_id = layer_id + 1
|
|
index_prefetcher = getattr(
|
|
forward_batch, "cp_shared_kv_index_prefetcher", None
|
|
)
|
|
if (
|
|
cp_shared_kv_mla_prefetch_log_enabled()
|
|
and cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id)
|
|
):
|
|
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
|
|
seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
|
|
cp_shared_kv_mla_prefetch_log(
|
|
"index_start_request cp_rank=%s cp_size=%s layer=%s "
|
|
"next_layer=%s has_index=%s uses_cp_shared_kv=%s "
|
|
"seq_lens_cpu_len=%s",
|
|
getattr(layout, "cp_rank", None),
|
|
getattr(layout, "cp_size", None),
|
|
layer_id,
|
|
next_layer_id,
|
|
index_prefetcher is not None,
|
|
getattr(forward_batch, "uses_cp_shared_kv", None),
|
|
len(seq_lens_cpu) if seq_lens_cpu is not None else None,
|
|
)
|
|
if index_prefetcher is None:
|
|
return
|
|
if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id):
|
|
return
|
|
index_prefetcher.start_next_layer_prefix(
|
|
next_layer_id=next_layer_id,
|
|
token_to_kv_pool=forward_batch.token_to_kv_pool,
|
|
)
|
|
|
|
def _filter_shared_index_write(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
key: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
out_loc = forward_batch.out_cache_loc
|
|
if not forward_batch.uses_cp_shared_kv:
|
|
return out_loc, key
|
|
|
|
if key.shape[0] != out_loc.numel():
|
|
raise RuntimeError(
|
|
"CP shared KV index write shape mismatch: "
|
|
f"key_tokens={key.shape[0]} out_cache_loc={out_loc.numel()} "
|
|
f"forward_mode={forward_batch.forward_mode}"
|
|
)
|
|
|
|
assert forward_batch.cp_shared_kv_layout is not None
|
|
layout = forward_batch.cp_shared_kv_layout
|
|
owned_mask, physical_out_loc = filter_owned_logical_locs(
|
|
out_loc,
|
|
layout,
|
|
)
|
|
if cp_shared_kv_debug_enabled():
|
|
valid_mask = out_loc > 0
|
|
owned_count = int(owned_mask.sum().item())
|
|
valid_count = int(valid_mask.sum().item())
|
|
cp_shared_kv_debug_log(
|
|
"index_write_filter",
|
|
"NSA index shard write filter cp_rank=%s forward_mode=%s logical_out_loc=%s "
|
|
"valid=%s owned=%s physical_out_loc=%s key_ck=%s key_valid_ck=%s "
|
|
"key_owned_ck=%s",
|
|
layout.cp_rank,
|
|
forward_batch.forward_mode,
|
|
tensor_debug_summary(out_loc),
|
|
valid_count,
|
|
owned_count,
|
|
tensor_debug_summary(physical_out_loc),
|
|
tensor_debug_checksum(key),
|
|
tensor_debug_checksum(key[valid_mask].contiguous()),
|
|
tensor_debug_checksum(key[owned_mask].contiguous()),
|
|
)
|
|
if physical_out_loc.numel() == out_loc.numel():
|
|
return physical_out_loc, key
|
|
|
|
return physical_out_loc, key[owned_mask].contiguous()
|
|
|
|
def _can_reuse_current_index_kv(self, forward_batch: ForwardBatch) -> bool:
|
|
return (
|
|
cp_shared_kv_current_reuse_enabled()
|
|
and forward_batch.uses_cp_shared_kv
|
|
and self.nsa_enable_prefill_cp
|
|
and forward_batch.nsa_cp_metadata is not None
|
|
and is_nsa_prefill_cp_in_seq_split()
|
|
and is_current_only_extend_batch(forward_batch)
|
|
and forward_batch.hisparse_coordinator is None
|
|
and _is_cuda
|
|
and not _is_fp8_fnuz
|
|
)
|
|
|
|
@contextlib.contextmanager
|
|
def _with_real_sm_count(self):
|
|
# When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch
|
|
# request to receive the PP proxy tensor or output from the previous stage, occupying one SM resource.
|
|
# Model execution runs in parallel with the recv operation, so the SMs available to the indexer must be reduced
|
|
# by 1. Currently, the last rank starts the send result + recv request only after waiting for execution results.
|
|
if self.logits_with_pp_recv:
|
|
pp_recv_sm_count = 1
|
|
with deep_gemm_wrapper.configure_deep_gemm_num_sms(
|
|
self.sm_count - pp_recv_sm_count
|
|
):
|
|
yield
|
|
else:
|
|
yield
|
|
|
|
def _weights_proj_bf16_in_fp32_out(self, x: torch.Tensor) -> torch.Tensor:
|
|
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
|
weight = self.weights_proj.weight
|
|
out = torch.empty(
|
|
(x.shape[0], weight.shape[0]),
|
|
dtype=torch.float32,
|
|
device=x.device,
|
|
)
|
|
deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, weight, out)
|
|
return out
|
|
|
|
if _is_hip:
|
|
x = x.to(self.weights_proj.weight.dtype)
|
|
weights, _ = self.weights_proj(x)
|
|
return weights.float()
|
|
|
|
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
|
|
def _project_and_scale_head_gates(self, x: torch.Tensor):
|
|
weights = self._weights_proj_bf16_in_fp32_out(x)
|
|
weights = weights * self.n_heads**-0.5
|
|
return weights
|
|
|
|
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
|
|
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
|
|
weights = self._weights_proj_bf16_in_fp32_out(x)
|
|
weights = weights * self.n_heads**-0.5
|
|
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
|
|
return weights
|
|
|
|
def _get_q_k_bf16(
|
|
self,
|
|
q_lora: torch.Tensor,
|
|
x: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
enable_dual_stream: bool,
|
|
forward_batch: ForwardBatch,
|
|
):
|
|
if enable_dual_stream:
|
|
current_stream = torch.cuda.current_stream()
|
|
self.alt_stream.wait_stream(current_stream)
|
|
|
|
with deep_gemm_wrapper.configure_deep_gemm_num_sms(
|
|
self.half_device_sm_count
|
|
):
|
|
query, _ = self.wq_b(q_lora)
|
|
query = rearrange(query, "l (h d) -> l h d", d=self.head_dim)
|
|
q_rope, _ = torch.split(
|
|
query,
|
|
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
|
|
dim=-1,
|
|
)
|
|
with torch.cuda.stream(self.alt_stream):
|
|
# TODO we should also put DeepGEMM half SM here?
|
|
key, _ = self.wk(x)
|
|
key = self.k_norm(key)
|
|
|
|
k_rope, _ = torch.split(
|
|
key,
|
|
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
|
|
dim=-1,
|
|
)
|
|
|
|
current_stream.wait_stream(self.alt_stream)
|
|
else:
|
|
query, _ = self.wq_b(q_lora)
|
|
query = rearrange(query, "l (h d) -> l h d", d=self.head_dim)
|
|
q_rope, _ = torch.split(
|
|
query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
|
|
)
|
|
key, _ = self.wk(x)
|
|
key = self.k_norm(key)
|
|
k_rope, _ = torch.split(
|
|
key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
|
|
)
|
|
|
|
q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope)
|
|
|
|
query[..., : self.rope_head_dim] = q_rope.clone()
|
|
key[..., : self.rope_head_dim] = k_rope.clone()
|
|
|
|
if enable_dual_stream:
|
|
current_stream = torch.cuda.current_stream()
|
|
self.alt_stream.wait_stream(current_stream)
|
|
query = rotate_activation(query)
|
|
|
|
with torch.cuda.stream(self.alt_stream):
|
|
key = rotate_activation(key)
|
|
current_stream.wait_stream(self.alt_stream)
|
|
else:
|
|
query = rotate_activation(query)
|
|
key = rotate_activation(key)
|
|
|
|
local_key = key
|
|
# allgather+rerrange
|
|
if forward_batch.nsa_cp_metadata is not None and self.nsa_enable_prefill_cp:
|
|
key = cp_all_gather_rerange_output(
|
|
key.contiguous(),
|
|
self.cp_size,
|
|
forward_batch,
|
|
torch.cuda.current_stream(),
|
|
)
|
|
return query, key, local_key
|
|
|
|
def _get_k_bf16(
|
|
self,
|
|
x: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
enable_dual_stream: bool,
|
|
):
|
|
# Compute only key, skip query
|
|
key, _ = self.wk(x)
|
|
key = self.k_norm(key)
|
|
k_rope, _ = torch.split(
|
|
key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
|
|
)
|
|
|
|
_, k_rope = self.rotary_emb(positions, k_rope, k_rope)
|
|
key[..., : self.rope_head_dim] = k_rope.clone()
|
|
key = rotate_activation(key)
|
|
|
|
return key
|
|
|
|
def _get_topk_paged(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
q_fp8: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
metadata: BaseIndexerMetadata,
|
|
) -> torch.Tensor:
|
|
if TYPE_CHECKING:
|
|
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
|
|
|
|
page_size = forward_batch.token_to_kv_pool.page_size
|
|
# NOTE(dark): blocksize = 64 is hardcoded in deep_gemm
|
|
if _is_hip:
|
|
assert page_size == 1, "only support page size 1"
|
|
block_tables = metadata.get_page_table_1()
|
|
else:
|
|
assert page_size == 64, "only support page size 64"
|
|
# NOTE(dark): this support extend/decode/decode+graph
|
|
block_tables = metadata.get_page_table_64()
|
|
|
|
max_seq_len = block_tables.shape[1] * page_size
|
|
kv_cache_fp8, block_tables = self._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id,
|
|
block_tables,
|
|
)
|
|
|
|
blocksize = page_size
|
|
if (
|
|
forward_batch.forward_mode.is_target_verify()
|
|
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
|
|
):
|
|
seqlens_32 = metadata.get_seqlens_expanded()
|
|
else:
|
|
seqlens_32 = metadata.get_seqlens_int32()
|
|
# Reuse pre-computed schedule metadata if available (from init_forward_metadata),
|
|
# otherwise fall back to computing it here.
|
|
schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None)
|
|
if _is_cuda:
|
|
if schedule_metadata is None:
|
|
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
|
|
seqlens_32, blocksize, self.sm_count
|
|
)
|
|
|
|
assert len(q_fp8.shape) == 3
|
|
q_fp8 = q_fp8.unsqueeze(1) # the next_n dim is 1 now
|
|
assert len(kv_cache_fp8.shape) == 2
|
|
block_kv = 1 if _is_hip else 64
|
|
num_heads_kv = 1
|
|
head_dim_with_sf = 132
|
|
if _is_hip:
|
|
kv_cache_fp8 = kv_cache_fp8.view(
|
|
-1, block_kv, num_heads_kv, head_dim_with_sf
|
|
)
|
|
else:
|
|
kv_cache_fp8 = kv_cache_fp8.view(
|
|
kv_cache_fp8.shape[0], block_kv, num_heads_kv, head_dim_with_sf
|
|
)
|
|
assert len(weights.shape) == 3
|
|
weights = weights.squeeze(2)
|
|
|
|
# When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states,
|
|
# and it is necessary to extract the actual q length.
|
|
q_offset = sum(metadata.get_nsa_extend_len_cpu())
|
|
if _is_hip:
|
|
from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
|
|
|
|
batch_size, next_n, heads, _ = q_fp8.shape
|
|
logits = torch.full(
|
|
(batch_size * next_n, max_seq_len),
|
|
float("-inf"),
|
|
device=q_fp8.device,
|
|
dtype=torch.float32,
|
|
)
|
|
deepgemm_fp8_paged_mqa_logits(
|
|
q_fp8,
|
|
kv_cache_fp8,
|
|
weights,
|
|
logits,
|
|
seqlens_32,
|
|
block_tables,
|
|
max_seq_len,
|
|
Preshuffle=False,
|
|
KVBlockSize=block_kv,
|
|
ChunkK=128,
|
|
TotalCuCount=256,
|
|
WavePerEU=5,
|
|
)
|
|
else:
|
|
logits = deep_gemm.fp8_paged_mqa_logits(
|
|
q_fp8[:q_offset],
|
|
kv_cache_fp8,
|
|
weights[:q_offset],
|
|
seqlens_32,
|
|
block_tables,
|
|
schedule_metadata,
|
|
max_seq_len,
|
|
clean_logits=False,
|
|
)
|
|
|
|
# NOTE(dark): logits should be cleaned in topk_transform
|
|
topk_result = metadata.topk_transform(logits, self.index_topk)
|
|
# Restore possible padding exist in the hidden states.
|
|
if not _is_hip and q_offset < q_fp8.shape[0]:
|
|
pad_len = q_fp8.shape[0] - q_offset
|
|
padding = torch.full(
|
|
(pad_len, topk_result.shape[1]),
|
|
-1,
|
|
dtype=topk_result.dtype,
|
|
device=topk_result.device,
|
|
)
|
|
topk_result = torch.cat([topk_result, padding], dim=0)
|
|
return topk_result
|
|
|
|
def _should_chunk_mqa_logits(
|
|
self, num_q: int, num_k: int, device: torch.device
|
|
) -> Tuple[bool, int]:
|
|
"""
|
|
Detect whether we need to chunk the MQA logits computation to avoid OOM
|
|
Return: (need_chunk, free_mem)
|
|
"""
|
|
# Quick static check for normal batches
|
|
if num_q * num_k < 8_000_000: # 8M elements ≈ 32MB logits
|
|
return False, 0
|
|
|
|
free_mem, total_mem = torch.cuda.mem_get_info(device)
|
|
bytes_per_elem = 4 # float32
|
|
logits_bytes = num_q * num_k * bytes_per_elem
|
|
|
|
# Logits should not exceed 50% of free memory or 30% of total memory
|
|
need_chunk = (logits_bytes * 2 > free_mem) or (logits_bytes > total_mem * 0.3)
|
|
return need_chunk, free_mem
|
|
|
|
def _get_topk_ragged(
|
|
self,
|
|
enable_dual_stream: bool,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
q_fp8: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
metadata: BaseIndexerMetadata,
|
|
) -> torch.Tensor:
|
|
if TYPE_CHECKING:
|
|
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
|
|
|
|
assert forward_batch.forward_mode.is_extend_without_speculative()
|
|
|
|
page_size = forward_batch.token_to_kv_pool.page_size
|
|
if _is_hip:
|
|
assert page_size == 1, "only support page size 1"
|
|
else:
|
|
assert page_size == 64, "only support page size 64"
|
|
|
|
assert len(weights.shape) == 3
|
|
assert (
|
|
forward_batch.seq_lens_cpu is not None
|
|
and forward_batch.extend_seq_lens_cpu is not None
|
|
)
|
|
weights = weights.squeeze(-1)
|
|
|
|
if _is_hip:
|
|
block_tables = metadata.get_page_table_1()
|
|
else:
|
|
block_tables = metadata.get_page_table_64()
|
|
|
|
assert (
|
|
forward_batch.seq_lens_cpu is not None
|
|
and forward_batch.extend_seq_lens_cpu is not None
|
|
)
|
|
|
|
batch_size = len(block_tables)
|
|
token_nums, _, _ = q_fp8.shape
|
|
device = q_fp8.device
|
|
|
|
topk_result = torch.full(
|
|
(token_nums, self.index_topk), -1, device=device, dtype=torch.int32
|
|
)
|
|
if batch_size == 0:
|
|
return topk_result
|
|
|
|
ks, ke = metadata.get_indexer_kvcache_range()
|
|
|
|
indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu()
|
|
seq_len_sum = torch.sum(indexer_seq_lens_cpu).item()
|
|
max_seq_len = torch.max(indexer_seq_lens_cpu).item()
|
|
index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id,
|
|
block_tables,
|
|
)
|
|
k_fp8, k_scale = index_buf_accessor.GetKAndS.execute(
|
|
forward_batch.token_to_kv_pool,
|
|
index_buffer,
|
|
page_indices=block_tables,
|
|
seq_len_tensor=metadata.get_indexer_seq_len(),
|
|
seq_len_sum=seq_len_sum,
|
|
max_seq_len=max_seq_len,
|
|
)
|
|
if _is_fp8_fnuz:
|
|
k_fp8 = k_fp8.view(torch.float8_e4m3fnuz)
|
|
else:
|
|
k_fp8 = k_fp8.view(torch.float8_e4m3fn)
|
|
|
|
k_scale = k_scale.view(torch.float32).squeeze(-1)
|
|
kv_fp8 = (k_fp8, k_scale)
|
|
|
|
# Check if we need to chunk to avoid OOM
|
|
seq_lens_expanded = metadata.get_seqlens_expanded()
|
|
token_to_batch_idx = metadata.get_token_to_batch_idx()
|
|
q_offset = ks.shape[0]
|
|
k_offset = k_fp8.shape[0]
|
|
need_chunk, free_mem = self._should_chunk_mqa_logits(q_offset, k_offset, device)
|
|
|
|
if not need_chunk:
|
|
assert q_fp8[:q_offset].shape[0] != 0
|
|
with self._with_real_sm_count():
|
|
if _is_hip:
|
|
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
|
|
|
|
kv, scale = kv_fp8
|
|
logits = fp8_mqa_logits(
|
|
q_fp8[:q_offset], kv, scale, weights[:q_offset], ks, ke
|
|
)
|
|
else:
|
|
logits = deep_gemm.fp8_mqa_logits(
|
|
q_fp8[:q_offset],
|
|
kv_fp8,
|
|
weights[:q_offset],
|
|
ks,
|
|
ke,
|
|
clean_logits=False,
|
|
)
|
|
assert logits.shape[0] == len(seq_lens_expanded)
|
|
assert logits.shape[1] == k_offset
|
|
|
|
raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks)
|
|
topk_result[:q_offset] = raw_topk_result
|
|
return topk_result
|
|
|
|
# Chunk path
|
|
bytes_per_elem = 4 # float32
|
|
bytes_per_row = k_offset * bytes_per_elem
|
|
# Reserve 50% of free memory for logits
|
|
max_rows = max(1, int((free_mem * 0.5) // max(bytes_per_row, 1)))
|
|
max_rows = min(max_rows, q_offset)
|
|
|
|
global_topk_offset = metadata.attn_metadata.topk_indices_offset
|
|
|
|
assert (
|
|
seq_lens_expanded.shape[0] == q_offset
|
|
), f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}"
|
|
if global_topk_offset is not None:
|
|
assert (
|
|
global_topk_offset.shape[0] >= q_offset
|
|
), f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}"
|
|
|
|
start = 0
|
|
while start < q_offset:
|
|
end = min(start + max_rows, q_offset)
|
|
|
|
with self._with_real_sm_count():
|
|
if _is_hip:
|
|
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
|
|
|
|
kv, scale = kv_fp8
|
|
logits_chunk = fp8_mqa_logits(
|
|
q_fp8[start:end],
|
|
kv,
|
|
scale,
|
|
weights[start:end],
|
|
ks[start:end],
|
|
ke[start:end],
|
|
)
|
|
else:
|
|
logits_chunk = deep_gemm.fp8_mqa_logits(
|
|
q_fp8[start:end],
|
|
kv_fp8,
|
|
weights[start:end],
|
|
ks[start:end],
|
|
ke[start:end],
|
|
clean_logits=False,
|
|
)
|
|
|
|
lengths_chunk = seq_lens_expanded[start:end]
|
|
|
|
# RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk
|
|
if global_topk_offset is not None:
|
|
# RAGGED path
|
|
topk_offset_chunk = global_topk_offset[start:end]
|
|
cu_seqlens_q_chunk = None
|
|
batch_idx_chunk = None
|
|
else:
|
|
# PAGED path: treat each token as a length-1 sequence
|
|
topk_offset_chunk = None
|
|
B_chunk = logits_chunk.shape[0]
|
|
cu_seqlens_q_chunk = torch.ones(
|
|
B_chunk, dtype=torch.int32, device=device
|
|
)
|
|
batch_idx_chunk = token_to_batch_idx[start:end]
|
|
|
|
raw_topk_chunk = metadata.topk_transform(
|
|
logits_chunk,
|
|
self.index_topk,
|
|
ks=ks[start:end],
|
|
cu_seqlens_q=cu_seqlens_q_chunk,
|
|
ke_offset=lengths_chunk,
|
|
batch_idx_list=batch_idx_chunk,
|
|
topk_indices_offset_override=topk_offset_chunk,
|
|
)
|
|
topk_result[start:end] = raw_topk_chunk
|
|
start = end
|
|
|
|
return topk_result
|
|
|
|
def _forward_cuda_k_only(
|
|
self,
|
|
x: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
act_quant,
|
|
enable_dual_stream: bool,
|
|
metadata: BaseIndexerMetadata,
|
|
return_indices: bool = True,
|
|
) -> Optional[torch.Tensor]:
|
|
assert forward_batch.forward_mode.is_extend_without_speculative()
|
|
x_meta = x[0] if isinstance(x, tuple) else x
|
|
|
|
# Fast path: only compute and store k cache, skip all q and weights ops
|
|
key = self._get_k_bf16(x, positions, enable_dual_stream)
|
|
|
|
if not forward_batch.out_cache_loc.is_contiguous():
|
|
forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
|
|
|
|
self._store_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
key=key,
|
|
act_quant=act_quant,
|
|
)
|
|
|
|
# MHA doesn't need topk_indices
|
|
if not return_indices:
|
|
return None
|
|
|
|
# MLA: use dummy logits with topk kernel's fast path to generate indices
|
|
# When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...]
|
|
seq_lens_expanded = metadata.get_seqlens_expanded()
|
|
dummy_logits = torch.zeros(
|
|
seq_lens_expanded.shape[0],
|
|
self.index_topk,
|
|
dtype=torch.float32,
|
|
device=x_meta.device,
|
|
)
|
|
return metadata.topk_transform(dummy_logits, self.index_topk)
|
|
|
|
def _get_topk_ragged_with_cp(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
q_fp8: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
metadata: BaseIndexerMetadata,
|
|
kv_len: int,
|
|
actual_seq_q: int,
|
|
cp_index: List[Tuple[int, int, int]] = None,
|
|
current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
|
shared_index_buffer: Optional[torch.Tensor] = None,
|
|
shared_block_tables: Optional[torch.Tensor] = None,
|
|
actual_seq_q_tensor: Optional[torch.Tensor] = None,
|
|
actual_seq_q_cu_tensor: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
if TYPE_CHECKING:
|
|
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
|
|
|
|
page_size = forward_batch.token_to_kv_pool.page_size
|
|
assert page_size == 64, "only support page size 64"
|
|
assert len(weights.shape) == 3
|
|
weights = weights.squeeze(-1)
|
|
k_fp8_list = []
|
|
k_scale_list = []
|
|
ks_list = []
|
|
ke_offset_list = []
|
|
offset = 0
|
|
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 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:
|
|
raise RuntimeError(
|
|
"shared index buffer and block tables must be provided together"
|
|
)
|
|
index_buffer = shared_index_buffer
|
|
block_tables = shared_block_tables
|
|
else:
|
|
block_tables = metadata.get_page_table_64()
|
|
index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id,
|
|
block_tables,
|
|
)
|
|
else:
|
|
index_buffer = None
|
|
if cp_shared_kv_debug_enabled():
|
|
cp_shared_kv_debug_log(
|
|
"index_current_reuse",
|
|
"NSA index current reuse cp_rank=%s layer=%s kv_len=%s actual_seq_q=%s k_ck=%s s_ck=%s",
|
|
forward_batch.cp_shared_kv_layout.cp_rank
|
|
if forward_batch.cp_shared_kv_layout is not None
|
|
else None,
|
|
layer_id,
|
|
kv_len,
|
|
actual_seq_q,
|
|
tensor_debug_checksum(current_index_kv[0]),
|
|
tensor_debug_checksum(current_index_kv[1]),
|
|
)
|
|
|
|
assert (
|
|
forward_batch.seq_lens_cpu is not None
|
|
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:
|
|
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],
|
|
)
|
|
|
|
extend_seq_len = end_seq_position - start_seq_position
|
|
ks = torch.full(
|
|
(extend_seq_len,), offset, dtype=torch.int32, device="cuda"
|
|
)
|
|
k_fp8_list.append(k_fp8)
|
|
k_scale_list.append(k_scale)
|
|
ks_list.append(ks)
|
|
ke_offset = torch.arange(
|
|
start_seq_position + 1,
|
|
end_seq_position + 1,
|
|
dtype=torch.int32,
|
|
device="cuda",
|
|
)
|
|
ke_offset_list.append(ke_offset)
|
|
actual_seq_q = torch.tensor(
|
|
[extend_seq_len], dtype=torch.int32, device="cuda"
|
|
)
|
|
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)
|
|
kv_fp8 = (k_fp8, k_scale)
|
|
ks = torch.cat(ks_list, dim=0)
|
|
ke_offset = torch.cat(ke_offset_list, dim=0)
|
|
ke = ks + ke_offset
|
|
actual_seq_q = torch.cat(actual_seq_q_list, dim=0)
|
|
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,
|
|
)
|
|
else:
|
|
cp_kv_end = (
|
|
forward_batch.seq_lens_cpu[0].item()
|
|
- forward_batch.extend_seq_lens_cpu[0]
|
|
+ kv_len
|
|
)
|
|
page_table_1 = metadata.get_page_table_1()
|
|
logical_kv_limit = min(
|
|
int(forward_batch.seq_lens_cpu[0].item()),
|
|
int(page_table_1.shape[1]),
|
|
)
|
|
valid_q_count = _compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end=cp_kv_end,
|
|
actual_seq_q=actual_seq_q,
|
|
logical_kv_limit=logical_kv_limit,
|
|
)
|
|
if valid_q_count == 0:
|
|
return torch.full(
|
|
(actual_seq_q, self.index_topk),
|
|
-1,
|
|
dtype=torch.int32,
|
|
device=q_fp8.device,
|
|
)
|
|
|
|
ke_start = cp_kv_end - actual_seq_q + 1
|
|
|
|
q_fp8 = q_fp8[:valid_q_count]
|
|
weights = weights[:valid_q_count]
|
|
kv_len = min(cp_kv_end, logical_kv_limit)
|
|
if current_index_kv is None:
|
|
assert index_buffer is not None
|
|
tai_prepared = try_tai_prepare_cp_mqa_index(
|
|
index_buffer=index_buffer,
|
|
page_indices=block_tables[0],
|
|
kv_len=kv_len,
|
|
valid_q_count=valid_q_count,
|
|
ke_start=ke_start,
|
|
page_size=page_size,
|
|
index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
|
|
)
|
|
if tai_prepared is not None:
|
|
k_fp8_u8, k_scale, ks, ke_offset = tai_prepared
|
|
k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
|
|
else:
|
|
tai_range = try_tai_prepare_cp_mqa_range(
|
|
valid_q_count=valid_q_count,
|
|
ke_start=ke_start,
|
|
device=q_fp8.device,
|
|
)
|
|
if tai_range is not None:
|
|
ks, ke_offset = tai_range
|
|
else:
|
|
ke_offset = torch.arange(
|
|
ke_start,
|
|
ke_start + valid_q_count,
|
|
dtype=torch.int32,
|
|
device=q_fp8.device,
|
|
)
|
|
ks = torch.zeros_like(ke_offset)
|
|
k_fp8 = index_buf_accessor.GetK.execute(
|
|
forward_batch.token_to_kv_pool,
|
|
index_buffer,
|
|
seq_len=kv_len,
|
|
page_indices=block_tables[0],
|
|
)
|
|
k_scale = index_buf_accessor.GetS.execute(
|
|
forward_batch.token_to_kv_pool,
|
|
index_buffer,
|
|
seq_len=kv_len,
|
|
page_indices=block_tables[0],
|
|
)
|
|
|
|
k_fp8 = k_fp8.view(torch.float8_e4m3fn)
|
|
k_scale = k_scale.view(torch.float32).squeeze(-1)
|
|
else:
|
|
tai_range = try_tai_prepare_cp_mqa_range(
|
|
valid_q_count=valid_q_count,
|
|
ke_start=ke_start,
|
|
device=q_fp8.device,
|
|
)
|
|
if tai_range is not None:
|
|
ks, ke_offset = tai_range
|
|
else:
|
|
ke_offset = torch.arange(
|
|
ke_start,
|
|
ke_start + valid_q_count,
|
|
dtype=torch.int32,
|
|
device=q_fp8.device,
|
|
)
|
|
ks = torch.zeros_like(ke_offset)
|
|
k_fp8, k_scale = current_index_kv
|
|
k_fp8 = k_fp8[:kv_len].contiguous()
|
|
k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous()
|
|
kv_fp8 = (k_fp8, k_scale)
|
|
ke = 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_indices_offset_override = None
|
|
cu_seqlens_q_topk_override = None
|
|
if (
|
|
getattr(getattr(metadata, "topk_transform_method", None), "name", None)
|
|
== "RAGGED"
|
|
):
|
|
# For this single CP segment the old
|
|
# compute_cu_seqlens([valid_q_count]) + repeat_interleave path
|
|
# produced a zero offset per query. Reuse `ks` and avoid the
|
|
# post-MQA metadata kernels entirely.
|
|
topk_indices_offset_override = ks
|
|
actual_seq_q_tensor = None
|
|
elif valid_q_count == actual_seq_q and actual_seq_q_cu_tensor is not None:
|
|
cu_seqlens_q_topk_override = actual_seq_q_cu_tensor
|
|
elif actual_seq_q_tensor is None or valid_q_count != actual_seq_q:
|
|
actual_seq_q_tensor = ke_offset.new_full((1,), valid_q_count)
|
|
cu_seqlens_q_topk_override = ke_offset.new_empty((2,))
|
|
cu_seqlens_q_topk_override[0] = 0
|
|
cu_seqlens_q_topk_override[1] = actual_seq_q_tensor.reshape(-1)[0]
|
|
elif actual_seq_q_tensor.ndim == 0:
|
|
actual_seq_q_tensor = actual_seq_q_tensor.reshape(1)
|
|
valid_topk_result = metadata.topk_transform(
|
|
logits,
|
|
self.index_topk,
|
|
ks=ks,
|
|
cu_seqlens_q=actual_seq_q_tensor,
|
|
ke_offset=ke_offset,
|
|
topk_indices_offset_override=topk_indices_offset_override,
|
|
cu_seqlens_q_topk_override=cu_seqlens_q_topk_override,
|
|
)
|
|
if valid_q_count == actual_seq_q:
|
|
topk_result = valid_topk_result
|
|
else:
|
|
topk_result = torch.full(
|
|
(actual_seq_q, self.index_topk),
|
|
-1,
|
|
dtype=torch.int32,
|
|
device=q_fp8.device,
|
|
)
|
|
topk_result[:valid_q_count] = valid_topk_result
|
|
|
|
return topk_result
|
|
|
|
def _get_topk_in_seq_cp_pair(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
q_fp8: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
metadata: BaseIndexerMetadata,
|
|
current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
|
) -> torch.Tensor:
|
|
assert forward_batch.nsa_cp_metadata is not None
|
|
kv_len_prev = forward_batch.nsa_cp_metadata.kv_len_prev
|
|
kv_len_next = forward_batch.nsa_cp_metadata.kv_len_next
|
|
actual_seq_q_prev = forward_batch.nsa_cp_metadata.actual_seq_q_prev
|
|
actual_seq_q_next = forward_batch.nsa_cp_metadata.actual_seq_q_next
|
|
|
|
# TODO support mutil-batch
|
|
# cp_batch_seq_index_prev = forward_batch.nsa_cp_metadata["cp_batch_seq_index_prev"]
|
|
# cp_batch_seq_index_next = forward_batch.nsa_cp_metadata["cp_batch_seq_index_next"]
|
|
q_fp8_prev, q_fp8_next = split_in_seq_cp_local_pair(
|
|
q_fp8,
|
|
actual_seq_q_prev,
|
|
actual_seq_q_next,
|
|
name="q_fp8",
|
|
)
|
|
weights_prev, weights_next = split_in_seq_cp_local_pair(
|
|
weights,
|
|
actual_seq_q_prev,
|
|
actual_seq_q_next,
|
|
name="weights",
|
|
)
|
|
|
|
shared_index_buffer = None
|
|
shared_block_tables = None
|
|
if current_index_kv is None:
|
|
shared_block_tables = metadata.get_page_table_64()
|
|
shared_index_buffer, shared_block_tables = (
|
|
self._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id,
|
|
shared_block_tables,
|
|
)
|
|
)
|
|
|
|
topk_result_prev = self._get_topk_ragged_with_cp(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8_prev,
|
|
weights_prev,
|
|
metadata,
|
|
kv_len_prev,
|
|
actual_seq_q_prev,
|
|
current_index_kv=current_index_kv,
|
|
shared_index_buffer=shared_index_buffer,
|
|
shared_block_tables=shared_block_tables,
|
|
actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor,
|
|
actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_cu_tensor,
|
|
)
|
|
|
|
topk_result_next = self._get_topk_ragged_with_cp(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8_next,
|
|
weights_next,
|
|
metadata,
|
|
kv_len_next,
|
|
actual_seq_q_next,
|
|
current_index_kv=current_index_kv,
|
|
shared_index_buffer=shared_index_buffer,
|
|
shared_block_tables=shared_block_tables,
|
|
actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor,
|
|
actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_cu_tensor,
|
|
)
|
|
return torch.cat([topk_result_prev, topk_result_next], dim=0)
|
|
|
|
def forward_indexer(
|
|
self,
|
|
q_fp8: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
topk: int,
|
|
layer_id: int,
|
|
) -> Optional[torch.Tensor]:
|
|
if not _is_npu:
|
|
from sglang.srt.layers.attention.nsa.tilelang_kernel import fp8_index
|
|
|
|
page_size = forward_batch.token_to_kv_pool.page_size
|
|
assert page_size == 64, "only support page size 64"
|
|
|
|
assert len(weights.shape) == 3
|
|
weights = weights.squeeze(-1)
|
|
|
|
# logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)
|
|
k_fp8_list = []
|
|
k_scale_list = []
|
|
|
|
topk_indices_list = []
|
|
|
|
block_tables = forward_batch.req_to_token_pool.req_to_token[
|
|
forward_batch.req_pool_indices, :
|
|
]
|
|
strided_indices = torch.arange(
|
|
0, block_tables.shape[-1], page_size, device="cuda"
|
|
)
|
|
block_tables = block_tables[:, strided_indices] // page_size
|
|
index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id,
|
|
block_tables,
|
|
)
|
|
|
|
q_len_start = 0
|
|
|
|
for i in range(forward_batch.batch_size):
|
|
seq_len = forward_batch.seq_lens[i].item()
|
|
q_len = (
|
|
forward_batch.extend_seq_lens_cpu[i]
|
|
if forward_batch.forward_mode.is_extend()
|
|
else 1
|
|
)
|
|
q_len_end = q_len_start + q_len
|
|
|
|
q_fp8_partial = q_fp8[q_len_start:q_len_end]
|
|
q_fp8_partial = q_fp8_partial.unsqueeze(0).contiguous()
|
|
|
|
weights_partial = weights[q_len_start:q_len_end]
|
|
weights_partial = weights_partial.squeeze(-1).unsqueeze(0).contiguous()
|
|
|
|
k_fp8 = index_buf_accessor.GetK.execute(
|
|
forward_batch.token_to_kv_pool,
|
|
index_buffer,
|
|
seq_len=seq_len,
|
|
page_indices=block_tables[i],
|
|
)
|
|
k_scale = index_buf_accessor.GetS.execute(
|
|
forward_batch.token_to_kv_pool,
|
|
index_buffer,
|
|
seq_len=seq_len,
|
|
page_indices=block_tables[i],
|
|
)
|
|
|
|
k_fp8 = k_fp8.view(torch.float8_e4m3fn).unsqueeze(0).contiguous()
|
|
k_scale = k_scale.view(torch.float32).squeeze(-1).unsqueeze(0).contiguous()
|
|
|
|
index_score = fp8_index(
|
|
q_fp8_partial,
|
|
weights_partial,
|
|
k_fp8,
|
|
k_scale,
|
|
)
|
|
end_pos = seq_len
|
|
topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0)
|
|
|
|
pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1]
|
|
topk_indices = torch.nn.functional.pad(
|
|
topk_indices, (0, pad_len), "constant", -1
|
|
)
|
|
|
|
topk_indices_list.append(topk_indices)
|
|
|
|
q_len_start = q_len_end
|
|
|
|
topk_indices = torch.cat(topk_indices_list, dim=0)
|
|
return topk_indices
|
|
|
|
def _store_index_k_cache(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
key: torch.Tensor,
|
|
*,
|
|
act_quant=None, # fallback only
|
|
out_loc_override: Optional[torch.Tensor] = None,
|
|
) -> None:
|
|
"""
|
|
Store NSA indexer K cache for current step.
|
|
|
|
Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size)
|
|
Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...)
|
|
"""
|
|
if out_loc_override is None:
|
|
out_loc, key = self._filter_shared_index_write(forward_batch, key)
|
|
else:
|
|
out_loc = out_loc_override
|
|
if out_loc.numel() == 0:
|
|
return
|
|
if not out_loc.is_contiguous():
|
|
out_loc = out_loc.contiguous()
|
|
|
|
# Fast path: JIT fused store (CUDA, page_size=64, non-fnuz)
|
|
if (
|
|
_is_cuda
|
|
and (not _is_fp8_fnuz)
|
|
and can_use_nsa_fused_store(
|
|
key.dtype,
|
|
out_loc.dtype,
|
|
forward_batch.token_to_kv_pool.page_size,
|
|
)
|
|
):
|
|
# NOTE: wrapper already normalizes shape/contiguity and asserts dtypes.
|
|
buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
|
|
layer_id=layer_id
|
|
)
|
|
fused_store_index_k_cache(
|
|
key,
|
|
buf,
|
|
out_loc,
|
|
forward_batch.token_to_kv_pool.page_size,
|
|
)
|
|
return
|
|
|
|
# Fallback: original path
|
|
assert act_quant is not None
|
|
k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
|
|
|
|
forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
|
|
layer_id=layer_id,
|
|
loc=out_loc,
|
|
index_k=k_fp8,
|
|
index_k_scale=k_scale,
|
|
)
|
|
|
|
def _store_cp_shared_local_index_k_cache(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
local_key: torch.Tensor,
|
|
*,
|
|
act_quant,
|
|
) -> bool:
|
|
if not nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
|
|
return False
|
|
|
|
local_out_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
if local_out_loc is None:
|
|
return False
|
|
if local_key.shape[0] != local_out_loc.numel():
|
|
log_cp_shared_kv_direct_write_fallback(
|
|
"index_local_shape_mismatch",
|
|
"NSA index local key token count does not match local out_cache_loc: "
|
|
"local_key=%s local_out_cache_loc=%s layer_id=%s",
|
|
local_key.shape[0],
|
|
local_out_loc.numel(),
|
|
layer_id,
|
|
)
|
|
return False
|
|
if local_out_loc.numel() == 0:
|
|
return True
|
|
|
|
physical_out_loc = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
|
if physical_out_loc is None:
|
|
return False
|
|
log_cp_draft_shared_kv_debug(
|
|
"index_write",
|
|
"index_write layer=%s tokens=%s physical_tokens=%s pool=%s key_shape=%s",
|
|
layer_id,
|
|
local_out_loc.numel(),
|
|
physical_out_loc.numel(),
|
|
forward_batch.token_to_kv_pool.__class__.__name__,
|
|
tuple(local_key.shape),
|
|
)
|
|
self._store_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
key=local_key,
|
|
act_quant=act_quant,
|
|
out_loc_override=physical_out_loc,
|
|
)
|
|
return True
|
|
|
|
def forward_cuda(
|
|
self,
|
|
x: torch.Tensor,
|
|
q_lora: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
return_indices: bool = True,
|
|
) -> Optional[torch.Tensor]:
|
|
if _is_hip:
|
|
from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant
|
|
elif not _is_npu:
|
|
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
|
|
|
|
if TYPE_CHECKING:
|
|
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
|
|
|
|
# When upstream uses fused FP8 RMSNorm+quant, activations may be passed as
|
|
# a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries.
|
|
x_meta = x[0] if isinstance(x, tuple) else x
|
|
|
|
metadata = forward_batch.attn_backend.get_indexer_metadata(
|
|
layer_id, forward_batch
|
|
)
|
|
|
|
enable_dual_stream = (
|
|
self.alt_stream is not None
|
|
and get_is_capture_mode()
|
|
and q_lora.shape[0] > 0
|
|
and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
|
|
)
|
|
|
|
# skip NSA if attention backend choose to skip this batch
|
|
if metadata is None:
|
|
return None
|
|
|
|
# Determine if should skip topk based on sequence length
|
|
# We can only skip the logits computation if cuda graph is not involved
|
|
skip_logits_computation = False
|
|
if forward_batch.forward_mode.is_extend_without_speculative():
|
|
if forward_batch.seq_lens_cpu is not None:
|
|
max_kv_len = forward_batch.seq_lens_cpu.max().item()
|
|
skip_logits_computation = max_kv_len <= self.index_topk
|
|
|
|
# Optimization: fast path when skipping topk computation
|
|
if skip_logits_computation and (not self.nsa_enable_prefill_cp):
|
|
return self._forward_cuda_k_only(
|
|
x,
|
|
positions,
|
|
forward_batch,
|
|
layer_id,
|
|
act_quant,
|
|
enable_dual_stream,
|
|
metadata,
|
|
return_indices,
|
|
)
|
|
|
|
if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
|
|
current_stream = torch.cuda.current_stream()
|
|
self.alt_stream.wait_stream(current_stream)
|
|
weights = self._project_and_scale_head_gates(x)
|
|
query, key, local_key = self._get_q_k_bf16(
|
|
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
|
|
)
|
|
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
|
|
with torch.cuda.stream(self.alt_stream):
|
|
if not self._store_cp_shared_local_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
local_key=local_key,
|
|
act_quant=act_quant,
|
|
):
|
|
self._store_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
key=key,
|
|
act_quant=act_quant,
|
|
)
|
|
current_stream.wait_stream(self.alt_stream)
|
|
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
|
|
else:
|
|
query, key, local_key = self._get_q_k_bf16(
|
|
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
|
|
)
|
|
|
|
if enable_dual_stream:
|
|
current_stream = torch.cuda.current_stream()
|
|
self.alt_stream.wait_stream(current_stream)
|
|
|
|
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
|
|
with torch.cuda.stream(self.alt_stream):
|
|
if not self._store_cp_shared_local_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
local_key=local_key,
|
|
act_quant=act_quant,
|
|
):
|
|
self._store_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
key=key,
|
|
act_quant=act_quant,
|
|
)
|
|
current_stream.wait_stream(self.alt_stream)
|
|
else:
|
|
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
|
|
if not self._store_cp_shared_local_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
local_key=local_key,
|
|
act_quant=act_quant,
|
|
):
|
|
self._store_index_k_cache(
|
|
forward_batch=forward_batch,
|
|
layer_id=layer_id,
|
|
key=key,
|
|
act_quant=act_quant,
|
|
)
|
|
|
|
# `_get_logits_head_gate` expects a Tensor. For tuple activations, dequantize
|
|
# to a float tensor here (callsite), keeping `_get_logits_head_gate` backend-agnostic.
|
|
if isinstance(x, tuple):
|
|
assert len(x) in (
|
|
2,
|
|
3,
|
|
), "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted"
|
|
x_q, x_s = x[0], x[1]
|
|
if (
|
|
x_s is not None
|
|
and x_q.dim() == 2
|
|
and x_s.dim() == 2
|
|
and x_q.shape[0] == x_s.shape[0]
|
|
):
|
|
m, n = x_q.shape
|
|
ng = x_s.shape[1]
|
|
if ng > 0 and n % ng == 0:
|
|
group = n // ng
|
|
x_for_gate = (
|
|
x_q.to(torch.float32)
|
|
.view(m, ng, group)
|
|
.mul_(x_s.to(torch.float32).unsqueeze(-1))
|
|
.view(m, n)
|
|
.to(torch.bfloat16)
|
|
)
|
|
else:
|
|
x_for_gate = x_q.to(torch.bfloat16)
|
|
else:
|
|
x_for_gate = x_q.to(torch.bfloat16)
|
|
else:
|
|
x_for_gate = x
|
|
|
|
weights = self._get_logits_head_gate(x_for_gate, q_scale)
|
|
|
|
current_index_kv = None
|
|
if self._can_reuse_current_index_kv(forward_batch):
|
|
if key.shape[0] == forward_batch.out_cache_loc.numel():
|
|
current_k_fp8, current_k_scale = act_quant(
|
|
key.contiguous(), self.block_size, self.scale_fmt
|
|
)
|
|
current_index_kv = (
|
|
current_k_fp8.contiguous(),
|
|
current_k_scale.contiguous(),
|
|
)
|
|
if cp_shared_kv_debug_enabled():
|
|
cp_shared_kv_debug_log(
|
|
"index_current_reuse_prepare",
|
|
"NSA index current reuse prepared cp_rank=%s layer=%s locs=%s k_ck=%s s_ck=%s",
|
|
forward_batch.cp_shared_kv_layout.cp_rank
|
|
if forward_batch.cp_shared_kv_layout is not None
|
|
else None,
|
|
layer_id,
|
|
tensor_debug_summary(forward_batch.out_cache_loc),
|
|
tensor_debug_checksum(current_index_kv[0]),
|
|
tensor_debug_checksum(current_index_kv[1]),
|
|
)
|
|
elif cp_shared_kv_debug_enabled():
|
|
raise RuntimeError(
|
|
"CP shared KV current index reuse shape mismatch: "
|
|
f"key_tokens={key.shape[0]} out_cache_loc={forward_batch.out_cache_loc.numel()}"
|
|
)
|
|
|
|
if _is_cuda or _is_hip:
|
|
assert forward_batch.seq_lens_cpu is not None
|
|
if len(forward_batch.seq_lens_cpu) == 0:
|
|
if (
|
|
cp_shared_kv_mla_prefetch_log_enabled()
|
|
and cp_shared_kv_mla_prefetch_should_log_layer(layer_id)
|
|
):
|
|
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
|
|
index_prefetcher = getattr(
|
|
forward_batch, "cp_shared_kv_index_prefetcher", None
|
|
)
|
|
cp_shared_kv_mla_prefetch_log(
|
|
"index_forward_early_return_empty_seq_lens cp_rank=%s "
|
|
"cp_size=%s layer=%s has_index=%s uses_cp_shared_kv=%s "
|
|
"x_meta_shape=%s",
|
|
getattr(layout, "cp_rank", None),
|
|
getattr(layout, "cp_size", None),
|
|
layer_id,
|
|
index_prefetcher is not None,
|
|
getattr(forward_batch, "uses_cp_shared_kv", None),
|
|
tuple(x_meta.shape),
|
|
)
|
|
# this seems b/c max-pad, no worries?
|
|
# if x.shape[0] != 0:
|
|
# print(
|
|
# "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result"
|
|
# )
|
|
return torch.full(
|
|
(x_meta.shape[0], self.index_topk),
|
|
-1,
|
|
dtype=torch.int,
|
|
device=x_meta.device,
|
|
)
|
|
|
|
if (
|
|
forward_batch.forward_mode.is_decode_or_idle()
|
|
or forward_batch.forward_mode.is_target_verify()
|
|
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
|
|
):
|
|
topk_result = self._get_topk_paged(
|
|
forward_batch, layer_id, q_fp8, weights, metadata
|
|
)
|
|
else:
|
|
if (
|
|
forward_batch.nsa_cp_metadata is not None
|
|
and is_nsa_prefill_cp_in_seq_split()
|
|
):
|
|
topk_result = self._get_topk_in_seq_cp_pair(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
else:
|
|
topk_result = self._get_topk_ragged(
|
|
enable_dual_stream,
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
)
|
|
else:
|
|
topk_result = self.forward_indexer(
|
|
q_fp8.contiguous(),
|
|
weights,
|
|
forward_batch,
|
|
topk=self.index_topk,
|
|
layer_id=layer_id,
|
|
)
|
|
self._maybe_start_next_layer_index_prefetch(forward_batch, layer_id)
|
|
return topk_result
|
|
|
|
def forward_npu(
|
|
self,
|
|
x: torch.Tensor,
|
|
q_lora: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
layer_id: int,
|
|
layer_scatter_modes=None,
|
|
dynamic_scale: torch.Tensor = None,
|
|
) -> torch.Tensor:
|
|
if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None:
|
|
actual_seq_lengths_kv = forward_batch.attn_backend.forward_metadata.seq_lens
|
|
else:
|
|
actual_seq_lengths_kv = (
|
|
forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int
|
|
)
|
|
is_prefill = (
|
|
forward_batch.forward_mode.is_extend()
|
|
and not forward_batch.forward_mode.is_draft_extend_v2()
|
|
and not forward_batch.forward_mode.is_target_verify()
|
|
and not forward_batch.forward_mode.is_draft_extend()
|
|
)
|
|
|
|
cos_sin = self.rotary_emb.cos_sin_cache[positions]
|
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
|
cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
|
|
sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
|
|
|
|
bs = q_lora.shape[0]
|
|
if self.alt_stream is not None:
|
|
self.alt_stream.wait_stream(torch.npu.current_stream())
|
|
with torch.npu.stream(self.alt_stream):
|
|
q_lora = (
|
|
(q_lora, dynamic_scale) if dynamic_scale is not None else q_lora
|
|
)
|
|
q = self.wq_b(q_lora)[
|
|
0
|
|
] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
|
|
wq_b_event = self.alt_stream.record_event()
|
|
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
|
|
q_pe, q_nope = torch.split(
|
|
q,
|
|
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
|
|
dim=-1,
|
|
) # [bs, 64, 64 + 64]
|
|
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim)
|
|
q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view(
|
|
bs, self.n_heads, self.rope_head_dim
|
|
) # [bs, n, d]
|
|
q = torch.cat([q_pe, q_nope], dim=-1)
|
|
q.record_stream(self.alt_stream)
|
|
q_rope_event = self.alt_stream.record_event()
|
|
else:
|
|
q_lora = (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora
|
|
q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
|
|
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
|
|
q_pe, q_nope = torch.split(
|
|
q,
|
|
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
|
|
dim=-1,
|
|
) # [bs, 64, 64 + 64]
|
|
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim)
|
|
q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view(
|
|
bs, self.n_heads, self.rope_head_dim
|
|
) # [bs, n, d]
|
|
q = torch.cat([q_pe, q_nope], dim=-1)
|
|
|
|
if envs.SGLANG_NPU_USE_MULTI_STREAM.get():
|
|
indexer_weight_stream = get_indexer_weight_stream()
|
|
indexer_weight_stream.wait_stream(torch.npu.current_stream())
|
|
with torch.npu.stream(indexer_weight_stream):
|
|
x = x.view(-1, self.hidden_size)
|
|
weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
|
|
weights.record_stream(indexer_weight_stream)
|
|
weights_event = indexer_weight_stream.record_event()
|
|
else:
|
|
x = x.view(-1, self.hidden_size)
|
|
weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
|
|
|
|
k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128]
|
|
k = self.k_norm(k_proj)
|
|
if (
|
|
_use_ag_after_qlora
|
|
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
|
|
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
|
|
):
|
|
k = scattered_to_tp_attn_full(k, forward_batch)
|
|
k_pe, k_nope = torch.split(
|
|
k,
|
|
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
|
|
dim=-1,
|
|
) # [bs, 64 + 64]
|
|
|
|
k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim)
|
|
k_pe = torch.ops.npu.npu_rotary_mul(k_pe, cos, sin).view(
|
|
bs, 1, self.rope_head_dim
|
|
) # [bs, 1, d]
|
|
k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128]
|
|
|
|
if (
|
|
is_prefill
|
|
and self.nsa_enable_prefill_cp
|
|
and forward_batch.nsa_cp_metadata is not None
|
|
):
|
|
k = cp_all_gather_rerange_output(
|
|
k.contiguous().view(-1, self.head_dim),
|
|
self.cp_size,
|
|
forward_batch,
|
|
torch.npu.current_stream(),
|
|
)
|
|
|
|
forward_batch.token_to_kv_pool.set_index_k_buffer(
|
|
layer_id, forward_batch.out_cache_loc, k
|
|
)
|
|
if is_prefill:
|
|
if self.nsa_enable_prefill_cp and forward_batch.nsa_cp_metadata is not None:
|
|
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q = (
|
|
forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor,
|
|
forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor,
|
|
)
|
|
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = (
|
|
forward_batch.nsa_cp_metadata.kv_len_prev_tensor,
|
|
forward_batch.nsa_cp_metadata.kv_len_next_tensor,
|
|
)
|
|
actual_seq_lengths_q = (
|
|
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q
|
|
)
|
|
actual_seq_lengths_kv = (
|
|
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv
|
|
)
|
|
else:
|
|
actual_seq_lengths_kv = forward_batch.seq_lens
|
|
actual_seq_lengths_q = forward_batch.extend_seq_lens.cumsum(dim=0)
|
|
else:
|
|
if forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q is None:
|
|
if (
|
|
forward_batch.forward_mode.is_draft_extend_v2()
|
|
or forward_batch.forward_mode.is_target_verify()
|
|
or forward_batch.forward_mode.is_draft_extend()
|
|
):
|
|
num_draft_tokens = (
|
|
forward_batch.attn_backend.speculative_num_draft_tokens
|
|
)
|
|
actual_seq_lengths_q = torch.arange(
|
|
num_draft_tokens,
|
|
num_draft_tokens + bs,
|
|
num_draft_tokens,
|
|
dtype=torch.int32,
|
|
device=k.device,
|
|
)
|
|
else:
|
|
actual_seq_lengths_q = torch.tensor(
|
|
[1 + i * 1 for i in range(bs)],
|
|
dtype=torch.int32,
|
|
device=k.device,
|
|
)
|
|
else:
|
|
actual_seq_lengths_q = (
|
|
forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q
|
|
)
|
|
|
|
past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id)
|
|
|
|
if self.alt_stream is not None:
|
|
torch.npu.current_stream().wait_event(q_rope_event)
|
|
if envs.SGLANG_NPU_USE_MULTI_STREAM.get():
|
|
torch.npu.current_stream().wait_event(weights_event)
|
|
if (
|
|
_use_ag_after_qlora
|
|
and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
|
|
and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
|
|
):
|
|
weights = scattered_to_tp_attn_full(weights, forward_batch)
|
|
block_table = forward_batch.attn_backend.forward_metadata.block_tables
|
|
if (
|
|
is_prefill
|
|
and self.nsa_enable_prefill_cp
|
|
and forward_batch.nsa_cp_metadata is not None
|
|
):
|
|
block_table = block_table[: actual_seq_lengths_q[0].numel()]
|
|
topk_indices = self.do_npu_cp_balance_indexer(
|
|
q.view(-1, self.n_heads, self.head_dim),
|
|
past_key_states,
|
|
weights,
|
|
actual_seq_lengths_q,
|
|
actual_seq_lengths_kv,
|
|
block_table,
|
|
)
|
|
return topk_indices
|
|
else:
|
|
block_table = (
|
|
block_table[: actual_seq_lengths_q.size()[0]]
|
|
if is_prefill
|
|
else block_table
|
|
)
|
|
|
|
topk_indices = torch_npu.npu_lightning_indexer(
|
|
query=q.view(-1, self.n_heads, self.head_dim),
|
|
key=past_key_states,
|
|
weights=weights,
|
|
actual_seq_lengths_query=actual_seq_lengths_q.to(torch.int32),
|
|
actual_seq_lengths_key=actual_seq_lengths_kv.to(k.device).to(
|
|
torch.int32
|
|
),
|
|
block_table=block_table,
|
|
layout_query="TND",
|
|
layout_key="PA_BSND",
|
|
sparse_count=self.index_topk,
|
|
sparse_mode=3,
|
|
)
|
|
return topk_indices[0]
|
|
|
|
def do_npu_cp_balance_indexer(
|
|
self,
|
|
q,
|
|
past_key_states,
|
|
indexer_weights,
|
|
actual_seq_lengths_q,
|
|
actual_seq_lengths_kv,
|
|
block_table,
|
|
):
|
|
q_prev, q_next = torch.split(q, (q.size(0) + 1) // 2, dim=0)
|
|
weights_prev, weights_next = None, None
|
|
if indexer_weights is not None:
|
|
weights_prev, weights_next = torch.split(
|
|
indexer_weights, (indexer_weights.size(0) + 1) // 2, dim=0
|
|
)
|
|
weights_prev = weights_prev.contiguous().view(-1, weights_prev.shape[-1])
|
|
weights_next = weights_next.contiguous().view(-1, weights_next.shape[-1])
|
|
|
|
actual_seq_lengths_q_prev, actual_seq_lengths_q_next = actual_seq_lengths_q
|
|
actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv
|
|
|
|
topk_indices_prev = torch_npu.npu_lightning_indexer(
|
|
query=q_prev,
|
|
key=past_key_states,
|
|
weights=weights_prev,
|
|
actual_seq_lengths_query=actual_seq_lengths_q_prev.to(
|
|
device=q.device, dtype=torch.int32
|
|
),
|
|
actual_seq_lengths_key=actual_seq_lengths_kv_prev.to(
|
|
device=q.device, dtype=torch.int32
|
|
),
|
|
block_table=block_table,
|
|
layout_query="TND",
|
|
layout_key="PA_BSND",
|
|
sparse_count=self.index_topk,
|
|
sparse_mode=3,
|
|
)
|
|
topk_indices_next = torch_npu.npu_lightning_indexer(
|
|
query=q_next,
|
|
key=past_key_states,
|
|
weights=weights_next,
|
|
actual_seq_lengths_query=actual_seq_lengths_q_next.to(
|
|
device=q.device, dtype=torch.int32
|
|
),
|
|
actual_seq_lengths_key=actual_seq_lengths_kv_next.to(
|
|
device=q.device, dtype=torch.int32
|
|
),
|
|
block_table=block_table,
|
|
layout_query="TND",
|
|
layout_key="PA_BSND",
|
|
sparse_count=self.index_topk,
|
|
sparse_mode=3,
|
|
)
|
|
return topk_indices_prev[0], topk_indices_next[0]
|
|
|
|
|
|
def scattered_to_tp_attn_full(
|
|
hidden_states: torch.Tensor,
|
|
forward_batch,
|
|
) -> torch.Tensor:
|
|
hidden_states, local_hidden_states = (
|
|
torch.empty(
|
|
(forward_batch.input_ids.shape[0], hidden_states.shape[1]),
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
),
|
|
hidden_states,
|
|
)
|
|
attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous())
|
|
return hidden_states
|