From aa27a444f6493708fee45b2cdbf3a6bd8a598217 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Fri, 8 May 2026 17:23:38 +0800
Subject: [PATCH] Reuse TAI range prepare for CP MQA current path

The existing CP MQA prepare integration only used the TAI GetK/GetS+range kernel when current_index_kv was absent. Current-index reuse and full-kernel fallback still launched torch arange and zeros_like before fp8_mqa_logits. Route both cases through the new tai-kernel range-only API while preserving the torch fallback when the kernel is disabled or unavailable.

Constraint: SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE remains the gate for this optimization.

Rejected: Force the full prepare kernel in current_index_kv path | it would redo index KV gathering that current reuse intentionally avoids.

Confidence: high

Scope-risk: narrow

Directive: Do not remove the torch fallback; mixed deployments may run without the updated tai-kernel package.

Related: tai-kernel 34cb7a8

Tested: Remote container g0034 docker py_compile for cp_shared_kv_runtime.py and nsa_indexer.py; tai-kernel range unit tests passed remotely (4 passed).

Not-tested: Full GLM5 prefill/decode server profile after this exact commit.
---
 .../attention/nsa/cp_shared_kv_runtime.py     | 39 +++++++++++++++++++
 .../srt/layers/attention/nsa/nsa_indexer.py   | 37 +++++++++++++-----
 2 files changed, 66 insertions(+), 10 deletions(-)

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