perf(nsa): parallelize long ragged index copies

lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py

@@ -20,6 +20,7 @@ def _gen_nsa_ks_ke(
     ragged_mem_index,
     BLOCK_REQ: tl.constexpr,
     BLOCK_SEQ_SPLIT: tl.constexpr,
+    COPY_RAGGED: tl.constexpr = True,
 ):
     cur_index = tl.program_id(0)
     # 只处理最后一个同样req_idx的req进行处理，代表seq_len最长的那个。
@@ -65,24 +66,57 @@ def _gen_nsa_ks_ke(
             mask=mask,
         )
 
-    for block_index in range(tl.cdiv(cur_total_len, BLOCK_SEQ_SPLIT)):
-        block_start = block_index * BLOCK_SEQ_SPLIT
-        block_end = min(cur_total_len, (block_index + 1) * BLOCK_SEQ_SPLIT)
-        mask = (block_start + tl.arange(0, BLOCK_SEQ_SPLIT)) < block_end
-
-        src_mem_index_ptr = (
-            req_to_token_index
-            + strided_req_to_token_index_b * cur_req_idx
-            + block_start
-            + tl.arange(0, BLOCK_SEQ_SPLIT)
-        )
-        src_mem_index = tl.load(src_mem_index_ptr, mask=mask, other=-1)
-        tl.store(
-            ragged_mem_index + pre_sum_seq_len + block_start + tl.arange(0, BLOCK_SEQ_SPLIT), src_mem_index, mask=mask
-        )
+    if COPY_RAGGED:
+        for block_index in range(tl.cdiv(cur_total_len, BLOCK_SEQ_SPLIT)):
+            block_start = block_index * BLOCK_SEQ_SPLIT
+            block_end = min(cur_total_len, (block_index + 1) * BLOCK_SEQ_SPLIT)
+            mask = (block_start + tl.arange(0, BLOCK_SEQ_SPLIT)) < block_end
+
+            src_mem_index_ptr = (
+                req_to_token_index
+                + strided_req_to_token_index_b * cur_req_idx
+                + block_start
+                + tl.arange(0, BLOCK_SEQ_SPLIT)
+            )
+            src_mem_index = tl.load(src_mem_index_ptr, mask=mask, other=-1)
+            tl.store(
+                ragged_mem_index + pre_sum_seq_len + block_start + tl.arange(0, BLOCK_SEQ_SPLIT),
+                src_mem_index,
+                mask=mask,
+            )
     return
 
 
