diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 150f4fbb9..602c93f95 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -247,6 +247,16 @@ def _load_tai_index_mqa_prepare_kernel(): return None +@lru_cache(maxsize=1) +def _load_tai_index_mqa_range_kernel(): + try: + from tai_kernel.nsa_prefill import prepare_cp_mqa_range + + return prepare_cp_mqa_range + except Exception: + return None + + def _log_tai_materialize_fallback( key: str, message: str, @@ -322,6 +332,35 @@ def try_tai_prepare_cp_mqa_index( return None +def try_tai_prepare_cp_mqa_range( + *, + valid_q_count: int, + ke_start: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Try the TAI fused MQA range prepare path. + + This is used when the caller already has dense/current index KV and only + needs `ks`/`ke_offset`, so the full GetK/GetS fused kernel is not applicable. + """ + + if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + return None + + kernel = _load_tai_index_mqa_range_kernel() + if kernel is None: + return None + + try: + return kernel( + valid_q_count=int(valid_q_count), + ke_start=int(ke_start), + device=device, + ) + except Exception: + return None + + def try_tai_fused_mla_store( *, token_to_kv_pool, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 659e8c02c..b6fff7700 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -25,6 +25,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( 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 @@ -1117,12 +1118,21 @@ class Indexer(MultiPlatformOp): k_fp8_u8, k_scale, ks, ke_offset = tai_prepared k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn) else: - ke_offset = torch.arange( - ke_start, - ke_start + valid_q_count, - dtype=torch.int32, + 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, @@ -1138,18 +1148,25 @@ class Indexer(MultiPlatformOp): k_fp8 = k_fp8.view(torch.float8_e4m3fn) k_scale = k_scale.view(torch.float32).squeeze(-1) - ks = torch.zeros_like(ke_offset) else: - ke_offset = torch.arange( - ke_start, - ke_start + valid_q_count, - dtype=torch.int32, + 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() - ks = torch.zeros_like(ke_offset) kv_fp8 = (k_fp8, k_scale) ke = ke_offset