diff --git a/lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py b/lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py index 44c64c4d3c..48cf4ff342 100644 --- a/lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py +++ b/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 diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py b/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py index ba168f4a2f..c5e16d7c59 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py +++ b/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()