+@triton.jit
+def _copy_nsa_ragged_mem_index(
+    b_seq_len,
+    b_req_idx,
+    b_same_req_mark,
+    req_to_token_index,
+    ragged_mem_index,
+    strided_req_to_token_index_b: tl.constexpr,
+    BLOCK_REQ: tl.constexpr,
+    BLOCK_SEQ_SPLIT: tl.constexpr,
+):
+    cur_index = tl.program_id(0)
+    block_index = tl.program_id(1)
+    req_mark = tl.load(b_same_req_mark + cur_index)
+    cur_total_len = tl.load(b_seq_len + cur_index)
+    if req_mark == 0 or block_index * BLOCK_SEQ_SPLIT >= cur_total_len:
+        return
+    off = tl.arange(0, BLOCK_REQ)
+    pre_req_marks = tl.load(b_same_req_mark + off, off < cur_index, other=0)
+    pre_seq_lens = tl.load(b_seq_len + off, (off < cur_index) & (pre_req_marks != 0), other=0)
+    pre_sum_seq_len = tl.sum(pre_seq_lens)
+    cur_req_idx = tl.load(b_req_idx + cur_index)
+    block_off = block_index * BLOCK_SEQ_SPLIT + tl.arange(0, BLOCK_SEQ_SPLIT)
+    mask = block_off < cur_total_len
+    src_mem_index = tl.load(
+        req_to_token_index + cur_req_idx * strided_req_to_token_index_b + block_off, mask=mask, other=-1
+    )
+    tl.store(ragged_mem_index + pre_sum_seq_len + block_off, src_mem_index, mask=mask)
+
+
 @torch.no_grad()
 def gen_nsa_ks_ke(
     b_seq_len: torch.Tensor,
@@ -92,15 +126,28 @@ def gen_nsa_ks_ke(
     q_token_num: int,
     ragged_mem_index: torch.Tensor,
     hold_req_idx: int = -1,
+    max_seq_len: Optional[int] = None,
 ):
     """
     hold_req_idx 这是一个特殊req idx，主要是用于padding 请求数量时使用，所以其处理存在特殊性。
+    max_seq_len 是 b_seq_len 的 CPU 上界，长序列可据此并行复制 ragged 索引。
     """
     batch_size = b_seq_len.shape[0]
     ks = torch.empty((q_token_num,), dtype=torch.int32, device=b_seq_len.device)
     ke = torch.empty((q_token_num,), dtype=torch.int32, device=b_seq_len.device)
     lengths = torch.empty((q_token_num,), dtype=torch.int32, device=b_seq_len.device)
     b_same_req_mark = gen_same_req_mark(b_req_idx, hold_req_idx=hold_req_idx)
+    # Long GLM index tables otherwise traverse the entire prefix on one CTA.
+    # init_att_state runs outside the model graph, so this bound is refreshed
+    # for each decode call. Keep the established short/unspecified path.
+    if max_seq_len is None and req_to_token_index.shape[1] >= 32768 and req_to_token_index.stride(1) == 1:
+        # Long contiguous request tables are the GLM-5.3 layout. Infer the
+        # static copy bound here so model-specific attention wrappers do not
+        # need to duplicate this dispatch plumbing.
+        max_seq_len = req_to_token_index.shape[1]
+    use_parallel_ragged_copy = max_seq_len is not None and max_seq_len >= 32768
+    if use_parallel_ragged_copy:
+        assert req_to_token_index.stride(1) == 1
 
     _gen_nsa_ks_ke[(batch_size,)](
         b_seq_len=b_seq_len,
@@ -116,7 +163,19 @@ def gen_nsa_ks_ke(
         ragged_mem_index=ragged_mem_index,
         BLOCK_REQ=triton.next_power_of_2(batch_size),
         BLOCK_SEQ_SPLIT=256,
+        COPY_RAGGED=not use_parallel_ragged_copy,
     )
+    if use_parallel_ragged_copy:
+        _copy_nsa_ragged_mem_index[(batch_size, triton.cdiv(max_seq_len, 4096))](
+            b_seq_len=b_seq_len,
+            b_req_idx=b_req_idx,
+            b_same_req_mark=b_same_req_mark,
+            req_to_token_index=req_to_token_index,
+            ragged_mem_index=ragged_mem_index,
+            strided_req_to_token_index_b=req_to_token_index.stride(0),
+            BLOCK_REQ=triton.next_power_of_2(batch_size),
+            BLOCK_SEQ_SPLIT=4096,
+        )
     return ks, ke, lengths
 
 

unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py

@@ -1,6 +1,9 @@
 import torch
 import pytest
-from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke, gen_same_req_mark
+from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import (
+    gen_nsa_ks_ke,
+    gen_same_req_mark,
+)
 
 
 def test_gen_nsa_ks_ke_basic():
@@ -33,14 +36,20 @@ def test_gen_nsa_ks_ke_basic():
     ragged_mem_index = torch.empty_like(req_to_token_index.view(-1))
 
     ks, ke, lengths = gen_nsa_ks_ke(
-        b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index
+        b_seq_len,
+        b_q_seq_len,
+        b_req_idx,
+        req_to_token_index,
+        q_token_num,
+        ragged_mem_index,
     )
 
     assert torch.equal(ks, torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32, device="cuda"))
     assert torch.equal(ke, torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32, device="cuda"))
     assert torch.equal(lengths, ke - ks)
     assert torch.equal(
-        ragged_mem_index[0:10], torch.tensor([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=torch.int32, device="cuda")
+        ragged_mem_index[0:10],
+        torch.tensor([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=torch.int32, device="cuda"),
     )
     return
 
@@ -55,7 +64,12 @@ def test_gen_nsa_ks_ke_batch():
     ragged_mem_index = torch.empty_like(req_to_token_index.view(-1))
 
     ks, ke, lengths = gen_nsa_ks_ke(
-        b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index
+        b_seq_len,
+        b_q_seq_len,
+        b_req_idx,
+        req_to_token_index,
+        q_token_num,
+        ragged_mem_index,
     )
 
     assert torch.equal(
@@ -83,5 +97,131 @@ def test_gen_same_req_mark():
     assert torch.equal(same_req_mark, expected_same_req_mark)
 
 
+@pytest.mark.parametrize("context_length", [4096, 32768, 131072, 524288, 1000000])
+@pytest.mark.parametrize("verify_width", [1, 4])
+def test_parallel_ragged_indices_preserve_mtp_and_hold_rows(context_length, verify_width):
+    # Include recycled/nonsequential slots and independent HOLD padding rows.
+    hold = 3
+    req = torch.tensor(
+        [2] * verify_width + [0] * verify_width + [hold] * 4,
+        device="cuda",
+        dtype=torch.int32,
+    )
+    seq = torch.tensor(
+        list(range(context_length - verify_width + 1, context_length + 1))
+        + list(range(4096 - verify_width + 1, 4097))
+        + [2] * 4,
+        device="cuda",
+        dtype=torch.int32,
+    )
+    q = torch.ones_like(req)
+    table = torch.arange(4 * context_length, device="cuda", dtype=torch.int32).reshape(4, context_length)
+    capacity = context_length + 4096 + 8
+    reference_ragged = torch.full((capacity,), -1, device="cuda", dtype=torch.int32)
+    parallel_ragged = reference_ragged.clone()
+    reference = gen_nsa_ks_ke(seq, q, req, table, len(req), reference_ragged, hold_req_idx=hold)
+    actual = gen_nsa_ks_ke(
+        seq,
+        q,
+        req,
+        table,
+        len(req),
+        parallel_ragged,
+        hold_req_idx=hold,
+        max_seq_len=context_length,
+    )
+    for observed, expected in zip(actual, reference):
+        assert torch.equal(observed, expected)
+    assert torch.equal(parallel_ragged, reference_ragged)
+    assert torch.equal(parallel_ragged[:context_length], table[2, :context_length])
+    assert torch.equal(parallel_ragged[context_length : context_length + 4096], table[0, :4096])
+
+
+def test_parallel_ragged_metadata_matches_prefill_and_recycled_layout():
+    table = torch.arange(4 * 131072, device="cuda", dtype=torch.int32).reshape(4, -1)
+    q = torch.tensor([8192, 256], device="cuda", dtype=torch.int32)
+    req = torch.tensor([2, 0], device="cuda", dtype=torch.int32)
+    for lengths in ([131072, 32768], [32768, 65536], [65536, 131072]):
+        seq = torch.tensor(lengths, device="cuda", dtype=torch.int32)
+        ragged = torch.empty(sum(lengths), device="cuda", dtype=torch.int32)
+        ks, ke, actual_lengths = gen_nsa_ks_ke(seq, q, req, table, 8192 + 256, ragged, max_seq_len=max(lengths))
+        expected_lengths = torch.cat(
+            (
+                torch.arange(
+                    lengths[0] - 8192 + 1,
+                    lengths[0] + 1,
+                    device="cuda",
+                    dtype=torch.int32,
+                ),
+                torch.arange(
+                    lengths[1] - 256 + 1,
+                    lengths[1] + 1,
+                    device="cuda",
+                    dtype=torch.int32,
+                ),
+            )
+        )
+        expected_ks = torch.cat(
+            (
+                torch.zeros(8192, device="cuda", dtype=torch.int32),
+                torch.full((256,), lengths[0], device="cuda", dtype=torch.int32),
+            )
+        )
+        assert torch.equal(actual_lengths, expected_lengths)
+        assert torch.equal(ks, expected_ks)
+        assert torch.equal(ke, expected_ks + expected_lengths)
+        assert torch.equal(ragged, torch.cat((table[2, : lengths[0]], table[0, : lengths[1]])))
+
+
+def test_parallel_ragged_indices_cuda_graph_replay():
+    max_seq_len = 65536
+    # Preserve a nontrivial request-row stride while keeping token rows contiguous.
+    req_to_token_index = torch.arange(4 * (max_seq_len + 7), device="cuda", dtype=torch.int32).reshape(4, -1)
+    req_to_token_index = req_to_token_index[:, :max_seq_len]
+    b_req_idx = torch.tensor([2, 2, 0, 0, 3, 3], dtype=torch.int32, device="cuda")
+    b_seq_len = torch.tensor([65535, 65536, 4095, 4096, 2, 2], dtype=torch.int32, device="cuda")
+    b_q_seq_len = torch.ones_like(b_seq_len)
+    ragged_mem_index = torch.full((2 * max_seq_len + 4,), -1, dtype=torch.int32, device="cuda")
+
+    def run():
+        return gen_nsa_ks_ke(
+            b_seq_len,
+            b_q_seq_len,
+            b_req_idx,
+            req_to_token_index,
+            b_req_idx.numel(),
+            ragged_mem_index,
+            hold_req_idx=3,
+        )
+
+    run()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = run()
+
+    for req_indices, seq_lens in [
+        ([2, 2, 0, 0, 3, 3], [65535, 65536, 4095, 4096, 2, 2]),
+        ([0, 0, 1, 1, 3, 3], [32767, 32768, 65534, 65535, 2, 2]),
+        ([1, 1, 2, 2, 3, 3], [3, 4, 4096, 4097, 2, 2]),
+    ]:
+        b_req_idx.copy_(torch.tensor(req_indices, dtype=torch.int32, device="cuda"))
+        b_seq_len.copy_(torch.tensor(seq_lens, dtype=torch.int32, device="cuda"))
+        ragged_mem_index.fill_(-1)
+        graph.replay()
+        expected_ragged_mem_index = torch.full_like(ragged_mem_index, -1)
+        expected = gen_nsa_ks_ke(
+            b_seq_len,
+            b_q_seq_len,
+            b_req_idx,
+            req_to_token_index,
+            b_req_idx.numel(),
+            expected_ragged_mem_index,
+            hold_req_idx=3,
+        )
+        for output, reference in zip(actual, expected):
+            torch.testing.assert_close(output, reference, rtol=0, atol=0)
+        torch.testing.assert_close(ragged_mem_index, expected_ragged_mem_index, rtol=0, atol=0)
+
+
 if __name__ == "__main__":
     pytest.main()