From 37a9d4c5ea3f497ea95f01a93f0824572e011d61 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:49:48 +0800 Subject: [PATCH 01/10] perf(nsa): parallelize long ragged index copies --- .../basemodel/triton_kernel/gen_nsa_ks_ke.py | 89 +++++++++-- .../triton_kernel/test_gen_nsa_ks_ke.py | 148 +++++++++++++++++- 2 files changed, 218 insertions(+), 19 deletions(-) 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() From fb559a86f775281e5932f1cbb4c8915c4f9b83d8 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 07:24:52 +0000 Subject: [PATCH 02/10] perf(glm): eliminate redundant ragged index preparation --- .../basemodel/attention/nsa/fp8_glm5_next.py | 53 +++- .../basemodel/attention/nsa/glm5_next.py | 20 +- .../basemodel/triton_kernel/gen_nsa_ks_ke.py | 89 ++----- lightllm/models/glm5_next/indexer.py | 16 +- .../context_flashattention_nopad_fp8.py | 176 +++++++++---- .../models/glm5_next/triton_kernel/kpool.py | 49 ++-- .../triton_kernel/test_gen_nsa_ks_ke.py | 148 +---------- .../glm5_next/test_direct_nsa_indices.py | 238 ++++++++++++++++++ unit_tests/models/glm5_next/test_fp8_kv.py | 9 + unit_tests/models/glm5_next/test_kernels.py | 19 +- unit_tests/models/glm5_next/test_mtp.py | 6 +- unit_tests/models/glm5_next/test_pd_cache.py | 10 +- 12 files changed, 518 insertions(+), 315 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_direct_nsa_indices.py diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index d245b658c6..cd8e43f8cf 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -3,15 +3,15 @@ import torch -from .fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend, NsaFlashMlaFp8SparseDecodeAttState +from .fp8_flashmla_sparse import NsaFlashMlaFp8SparseDecodeAttState from .glm5_next import Glm5NextSparsePrefillState +from ..base_att import BaseAttBackend -class Fp8Glm5NextSparseAttBackend(NsaFlashMlaFp8SparseAttBackend): +class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - # Match the backend's existing microbatch-separated ragged buffers. - self._prefill_workspaces = [{"smallq": None, "generic": None} for _ in self.ragged_mem_buffers] + self._prefill_workspaces = [{"smallq": None, "generic": None} for _ in range(2)] def get_smallq_workspace(self, infer_state, q): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( @@ -39,7 +39,7 @@ def get_smallq_workspace(self, infer_state, q): "selected_kv": selected["selected_kv"][: batch * slots], } - def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices): + def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices, use_request_table=False): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( allocate_generic_prefill_workspace, ) @@ -52,12 +52,16 @@ def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices) or entry["buffers"]["chunk_indices"].shape[0] < rows or entry["buffers"]["chunk_indices"].shape[1] != columns or entry["buffers"]["chunk_indices"].dtype != topk_indices.dtype - or entry["buffers"]["prefill_row_table"].numel() < packed_kv.shape[0] + or entry["use_request_table"] != use_request_table + or (not use_request_table and entry["buffers"]["prefill_row_table"].numel() < packed_kv.shape[0]) ): capacity = max(self.model.batch_max_tokens or 0, rows) entry = { - "buffers": allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=capacity), + "buffers": allocate_generic_prefill_workspace( + q, packed_kv, topk_indices, query_capacity=capacity, use_request_table=use_request_table + ), "owner": None, + "use_request_table": use_request_table, } slot["generic"] = entry mem_index = infer_state.mem_index @@ -65,7 +69,8 @@ def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices) if entry["owner"] is None or entry["owner"]() is not infer_state or entry["binding"] != binding: buffers = entry["buffers"] # Capacity is persistent; the current batch's mapping is not. - buffers["prefill_row_table"].fill_(-1) + if not use_request_table: + buffers["prefill_row_table"].fill_(-1) entry["workspace"] = { "prefill_row_table": buffers["prefill_row_table"][: packed_kv.shape[0]], "chunk_indices": buffers["chunk_indices"][:rows], @@ -97,6 +102,20 @@ def _nsa_prefill_att(self, q, kv, att_control): ) params = att_control.nsa_prefill_dict + request_layout = {} + index_layout = self.ragged_mem_index + if index_layout is None: + state = self.infer_state + index_layout = params["topk_mem_indices"] + request_layout = { + "topk_mem_indices": index_layout, + "req_to_token_index": state.req_manager.req_to_token_indexs, + "b_req_idx": state.b_req_idx, + "cu_kv_seq_len": state.b1_cu_kv_seq_len, + "total_kv_tokens": state.total_token_num, + "cu_q_seq_len": state.b1_cu_q_seq_len, + } + total_kv_tokens = request_layout.get("total_kv_tokens") # Only H64 uses the small-Q workspace paths. if q.shape[1] != 64: return context_attention_fwd_fp8( @@ -108,6 +127,7 @@ def _nsa_prefill_att(self, q, kv, att_control): self.ks, params["topk_indices"], params["softmax_scale"], + **request_layout, ) from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( @@ -120,20 +140,22 @@ def _nsa_prefill_att(self, q, kv, att_control): if can_use_direct_smallq( q, kv, - self.ragged_mem_index, + index_layout, self.infer_state.mem_index, params["prefill_cache_kv"], params["topk_indices"], + total_kv_tokens, ) else None ) selected_eligible = direct_workspace is None and can_use_selected_prefill( q, kv, - self.ragged_mem_index, + index_layout, self.infer_state.mem_index, params["prefill_cache_kv"], params["topk_indices"], + total_kv_tokens, ) workspace = self.backend.get_smallq_workspace(self.infer_state, q) if selected_eligible else None generic_workspace = None @@ -145,7 +167,7 @@ def _nsa_prefill_att(self, q, kv, att_control): ) if not graph_guard and not torch.cuda.is_current_stream_capturing(): generic_workspace = self.backend.get_generic_prefill_workspace( - self.infer_state, q, kv, params["topk_indices"] + self.infer_state, q, kv, params["topk_indices"], use_request_table=bool(request_layout) ) return context_attention_fwd_fp8( q, @@ -159,12 +181,21 @@ def _nsa_prefill_att(self, q, kv, att_control): workspace, direct_workspace, generic_workspace, + **request_layout, ) @dataclasses.dataclass class Fp8Glm5NextSparseDecodeState(NsaFlashMlaFp8SparseDecodeAttState): decode_topk_bound: int = None + query_req: torch.Tensor = None + + def init_state(self): + import flash_mla + + self.lengths = self.infer_state.b_seq_len + self.query_req = self.infer_state.b_req_idx + self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() def _nsa_decode_att(self, q, packed_kv, att_control): import flash_mla diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py index bccb6a4940..aeb895e9e8 100644 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -3,14 +3,15 @@ import torch +from ..base_att import BaseAttBackend + from .flashmla_sparse import ( - NsaFlashMlaSparseAttBackend, NsaFlashMlaSparsePrefillAttState, NsaFlashMlaSparseDecodeAttState, ) -class Glm5NextSparseAttBackend(NsaFlashMlaSparseAttBackend): +class Glm5NextSparseAttBackend(BaseAttBackend): def create_att_prefill_state(self, infer_state): return Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -20,6 +21,16 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Glm5NextSparsePrefillState(NsaFlashMlaSparsePrefillAttState): + query_req: torch.Tensor = None + + def init_state(self): + from lightllm.models.glm5_next.triton_kernel.kpool import prefill_nsa_params + + state = self.infer_state + self.lengths, self.query_req, self.ks = prefill_nsa_params( + state.position_ids, state.b_req_idx, state.b1_cu_q_seq_len, state.b1_cu_kv_seq_len + ) + def _nsa_prefill_att(self, q, kv, att_control): from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -43,8 +54,11 @@ def _nsa_prefill_att(self, q, kv, att_control): @dataclasses.dataclass class Glm5NextSparseDecodeState(NsaFlashMlaSparseDecodeAttState): + query_req: torch.Tensor = None + def init_state(self): - super().init_state() + self.lengths = self.infer_state.b_seq_len + self.query_req = self.infer_state.b_req_idx pool = self.backend.model.config["index_kpool"] topk = self.backend.model.config["index_topk"] self.nsa_cache_seqlens = ( 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 48cf4ff342..44c64c4d3c 100644 --- a/lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py +++ b/lightllm/common/basemodel/triton_kernel/gen_nsa_ks_ke.py @@ -20,7 +20,6 @@ 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最长的那个。 @@ -66,57 +65,24 @@ def _gen_nsa_ks_ke( 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, - ) + 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, @@ -126,28 +92,15 @@ 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, @@ -163,19 +116,7 @@ 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/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 5df510e5de..3a13662e32 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -57,8 +57,7 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig packed_buffer=packed_buffer, ape=layer_weight.index_kpool_compress_ape.weight, lengths=att_state.lengths, - starts=att_state.ks, - ragged=att_state.ragged_mem_index, + mem_index=infer_state.mem_index, req_idx=infer_state.b_req_idx, cu_q_lens=infer_state.b1_cu_q_seq_len, seq_lens=infer_state.b_seq_len, @@ -67,7 +66,14 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig ) if infer_state.max_kv_seq_len <= self.topk: - return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True) + return expand_topk( + None, + att_state.lengths, + att_state.query_req, + infer_state.req_manager.req_to_token_indexs, + self.topk, + dense=True, + ) # The small indexer is replicated: no all-gather of query heads and # identical pool selection on every TP rank. @@ -82,7 +88,9 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig groups = self._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) else: groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) - return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + return expand_topk( + groups, att_state.lengths, att_state.query_req, infer_state.req_manager.req_to_token_indexs, self.topk + ) def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): keys = gather_pools( diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index e9904e5054..a6782aa439 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -34,6 +34,7 @@ def _direct_smallq_packed_kernel( BLOCKS_: tl.constexpr, SLOTS_: tl.constexpr, softmax_scale: tl.constexpr, + PHYSICAL_INDICES: tl.constexpr, ): qid = tl.program_id(0) split = tl.program_id(1) @@ -57,7 +58,8 @@ def _direct_smallq_packed_kernel( acc = tl.zeros((BLOCK_H_, 512), tl.float32) m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) denominator = tl.zeros((BLOCK_H_,), tl.float32) - start = tl.load(KVStarts + qid).to(tl.int64) + if not PHYSICAL_INDICES: + start = tl.load(KVStarts + qid).to(tl.int64) for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): n = tl.arange(0, 64) slot = split * SLOTS_ + b * 64 + n @@ -66,9 +68,13 @@ def _direct_smallq_packed_kernel( slot < 2176, -1, ).to(tl.int64) - ragged_pos = start + top - valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) - mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + if PHYSICAL_INDICES: + valid = (top >= 0) & (top < ragged_n) + mem = tl.where(valid, top, 0) + else: + ragged_pos = start + top + valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) current_row = tl.full((64,), -1, tl.int32) for j in tl.static_range(PREFILL_N): current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) @@ -149,7 +155,9 @@ def reduce_direct_smallq_kernel( tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) -def can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): +def can_use_direct_smallq( + q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens=None +): return ( q.ndim == 3 and 1 <= q.shape[0] <= 8 @@ -159,7 +167,7 @@ def can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, pre and packed_kv.shape[1] == 1 and packed_kv.shape[2] >= 584 and ragged_mem_index.dtype == torch.int32 - and ragged_mem_index.numel() >= 65536 + and (ragged_mem_index.numel() if total_kv_tokens is None else total_kv_tokens) >= 65536 and prefill_mem_index.dtype == torch.int32 and prefill_mem_index.numel() == q.shape[0] and prefill_kv.dtype == torch.bfloat16 @@ -191,6 +199,7 @@ def _context_attention_fwd_fp8_direct_smallq( topk_indices, softmax_scale, direct_workspace=None, + topk_mem_indices=None, ): qn = q.shape[0] heads = q.shape[1] @@ -209,6 +218,7 @@ def _context_attention_fwd_fp8_direct_smallq( partial_lse = workspace["partial_lse"] output = workspace["output"] launch_grid = (qn, splits, triton.cdiv(heads, block_h)) + indices = topk_indices if topk_mem_indices is None else topk_mem_indices _direct_smallq_packed_kernel[launch_grid]( q, packed_kv, @@ -216,18 +226,18 @@ def _context_attention_fwd_fp8_direct_smallq( prefill_mem_index, prefill_kv, kv_starts, - topk_indices, + indices, partial_o, partial_lse, - ragged_mem_index.numel(), + ragged_mem_index.numel() if topk_mem_indices is None else packed_kv.shape[0], q.stride(0), q.stride(1), q.stride(2), packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - topk_indices.stride(0), - topk_indices.stride(1), + indices.stride(0), + indices.stride(1), HEADS_=heads, BLOCK_H_=block_h, PREFILL_N=qn, @@ -235,6 +245,7 @@ def _context_attention_fwd_fp8_direct_smallq( BLOCKS_=blocks, SLOTS_=slots, softmax_scale=softmax_scale, + PHYSICAL_INDICES=topk_mem_indices is not None, num_warps=8, ) reduce_direct_smallq_kernel[(qn, heads, 4)]( @@ -263,18 +274,51 @@ def _gather_prefill_kv_kernel( PrefillRowTable, PrefillKV, Output, + ReqToToken, + ReqIdx, + CuKV, + CuQ, + chunk_start, num_tokens, packed_stride_s: tl.constexpr, prefill_stride_s: tl.constexpr, prefill_stride_d: tl.constexpr, + REQ_STRIDE: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + BATCH: tl.constexpr, + USE_REQ_TABLE: tl.constexpr, BLOCK_ROWS: tl.constexpr, ): - row = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + block = tl.program_id(0) + if USE_REQ_TABLE: + lanes = tl.arange(0, triton.next_power_of_2(BATCH)) + ends = tl.load(CuKV + lanes + 1, lanes < BATCH, 0) + ends = tl.minimum(tl.maximum(ends - chunk_start, 0), num_tokens) + block_ends = ends // BLOCK_ROWS + lanes + 1 + batch = tl.sum(((block >= block_ends) & (lanes < BATCH)).to(tl.int32), 0) + start = tl.load(CuKV + batch) + copy_start = tl.minimum(tl.maximum(start - chunk_start, 0), num_tokens) + copy_end = tl.minimum(tl.maximum(tl.load(CuKV + batch + 1) - chunk_start, 0), num_tokens) + local_block = block - (copy_start // BLOCK_ROWS + batch) + row = copy_start + local_block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + live = row < copy_end + req = tl.load(ReqIdx + batch).to(tl.int64) + position = chunk_start + row - start + q_start = tl.load(CuQ + batch) + q_len = tl.load(CuQ + batch + 1) - q_start + cache_len = tl.load(CuKV + batch + 1) - start - q_len + prefill_row = tl.where(position >= cache_len, q_start + position - cache_len, -1) + mem_index = tl.load(ReqToToken + req * REQ_STRIDE + position * TOKEN_STRIDE, live & (prefill_row < 0), 0).to( + tl.int64 + ) + else: + row = block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + live = row < num_tokens + mem_index = tl.load(MemIndex + row, live, 0).to(tl.int64) + prefill_row = tl.load(PrefillRowTable + mem_index, live, -1) dim = tl.arange(0, 512) - mem_index = tl.load(MemIndex + row, row < num_tokens, 0).to(tl.int64) - prefill_row = tl.load(PrefillRowTable + mem_index, row < num_tokens, -1) - cached = (row < num_tokens) & (prefill_row < 0) - current = (row < num_tokens) & (prefill_row >= 0) + cached = live & (prefill_row < 0) + current = live & (prefill_row >= 0) packed_row = PackedKV + mem_index[:, None] * packed_stride_s quantized = tl.load( @@ -298,7 +342,7 @@ def _gather_prefill_kv_kernel( tl.store( Output + row[:, None] * 512 + dim[None, :], tl.where(current[:, None], current_values, values), - row[:, None] < num_tokens, + live[:, None], ) @@ -387,13 +431,18 @@ def _selected_prefill_compact_kernel( topk_stride_d: tl.constexpr, BATCH: tl.constexpr, BLOCK: tl.constexpr, + PHYSICAL_INDICES: tl.constexpr, ): qid = tl.program_id(0) lane = tl.arange(0, BLOCK) local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) - ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local - valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) - physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + if PHYSICAL_INDICES: + valid = (lane < 2176) & (local >= 0) + physical = local + else: + ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local + valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) valid = valid & (physical >= 0) & (physical < packed_rows) rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) @@ -466,7 +515,9 @@ def _selected_prefill_gather_kernel( tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) -def can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): +def can_use_selected_prefill( + q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens=None +): return ( q.ndim == 3 and 2 <= q.shape[0] <= 8 @@ -479,7 +530,7 @@ def can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, and packed_kv.device == q.device and ragged_mem_index.dtype == torch.int32 and ragged_mem_index.device == q.device - and ragged_mem_index.numel() >= 65536 + and (ragged_mem_index.numel() if total_kv_tokens is None else total_kv_tokens) >= 65536 and prefill_mem_index.dtype == torch.int32 and prefill_mem_index.device == q.device and prefill_mem_index.shape == (q.shape[0],) @@ -504,11 +555,13 @@ def allocate_selected_prefill_workspace(q, batch_capacity=None): } -def allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None): +def allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None, use_request_table=False): """Allocate generic prefill metadata up to the backend's planned query capacity.""" capacity = q.shape[0] if query_capacity is None else query_capacity return { - "prefill_row_table": torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device), + "prefill_row_table": torch.full( + (0 if use_request_table else packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device + ), "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), @@ -527,26 +580,29 @@ def _context_attention_fwd_fp8_selected_prefill( topk_indices, softmax_scale, selected_workspace=None, + topk_mem_indices=None, ): from sgl_kernel.flash_attn import flash_attn_with_kvcache batch = q.shape[0] workspace = selected_workspace or allocate_selected_prefill_workspace(q) batch_pad = triton.next_power_of_2(batch) + indices = topk_indices if topk_mem_indices is None else topk_mem_indices _selected_prefill_compact_kernel[(batch,)]( - topk_indices, + indices, kv_starts, ragged_mem_index, workspace["compact"], workspace["lengths"], workspace["page_table"], workspace["cu_q"], - ragged_mem_index.numel(), + ragged_mem_index.numel() if ragged_mem_index is not None else 0, packed_kv.shape[0], - topk_indices.stride(0), - topk_indices.stride(1), + indices.stride(0), + indices.stride(1), BATCH=batch, BLOCK=4096, + PHYSICAL_INDICES=topk_mem_indices is not None, num_warps=8, ) _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( @@ -593,11 +649,22 @@ def context_attention_fwd_fp8( selected_workspace=None, direct_workspace=None, generic_workspace=None, + topk_mem_indices=None, + req_to_token_index=None, + b_req_idx=None, + cu_kv_seq_len=None, + total_kv_tokens=None, + cu_q_seq_len=None, ): - """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV.""" + """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV. + + GLM passes physical top-k slots and the existing request table/prefix sums, + leaving ragged_mem_index unset. The ragged form supports standalone callers. + """ # Direct packed attention is intentionally restricted to the validated # small-Q GLM-5.3 shape; all other cases retain the bounded-gather path. - if can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + index_layout = ragged_mem_index if ragged_mem_index is not None else topk_mem_indices + if can_use_direct_smallq(q, packed_kv, index_layout, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens): return _context_attention_fwd_fp8_direct_smallq( q, packed_kv, @@ -608,8 +675,11 @@ def context_attention_fwd_fp8( topk_indices, softmax_scale, direct_workspace, + topk_mem_indices, ) - if can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + if can_use_selected_prefill( + q, packed_kv, index_layout, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens + ): return _context_attention_fwd_fp8_selected_prefill( q, packed_kv, @@ -620,17 +690,22 @@ def context_attention_fwd_fp8( topk_indices, softmax_scale, selected_workspace, + topk_mem_indices, ) from sgl_kernel.flash_attn import flash_attn_with_kvcache num_q_tokens, num_heads, _ = q.shape - gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, ragged_mem_index.numel()) + use_req_table = ragged_mem_index is None + num_kv_tokens = total_kv_tokens if use_req_table else ragged_mem_index.numel() + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, num_kv_tokens) if generic_workspace is None: # Build the row table before gathering KV and allocating compaction metadata. - prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) - _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( - prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 - ) + prefill_row_table = None + if not use_req_table: + prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) # Keep this potentially 2 GiB buffer local to a DSA layer. gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) chunk_indices = torch.empty_like(topk_indices) @@ -641,7 +716,7 @@ def context_attention_fwd_fp8( chunk_indices = generic_workspace["chunk_indices"] chunk_lengths = generic_workspace["chunk_lengths"] cu_seqlens_q = generic_workspace["cu_seqlens_q"] - if not generic_workspace["row_table_built"]: + if not use_req_table and not generic_workspace["row_table_built"]: _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 ) @@ -649,26 +724,41 @@ def context_attention_fwd_fp8( # Keep this potentially 2 GiB buffer local to a DSA layer. gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) output = lse = None - for chunk_start in range(0, ragged_mem_index.numel(), gather_chunk_size): - num_chunk_tokens = min(gather_chunk_size, ragged_mem_index.numel() - chunk_start) - _gather_prefill_kv_kernel[(triton.cdiv(num_chunk_tokens, 4),)]( + for chunk_start in range(0, num_kv_tokens, gather_chunk_size): + num_chunk_tokens = min(gather_chunk_size, num_kv_tokens - chunk_start) + if use_req_table: + gather_rows = 4 + gather_grid = (num_chunk_tokens // gather_rows + b_req_idx.numel(),) + else: + gather_rows = 4 + gather_grid = (triton.cdiv(num_chunk_tokens, 4),) + _gather_prefill_kv_kernel[gather_grid]( packed_kv, - ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], + None if use_req_table else ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], prefill_row_table, prefill_kv, gathered_kv, + req_to_token_index, + b_req_idx, + cu_kv_seq_len, + cu_q_seq_len, + chunk_start, num_chunk_tokens, packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - BLOCK_ROWS=4, + REQ_STRIDE=req_to_token_index.stride(0) if use_req_table else 0, + TOKEN_STRIDE=req_to_token_index.stride(1) if use_req_table else 0, + BATCH=b_req_idx.numel() if use_req_table else 1, + USE_REQ_TABLE=use_req_table, + BLOCK_ROWS=gather_rows, ) chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) if ( num_q_tokens > 8 and num_heads == 64 - and ragged_mem_index.numel() > num_q_tokens - and ragged_mem_index.numel() <= KV_GATHER_CHUNK_SIZE + and num_kv_tokens > num_q_tokens + and num_kv_tokens <= KV_GATHER_CHUNK_SIZE and q.dtype == torch.bfloat16 and q.is_contiguous() and packed_kv.dtype == torch.uint8 diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index e08d1e456b..dd2f4c0216 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -28,6 +28,25 @@ def _get_query_block(CuQLens, BATCH, BLOCK: tl.constexpr): return batch, q_start, q_end, local_block +@triton.jit +def _prefill_nsa_params(PositionIds, ReqIdx, CuQ, CuKV, Lengths, QueryReq, KVStarts, BATCH, BLOCK: tl.constexpr): + batch, q_start, q_end, block = _get_query_block(CuQ, BATCH, BLOCK) + rows = q_start + block * BLOCK + tl.arange(0, BLOCK) + valid = rows < q_end + tl.store(Lengths + rows, tl.load(PositionIds + rows, valid, 0) + 1, valid) + tl.store(QueryReq + rows, tl.load(ReqIdx + batch), valid) + tl.store(KVStarts + rows, tl.load(CuKV + batch), valid) + + +def prefill_nsa_params(position_ids, req_idx, cu_q_lens, cu_kv_lens): + """Expand existing request metadata to query rows without visiting historical KV.""" + lengths, query_req, kv_starts = [torch.empty_like(position_ids) for _ in range(3)] + _prefill_nsa_params[(position_ids.numel() // 256 + req_idx.numel(),)]( + position_ids, req_idx, cu_q_lens, cu_kv_lens, lengths, query_req, kv_starts, req_idx.numel(), BLOCK=256 + ) + return lengths, query_req, kv_starts + + @triton.jit def _compress_pools( Raw, @@ -35,8 +54,7 @@ def _compress_pools( Packed, Ape, Lengths, - Starts, - Ragged, + MemIndex, ReqIdx, CuQLens, MtpIndex, @@ -67,7 +85,6 @@ def _compress_pools( valid = row < q_end req = tl.load(ReqIdx + batch) if valid and req != HOLD_REQ: - start = tl.load(Starts + row) pool = tl.arange(0, 4) cols = tl.arange(0, 128) raw_rows = row - 3 + pool @@ -94,7 +111,8 @@ def _compress_pools( key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) - loc = tl.load(Ragged + start + length - 1).to(tl.int64) + # A completed pool lives at its closing token, which is a current query row. + loc = tl.load(MemIndex + row).to(tl.int64) dest = Packed + loc * PACKED_STRIDE tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) @@ -135,7 +153,7 @@ def _save_pool_tails( def compress_pools( - raw, tail, packed_buffer, ape, lengths, starts, ragged, req_idx, cu_q_lens, seq_lens, max_q_len, mtp_index=None + raw, tail, packed_buffer, ape, lengths, mem_index, req_idx, cu_q_lens, seq_lens, max_q_len, mtp_index=None ): batch = req_idx.numel() single_query = max_q_len == 1 and lengths.numel() == batch @@ -146,8 +164,7 @@ def compress_pools( packed_buffer, ape, lengths, - starts, - ragged, + mem_index, req_idx, cu_q_lens, mtp_index, @@ -384,10 +401,12 @@ def gather_paged_pools(packed_buffer, req_to_token_indexs, b_req_idx, pool_lengt def _expand_topk( Groups, Lengths, - Starts, - Ragged, + QueryReq, + ReqToToken, Out, Relative, + REQ_STRIDE: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, TOPK: tl.constexpr, WIDTH: tl.constexpr, DENSE: tl.constexpr, @@ -396,7 +415,7 @@ def _expand_topk( row = tl.program_id(0) lane = tl.arange(0, BLOCK) length = tl.load(Lengths + row) - start = tl.load(Starts + row) + req = tl.load(QueryReq + row).to(tl.int64) if DENSE: token = lane valid = lane < length @@ -405,22 +424,24 @@ def _expand_topk( group = tl.load(Groups + row * (TOPK // 4) + lane // 4, lane < closed_tokens, -1) token = tl.where(lane < closed_tokens, group * 4 + lane % 4, length // 4 * 4 + lane - closed_tokens) valid = (lane < closed_tokens + length % 4) & (token >= 0) - mem = tl.load(Ragged + start + token, valid & (lane < WIDTH), -1) + mem = tl.load(ReqToToken + req * REQ_STRIDE + token * TOKEN_STRIDE, valid & (lane < WIDTH), -1) tl.store(Out + row * WIDTH + lane, mem, lane < WIDTH) tl.store(Relative + row * WIDTH + lane, tl.where(valid, token, -1), lane < WIDTH) -def expand_topk(groups, lengths, starts, ragged, topk, dense=False): +def expand_topk(groups, lengths, query_req, req_to_token, topk, dense=False): width = triton.cdiv(topk + 3, 128) * 128 out = torch.empty((lengths.numel(), width), dtype=torch.int32, device=lengths.device) relative = torch.empty_like(out) _expand_topk[(lengths.numel(),)]( groups, lengths, - starts, - ragged, + query_req, + req_to_token, out, relative, + req_to_token.stride(0), + req_to_token.stride(1), topk, width, dense, 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 c5e16d7c59..ba168f4a2f 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,9 +1,6 @@ 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(): @@ -36,20 +33,14 @@ 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 @@ -64,12 +55,7 @@ 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( @@ -97,131 +83,5 @@ 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() diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py new file mode 100644 index 0000000000..732bfa10a8 --- /dev/null +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -0,0 +1,238 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.nsa.glm5_next import ( + Glm5NextSparseAttBackend, + Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + Fp8Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv +from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill +from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def tensor(values): + return torch.tensor(values, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("counts", [(3, 0, 1, 2), (257, 0, 1, 513), (1,) * 129]) +def test_prefill_metadata_reuses_existing_prefixes_during_graph_replay(counts): + batch, queries = len(counts), sum(counts) + positions = torch.empty(queries, device="cuda", dtype=torch.int32) + req = torch.empty(batch, device="cuda", dtype=torch.int32) + cu_q, cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32), torch.empty( + batch + 1, device="cuda", dtype=torch.int32 + ) + + def update(counts, phase): + q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] + req_values = [(batch - i + phase) % (batch + 1) for i in range(batch)] + for i, count in enumerate(counts): + prefix = 1000000 + i * 7 + phase + visible.extend(range(prefix + 1, prefix + count + 1)) + starts.extend([kv_prefix[-1]] * count) + requests.extend([req_values[i]] * count) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + prefix + count) + positions.copy_(tensor(visible) - 1) + req.copy_(tensor(req_values)) + cu_q.copy_(tensor(q_prefix)) + cu_kv.copy_(tensor(kv_prefix)) + return tensor(visible), tensor(requests), tensor(starts) + + expected = update(counts, 0) + actual = prefill_nsa_params(positions, req, cu_q, cu_kv) + for observed, reference in zip(actual, expected): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = prefill_nsa_params(positions, req, cu_q, cu_kv) + expected = update(counts[::-1], 11) + graph.replay() + for observed, reference in zip(actual, expected): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize("dense", [False, True]) +def test_topk_direct_request_lookup_preserves_mtp_hold_and_strided_tables(dense): + table = torch.arange(8 * 70, device="cuda", dtype=torch.int32).reshape(8, 70)[::2, ::2] + req = tensor([2, 2, 0, 3, 3]) + lengths = tensor([5, 7, 9, 2, 2] if dense else [17, 18, 9, 2, 2]) + groups = tensor([[3, 1], [1, 0], [0, 1], [0, 0], [0, 0]]) + + def check(actual): + mem, relative = actual + expected_mem, expected_relative = torch.full_like(mem, -1), torch.full_like(relative, -1) + for i, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): + tokens = ( + list(range(length)) + if dense + else [group * 4 + j for group in groups[i].tolist()[: min(length // 4, 2)] for j in range(4)] + + list(range(length // 4 * 4, length)) + ) + expected_relative[i, : len(tokens)] = tensor(tokens) + expected_mem[i, : len(tokens)] = table[request, tensor(tokens).long()] + torch.testing.assert_close(mem, expected_mem, rtol=0, atol=0) + torch.testing.assert_close(relative, expected_relative, rtol=0, atol=0) + + check(expand_topk(groups, lengths, req, table, 8, dense=dense)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = expand_topk(groups, lengths, req, table, 8, dense=dense) + req.copy_(tensor([0, 0, 2, 3, 3])) + table.add_(10000) + lengths[:2].add_(1) + graph.replay() + check(actual) + + +@pytest.mark.parametrize("fp8", [False, True]) +def test_decode_aliases_infer_state_without_ragged_allocation(monkeypatch, fp8): + import lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke as legacy + + def forbidden(*args, **kwargs): + raise AssertionError("GLM must not generate or copy a full ragged index") + + monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) + model = type("Model", (), {"config": {"index_kpool": 4, "index_topk": 2048}})() + backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) + assert not hasattr(backend, "ragged_mem_buffers") + infer = InferStateInfo() + infer.b_req_idx = tensor([2, 2, 0, 3, 3]) + infer.b_seq_len = tensor([999999, 1000000, 65536, 2, 2]) + cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState + state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + state.init_state() + assert state.lengths is infer.b_seq_len + assert state.query_req is infer.b_req_idx + assert state.ragged_mem_index is None and state.ks is None and state.ke is None + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + state.init_state() + infer.b_seq_len[:3].sub_(100) + infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) + graph.replay() + assert state.lengths is infer.b_seq_len and state.query_req is infer.b_req_idx + + +@pytest.mark.parametrize( + "queries,heads,total_kv,chunk_size", + [ + (1, 64, 65536, 2097152), + (2, 64, 65536, 2097152), + (3, 8, 65536, 2097152), + (9, 64, 512, 2097152), + (9, 64, 131073, 2097152), + (4, 16, 40, 3), + (4, 64, 40, 7), + ], +) +def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + + class ForbiddenRowTable: + def __getitem__(self, grid): + raise AssertionError("Request-layout prefill must not build a physical-slot row table") + + monkeypatch.setattr(fp8_prefill, "_build_prefill_row_table_kernel", ForbiddenRowTable()) + torch.manual_seed(1685) + batch = min(queries, 2) + counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] + seqs = [total_kv] if batch == 1 else [total_kv // 2, total_kv - total_kv // 2] + req = tensor([2, 0][:batch]) + table = torch.full((4, 2 * (max(seqs) + 7)), -1, device="cuda", dtype=torch.int32)[:, ::2] + physical = torch.randperm(total_kv + 32, device="cuda", dtype=torch.int32)[:total_kv] + offset = 0 + for request, length in zip(req.tolist(), seqs): + table[request, :length] = physical[offset : offset + length] + offset += length + q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) + source = torch.randn(total_kv + 32, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.zeros(total_kv + 32, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, torch.arange(total_kv + 32, device="cuda", dtype=torch.int32), packed) + prefill_kv = torch.randn(queries, 1, 512, device="cuda", dtype=torch.bfloat16) + visible, query_req, current_positions, q_prefix, kv_prefix = [], [], [], [0], [0] + for request, length, count in zip(req.tolist(), seqs, counts): + visible.extend(range(length - count + 1, length + 1)) + query_req.extend([request] * count) + current_positions.extend(range(length - count, length)) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + length) + query_req, current_positions = tensor(query_req), tensor(current_positions) + mem_index = table[query_req.long(), current_positions.long()] + topk = torch.full((queries, 2176), -1, device="cuda", dtype=torch.int32) + for i, length in enumerate(visible): + topk[i, :7] = tensor([0, length - 1, 0, 1, length // 2, length - 1, -1]) + topk_mem = torch.empty_like(topk) + + def refresh_physical(): + topk_mem.copy_(table[query_req.long()[:, None], topk.clamp_min(0).long()]) + topk_mem.masked_fill_(topk < 0, -1) + mem_index.copy_(table[query_req.long(), current_positions.long()]) + + refresh_physical() + infer = InferStateInfo() + infer.mem_index, infer.b_req_idx = mem_index, req + infer.position_ids = tensor(visible) - 1 + infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) + infer.total_token_num = total_kv + infer.req_manager = SimpleNamespace(req_to_token_indexs=table) + model = type("Model", (), {"batch_max_tokens": queries, "prefill_graph": None})() + backend = Fp8Glm5NextSparseAttBackend(model) + state = backend.create_att_prefill_state(infer) + state.init_state() + assert state.ragged_mem_index is None and not hasattr(backend, "ragged_mem_buffers") + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk, + "topk_mem_indices": topk_mem, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 512 ** -0.5, + "kv_lora_rank": 512, + }, + ) + + def run(): + return state.prefill_att(q, packed[..., :584], None, control) + + def check(output): + scales = packed[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) + values = packed[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) * scales + kv = torch.cat((values.reshape(-1, 1, 448).bfloat16(), packed[..., 448:576].view(torch.bfloat16)), -1) + kv[mem_index.long()] = prefill_kv + expected = [] + for i in range(queries): + keys = kv[topk_mem[i][topk_mem[i] >= 0].long(), 0].float() + expected.append((q[i].float() @ keys.T * 512 ** -0.5).softmax(-1) @ keys) + torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) + + check(run()) + generic = backend._prefill_workspaces[0]["generic"] + if generic is not None: + assert generic["buffers"]["prefill_row_table"].numel() == 0 + assert not generic["workspace"]["row_table_built"] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + graph.replay() + check(output) + q.normal_() + prefill_kv.normal_() + for request, length in zip(req.tolist(), seqs): + table[request, :length].copy_(table[request, :length].roll(1)) + topk[0].fill_(-1) + refresh_physical() + graph.replay() + check(output) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 8b4e7aba69..d47d9f808f 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -152,10 +152,19 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid row_table, prefill_kv, gathered_kv, + None, + None, + None, + None, + 0, indices.numel(), packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), + REQ_STRIDE=0, + TOKEN_STRIDE=0, + BATCH=1, + USE_REQ_TABLE=False, BLOCK_ROWS=4, ) expected = _dequantize_kv_cache(packed_kv) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 0ea5551fff..c382930a82 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -211,8 +211,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): packed, ape, tensor(lengths), - tensor(starts), - tensor(locations), + tensor(locations)[tensor(starts) + tensor(lengths) - 1], tensor([r for r, _, _ in chunks]), tensor(cu_q_lens), tensor([end for _, _, end in chunks]), @@ -248,7 +247,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) starts = torch.tensor([0, 9], device="cuda", dtype=torch.int32) groups = torch.tensor([[1, 0], [0, -1]], device="cuda", dtype=torch.int32) - out, relative = expand_topk(groups, lengths, starts, ragged, topk=8) + out, relative = expand_topk(groups, lengths, req_idx, table, topk=8) for batch, expected in enumerate([[4, 5, 6, 7, 0, 1, 2, 3, 8], list(range(7))]): assert relative[batch, : len(expected)].tolist() == expected assert ( @@ -270,12 +269,11 @@ def test_kpool_decode_cuda_graph_and_padding(): raw = torch.zeros(4, 256, device="cuda", dtype=torch.bfloat16) req_idx = torch.tensor([1, 3, 4, 4], device="cuda", dtype=torch.int32) lengths = torch.ones(4, device="cuda", dtype=torch.int32) - starts = torch.arange(4, device="cuda", dtype=torch.int32) * 8 - ragged = table.flatten().clone() + mem_index = torch.zeros(4, device="cuda", dtype=torch.int32) cu_q_lens = torch.arange(5, device="cuda", dtype=torch.int32) def run(): - compress_pools(raw, tail, packed, ape, lengths, starts, ragged, req_idx, cu_q_lens, lengths, max_q_len=1) + compress_pools(raw, tail, packed, ape, lengths, mem_index, req_idx, cu_q_lens, lengths, max_q_len=1) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -291,7 +289,7 @@ def run(): req_idx[:2] = torch.tensor([1 if i == 0 else 3 for i in order], device="cuda") raw[:2] = source[order, pos] lengths.fill_(pos + 1) - ragged[:16] = table[order].flatten() + mem_index[:2] = table[order, pos] graph.replay() for batch, req in enumerate([1, 3]): pending = (pos + 1) % 4 @@ -404,7 +402,7 @@ def test_kpool_packed_kv_above_two_gib(): zero = torch.zeros(4, device="cuda", dtype=torch.int32) ragged = torch.tensor([0, 1, 2, loc], device="cuda", dtype=torch.int32) cu_q_lens = torch.tensor([0, 4], device="cuda", dtype=torch.int32) - compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) + compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) @@ -614,8 +612,7 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): ) state = SimpleNamespace( lengths=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), - ks=torch.zeros(tokens, device="cuda", dtype=torch.int32), - ragged_mem_index=ragged, + query_req=torch.zeros(tokens, device="cuda", dtype=torch.int32), ) indexer = Glm5NextNsaInfer( 0, @@ -641,7 +638,7 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): infer.max_kv_seq_len = pos + 1 infer.b_seq_len.fill_(pos + 1) state.lengths = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) - state.ks = torch.zeros_like(state.lengths) + state.query_req = infer.b_req_idx _, decoded = indexer._get_indices(hidden[pos : pos + 1], hidden[pos : pos + 1], infer, state, weights) valid = 2048 + (pos + 1) % 4 assert set(decoded[0, :valid].tolist()) == set(full[pos, :valid].tolist()) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 82e029c5dd..354521250f 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -91,8 +91,7 @@ def prefill(values, destination, tail): destination, ape, torch.arange(1, count + 1, device="cuda", dtype=torch.int32), - torch.zeros(count, device="cuda", dtype=torch.int32), - ragged, + ragged[:count], torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([0, count], device="cuda", dtype=torch.int32), torch.tensor([count], device="cuda", dtype=torch.int32), @@ -110,8 +109,7 @@ def prefill(values, destination, tail): packed, ape, lengths, - torch.zeros(count, device="cuda", dtype=torch.int32), - ragged, + ragged[start : start + count], torch.zeros(count, device="cuda", dtype=torch.int32), torch.arange(count + 1, device="cuda", dtype=torch.int32), lengths, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index ff93312516..1e14047c24 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -267,15 +267,13 @@ def test_pd_kv_and_runtime_state_roundtrip( outputs = [] for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): new_indexes = [i for i in range(32) if i not in indexes][:count] - all_indexes = indexes + new_indexes compress_pools( raw=raw, tail=mem.tail.buffer[0], packed_buffer=mem.full.get_layer_view(3).indexer_k, ape=ape, lengths=tensor(range(length + 1, length + count + 1)), - starts=tensor([0] * count), - ragged=tensor(all_indexes), + mem_index=tensor(new_indexes), req_idx=tensor([req_idx]), cu_q_lens=tensor([0, count]), seq_lens=tensor([length + count]), @@ -357,8 +355,7 @@ def prefill(values, destination, tail): device="cuda", dtype=torch.int32, ), - starts=torch.zeros(count, device="cuda", dtype=torch.int32), - ragged=ragged, + mem_index=ragged[:count], req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), cu_q_lens=torch.tensor([0, count], device="cuda", dtype=torch.int32), seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), @@ -412,8 +409,7 @@ def prefill(values, destination, tail): packed_buffer=packed, ape=ape, lengths=lengths, - starts=torch.zeros(count, device="cuda", dtype=torch.int32), - ragged=ragged, + mem_index=ragged[start : start + count], req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), cu_q_lens=torch.arange(count + 1, device="cuda", dtype=torch.int32), seq_lens=lengths, From 4106af4b4cf7313dc833a62191fc6d9639c8fa72 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 08:47:50 +0000 Subject: [PATCH 03/10] refactor(glm): remove legacy ragged prefill compatibility --- .../basemodel/attention/nsa/fp8_glm5_next.py | 166 ++-- .../basemodel/attention/nsa/glm5_next.py | 24 +- .../context_flashattention_nopad_fp8.py | 288 ++----- .../glm5_next/test_direct_nsa_indices.py | 12 +- unit_tests/models/glm5_next/test_fp8_kv.py | 732 +++++------------- 5 files changed, 364 insertions(+), 858 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index cd8e43f8cf..2d1205c714 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -3,8 +3,7 @@ import torch -from .fp8_flashmla_sparse import NsaFlashMlaFp8SparseDecodeAttState -from .glm5_next import Glm5NextSparsePrefillState +from .glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState from ..base_att import BaseAttBackend @@ -39,48 +38,29 @@ def get_smallq_workspace(self, infer_state, q): "selected_kv": selected["selected_kv"][: batch * slots], } - def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices, use_request_table=False): + def get_generic_prefill_workspace(self, infer_state, q, topk_indices): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( allocate_generic_prefill_workspace, ) slot = self._prefill_workspaces[infer_state.microbatch_index] - entry = slot["generic"] + buffers = slot["generic"] rows, columns = topk_indices.shape if ( - entry is None - or entry["buffers"]["chunk_indices"].shape[0] < rows - or entry["buffers"]["chunk_indices"].shape[1] != columns - or entry["buffers"]["chunk_indices"].dtype != topk_indices.dtype - or entry["use_request_table"] != use_request_table - or (not use_request_table and entry["buffers"]["prefill_row_table"].numel() < packed_kv.shape[0]) + buffers is None + or buffers["chunk_indices"].shape[0] < rows + or buffers["chunk_indices"].shape[1] != columns + or buffers["chunk_indices"].dtype != topk_indices.dtype ): - capacity = max(self.model.batch_max_tokens or 0, rows) - entry = { - "buffers": allocate_generic_prefill_workspace( - q, packed_kv, topk_indices, query_capacity=capacity, use_request_table=use_request_table - ), - "owner": None, - "use_request_table": use_request_table, - } - slot["generic"] = entry - mem_index = infer_state.mem_index - binding = (id(mem_index), mem_index.data_ptr(), tuple(topk_indices.shape), packed_kv.shape[0]) - if entry["owner"] is None or entry["owner"]() is not infer_state or entry["binding"] != binding: - buffers = entry["buffers"] - # Capacity is persistent; the current batch's mapping is not. - if not use_request_table: - buffers["prefill_row_table"].fill_(-1) - entry["workspace"] = { - "prefill_row_table": buffers["prefill_row_table"][: packed_kv.shape[0]], - "chunk_indices": buffers["chunk_indices"][:rows], - "chunk_lengths": buffers["chunk_lengths"][:rows], - "cu_seqlens_q": buffers["cu_seqlens_q"][: rows + 1], - "row_table_built": False, - } - entry["owner"] = weakref.ref(infer_state) - entry["binding"] = binding - return entry["workspace"] + buffers = allocate_generic_prefill_workspace( + q, topk_indices, query_capacity=max(self.model.batch_max_tokens or 0, rows) + ) + slot["generic"] = buffers + return { + "chunk_indices": buffers["chunk_indices"][:rows], + "chunk_lengths": buffers["chunk_lengths"][:rows], + "cu_seqlens_q": buffers["cu_seqlens_q"][: rows + 1], + } def create_att_prefill_state(self, infer_state): return Fp8Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -99,96 +79,52 @@ class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): def _nsa_prefill_att(self, q, kv, att_control): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( context_attention_fwd_fp8, - ) - - params = att_control.nsa_prefill_dict - request_layout = {} - index_layout = self.ragged_mem_index - if index_layout is None: - state = self.infer_state - index_layout = params["topk_mem_indices"] - request_layout = { - "topk_mem_indices": index_layout, - "req_to_token_index": state.req_manager.req_to_token_indexs, - "b_req_idx": state.b_req_idx, - "cu_kv_seq_len": state.b1_cu_kv_seq_len, - "total_kv_tokens": state.total_token_num, - "cu_q_seq_len": state.b1_cu_q_seq_len, - } - total_kv_tokens = request_layout.get("total_kv_tokens") - # Only H64 uses the small-Q workspace paths. - if q.shape[1] != 64: - return context_attention_fwd_fp8( - q, - kv, - self.ragged_mem_index, - self.infer_state.mem_index, - params["prefill_cache_kv"], - self.ks, - params["topk_indices"], - params["softmax_scale"], - **request_layout, - ) - - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( can_use_direct_smallq, can_use_selected_prefill, ) - direct_workspace = ( - self.backend.get_smallq_workspace(self.infer_state, q) - if can_use_direct_smallq( - q, - kv, - index_layout, - self.infer_state.mem_index, - params["prefill_cache_kv"], - params["topk_indices"], - total_kv_tokens, - ) - else None - ) - selected_eligible = direct_workspace is None and can_use_selected_prefill( - q, - kv, - index_layout, - self.infer_state.mem_index, - params["prefill_cache_kv"], - params["topk_indices"], - total_kv_tokens, - ) - workspace = self.backend.get_smallq_workspace(self.infer_state, q) if selected_eligible else None - generic_workspace = None - if direct_workspace is None and not selected_eligible: - graph_guard = ( - self.infer_state.is_cuda_graph_capture - or self.backend.model.prefill_graph is not None - or hasattr(self.infer_state, "prefill_cuda_graph_exe_list") - ) - if not graph_guard and not torch.cuda.is_current_stream_capturing(): - generic_workspace = self.backend.get_generic_prefill_workspace( - self.infer_state, q, kv, params["topk_indices"], use_request_table=bool(request_layout) + params = att_control.nsa_prefill_dict + state = self.infer_state + topk_mem_indices = params["topk_mem_indices"] + direct_workspace = selected_workspace = generic_workspace = None + if q.shape[1] == 64: + eligibility = (q, kv, state.mem_index, params["prefill_cache_kv"], topk_mem_indices, state.total_token_num) + if can_use_direct_smallq(*eligibility): + direct_workspace = self.backend.get_smallq_workspace(state, q) + elif can_use_selected_prefill(*eligibility): + selected_workspace = self.backend.get_smallq_workspace(state, q) + else: + graph_guard = ( + state.is_cuda_graph_capture + or self.backend.model.prefill_graph is not None + or hasattr(state, "prefill_cuda_graph_exe_list") ) + if not graph_guard and not torch.cuda.is_current_stream_capturing(): + generic_workspace = self.backend.get_generic_prefill_workspace(state, q, params["topk_indices"]) return context_attention_fwd_fp8( - q, - kv, - self.ragged_mem_index, - self.infer_state.mem_index, - params["prefill_cache_kv"], - self.ks, - params["topk_indices"], - params["softmax_scale"], - workspace, - direct_workspace, - generic_workspace, - **request_layout, + q=q, + packed_kv=kv, + prefill_mem_index=state.mem_index, + prefill_kv=params["prefill_cache_kv"], + kv_starts=self.ks, + topk_indices=params["topk_indices"], + topk_mem_indices=topk_mem_indices, + req_to_token_index=state.req_manager.req_to_token_indexs, + b_req_idx=state.b_req_idx, + cu_kv_seq_len=state.b1_cu_kv_seq_len, + cu_q_seq_len=state.b1_cu_q_seq_len, + total_kv_tokens=state.total_token_num, + softmax_scale=params["softmax_scale"], + selected_workspace=selected_workspace, + direct_workspace=direct_workspace, + generic_workspace=generic_workspace, ) @dataclasses.dataclass -class Fp8Glm5NextSparseDecodeState(NsaFlashMlaFp8SparseDecodeAttState): +class Fp8Glm5NextSparseDecodeState(Glm5NextSparseDecodeState): decode_topk_bound: int = None - query_req: torch.Tensor = None + flashmla_sched_meta: object = None def init_state(self): import flash_mla diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py index aeb895e9e8..0e9b1b7900 100644 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -3,12 +3,7 @@ import torch -from ..base_att import BaseAttBackend - -from .flashmla_sparse import ( - NsaFlashMlaSparsePrefillAttState, - NsaFlashMlaSparseDecodeAttState, -) +from ..base_att import AttControl, BaseAttBackend, BasePrefillAttState, BaseDecodeAttState class Glm5NextSparseAttBackend(BaseAttBackend): @@ -20,8 +15,10 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass -class Glm5NextSparsePrefillState(NsaFlashMlaSparsePrefillAttState): +class Glm5NextSparsePrefillState(BasePrefillAttState): + lengths: torch.Tensor = None query_req: torch.Tensor = None + ks: torch.Tensor = None def init_state(self): from lightllm.models.glm5_next.triton_kernel.kpool import prefill_nsa_params @@ -31,6 +28,10 @@ def init_state(self): state.position_ids, state.b_req_idx, state.b1_cu_q_seq_len, state.b1_cu_kv_seq_len ) + def prefill_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): + assert att_control.nsa_prefill and att_control.nsa_prefill_dict is not None + return self._nsa_prefill_att(q=q, kv=k, att_control=att_control) + def _nsa_prefill_att(self, q, kv, att_control): from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -53,8 +54,11 @@ def _nsa_prefill_att(self, q, kv, att_control): @dataclasses.dataclass -class Glm5NextSparseDecodeState(NsaFlashMlaSparseDecodeAttState): +class Glm5NextSparseDecodeState(BaseDecodeAttState): + lengths: torch.Tensor = None query_req: torch.Tensor = None + nsa_cache_seqlens: torch.Tensor = None + nsa_cu_seqlens_k_new: torch.Tensor = None def init_state(self): self.lengths = self.infer_state.b_seq_len @@ -66,6 +70,10 @@ def init_state(self): ) self.nsa_cu_seqlens_k_new = torch.nn.functional.pad(self.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0)) + def decode_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): + assert att_control.nsa_decode and att_control.nsa_decode_dict is not None + return self._nsa_decode_att(q, k, att_control) + def _nsa_decode_att(self, q, kv, att_control): from sgl_kernel.flash_attn import flash_attn_with_kvcache diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index a6782aa439..f87bd3a235 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -11,14 +11,12 @@ def _direct_smallq_packed_kernel( Q, PackedKV, - RaggedMemIndex, PrefillMemIndex, PrefillKV, - KVStarts, Topk, PartialO, PartialLse, - ragged_n, + packed_rows, q_stride_s: tl.constexpr, q_stride_h: tl.constexpr, q_stride_d: tl.constexpr, @@ -34,7 +32,6 @@ def _direct_smallq_packed_kernel( BLOCKS_: tl.constexpr, SLOTS_: tl.constexpr, softmax_scale: tl.constexpr, - PHYSICAL_INDICES: tl.constexpr, ): qid = tl.program_id(0) split = tl.program_id(1) @@ -58,8 +55,6 @@ def _direct_smallq_packed_kernel( acc = tl.zeros((BLOCK_H_, 512), tl.float32) m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) denominator = tl.zeros((BLOCK_H_,), tl.float32) - if not PHYSICAL_INDICES: - start = tl.load(KVStarts + qid).to(tl.int64) for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): n = tl.arange(0, 64) slot = split * SLOTS_ + b * 64 + n @@ -68,13 +63,8 @@ def _direct_smallq_packed_kernel( slot < 2176, -1, ).to(tl.int64) - if PHYSICAL_INDICES: - valid = (top >= 0) & (top < ragged_n) - mem = tl.where(valid, top, 0) - else: - ragged_pos = start + top - valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) - mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + valid = (top >= 0) & (top < packed_rows) + mem = tl.where(valid, top, 0) current_row = tl.full((64,), -1, tl.int32) for j in tl.static_range(PREFILL_N): current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) @@ -155,9 +145,7 @@ def reduce_direct_smallq_kernel( tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) -def can_use_direct_smallq( - q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens=None -): +def can_use_direct_smallq(q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens): return ( q.ndim == 3 and 1 <= q.shape[0] <= 8 @@ -166,14 +154,13 @@ def can_use_direct_smallq( and packed_kv.ndim == 3 and packed_kv.shape[1] == 1 and packed_kv.shape[2] >= 584 - and ragged_mem_index.dtype == torch.int32 - and (ragged_mem_index.numel() if total_kv_tokens is None else total_kv_tokens) >= 65536 + and total_kv_tokens >= 65536 and prefill_mem_index.dtype == torch.int32 and prefill_mem_index.numel() == q.shape[0] and prefill_kv.dtype == torch.bfloat16 and prefill_kv.shape == (q.shape[0], 1, 512) - and topk_indices.dtype == torch.int32 - and topk_indices.shape == (q.shape[0], 2176) + and topk_mem_indices.dtype == torch.int32 + and topk_mem_indices.shape == (q.shape[0], 2176) ) @@ -190,16 +177,7 @@ def allocate_direct_smallq_workspace(q): @torch.no_grad() def _context_attention_fwd_fp8_direct_smallq( - q, - packed_kv, - ragged_mem_index, - prefill_mem_index, - prefill_kv, - kv_starts, - topk_indices, - softmax_scale, - direct_workspace=None, - topk_mem_indices=None, + q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, direct_workspace=None ): qn = q.shape[0] heads = q.shape[1] @@ -218,26 +196,23 @@ def _context_attention_fwd_fp8_direct_smallq( partial_lse = workspace["partial_lse"] output = workspace["output"] launch_grid = (qn, splits, triton.cdiv(heads, block_h)) - indices = topk_indices if topk_mem_indices is None else topk_mem_indices _direct_smallq_packed_kernel[launch_grid]( q, packed_kv, - ragged_mem_index, prefill_mem_index, prefill_kv, - kv_starts, - indices, + topk_mem_indices, partial_o, partial_lse, - ragged_mem_index.numel() if topk_mem_indices is None else packed_kv.shape[0], + packed_kv.shape[0], q.stride(0), q.stride(1), q.stride(2), packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - indices.stride(0), - indices.stride(1), + topk_mem_indices.stride(0), + topk_mem_indices.stride(1), HEADS_=heads, BLOCK_H_=block_h, PREFILL_N=qn, @@ -245,7 +220,6 @@ def _context_attention_fwd_fp8_direct_smallq( BLOCKS_=blocks, SLOTS_=slots, softmax_scale=softmax_scale, - PHYSICAL_INDICES=topk_mem_indices is not None, num_warps=8, ) reduce_direct_smallq_kernel[(qn, heads, 4)]( @@ -260,18 +234,9 @@ def _context_attention_fwd_fp8_direct_smallq( return output -@triton.jit -def _build_prefill_row_table_kernel(PrefillMemIndex, RowTable, num_tokens, BLOCK: tl.constexpr): - row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mem_index = tl.load(PrefillMemIndex + row, row < num_tokens, 0) - tl.store(RowTable + mem_index, row, row < num_tokens) - - @triton.jit def _gather_prefill_kv_kernel( PackedKV, - MemIndex, - PrefillRowTable, PrefillKV, Output, ReqToToken, @@ -286,36 +251,29 @@ def _gather_prefill_kv_kernel( REQ_STRIDE: tl.constexpr, TOKEN_STRIDE: tl.constexpr, BATCH: tl.constexpr, - USE_REQ_TABLE: tl.constexpr, BLOCK_ROWS: tl.constexpr, ): block = tl.program_id(0) - if USE_REQ_TABLE: - lanes = tl.arange(0, triton.next_power_of_2(BATCH)) - ends = tl.load(CuKV + lanes + 1, lanes < BATCH, 0) - ends = tl.minimum(tl.maximum(ends - chunk_start, 0), num_tokens) - block_ends = ends // BLOCK_ROWS + lanes + 1 - batch = tl.sum(((block >= block_ends) & (lanes < BATCH)).to(tl.int32), 0) - start = tl.load(CuKV + batch) - copy_start = tl.minimum(tl.maximum(start - chunk_start, 0), num_tokens) - copy_end = tl.minimum(tl.maximum(tl.load(CuKV + batch + 1) - chunk_start, 0), num_tokens) - local_block = block - (copy_start // BLOCK_ROWS + batch) - row = copy_start + local_block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) - live = row < copy_end - req = tl.load(ReqIdx + batch).to(tl.int64) - position = chunk_start + row - start - q_start = tl.load(CuQ + batch) - q_len = tl.load(CuQ + batch + 1) - q_start - cache_len = tl.load(CuKV + batch + 1) - start - q_len - prefill_row = tl.where(position >= cache_len, q_start + position - cache_len, -1) - mem_index = tl.load(ReqToToken + req * REQ_STRIDE + position * TOKEN_STRIDE, live & (prefill_row < 0), 0).to( - tl.int64 - ) - else: - row = block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) - live = row < num_tokens - mem_index = tl.load(MemIndex + row, live, 0).to(tl.int64) - prefill_row = tl.load(PrefillRowTable + mem_index, live, -1) + lanes = tl.arange(0, triton.next_power_of_2(BATCH)) + ends = tl.load(CuKV + lanes + 1, lanes < BATCH, 0) + ends = tl.minimum(tl.maximum(ends - chunk_start, 0), num_tokens) + block_ends = ends // BLOCK_ROWS + lanes + 1 + batch = tl.sum(((block >= block_ends) & (lanes < BATCH)).to(tl.int32), 0) + start = tl.load(CuKV + batch) + copy_start = tl.minimum(tl.maximum(start - chunk_start, 0), num_tokens) + copy_end = tl.minimum(tl.maximum(tl.load(CuKV + batch + 1) - chunk_start, 0), num_tokens) + local_block = block - (copy_start // BLOCK_ROWS + batch) + row = copy_start + local_block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + live = row < copy_end + req = tl.load(ReqIdx + batch).to(tl.int64) + position = chunk_start + row - start + q_start = tl.load(CuQ + batch) + q_len = tl.load(CuQ + batch + 1) - q_start + cache_len = tl.load(CuKV + batch + 1) - start - q_len + prefill_row = tl.where(position >= cache_len, q_start + position - cache_len, -1) + mem_index = tl.load(ReqToToken + req * REQ_STRIDE + position * TOKEN_STRIDE, live & (prefill_row < 0), 0).to( + tl.int64 + ) dim = tl.arange(0, 512) cached = live & (prefill_row < 0) current = live & (prefill_row >= 0) @@ -361,16 +319,16 @@ def _compact_topk_indices_kernel( row = tl.program_id(0) lane = tl.arange(0, BLOCK) index = tl.load(TopkIndices + row * TOPK + lane, lane < TOPK, -1) - ragged_index = index + tl.load(KVStarts + row) - valid = (lane < TOPK) & (index >= 0) & (ragged_index >= chunk_start) & (ragged_index < chunk_end) + kv_index = index + tl.load(KVStarts + row) + valid = (lane < TOPK) & (index >= 0) & (kv_index >= chunk_start) & (kv_index < chunk_end) if COMPACT: compact_index = tl.cumsum(valid.to(tl.int32)) - 1 - tl.store(ChunkIndices + row * TOPK + compact_index, ragged_index - chunk_start, valid) + tl.store(ChunkIndices + row * TOPK + compact_index, kv_index - chunk_start, valid) tl.store(ChunkLengths + row, tl.sum(valid.to(tl.int32))) else: tl.store( ChunkIndices + row * TOPK + lane, - tl.where(valid, ragged_index - chunk_start, -1), + tl.where(valid, kv_index - chunk_start, -1), lane < TOPK, ) @@ -419,30 +377,21 @@ def _merge_att_output_kernel( @triton.jit def _selected_prefill_compact_kernel( Topk, - KVStarts, - RaggedMemIndex, CompactPhysical, Lengths, PageTable, CuQ, - ragged_n, packed_rows: tl.constexpr, topk_stride_s: tl.constexpr, topk_stride_d: tl.constexpr, BATCH: tl.constexpr, BLOCK: tl.constexpr, - PHYSICAL_INDICES: tl.constexpr, ): qid = tl.program_id(0) lane = tl.arange(0, BLOCK) local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) - if PHYSICAL_INDICES: - valid = (lane < 2176) & (local >= 0) - physical = local - else: - ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local - valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) - physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + valid = (lane < 2176) & (local >= 0) + physical = local valid = valid & (physical >= 0) & (physical < packed_rows) rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) @@ -515,9 +464,7 @@ def _selected_prefill_gather_kernel( tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) -def can_use_selected_prefill( - q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens=None -): +def can_use_selected_prefill(q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens): return ( q.ndim == 3 and 2 <= q.shape[0] <= 8 @@ -528,18 +475,16 @@ def can_use_selected_prefill( and packed_kv.shape[1:] == (1, 584) and packed_kv.dtype == torch.uint8 and packed_kv.device == q.device - and ragged_mem_index.dtype == torch.int32 - and ragged_mem_index.device == q.device - and (ragged_mem_index.numel() if total_kv_tokens is None else total_kv_tokens) >= 65536 + and total_kv_tokens >= 65536 and prefill_mem_index.dtype == torch.int32 and prefill_mem_index.device == q.device and prefill_mem_index.shape == (q.shape[0],) and prefill_kv.dtype == torch.bfloat16 and prefill_kv.device == q.device and prefill_kv.shape == (q.shape[0], 1, 512) - and topk_indices.dtype == torch.int32 - and topk_indices.device == q.device - and topk_indices.shape == (q.shape[0], 2176) + and topk_mem_indices.dtype == torch.int32 + and topk_mem_indices.device == q.device + and topk_mem_indices.shape == (q.shape[0], 2176) ) @@ -555,54 +500,35 @@ def allocate_selected_prefill_workspace(q, batch_capacity=None): } -def allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None, use_request_table=False): - """Allocate generic prefill metadata up to the backend's planned query capacity.""" +def allocate_generic_prefill_workspace(q, topk_indices, query_capacity=None): capacity = q.shape[0] if query_capacity is None else query_capacity return { - "prefill_row_table": torch.full( - (0 if use_request_table else packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device - ), "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), - "row_table_built": False, } @torch.no_grad() def _context_attention_fwd_fp8_selected_prefill( - q, - packed_kv, - ragged_mem_index, - prefill_mem_index, - prefill_kv, - kv_starts, - topk_indices, - softmax_scale, - selected_workspace=None, - topk_mem_indices=None, + q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, selected_workspace=None ): from sgl_kernel.flash_attn import flash_attn_with_kvcache batch = q.shape[0] workspace = selected_workspace or allocate_selected_prefill_workspace(q) batch_pad = triton.next_power_of_2(batch) - indices = topk_indices if topk_mem_indices is None else topk_mem_indices _selected_prefill_compact_kernel[(batch,)]( - indices, - kv_starts, - ragged_mem_index, + topk_mem_indices, workspace["compact"], workspace["lengths"], workspace["page_table"], workspace["cu_q"], - ragged_mem_index.numel() if ragged_mem_index is not None else 0, packed_kv.shape[0], - indices.stride(0), - indices.stride(1), + topk_mem_indices.stride(0), + topk_mem_indices.stride(1), BATCH=batch, BLOCK=4096, - PHYSICAL_INDICES=topk_mem_indices is not None, num_warps=8, ) _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( @@ -638,104 +564,49 @@ def _context_attention_fwd_fp8_selected_prefill( @torch.no_grad() def context_attention_fwd_fp8( - q: torch.Tensor, - packed_kv: torch.Tensor, - ragged_mem_index: torch.Tensor, - prefill_mem_index: torch.Tensor, - prefill_kv: torch.Tensor, - kv_starts: torch.Tensor, - topk_indices: torch.Tensor, - softmax_scale: float, + q, + packed_kv, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + topk_mem_indices, + req_to_token_index, + b_req_idx, + cu_kv_seq_len, + cu_q_seq_len, + total_kv_tokens, + softmax_scale, selected_workspace=None, direct_workspace=None, generic_workspace=None, - topk_mem_indices=None, - req_to_token_index=None, - b_req_idx=None, - cu_kv_seq_len=None, - total_kv_tokens=None, - cu_q_seq_len=None, ): - """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV. - - GLM passes physical top-k slots and the existing request table/prefix sums, - leaving ragged_mem_index unset. The ragged form supports standalone callers. - """ - # Direct packed attention is intentionally restricted to the validated - # small-Q GLM-5.3 shape; all other cases retain the bounded-gather path. - index_layout = ragged_mem_index if ragged_mem_index is not None else topk_mem_indices - if can_use_direct_smallq(q, packed_kv, index_layout, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens): + """Run NoPE sparse prefill from request mappings and physical top-k slots.""" + eligibility = (q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens) + if can_use_direct_smallq(*eligibility): return _context_attention_fwd_fp8_direct_smallq( - q, - packed_kv, - ragged_mem_index, - prefill_mem_index, - prefill_kv, - kv_starts, - topk_indices, - softmax_scale, - direct_workspace, - topk_mem_indices, + q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, direct_workspace ) - if can_use_selected_prefill( - q, packed_kv, index_layout, prefill_mem_index, prefill_kv, topk_indices, total_kv_tokens - ): + if can_use_selected_prefill(*eligibility): return _context_attention_fwd_fp8_selected_prefill( - q, - packed_kv, - ragged_mem_index, - prefill_mem_index, - prefill_kv, - kv_starts, - topk_indices, - softmax_scale, - selected_workspace, - topk_mem_indices, + q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, selected_workspace ) from sgl_kernel.flash_attn import flash_attn_with_kvcache num_q_tokens, num_heads, _ = q.shape - use_req_table = ragged_mem_index is None - num_kv_tokens = total_kv_tokens if use_req_table else ragged_mem_index.numel() + num_kv_tokens = total_kv_tokens gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, num_kv_tokens) - if generic_workspace is None: - # Build the row table before gathering KV and allocating compaction metadata. - prefill_row_table = None - if not use_req_table: - prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) - _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( - prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 - ) - # Keep this potentially 2 GiB buffer local to a DSA layer. - gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) - chunk_indices = torch.empty_like(topk_indices) - chunk_lengths = torch.empty(num_q_tokens, dtype=torch.int32, device=q.device) - cu_seqlens_q = torch.arange(num_q_tokens + 1, dtype=torch.int32, device=q.device) - else: - prefill_row_table = generic_workspace["prefill_row_table"] - chunk_indices = generic_workspace["chunk_indices"] - chunk_lengths = generic_workspace["chunk_lengths"] - cu_seqlens_q = generic_workspace["cu_seqlens_q"] - if not use_req_table and not generic_workspace["row_table_built"]: - _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( - prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 - ) - generic_workspace["row_table_built"] = True - # Keep this potentially 2 GiB buffer local to a DSA layer. - gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + workspace = generic_workspace or allocate_generic_prefill_workspace(q, topk_indices) + chunk_indices = workspace["chunk_indices"] + chunk_lengths = workspace["chunk_lengths"] + cu_seqlens_q = workspace["cu_seqlens_q"] + # Keep this potentially 2 GiB buffer local to a DSA layer. + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) output = lse = None for chunk_start in range(0, num_kv_tokens, gather_chunk_size): num_chunk_tokens = min(gather_chunk_size, num_kv_tokens - chunk_start) - if use_req_table: - gather_rows = 4 - gather_grid = (num_chunk_tokens // gather_rows + b_req_idx.numel(),) - else: - gather_rows = 4 - gather_grid = (triton.cdiv(num_chunk_tokens, 4),) - _gather_prefill_kv_kernel[gather_grid]( + _gather_prefill_kv_kernel[(num_chunk_tokens // 4 + b_req_idx.numel(),)]( packed_kv, - None if use_req_table else ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], - prefill_row_table, prefill_kv, gathered_kv, req_to_token_index, @@ -747,11 +618,10 @@ def context_attention_fwd_fp8( packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - REQ_STRIDE=req_to_token_index.stride(0) if use_req_table else 0, - TOKEN_STRIDE=req_to_token_index.stride(1) if use_req_table else 0, - BATCH=b_req_idx.numel() if use_req_table else 1, - USE_REQ_TABLE=use_req_table, - BLOCK_ROWS=gather_rows, + REQ_STRIDE=req_to_token_index.stride(0), + TOKEN_STRIDE=req_to_token_index.stride(1), + BATCH=b_req_idx.numel(), + BLOCK_ROWS=4, ) chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) if ( diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index 732bfa10a8..6b791b8e04 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -116,7 +116,7 @@ def forbidden(*args, **kwargs): state.init_state() assert state.lengths is infer.b_seq_len assert state.query_req is infer.b_req_idx - assert state.ragged_mem_index is None and state.ks is None and state.ke is None + assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): state.init_state() @@ -141,11 +141,6 @@ def forbidden(*args, **kwargs): def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) - class ForbiddenRowTable: - def __getitem__(self, grid): - raise AssertionError("Request-layout prefill must not build a physical-slot row table") - - monkeypatch.setattr(fp8_prefill, "_build_prefill_row_table_kernel", ForbiddenRowTable()) torch.manual_seed(1685) batch = min(queries, 2) counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] @@ -192,7 +187,7 @@ def refresh_physical(): backend = Fp8Glm5NextSparseAttBackend(model) state = backend.create_att_prefill_state(infer) state.init_state() - assert state.ragged_mem_index is None and not hasattr(backend, "ragged_mem_buffers") + assert not hasattr(state, "ragged_mem_index") and not hasattr(backend, "ragged_mem_buffers") control = AttControl( nsa_prefill=True, nsa_prefill_dict={ @@ -221,8 +216,7 @@ def check(output): check(run()) generic = backend._prefill_workspaces[0]["generic"] if generic is not None: - assert generic["buffers"]["prefill_row_table"].numel() == 0 - assert not generic["workspace"]["row_table_built"] + assert set(generic) == {"chunk_indices", "chunk_lengths", "cu_seqlens_q"} graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = run() diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index d47d9f808f..9993b2e162 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -24,7 +24,6 @@ from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - _build_prefill_row_table_kernel, _gather_prefill_kv_kernel, ) @@ -45,6 +44,99 @@ def _make_workspace_backend(batch_max_tokens): return backend +class _PrefillCase: + def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=None, softmax_scale=0.0625): + device = "cuda" + tensor = lambda values: torch.tensor(values, device=device, dtype=torch.int32) + self.seq_lens, self.q_lens = seq_lens, q_lens + self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) + total, queries = sum(seq_lens), sum(q_lens) + capacity = capacity or total + 17 + self.table = torch.full((2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32) + slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] + cu_q, cu_kv, query_req, positions = [0], [0], [], [] + for req, length, count in zip(self.requests, seq_lens, q_lens): + self.table[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] + query_req.extend([req] * count) + positions.extend(range(length - count, length)) + cu_q.append(cu_q[-1] + count) + cu_kv.append(cu_kv[-1] + length) + self.query_req, self.positions = tensor(query_req).long(), tensor(positions).long() + self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) + self.source = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) + self.packed = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) + self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) + write_fp8_e8m0_packed_kv(self.source, self.all_slots, self.packed) + self.prefill = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] + visible = tensor(positions) + 1 + self.topk = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) + self.topk.remainder_(visible[:, None]) + self.topk.masked_fill_(torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1) + self.topk[:, ::19] = -1 + self.topk[:, 0] = visible - 1 + self.topk[:, 1] = visible - 1 + self.physical = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] + self.infer = _make_infer_state( + req_manager=SimpleNamespace(req_to_token_indexs=self.table), + b_req_idx=tensor(self.requests), + b_seq_len=tensor(seq_lens), + b1_cu_q_seq_len=tensor(cu_q), + b1_cu_kv_seq_len=tensor(cu_kv), + position_ids=tensor(positions), + total_token_num=total, + mem_index=torch.empty(queries, device=device, dtype=torch.int32), + ) + self.backend = backend or _make_workspace_backend(queries) + _, self.state = self.new_state() + self.control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": self.topk, + "topk_mem_indices": self.physical, + "prefill_cache_kv": self.prefill, + "softmax_scale": softmax_scale, + "kv_lora_rank": 512, + }, + ) + + def new_state(self, **overrides): + if overrides: + infer = _make_infer_state(**vars(self.infer)) + vars(infer).update(overrides) + else: + infer = self.infer + state = self.backend.create_att_prefill_state(infer) + state.init_state() + return infer, state + + def refresh(self): + self.infer.mem_index.copy_(self.table[self.query_req, self.positions]) + self.physical.copy_(self.table[self.query_req[:, None], self.topk.clamp_min(0).long()]) + self.physical.masked_fill_(self.topk < 0, -1) + self.control.nsa_prefill_dict["topk_indices"] = self.topk + + def run(self, state=None): + self.refresh() + active = self.state if state is None else state + return active.prefill_att(self.q, self.packed[..., :584], None, self.control) + + def check(self, actual): + values = _dequantize_kv_cache(self.packed) + values[self.infer.mem_index.long()] = self.prefill + physical = self.table[self.query_req[:, None], self.topk.clamp_min(0).long()].masked_fill(self.topk < 0, -1) + expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + def mutate(self): + self.q.normal_() + self.prefill.normal_() + self.source.normal_() + write_fp8_e8m0_packed_kv(self.source, self.all_slots, self.packed) + for req, length in zip(self.requests, self.seq_lens): + self.table[req, :length].copy_(self.table[req, :length].roll(1)) + self.topk.copy_(self.topk.roll(1, dims=1)) + + @pytest.mark.parametrize("draft", [False, True]) @pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): @@ -141,35 +233,29 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) - indices = mem_index[[0, 2, 8, 2, 5, 1, 9, 3]] prefill_kv = source[8:].clone().add_(0.25) - row_table = torch.full((31,), -1, device="cuda", dtype=torch.int32) - _build_prefill_row_table_kernel[(1,)](mem_index[8:], row_table, 3, BLOCK=256) - gathered_kv = torch.empty((indices.numel(), 1, 512), device="cuda", dtype=torch.bfloat16) - _gather_prefill_kv_kernel[(triton.cdiv(indices.numel(), 4),)]( + gathered_kv = torch.empty((8, 1, 512), device="cuda", dtype=torch.bfloat16) + _gather_prefill_kv_kernel[(3,)]( packed_kv[..., :584], - indices, - row_table, prefill_kv, gathered_kv, - None, - None, - None, - None, - 0, - indices.numel(), + mem_index[None], + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([0, 11], device="cuda", dtype=torch.int32), + torch.tensor([0, 3], device="cuda", dtype=torch.int32), + 2, + 8, packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - REQ_STRIDE=0, - TOKEN_STRIDE=0, + REQ_STRIDE=11, + TOKEN_STRIDE=1, BATCH=1, - USE_REQ_TABLE=False, BLOCK_ROWS=4, ) expected = _dequantize_kv_cache(packed_kv) - expected[mem_index[8:]] = prefill_kv - assert torch.equal(gathered_kv, expected[indices.long()]) + expected[mem_index[8:].long()] = prefill_kv + assert torch.equal(gathered_kv, expected[mem_index[2:10].long()]) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @@ -224,68 +310,20 @@ def check(actual): def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) torch.manual_seed(53) - # Two requests whose physical cache locations are unrelated to batch order. - prefix_len = 4 if cached else 0 - seq_len = prefix_len + 3 - ragged_mem_index = torch.randperm(64, device="cuda")[: 2 * seq_len].int() - prefill_rows = torch.cat( - (torch.arange(prefix_len, seq_len), torch.arange(seq_len + prefix_len, 2 * seq_len)) - ).cuda() - prefill_mem_index = ragged_mem_index[prefill_rows] - source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) - prefill_kv = source[prefill_rows].clone().add_(0.25) - packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) - write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) - q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) - topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) - for row in range(6): - length = prefix_len + row % 3 + 1 - topk_indices[row, :length] = torch.arange(length, device="cuda") - topk_indices[0].fill_(-1) - topk_indices[1, :2] = prefix_len - topk_indices[3].fill_(-1) - topk_indices[3, 0] = prefix_len - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk_indices, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 512 ** -0.5, - "kv_lora_rank": 512, - }, - ) - infer_state = _make_infer_state(max_cache_len=prefix_len, mem_index=prefill_mem_index) - state = Fp8Glm5NextSparsePrefillState( - backend=_make_workspace_backend(q.shape[0]), - _infer_state_weak_ref=weakref.ref(infer_state), - ks=kv_start_offsets, - ragged_mem_index=ragged_mem_index, - ) - - def run(): - return state.prefill_att(q, packed_kv[..., :584], None, control) - - def check(out): - expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] - expected_kv[prefill_rows] = prefill_kv - indices = torch.where(topk_indices >= 0, topk_indices + kv_start_offsets[:, None], -1) - expected = _torch_sparse_attention(q, expected_kv, indices, control.nsa_prefill_dict["softmax_scale"]) - torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) - - check(run()) - graph, out = _capture_cuda_graph(run) + prefix = 4 if cached else 0 + case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) + case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + case.topk[0].fill_(-1) + case.topk[1, :2] = prefix + case.topk[3].fill_(-1) + case.topk[3, 0] = prefix + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) graph.replay() - check(out) - q.normal_() - source.normal_() - prefill_kv.copy_(source[prefill_rows] + 0.25) - ragged_mem_index.copy_(ragged_mem_index.roll(1)) - prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) - topk_indices.copy_(topk_indices.roll(1, dims=1)) - write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + case.check(output) + case.mutate() graph.replay() - check(out) + case.check(output) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @@ -293,111 +331,29 @@ def check(out): def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", 129) torch.manual_seed(53) - num_tokens = 1025 - source = torch.randn(num_tokens, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] - ragged_mem_index = torch.randperm(2048, device="cuda")[:num_tokens].int() - prefill_rows = torch.tensor([1, 15, 128, 129, 255, 257, 511, 512, 1024], device="cuda") - prefill_mem_index = ragged_mem_index[prefill_rows].long() - prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] - packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) - write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) - q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") - topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) - lanes = torch.arange(2176, device="cuda") - topk_indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) - topk_indices[1, 0] = 15 - topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") - topk_indices[3, :127] = prefill_rows[torch.randint(9, (127,), device="cuda")].int() - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk_indices, - "prefill_cache_kv": prefill_kv, - "softmax_scale": softmax_scale, - "kv_lora_rank": 512, - }, - ) - infer_state = _make_infer_state(mem_index=prefill_mem_index) - state = Fp8Glm5NextSparsePrefillState( - backend=_make_workspace_backend(q.shape[0]), - _infer_state_weak_ref=weakref.ref(infer_state), - ks=torch.zeros(9, device="cuda", dtype=torch.int32), - ragged_mem_index=ragged_mem_index, - ) - - def run(): - return state.prefill_att(q, packed_kv[..., :584], None, control) - - def check(actual): - expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] - expected_kv[prefill_rows] = prefill_kv - expected = _torch_sparse_attention(q, expected_kv, topk_indices, softmax_scale) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - check(run()) - graph, output = _capture_cuda_graph(run) + case.topk.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) + case.topk[1, 0] = 1016 + case.topk[2, :3] = torch.tensor([2, 3, 4], device="cuda") + case.topk[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) graph.replay() - check(output) - q.normal_() - source.normal_() - prefill_kv.normal_() - ragged_mem_index.copy_(ragged_mem_index.roll(1)) - prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) - write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) - topk_indices.copy_(topk_indices.roll(1, dims=1)) + case.check(output) + case.mutate() graph.replay() - check(output) - - -def _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples): - """Reference only a few H64 rows so the large-Q regression stays bounded.""" - expected = [] - for row, heads in samples: - local = topk[row] - local = local[local >= 0].long() - physical = ragged[(starts[row].long() + local)].long() - values = _dequantize_kv_cache(packed[physical])[:, 0].float() - current = physical[:, None] == prefill_mem[None, :].long() - has_current = current.any(dim=1) - current_row = current.int().argmax(dim=1) - values[has_current] = prefill_kv[current_row[has_current], 0].float() - expected.append(((q[row, heads].float() @ values.T * 0.0625).softmax(-1) @ values)) - return expected + case.check(output) @pytest.mark.parametrize("num_q", [16, 33]) -def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q): - """The native sparse prefill path maps top-k ragged positions, never physical IDs.""" +def test_fp8_h64_largeq_native_prefill_request_table_graph(monkeypatch, num_q): if torch.cuda.get_device_capability()[0] != 9: pytest.skip("Validated FlashMLA sparse prefill target is Hopper") from sgl_kernel.flash_mla import flash_mla_sparse_fwd torch.manual_seed(53100 + num_q) - capacity = 131072 - ragged_n = 65536 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) - ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] - q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) - starts = torch.where( - torch.arange(num_q, device="cuda", dtype=torch.int32) < num_q // 2, - torch.full((num_q,), 1024, device="cuda", dtype=torch.int32), - torch.full((num_q,), 32768, device="cuda", dtype=torch.int32), - ) - current_local = (torch.arange(num_q, device="cuda", dtype=torch.int32) * 37) % 2048 + 64 - current_positions = starts + current_local - prefill_mem = ragged[current_positions.long()].clone() - prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) - prefill_kv.add_(torch.arange(num_q, device="cuda", dtype=torch.bfloat16)[:, None, None]) - topk = torch.randint(0, 2176, (num_q, 2176), device="cuda", dtype=torch.int32) - topk[:, ::19] = -1 # Holes must survive the native position-preserving map. - topk[:, 0] = current_local - topk[:, 1] = current_local # Duplicates are valid selected keys. - topk[:, 2] = topk[:, 3] - assert q.is_contiguous() and topk.is_contiguous() - + case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], 64, capacity=131072) observed = [] def observed_sparse(*args, **kwargs): @@ -405,74 +361,37 @@ def observed_sparse(*args, **kwargs): return flash_mla_sparse_fwd(*args, **kwargs) monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", observed_sparse) - - def run(): - return fp8_prefill.context_attention_fwd_fp8( - q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 - ) - - samples = [(0, torch.tensor([0, 9, 63], device="cuda")), (num_q - 1, torch.tensor([3, 31, 62], device="cuda"))] - - def check(actual): - expected = _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples) - for (row, heads), value in zip(samples, expected): - torch.testing.assert_close(actual[row, heads].float(), value, atol=0.012, rtol=0.015) - - eager = run() - check(eager) - assert observed == [((num_q, 64, 512), (ragged_n, 1, 512), (num_q, 1, 2176))] - graph, output = _capture_cuda_graph(run) - q.normal_() - prefill_kv.normal_() - ragged.copy_(ragged.roll(1)) - prefill_mem.copy_(ragged[current_positions.long()]) - topk.copy_(topk.roll(1, dims=1)) - topk[:, 1] = current_local + case.check(case.run()) + assert observed == [((num_q, 64, 512), (65536, 1, 512), (num_q, 1, 2176))] + graph, output = _capture_cuda_graph(case.run) + case.mutate() graph.replay() - check(output) - topk.fill_(-1) + case.check(output) + case.topk.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 - assert len(observed) == 3 # eager, graph warmup, and capture; replay stays inside the CUDA graph. + assert len(observed) == 3 def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): - """H8 large-Q prefill must not enter the H64 native sparse fast path.""" import sgl_kernel.flash_attn as flash_attn - import sgl_kernel.flash_mla as flash_mla torch.manual_seed(53180) - capacity = 65536 - num_q = 16 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) - q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) - ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) - starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) - current_local = torch.arange(num_q, device="cuda", dtype=torch.int32) + 64 - prefill_mem = ragged[(starts + current_local).long()].clone() - prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) - topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) - topk[:, :64] = torch.arange(64, device="cuda", dtype=torch.int32) - topk[:, 0] = current_local - seen_flashattention = [] - original_flashattention = flash_attn.flash_attn_with_kvcache + case = _PrefillCase([65536], [16], 8, capacity=65536) + observed = [] + original = flash_attn.flash_attn_with_kvcache def observed_flashattention(*args, **kwargs): - seen_flashattention.append(True) - return original_flashattention(*args, **kwargs) + observed.append(True) + return original(*args, **kwargs) def unexpected_native(*args, **kwargs): raise AssertionError("H8 large-Q prefill entered the native H64 sparse path") monkeypatch.setattr(flash_attn, "flash_attn_with_kvcache", observed_flashattention) - monkeypatch.setattr(flash_mla, "flash_mla_sparse_fwd", unexpected_native) - output = fp8_prefill.context_attention_fwd_fp8( - q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 - ) - assert seen_flashattention - assert output.shape == q.shape + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_native) + case.check(case.run()) + assert observed @pytest.mark.parametrize( @@ -480,204 +399,59 @@ def unexpected_native(*args, **kwargs): [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], ) def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num_heads): - """Exercise eager and graph replay for no-rowtable, small-Q packed prefill.""" torch.manual_seed(53016 + num_q) - capacity = 65536 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - write_fp8_e8m0_packed_kv(source, mem, packed) - ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) - q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - # Requests deliberately have different spans; starts are not fixed-stride. - seq_lens = torch.tensor([127, 521, 1031, 2053, 4099, 8191, 12289, 16381], device="cuda", dtype=torch.int32)[:num_q] - starts = torch.nn.functional.pad(seq_lens.cumsum(0), (1, 0))[:-1] - current_rows = starts + seq_lens - 1 - prefill_mem = ragged[current_rows] - prefill_kv = torch.randn(num_q, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] - topk_backing = torch.full((num_q, 2176 * 2), -1, device="cuda", dtype=torch.int32) - topk = topk_backing[:, ::2] # Exercise the topk stride passed to the direct kernel. - topk[:, :3] = torch.tensor([3, 11, 11], device="cuda", dtype=torch.int32) - topk[:, 3] = seq_lens - 1 - - def run(): - return fp8_prefill.context_attention_fwd_fp8( - q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 - ) - - def check(actual): - assert actual.shape == q.shape and actual.is_contiguous() - expected_kv = _dequantize_kv_cache(packed)[ragged.long()] - expected_kv[current_rows.long()] = prefill_kv - physical = torch.where(topk >= 0, topk + starts[:, None], -1) - expected = _torch_sparse_attention(q, expected_kv, physical) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - actual = run() - check(actual) - assert torch.count_nonzero(actual) > 0 # Includes BF16 current KV in every row. - graph, output = _capture_cuda_graph(run) - graph.replay() - check(output) - q.normal_() - source.normal_() - prefill_kv.normal_() - ragged.copy_(ragged.roll(1)) - prefill_mem.copy_(ragged[current_rows]) - write_fp8_e8m0_packed_kv(source, mem, packed) - topk.copy_(topk.roll(1, dims=1)) + seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] + case = _PrefillCase(seqs, [1] * num_q, num_heads) + case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) + backing[:, ::2] = case.topk + case.topk = backing[:, ::2] + output = case.run() + case.check(output) + assert output.is_contiguous() and torch.count_nonzero(output) > 0 + graph, output = _capture_cuda_graph(case.run) graph.replay() - check(output) - topk.fill_(-1) # Empty graph replay must not read invalid packed pointers. + case.check(output) + case.mutate() graph.replay() - assert torch.count_nonzero(output) == 0 + case.check(output) def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): - """Q1 H64 direct workspace is owned by the backend and reused across infer states.""" torch.manual_seed(53150) - capacity = 65536 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - write_fp8_e8m0_packed_kv(source, mem, packed) - ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) - q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() - starts = torch.zeros(1, device="cuda", dtype=torch.int32) - current_rows = torch.tensor([2048], device="cuda", dtype=torch.int32) - prefill_mem = ragged[current_rows] - prefill_kv = torch.randn(1, 1, 512, device="cuda", dtype=torch.bfloat16) - topk = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) - topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) - topk[:, 2048] = current_rows - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, - ) - infer_state = _make_infer_state(mem_index=prefill_mem) - backend = _make_workspace_backend(q.shape[0]) - state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged - ) - - def expected(): - values = _dequantize_kv_cache(packed)[ragged.long()] - values[current_rows.long()] = prefill_kv - return _torch_sparse_attention(q, values, topk) - - def run(active_state=state): - return active_state.prefill_att(q, packed[..., :584], None, control) - - torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) - cache = backend._prefill_workspaces[0]["smallq"] - assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") - workspace = cache["direct"] - pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + case = _PrefillCase([65536], [1], 64) + case.check(case.run()) + cache = case.backend._prefill_workspaces[0]["smallq"] + pointers = {name: value.data_ptr() for name, value in cache["direct"].items()} for _ in range(10): - layer_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) - assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers - - fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) - fresh_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(fresh_infer_state), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(fresh_state).float(), expected(), atol=0.012, rtol=0.015) - fresh_workspace = backend._prefill_workspaces[0]["smallq"]["direct"] - assert fresh_workspace["partial_o"].data_ptr() == workspace["partial_o"].data_ptr() - - q.normal_() - source.normal_() - prefill_kv.normal_() - write_fp8_e8m0_packed_kv(source, mem, packed) - torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) - assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + infer, state = case.new_state() + case.check(case.run(state)) + assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers + fresh_infer, fresh = case.new_state(mem_index=case.infer.mem_index.clone()) + case.check(case.run(fresh)) + assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers + case.mutate() + case.check(case.run()) @pytest.mark.parametrize("num_q", [2, 3, 5, 8]) def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): - """H64 selected prefill rewrites backend workspace across DSA layers.""" torch.manual_seed(53200 + num_q) - capacity = 65536 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - write_fp8_e8m0_packed_kv(source, mem, packed) - ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) - q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 - current_rows = starts + 2048 - prefill_mem = ragged[current_rows] - prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) - topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) - topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) - topk[:, 2048] = 2048 - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, - ) - infer_state = _make_infer_state(mem_index=prefill_mem) - backend = _make_workspace_backend(q.shape[0]) - state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged - ) - - def expected(): - values = _dequantize_kv_cache(packed)[ragged.long()] - values[current_rows.long()] = prefill_kv - physical = torch.where(topk >= 0, topk + starts[:, None], -1) - return _torch_sparse_attention(q, values, physical) - - def run(active_state=state): - return active_state.prefill_att(q, packed[..., :584], None, control) - - eager = run() - torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) - cache = backend._prefill_workspaces[0]["smallq"] - assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") - workspace = cache["selected"] - pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} - # All 11 DSA layers share the backend allocation and rewrite its contents. + case = _PrefillCase([32768] * num_q, [1] * num_q, 64) + case.check(case.run()) + workspace = case.backend._prefill_workspaces[0]["smallq"]["selected"] + pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): - layer_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) - assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers - - fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) - fresh_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(fresh_infer_state), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(fresh_state).float(), expected(), atol=0.012, rtol=0.015) - fresh_workspace = backend._prefill_workspaces[0]["smallq"]["selected"] - assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() - - source.normal_() - write_fp8_e8m0_packed_kv(source, mem, packed) - prefill_kv.normal_() - q.normal_() - eager = run() - assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers - torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) - graph, output = _capture_cuda_graph(run) - q.normal_() - topk.copy_(topk.roll(1, dims=1)) + infer, state = case.new_state() + case.check(case.run(state)) + fresh_infer, fresh = case.new_state(mem_index=case.infer.mem_index.clone()) + case.check(case.run(fresh)) + assert {name: value.data_ptr() for name, value in workspace.items()} == pointers + graph, output = _capture_cuda_graph(case.run) + case.mutate() graph.replay() - torch.testing.assert_close(output.float(), expected(), atol=0.012, rtol=0.015) - topk.fill_(-1) + case.check(output) + case.topk.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 @@ -1158,132 +932,56 @@ def direct_must_not_run(*args, **kwargs): def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): - """Generic H64 backend buffers are reused, batch metadata resets, and graph falls back.""" torch.manual_seed(53450) - capacity, num_q = 65536, 13 - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - write_fp8_e8m0_packed_kv(source, mem, packed) - ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) - starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 - current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) - prefill_mem = ragged[(starts + current_local).long()].clone() - prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16).add_(2) - q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() - topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) - topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) - topk[:, 2048] = current_local - topk[:, 2049] = current_local - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, - ) - infer = _make_infer_state(mem_index=prefill_mem) - backend = _make_workspace_backend(2 * num_q) - - def expected(active_ragged=ragged): - values = _dequantize_kv_cache(packed)[active_ragged.long()] - values[(starts + current_local).long()] = prefill_kv - physical = torch.where(topk >= 0, topk + starts[:, None], -1) - return _torch_sparse_attention(q, values, physical) - - def run(active): - return active.prefill_att(q, packed[..., :584], None, control) - - state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) - cache = backend._prefill_workspaces[0] - assert not hasattr(infer, "_fp8_sparse_generic_prefill_workspace_cache") - workspace = cache["generic"]["workspace"] - pointers = {k: v.data_ptr() for k, v in workspace.items() if isinstance(v, torch.Tensor)} + backend = _make_workspace_backend(26) + case = _PrefillCase([32768, 16384, 16384], [5, 4, 4], 64, backend=backend) + case.check(case.run()) + workspace = backend._prefill_workspaces[0]["generic"] + assert set(workspace) == {"chunk_indices", "chunk_lengths", "cu_seqlens_q"} + pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): - layer = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(infer), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(layer).float(), expected(), atol=0.012, rtol=0.015) - assert workspace["row_table_built"] - assert { - k: v.data_ptr() for k, v in cache["generic"]["workspace"].items() if isinstance(v, torch.Tensor) - } == pointers - - # Change raw KV/top-k in place: reusable metadata must not cache either. - q.normal_() - prefill_kv.normal_() - topk.copy_(topk.roll(1, dims=1)) - topk[:, 2049] = current_local - torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) - - fresh_ragged = ragged.roll(1) - fresh_mem = fresh_ragged[(starts + current_local).long()].clone() - fresh = _make_infer_state(mem_index=fresh_mem) - fresh_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(fresh), ks=starts, ragged_mem_index=fresh_ragged - ) - torch.testing.assert_close(run(fresh_state).float(), expected(fresh_ragged), atol=0.012, rtol=0.015) - assert cache["generic"]["workspace"]["prefill_row_table"].data_ptr() == workspace["prefill_row_table"].data_ptr() - - graph_infer = _make_infer_state(mem_index=prefill_mem, prefill_cuda_graph_exe_list=[]) - graph_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(graph_infer), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(graph_state).float(), expected(), atol=0.012, rtol=0.015) - assert not hasattr(graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") - - explicit_graph_infer = _make_infer_state(mem_index=prefill_mem, is_cuda_graph_capture=True) - explicit_graph_state = Fp8Glm5NextSparsePrefillState( - backend=backend, _infer_state_weak_ref=weakref.ref(explicit_graph_infer), ks=starts, ragged_mem_index=ragged - ) - torch.testing.assert_close(run(explicit_graph_state).float(), expected(), atol=0.012, rtol=0.015) - assert not hasattr(explicit_graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + infer, state = case.new_state() + case.check(case.run(state)) + case.mutate() + case.check(case.run()) + fresh = _PrefillCase([16384, 32768, 16384], [3, 5, 5], 64, capacity=131072, backend=backend) + fresh.check(fresh.run()) + assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0]["generic"].items()} == pointers + for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): + infer, state = case.new_state(**overrides) + case.check(case.run(state)) + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) def test_fp8_generic_prefill_workspace_graph_query_is_generic_h64_only(monkeypatch): - """H8/direct/selected prefill never queries CUDA graph-capture state.""" - import lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 as fp8_prefill - - queries = [] + def forbidden(): + raise AssertionError("Non-generic prefill queried CUDA graph capture state") - def query_spy(): - queries.append(True) - raise AssertionError("non-generic path queried CUDA graph capture state") - - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", query_spy) - kv = torch.empty((1, 1, 584), dtype=torch.uint8) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", forbidden) + monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", lambda **kwargs: kwargs["q"]) params = { "prefill_cache_kv": torch.empty((1, 1, 512), dtype=torch.bfloat16), "topk_indices": torch.empty((1, 2176), dtype=torch.int32), + "topk_mem_indices": torch.empty((1, 2176), dtype=torch.int32), "softmax_scale": 0.0625, } control = AttControl(nsa_prefill=True, nsa_prefill_dict=params) - - def fake_context(*args): - return args[0] - - monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", fake_context) for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): monkeypatch.setattr(fp8_prefill, "can_use_direct_smallq", lambda *args, d=direct: d) monkeypatch.setattr(fp8_prefill, "can_use_selected_prefill", lambda *args, s=selected: s) - infer = _make_infer_state(mem_index=torch.empty(1, dtype=torch.int32)) + infer = _make_infer_state( + mem_index=torch.empty(1, dtype=torch.int32), + req_manager=SimpleNamespace(req_to_token_indexs=torch.empty((1, 1), dtype=torch.int32)), + ) backend = _make_workspace_backend(1) - backend.get_smallq_workspace = lambda active_infer, active_q, kind="direct" if direct else "selected": { - "kind": kind - } - state = Fp8Glm5NextSparsePrefillState( - backend=backend, - _infer_state_weak_ref=weakref.ref(infer), - ks=torch.zeros(1, dtype=torch.int32), - ragged_mem_index=torch.zeros(1, dtype=torch.int32), + backend.get_smallq_workspace = lambda *args: {} + state = Fp8Glm5NextSparsePrefillState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + assert ( + state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), None, control).shape[1] == heads ) - assert state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads - assert not queries def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): From c8466d6c95bcd61a87281c937104895086ff74d0 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 09:29:59 +0000 Subject: [PATCH 04/10] refactor(glm): align NSA names with inference state --- .../basemodel/attention/nsa/fp8_glm5_next.py | 101 +-- .../basemodel/attention/nsa/glm5_next.py | 61 +- lightllm/models/glm5_next/indexer.py | 138 ++-- .../context_flashattention_nopad_fp8.py | 589 +++++++++--------- .../models/glm5_next/triton_kernel/kpool.py | 457 ++++++++------ .../glm5_next/test_direct_nsa_indices.py | 8 +- unit_tests/models/glm5_next/test_fp8_kv.py | 120 ++-- .../models/glm5_next/test_indexer_topk.py | 2 +- unit_tests/models/glm5_next/test_kernels.py | 20 +- unit_tests/models/glm5_next/test_mtp.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 56 +- 11 files changed, 825 insertions(+), 729 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index 2d1205c714..129248fff0 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -22,20 +22,20 @@ def get_smallq_workspace(self, infer_state, q): if slot["smallq"] is None: slot["smallq"] = { "direct": allocate_direct_smallq_workspace(q[:1]), - "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), + "selected": allocate_selected_prefill_workspace(q, q_token_capacity=8), } if q.shape[0] == 1: return slot["smallq"]["direct"] - batch = q.shape[0] + q_token_num = q.shape[0] selected = slot["smallq"]["selected"] - slots = selected["compact"].shape[1] + topk_size = selected["compact_topk_mem_indices"].shape[1] return { - "compact": selected["compact"][:batch], - "lengths": selected["lengths"][:batch], - "page_table": selected["page_table"][:batch], - "cu_q": selected["cu_q"][: batch + 1], - "cu_k": selected["cu_k"][: batch + 1], - "selected_kv": selected["selected_kv"][: batch * slots], + "compact_topk_mem_indices": selected["compact_topk_mem_indices"][:q_token_num], + "cache_seqlens": selected["cache_seqlens"][:q_token_num], + "page_table": selected["page_table"][:q_token_num], + "cu_seqlens_q": selected["cu_seqlens_q"][: q_token_num + 1], + "cu_seqlens_k_new": selected["cu_seqlens_k_new"][: q_token_num + 1], + "selected_kv": selected["selected_kv"][: q_token_num * topk_size], } def get_generic_prefill_workspace(self, infer_state, q, topk_indices): @@ -45,21 +45,21 @@ def get_generic_prefill_workspace(self, infer_state, q, topk_indices): slot = self._prefill_workspaces[infer_state.microbatch_index] buffers = slot["generic"] - rows, columns = topk_indices.shape + q_token_num, topk_size = topk_indices.shape if ( buffers is None - or buffers["chunk_indices"].shape[0] < rows - or buffers["chunk_indices"].shape[1] != columns - or buffers["chunk_indices"].dtype != topk_indices.dtype + or buffers["chunk_topk_indices"].shape[0] < q_token_num + or buffers["chunk_topk_indices"].shape[1] != topk_size + or buffers["chunk_topk_indices"].dtype != topk_indices.dtype ): buffers = allocate_generic_prefill_workspace( - q, topk_indices, query_capacity=max(self.model.batch_max_tokens or 0, rows) + q, topk_indices, q_token_capacity=max(self.model.batch_max_tokens or 0, q_token_num) ) slot["generic"] = buffers return { - "chunk_indices": buffers["chunk_indices"][:rows], - "chunk_lengths": buffers["chunk_lengths"][:rows], - "cu_seqlens_q": buffers["cu_seqlens_q"][: rows + 1], + "chunk_topk_indices": buffers["chunk_topk_indices"][:q_token_num], + "chunk_cache_seqlens": buffers["chunk_cache_seqlens"][:q_token_num], + "cu_seqlens_q": buffers["cu_seqlens_q"][: q_token_num + 1], } def create_att_prefill_state(self, infer_state): @@ -76,45 +76,54 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): - def _nsa_prefill_att(self, q, kv, att_control): + def _nsa_prefill_att(self, q, packed_kv, att_control): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( context_attention_fwd_fp8, can_use_direct_smallq, can_use_selected_prefill, ) - params = att_control.nsa_prefill_dict - state = self.infer_state - topk_mem_indices = params["topk_mem_indices"] + nsa_prefill_dict = att_control.nsa_prefill_dict + infer_state = self.infer_state + topk_mem_indices = nsa_prefill_dict["topk_mem_indices"] direct_workspace = selected_workspace = generic_workspace = None if q.shape[1] == 64: - eligibility = (q, kv, state.mem_index, params["prefill_cache_kv"], topk_mem_indices, state.total_token_num) + eligibility = ( + q, + packed_kv, + infer_state.mem_index, + nsa_prefill_dict["prefill_cache_kv"], + topk_mem_indices, + infer_state.total_token_num, + ) if can_use_direct_smallq(*eligibility): - direct_workspace = self.backend.get_smallq_workspace(state, q) + direct_workspace = self.backend.get_smallq_workspace(infer_state, q) elif can_use_selected_prefill(*eligibility): - selected_workspace = self.backend.get_smallq_workspace(state, q) + selected_workspace = self.backend.get_smallq_workspace(infer_state, q) else: graph_guard = ( - state.is_cuda_graph_capture + infer_state.is_cuda_graph_capture or self.backend.model.prefill_graph is not None - or hasattr(state, "prefill_cuda_graph_exe_list") + or hasattr(infer_state, "prefill_cuda_graph_exe_list") ) if not graph_guard and not torch.cuda.is_current_stream_capturing(): - generic_workspace = self.backend.get_generic_prefill_workspace(state, q, params["topk_indices"]) + generic_workspace = self.backend.get_generic_prefill_workspace( + infer_state, q, nsa_prefill_dict["topk_indices"] + ) return context_attention_fwd_fp8( q=q, - packed_kv=kv, - prefill_mem_index=state.mem_index, - prefill_kv=params["prefill_cache_kv"], - kv_starts=self.ks, - topk_indices=params["topk_indices"], + packed_kv=packed_kv, + mem_index=infer_state.mem_index, + prefill_cache_kv=nsa_prefill_dict["prefill_cache_kv"], + q_kv_start_loc=self.q_kv_start_loc, + topk_indices=nsa_prefill_dict["topk_indices"], topk_mem_indices=topk_mem_indices, - req_to_token_index=state.req_manager.req_to_token_indexs, - b_req_idx=state.b_req_idx, - cu_kv_seq_len=state.b1_cu_kv_seq_len, - cu_q_seq_len=state.b1_cu_q_seq_len, - total_kv_tokens=state.total_token_num, - softmax_scale=params["softmax_scale"], + req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, + b_req_idx=infer_state.b_req_idx, + b1_cu_kv_seq_len=infer_state.b1_cu_kv_seq_len, + b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, + total_token_num=infer_state.total_token_num, + softmax_scale=nsa_prefill_dict["softmax_scale"], selected_workspace=selected_workspace, direct_workspace=direct_workspace, generic_workspace=generic_workspace, @@ -129,8 +138,8 @@ class Fp8Glm5NextSparseDecodeState(Glm5NextSparseDecodeState): def init_state(self): import flash_mla - self.lengths = self.infer_state.b_seq_len - self.query_req = self.infer_state.b_req_idx + self.q_kv_seq_len = self.infer_state.b_seq_len + self.q_req_idx = self.infer_state.b_req_idx self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() def _nsa_decode_att(self, q, packed_kv, att_control): @@ -139,8 +148,8 @@ def _nsa_decode_att(self, q, packed_kv, att_control): from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads - params = att_control.nsa_decode_dict - topk_mem_indices = params["topk_mem_indices"] + nsa_decode_dict = att_control.nsa_decode_dict + topk_mem_indices = nsa_decode_dict["topk_mem_indices"] if topk_mem_indices.ndim == 2: topk_mem_indices = topk_mem_indices.unsqueeze(1) assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" @@ -151,12 +160,12 @@ def _nsa_decode_att(self, q, packed_kv, att_control): q_nope, packed_kv, topk_mem_indices, - params["kv_lora_rank"], + nsa_decode_dict["kv_lora_rank"], self.decode_topk_bound, is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture, ): return fp8_sparse_decode.fp8_sparse_direct_decode( - q_nope, packed_kv, topk_mem_indices, params["softmax_scale"] + q_nope, packed_kv, topk_mem_indices, nsa_decode_dict["softmax_scale"] ) padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture) kv_cache = torch.as_strided( @@ -169,9 +178,9 @@ def _nsa_decode_att(self, q, packed_kv, att_control): k_cache=kv_cache, block_table=None, cache_seqlens=None, - head_dim_v=params["kv_lora_rank"], + head_dim_v=nsa_decode_dict["kv_lora_rank"], tile_scheduler_metadata=self.flashmla_sched_meta, - softmax_scale=params["softmax_scale"], + softmax_scale=nsa_decode_dict["softmax_scale"], causal=False, is_fp8_kvcache=True, indices=topk_mem_indices, diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py index 0e9b1b7900..6e33177dbe 100644 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -16,57 +16,60 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Glm5NextSparsePrefillState(BasePrefillAttState): - lengths: torch.Tensor = None - query_req: torch.Tensor = None - ks: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None def init_state(self): from lightllm.models.glm5_next.triton_kernel.kpool import prefill_nsa_params - state = self.infer_state - self.lengths, self.query_req, self.ks = prefill_nsa_params( - state.position_ids, state.b_req_idx, state.b1_cu_q_seq_len, state.b1_cu_kv_seq_len + infer_state = self.infer_state + self.q_kv_seq_len, self.q_req_idx, self.q_kv_start_loc = prefill_nsa_params( + infer_state.position_ids, infer_state.b_req_idx, infer_state.b1_cu_q_seq_len, infer_state.b1_cu_kv_seq_len ) def prefill_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): assert att_control.nsa_prefill and att_control.nsa_prefill_dict is not None - return self._nsa_prefill_att(q=q, kv=k, att_control=att_control) + return self._nsa_prefill_att(q, k, att_control) def _nsa_prefill_att(self, q, kv, att_control): from sgl_kernel.flash_mla import flash_mla_sparse_fwd - tokens, heads, dim = q.shape + q_token_num, q_head_num, head_dim = q.shape # FlashMLA uses 64-head tiles. - padded_heads = ((heads + 63) // 64) * 64 + padded_q_head_num = ((q_head_num + 63) // 64) * 64 padded_q = q - if padded_heads != heads: - padded_q = q.new_zeros((tokens, padded_heads, dim)) - padded_q[:, :heads] = q - params = att_control.nsa_prefill_dict - out, _, _ = flash_mla_sparse_fwd( + if padded_q_head_num != q_head_num: + padded_q = q.new_zeros((q_token_num, padded_q_head_num, head_dim)) + padded_q[:, :q_head_num] = q + nsa_prefill_dict = att_control.nsa_prefill_dict + output, _, _ = flash_mla_sparse_fwd( q=padded_q, kv=kv, - indices=params["topk_mem_indices"].unsqueeze(1), - sm_scale=params["softmax_scale"], - d_v=params["kv_lora_rank"], + indices=nsa_prefill_dict["topk_mem_indices"].unsqueeze(1), + sm_scale=nsa_prefill_dict["softmax_scale"], + d_v=nsa_prefill_dict["kv_lora_rank"], ) - return out[:, :heads] + return output[:, :q_head_num] @dataclasses.dataclass class Glm5NextSparseDecodeState(BaseDecodeAttState): - lengths: torch.Tensor = None - query_req: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None nsa_cache_seqlens: torch.Tensor = None nsa_cu_seqlens_k_new: torch.Tensor = None def init_state(self): - self.lengths = self.infer_state.b_seq_len - self.query_req = self.infer_state.b_req_idx - pool = self.backend.model.config["index_kpool"] - topk = self.backend.model.config["index_topk"] + self.q_kv_seq_len = self.infer_state.b_seq_len + self.q_req_idx = self.infer_state.b_req_idx + index_kpool = self.backend.model.config["index_kpool"] + index_topk = self.backend.model.config["index_topk"] self.nsa_cache_seqlens = ( - torch.minimum(self.lengths // pool * pool, torch.full_like(self.lengths, topk)) + self.lengths % pool + torch.minimum( + self.q_kv_seq_len // index_kpool * index_kpool, torch.full_like(self.q_kv_seq_len, index_topk) + ) + + self.q_kv_seq_len % index_kpool ) self.nsa_cu_seqlens_k_new = torch.nn.functional.pad(self.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0)) @@ -77,21 +80,21 @@ def decode_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): def _nsa_decode_att(self, q, kv, att_control): from sgl_kernel.flash_attn import flash_attn_with_kvcache - params = att_control.nsa_decode_dict + nsa_decode_dict = att_control.nsa_decode_dict q_nope, _ = q - kv_nope = kv.view(-1, 1, 1, params["kv_lora_rank"]) + kv_nope = kv.view(-1, 1, 1, nsa_decode_dict["kv_lora_rank"]) # NoPE has no RoPE Q/K tensors; only_qv uses q_nope and kv_nope directly. return flash_attn_with_kvcache( q=None, qv=q_nope, k_cache=None, v_cache=kv_nope, - page_table=params["topk_mem_indices"], + page_table=nsa_decode_dict["topk_mem_indices"], cache_seqlens=self.nsa_cache_seqlens, cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=params["softmax_scale"], + softmax_scale=nsa_decode_dict["softmax_scale"], causal=False, only_qv=True, ) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 3a13662e32..869cc13ea8 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -13,97 +13,108 @@ class Glm5NextNsaInfer: def __init__(self, layer_idx, network_config, tp_world_size): self.layer_idx = layer_idx - self.topk = network_config["index_topk"] - self.heads = network_config["index_n_heads"] - self.dim = network_config["index_head_dim"] - self.eps = network_config["rms_norm_eps"] - self.pool_size = network_config["index_kpool"] + self.index_topk = network_config["index_topk"] + self.index_n_heads = network_config["index_n_heads"] + self.index_head_dim = network_config["index_head_dim"] + self.rms_norm_eps = network_config["rms_norm_eps"] + self.index_kpool = network_config["index_kpool"] - def select_topk_indices(self, logits, lengths, indices): + def select_topk_indices(self, logits, q_pool_seq_len, topk_pool_indices): """Select row-relative indices, with valid entries before -1 padding.""" if HAS_VLLM and hasattr(vllm_ops, "top_k_per_row_decode"): # next_n=1 treats each query as an independent row, including prefill. # The decode entry splits long rows before merging their candidates. # Tested with vLLM 0.22.1; persistent_topk can drop candidates (#51782). vllm_ops.top_k_per_row_decode( - logits, 1, lengths, indices, logits.shape[0], logits.stride(0), logits.stride(1), indices.shape[1] + logits, + 1, + q_pool_seq_len, + topk_pool_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_pool_indices.shape[1], ) else: - positions = torch.arange(logits.shape[1], device=logits.device) - logits.masked_fill_(positions[None, :] >= lengths[:, None], -float("inf")) - selected = torch.topk(logits, indices.shape[1], dim=-1, sorted=True).indices - indices.copy_(selected.masked_fill(selected >= lengths[:, None], -1)) + pool_positions = torch.arange(logits.shape[1], device=logits.device) + logits.masked_fill_(pool_positions[None, :] >= q_pool_seq_len[:, None], -float("inf")) + selected = torch.topk(logits, topk_pool_indices.shape[1], dim=-1, sorted=True).indices + topk_pool_indices.copy_(selected.masked_fill(selected >= q_pool_seq_len[:, None], -1)) def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): - k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.eps) + k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.rms_norm_eps) gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) - raw = torch.cat((k, gate), -1) - role = infer_state.att_state_role - tail = infer_state.att_state_manager.get_state( - model_role=role, - state_type_filter=f"compressor_indexer_w{self.pool_size}_s{self.pool_size}", + k_and_gate = torch.cat((k, gate), -1) + att_state_role = infer_state.att_state_role + indexer_tail = infer_state.att_state_manager.get_state( + model_role=att_state_role, + state_type_filter=f"compressor_indexer_w{self.index_kpool}_s{self.index_kpool}", logical_layer_id=self.layer_idx, ).get_layer_view(logical_layer_id=self.layer_idx) - state = infer_state.att_state_manager.get_state( + full_att_state = infer_state.att_state_manager.get_state( model_role=infer_state.att_state_role, state_type_filter="mla", logical_layer_id=self.layer_idx, state_class=FullAttentionState, ) - packed_buffer = state.get_layer_view(logical_layer_id=self.layer_idx).indexer_k + indexer_k = full_att_state.get_layer_view(logical_layer_id=self.layer_idx).indexer_k compress_pools( - raw=raw, - tail=tail, - packed_buffer=packed_buffer, - ape=layer_weight.index_kpool_compress_ape.weight, - lengths=att_state.lengths, + k_and_gate=k_and_gate, + indexer_tail=indexer_tail, + indexer_k=indexer_k, + index_kpool_compress_ape=layer_weight.index_kpool_compress_ape.weight, + q_kv_seq_len=att_state.q_kv_seq_len, mem_index=infer_state.mem_index, - req_idx=infer_state.b_req_idx, - cu_q_lens=infer_state.b1_cu_q_seq_len, - seq_lens=infer_state.b_seq_len, - max_q_len=infer_state.max_q_seq_len, - mtp_index=None if infer_state.is_prefill else infer_state.b_mtp_index, + b_req_idx=infer_state.b_req_idx, + b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, + b_seq_len=infer_state.b_seq_len, + max_q_seq_len=infer_state.max_q_seq_len, + b_mtp_index=None if infer_state.is_prefill else infer_state.b_mtp_index, ) - if infer_state.max_kv_seq_len <= self.topk: + if infer_state.max_kv_seq_len <= self.index_topk: return expand_topk( None, - att_state.lengths, - att_state.query_req, + att_state.q_kv_seq_len, + att_state.q_req_idx, infer_state.req_manager.req_to_token_indexs, - self.topk, + self.index_topk, dense=True, ) # The small indexer is replicated: no all-gather of query heads and # identical pool selection on every TP rank. - q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.heads, self.dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim ** -0.5) + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.index_n_heads, self.index_head_dim) + q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.index_head_dim ** -0.5) weights = layer_weight.weights_proj_.mm(hidden_states.float()) - weights = weights * (self.heads ** -0.5 * self.dim ** -0.5) * q_scale.squeeze(-1) + weights = weights * (self.index_n_heads ** -0.5 * self.index_head_dim ** -0.5) * q_scale.squeeze(-1) max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. # Decode has few queries, so paged MQA also splits K across SMs for parallelism. if infer_state.is_prefill: - groups = self._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + topk_pool_indices = self._get_prefill_indices(q_fp8, weights, indexer_k, infer_state, att_state, max_pools) else: - groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + topk_pool_indices = self._get_decode_indices(q_fp8, weights, indexer_k, infer_state, att_state, max_pools) return expand_topk( - groups, att_state.lengths, att_state.query_req, infer_state.req_manager.req_to_token_indexs, self.topk + topk_pool_indices, + att_state.q_kv_seq_len, + att_state.q_req_idx, + infer_state.req_manager.req_to_token_indexs, + self.index_topk, ) - def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): - keys = gather_pools( - packed_buffer, + def _get_prefill_indices(self, q_fp8, weights, indexer_k, infer_state, att_state, max_pools): + k_fp8_and_scale = gather_pools( + indexer_k, infer_state.req_manager.req_to_token_indexs, infer_state.b_req_idx, infer_state.b_seq_len, max_pools, ) - starts, ends, lengths = get_pool_ranges( - att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools + ks, ke, q_pool_seq_len = get_pool_ranges( + att_state.q_kv_seq_len, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools ) - groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) + topk_pool_indices = torch.empty((q_fp8.shape[0], self.index_topk // 4), dtype=torch.int32, device=q_fp8.device) # Budget 64 MiB for FP32 logits. At 1M tokens this permits 64 queries per chunk, # limiting Q parallelism even when the prefill batch contains many tokens. chunk_size = max(1, min(q_fp8.shape[0], 16 * 1024 * 1024 // max_pools)) @@ -113,34 +124,41 @@ def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_s end = min(start + chunk_size, q_fp8.shape[0]) logits = deep_gemm.fp8_mqa_logits( q_fp8[start:end], - keys, + k_fp8_and_scale, weights[start:end], - starts[start:end], - ends[start:end], + ks[start:end], + ke[start:end], clean_logits=False, max_seqlen_k=max_pools, ) - self.select_topk_indices(logits, lengths[start:end], groups[start:end]) - return groups + self.select_topk_indices(logits, q_pool_seq_len[start:end], topk_pool_indices[start:end]) + return topk_pool_indices - def _get_decode_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): + def _get_decode_indices(self, q_fp8, weights, indexer_k, infer_state, att_state, max_pools): import deep_gemm - lengths = (att_state.lengths // 4).view(-1, 1) - pages, block_table = gather_paged_pools( - packed_buffer, + q_pool_seq_len = (att_state.q_kv_seq_len // 4).view(-1, 1) + indexer_k_pages, block_table = gather_paged_pools( + indexer_k, infer_state.req_manager.req_to_token_indexs, infer_state.b_req_idx, - lengths, + q_pool_seq_len, max_pools, ) - metadata = deep_gemm.get_paged_mqa_logits_metadata(lengths, 64, deep_gemm.get_num_sms()) + mqa_logits_metadata = deep_gemm.get_paged_mqa_logits_metadata(q_pool_seq_len, 64, deep_gemm.get_num_sms()) # Each MTP position has its own pool length; HOLD rows have no valid pools. # Fixed-width mtp_step=2 gives Q=[3 * num_requests, 1, heads, dim], so next_n stays 1. # All three verify positions are preserved despite the SM90 kernel's native next_n limit of 2. logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8.unsqueeze(1), pages, weights, lengths, block_table, metadata, max_pools, clean_logits=False + q_fp8.unsqueeze(1), + indexer_k_pages, + weights, + q_pool_seq_len, + block_table, + mqa_logits_metadata, + max_pools, + clean_logits=False, ) - groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) - self.select_topk_indices(logits, lengths.view(-1), groups) - return groups + topk_pool_indices = torch.empty((q_fp8.shape[0], self.index_topk // 4), dtype=torch.int32, device=q_fp8.device) + self.select_topk_indices(logits, q_pool_seq_len.view(-1), topk_pool_indices) + return topk_pool_indices diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index f87bd3a235..dcbd691260 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -9,68 +9,68 @@ @triton.jit def _direct_smallq_packed_kernel( - Q, - PackedKV, - PrefillMemIndex, - PrefillKV, - Topk, - PartialO, - PartialLse, - packed_rows, + q, + packed_kv, + mem_index, + prefill_cache_kv, + topk_mem_indices, + partial_output, + partial_softmax_lse, + kv_cache_capacity, q_stride_s: tl.constexpr, q_stride_h: tl.constexpr, q_stride_d: tl.constexpr, - packed_stride_s: tl.constexpr, - prefill_stride_s: tl.constexpr, - prefill_stride_d: tl.constexpr, - topk_stride_s: tl.constexpr, - topk_stride_d: tl.constexpr, - HEADS_: tl.constexpr, - BLOCK_H_: tl.constexpr, - PREFILL_N: tl.constexpr, - SPLITS_: tl.constexpr, - BLOCKS_: tl.constexpr, - SLOTS_: tl.constexpr, + packed_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_d: tl.constexpr, + topk_mem_indices_stride_s: tl.constexpr, + topk_mem_indices_stride_d: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_HEADS: tl.constexpr, + q_token_num: tl.constexpr, + NUM_SPLITS: tl.constexpr, + KV_BLOCKS_PER_SPLIT: tl.constexpr, + KV_SLOTS_PER_SPLIT: tl.constexpr, softmax_scale: tl.constexpr, ): qid = tl.program_id(0) split = tl.program_id(1) - if HEADS_ == 8: + if q_head_num == 8: # H8 fits in one fixed 16-lane tile; the head mask excludes inactive lanes. h = tl.arange(0, 16) else: - h = tl.program_id(2) * BLOCK_H_ + tl.arange(0, BLOCK_H_) + h = tl.program_id(2) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) d = tl.arange(0, 512) - hm = h < HEADS_ - q = tl.load( - Q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, + hm = h < q_head_num + q_block = tl.load( + q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, hm[:, None], 0.0, ).to(tl.bfloat16) - if HEADS_ == 8: + if q_head_num == 8: acc = tl.zeros((16, 512), tl.float32) m = tl.full((16,), -float("inf"), tl.float32) denominator = tl.zeros((16,), tl.float32) else: - acc = tl.zeros((BLOCK_H_, 512), tl.float32) - m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) - denominator = tl.zeros((BLOCK_H_,), tl.float32) - for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): + acc = tl.zeros((BLOCK_HEADS, 512), tl.float32) + m = tl.full((BLOCK_HEADS,), -float("inf"), tl.float32) + denominator = tl.zeros((BLOCK_HEADS,), tl.float32) + for b in tl.range(0, KV_BLOCKS_PER_SPLIT, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): n = tl.arange(0, 64) - slot = split * SLOTS_ + b * 64 + n + slot = split * KV_SLOTS_PER_SPLIT + b * 64 + n top = tl.load( - Topk + qid * topk_stride_s + slot * topk_stride_d, + topk_mem_indices + qid * topk_mem_indices_stride_s + slot * topk_mem_indices_stride_d, slot < 2176, -1, ).to(tl.int64) - valid = (top >= 0) & (top < packed_rows) - mem = tl.where(valid, top, 0) + valid = (top >= 0) & (top < kv_cache_capacity) + kv_mem_index = tl.where(valid, top, 0) current_row = tl.full((64,), -1, tl.int32) - for j in tl.static_range(PREFILL_N): - current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) + for j in tl.static_range(q_token_num): + current_row = tl.where(valid & (kv_mem_index == tl.load(mem_index + j).to(tl.int64)), j, current_row) cached = valid & (current_row < 0) current = valid & (current_row >= 0) - packed_row = PackedKV + mem[:, None] * packed_stride_s + packed_row = packed_kv + kv_mem_index[:, None] * packed_kv_stride_s fp8 = tl.load( packed_row.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], cached[:, None] & (d[None, :] < 448), @@ -89,12 +89,14 @@ def _direct_smallq_packed_kernel( ).to(tl.float32) cached_v = tl.where(d[None, :] < 448, fp8, tail) current_v = tl.load( - PrefillKV + current_row[:, None] * prefill_stride_s + d[None, :] * prefill_stride_d, + prefill_cache_kv + + current_row[:, None] * prefill_cache_kv_stride_s + + d[None, :] * prefill_cache_kv_stride_d, current[:, None], 0.0, ).to(tl.float32) kv = tl.where(current[:, None], current_v, cached_v).to(tl.bfloat16) - score = tl.dot(q, tl.trans(kv), out_dtype=tl.float32) * softmax_scale + score = tl.dot(q_block, tl.trans(kv), out_dtype=tl.float32) * softmax_scale score = tl.where(valid[None, :], score, -float("inf")) block_m = tl.max(score, axis=1) new_m = tl.maximum(m, block_m) @@ -106,12 +108,12 @@ def _direct_smallq_packed_kernel( has = denominator > 0.0 out = tl.where(has[:, None], acc / denominator[:, None], 0.0) tl.store( - PartialO + ((qid * SPLITS_ + split) * HEADS_ + h[:, None]) * 512 + d[None, :], + partial_output + ((qid * NUM_SPLITS + split) * q_head_num + h[:, None]) * 512 + d[None, :], out, hm[:, None], ) tl.store( - PartialLse + (qid * SPLITS_ + split) * HEADS_ + h, + partial_softmax_lse + (qid * NUM_SPLITS + split) * q_head_num + h, tl.where(has, m + tl.log(denominator), -float("inf")), hm, ) @@ -119,33 +121,33 @@ def _direct_smallq_packed_kernel( @triton.jit def reduce_direct_smallq_kernel( - PartialO, - PartialLse, - Out, - HEADS_: tl.constexpr, - SPLITS_: tl.constexpr, - SPAD_: tl.constexpr, + partial_output, + partial_softmax_lse, + output, + q_head_num: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, ): qid = tl.program_id(0) h = tl.program_id(1) d = tl.program_id(2) * 128 + tl.arange(0, 128) - s = tl.arange(0, SPAD_) - sm = s < SPLITS_ - lse = tl.load(PartialLse + (qid * SPLITS_ + s) * HEADS_ + h, sm, -float("inf")) - maximum = tl.max(lse, axis=0) + s = tl.arange(0, BLOCK_SPLITS) + sm = s < NUM_SPLITS + softmax_lse = tl.load(partial_softmax_lse + (qid * NUM_SPLITS + s) * q_head_num + h, sm, -float("inf")) + maximum = tl.max(softmax_lse, axis=0) safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) - weight = tl.where(sm & (lse > -float("inf")), tl.exp(lse - safe_maximum), 0.0) + weight = tl.where(sm & (softmax_lse > -float("inf")), tl.exp(softmax_lse - safe_maximum), 0.0) denominator = tl.sum(weight, axis=0) partial = tl.load( - PartialO + ((qid * SPLITS_ + s[:, None]) * HEADS_ + h) * 512 + d[None, :], + partial_output + ((qid * NUM_SPLITS + s[:, None]) * q_head_num + h) * 512 + d[None, :], sm[:, None], 0.0, ).to(tl.float32) value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) - tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) + tl.store(output + (qid * q_head_num + h) * 512 + d, value) -def can_use_direct_smallq(q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens): +def can_use_direct_smallq(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): return ( q.ndim == 3 and 1 <= q.shape[0] <= 8 @@ -154,22 +156,24 @@ def can_use_direct_smallq(q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_ and packed_kv.ndim == 3 and packed_kv.shape[1] == 1 and packed_kv.shape[2] >= 584 - and total_kv_tokens >= 65536 - and prefill_mem_index.dtype == torch.int32 - and prefill_mem_index.numel() == q.shape[0] - and prefill_kv.dtype == torch.bfloat16 - and prefill_kv.shape == (q.shape[0], 1, 512) + and total_token_num >= 65536 + and mem_index.dtype == torch.int32 + and mem_index.numel() == q.shape[0] + and prefill_cache_kv.dtype == torch.bfloat16 + and prefill_cache_kv.shape == (q.shape[0], 1, 512) and topk_mem_indices.dtype == torch.int32 and topk_mem_indices.shape == (q.shape[0], 2176) ) def allocate_direct_smallq_workspace(q): - qn, heads, _ = q.shape - splits = 34 if qn == 1 else 17 + q_token_num, q_head_num, _ = q.shape + num_splits = 34 if q_token_num == 1 else 17 return { - "partial_o": torch.empty((qn, splits, heads, q.shape[-1]), dtype=torch.float32, device=q.device), - "partial_lse": torch.empty((qn, splits, heads), dtype=torch.float32, device=q.device), + "partial_output": torch.empty( + (q_token_num, num_splits, q_head_num, q.shape[-1]), dtype=torch.float32, device=q.device + ), + "partial_softmax_lse": torch.empty((q_token_num, num_splits, q_head_num), dtype=torch.float32, device=q.device), # The reducer requires a contiguous [Q, H, 512] destination. "output": torch.empty(q.shape, dtype=q.dtype, device=q.device), } @@ -177,58 +181,58 @@ def allocate_direct_smallq_workspace(q): @torch.no_grad() def _context_attention_fwd_fp8_direct_smallq( - q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, direct_workspace=None + q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace=None ): - qn = q.shape[0] - heads = q.shape[1] + q_token_num = q.shape[0] + q_head_num = q.shape[1] # Q1 uses 34 64-key splits; H64 uses two H32 head tiles per split. - splits, blocks, slots = (34, 1, 64) if qn == 1 else (17, 2, 128) - block_h = 32 if heads == 64 else 16 - if heads == 8: + num_splits, kv_blocks_per_split, kv_slots_per_split = (34, 1, 64) if q_token_num == 1 else (17, 2, 128) + block_heads = 32 if q_head_num == 64 else 16 + if q_head_num == 8: # H8 uses local buffers because its direct path has no reusable workspace. - partial_o = torch.empty((qn, splits, 8, 512), dtype=torch.float32, device=q.device) - partial_lse = torch.empty((qn, splits, 8), dtype=torch.float32, device=q.device) + partial_output = torch.empty((q_token_num, num_splits, 8, 512), dtype=torch.float32, device=q.device) + partial_softmax_lse = torch.empty((q_token_num, num_splits, 8), dtype=torch.float32, device=q.device) output = torch.empty(q.shape, dtype=q.dtype, device=q.device) - launch_grid = (qn, splits) + launch_grid = (q_token_num, num_splits) else: workspace = direct_workspace or allocate_direct_smallq_workspace(q) - partial_o = workspace["partial_o"] - partial_lse = workspace["partial_lse"] + partial_output = workspace["partial_output"] + partial_softmax_lse = workspace["partial_softmax_lse"] output = workspace["output"] - launch_grid = (qn, splits, triton.cdiv(heads, block_h)) + launch_grid = (q_token_num, num_splits, triton.cdiv(q_head_num, block_heads)) _direct_smallq_packed_kernel[launch_grid]( q, packed_kv, - prefill_mem_index, - prefill_kv, + mem_index, + prefill_cache_kv, topk_mem_indices, - partial_o, - partial_lse, + partial_output, + partial_softmax_lse, packed_kv.shape[0], q.stride(0), q.stride(1), q.stride(2), packed_kv.stride(0), - prefill_kv.stride(0), - prefill_kv.stride(2), + prefill_cache_kv.stride(0), + prefill_cache_kv.stride(2), topk_mem_indices.stride(0), topk_mem_indices.stride(1), - HEADS_=heads, - BLOCK_H_=block_h, - PREFILL_N=qn, - SPLITS_=splits, - BLOCKS_=blocks, - SLOTS_=slots, + q_head_num=q_head_num, + BLOCK_HEADS=block_heads, + q_token_num=q_token_num, + NUM_SPLITS=num_splits, + KV_BLOCKS_PER_SPLIT=kv_blocks_per_split, + KV_SLOTS_PER_SPLIT=kv_slots_per_split, softmax_scale=softmax_scale, num_warps=8, ) - reduce_direct_smallq_kernel[(qn, heads, 4)]( - partial_o, - partial_lse, + reduce_direct_smallq_kernel[(q_token_num, q_head_num, 4)]( + partial_output, + partial_softmax_lse, output, - HEADS_=heads, - SPLITS_=splits, - SPAD_=triton.next_power_of_2(splits), + q_head_num=q_head_num, + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), num_warps=4, ) return output @@ -236,48 +240,50 @@ def _context_attention_fwd_fp8_direct_smallq( @triton.jit def _gather_prefill_kv_kernel( - PackedKV, - PrefillKV, - Output, - ReqToToken, - ReqIdx, - CuKV, - CuQ, + packed_kv, + prefill_cache_kv, + output, + req_to_token_indexs, + b_req_idx, + b1_cu_kv_seq_len, + b1_cu_q_seq_len, chunk_start, - num_tokens, - packed_stride_s: tl.constexpr, - prefill_stride_s: tl.constexpr, - prefill_stride_d: tl.constexpr, - REQ_STRIDE: tl.constexpr, - TOKEN_STRIDE: tl.constexpr, - BATCH: tl.constexpr, + chunk_token_num, + packed_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_d: tl.constexpr, + req_to_token_indexs_stride_b: tl.constexpr, + req_to_token_indexs_stride_s: tl.constexpr, + batch_size: tl.constexpr, BLOCK_ROWS: tl.constexpr, ): block = tl.program_id(0) - lanes = tl.arange(0, triton.next_power_of_2(BATCH)) - ends = tl.load(CuKV + lanes + 1, lanes < BATCH, 0) - ends = tl.minimum(tl.maximum(ends - chunk_start, 0), num_tokens) + lanes = tl.arange(0, triton.next_power_of_2(batch_size)) + ends = tl.load(b1_cu_kv_seq_len + lanes + 1, lanes < batch_size, 0) + ends = tl.minimum(tl.maximum(ends - chunk_start, 0), chunk_token_num) block_ends = ends // BLOCK_ROWS + lanes + 1 - batch = tl.sum(((block >= block_ends) & (lanes < BATCH)).to(tl.int32), 0) - start = tl.load(CuKV + batch) - copy_start = tl.minimum(tl.maximum(start - chunk_start, 0), num_tokens) - copy_end = tl.minimum(tl.maximum(tl.load(CuKV + batch + 1) - chunk_start, 0), num_tokens) + batch = tl.sum(((block >= block_ends) & (lanes < batch_size)).to(tl.int32), 0) + start = tl.load(b1_cu_kv_seq_len + batch) + copy_start = tl.minimum(tl.maximum(start - chunk_start, 0), chunk_token_num) + copy_end = tl.minimum(tl.maximum(tl.load(b1_cu_kv_seq_len + batch + 1) - chunk_start, 0), chunk_token_num) local_block = block - (copy_start // BLOCK_ROWS + batch) row = copy_start + local_block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) live = row < copy_end - req = tl.load(ReqIdx + batch).to(tl.int64) + req = tl.load(b_req_idx + batch).to(tl.int64) position = chunk_start + row - start - q_start = tl.load(CuQ + batch) - q_len = tl.load(CuQ + batch + 1) - q_start - cache_len = tl.load(CuKV + batch + 1) - start - q_len + q_start = tl.load(b1_cu_q_seq_len + batch) + q_len = tl.load(b1_cu_q_seq_len + batch + 1) - q_start + cache_len = tl.load(b1_cu_kv_seq_len + batch + 1) - start - q_len prefill_row = tl.where(position >= cache_len, q_start + position - cache_len, -1) - mem_index = tl.load(ReqToToken + req * REQ_STRIDE + position * TOKEN_STRIDE, live & (prefill_row < 0), 0).to( - tl.int64 - ) + kv_mem_index = tl.load( + req_to_token_indexs + req * req_to_token_indexs_stride_b + position * req_to_token_indexs_stride_s, + live & (prefill_row < 0), + 0, + ).to(tl.int64) dim = tl.arange(0, 512) cached = live & (prefill_row < 0) current = live & (prefill_row >= 0) - packed_row = PackedKV + mem_index[:, None] * packed_stride_s + packed_row = packed_kv + kv_mem_index[:, None] * packed_kv_stride_s quantized = tl.load( packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], @@ -293,12 +299,12 @@ def _gather_prefill_kv_kernel( ).to(tl.float32) values = tl.where(dim[None, :] < 448, values, tail) current_values = tl.load( - PrefillKV + prefill_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + prefill_cache_kv + prefill_row[:, None] * prefill_cache_kv_stride_s + dim[None, :] * prefill_cache_kv_stride_d, current[:, None], 0, ).to(tl.float32) tl.store( - Output + row[:, None] * 512 + dim[None, :], + output + row[:, None] * 512 + dim[None, :], tl.where(current[:, None], current_values, values), live[:, None], ) @@ -306,10 +312,10 @@ def _gather_prefill_kv_kernel( @triton.jit def _compact_topk_indices_kernel( - TopkIndices, - KVStarts, - ChunkIndices, - ChunkLengths, + topk_indices, + q_kv_start_loc, + chunk_topk_indices, + chunk_cache_seqlens, chunk_start, chunk_end, TOPK: tl.constexpr, @@ -318,16 +324,16 @@ def _compact_topk_indices_kernel( ): row = tl.program_id(0) lane = tl.arange(0, BLOCK) - index = tl.load(TopkIndices + row * TOPK + lane, lane < TOPK, -1) - kv_index = index + tl.load(KVStarts + row) + index = tl.load(topk_indices + row * TOPK + lane, lane < TOPK, -1) + kv_index = index + tl.load(q_kv_start_loc + row) valid = (lane < TOPK) & (index >= 0) & (kv_index >= chunk_start) & (kv_index < chunk_end) if COMPACT: compact_index = tl.cumsum(valid.to(tl.int32)) - 1 - tl.store(ChunkIndices + row * TOPK + compact_index, kv_index - chunk_start, valid) - tl.store(ChunkLengths + row, tl.sum(valid.to(tl.int32))) + tl.store(chunk_topk_indices + row * TOPK + compact_index, kv_index - chunk_start, valid) + tl.store(chunk_cache_seqlens + row, tl.sum(valid.to(tl.int32))) else: tl.store( - ChunkIndices + row * TOPK + lane, + chunk_topk_indices + row * TOPK + lane, tl.where(valid, kv_index - chunk_start, -1), lane < TOPK, ) @@ -335,103 +341,107 @@ def _compact_topk_indices_kernel( @triton.jit def _merge_att_output_kernel( - Output, - Lse, - ChunkOutput, - ChunkLse, + output, + softmax_lse, + chunk_output, + chunk_softmax_lse, chunk_lse_stride_s: tl.constexpr, chunk_lse_stride_h, - NUM_HEADS: tl.constexpr, - NUM_Q_TOKENS, + q_head_num: tl.constexpr, + q_token_num, BLOCK_HEADS: tl.constexpr, ): row = tl.program_id(0) head = tl.program_id(1) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) dim = tl.arange(0, 512) - previous_lse = tl.load(Lse + head * NUM_Q_TOKENS + row, head < NUM_HEADS, -float("inf")) - chunk_lse = tl.load( - ChunkLse + row * chunk_lse_stride_s + head * chunk_lse_stride_h, head < NUM_HEADS, -float("inf") + previous_lse = tl.load(softmax_lse + head * q_token_num + row, head < q_head_num, -float("inf")) + chunk_lse_block = tl.load( + chunk_softmax_lse + row * chunk_lse_stride_s + head * chunk_lse_stride_h, head < q_head_num, -float("inf") ) # FA3 uses natural-log LSE and returns +inf for an empty row. previous_lse = tl.where(previous_lse == float("inf"), -float("inf"), previous_lse) - chunk_lse = tl.where(chunk_lse == float("inf"), -float("inf"), chunk_lse) - maximum = tl.maximum(previous_lse, chunk_lse) + chunk_lse_block = tl.where(chunk_lse_block == float("inf"), -float("inf"), chunk_lse_block) + maximum = tl.maximum(previous_lse, chunk_lse_block) maximum = tl.where(maximum == -float("inf"), 0, maximum) previous_weight = tl.exp(previous_lse - maximum) - chunk_weight = tl.exp(chunk_lse - maximum) + chunk_weight = tl.exp(chunk_lse_block - maximum) denominator = previous_weight + chunk_weight - offset = (row * NUM_HEADS + head[:, None]) * 512 + dim[None, :] - previous = tl.load(Output + offset, (head[:, None] < NUM_HEADS) & (previous_weight[:, None] > 0), 0) - current = tl.load(ChunkOutput + offset, (head[:, None] < NUM_HEADS) & (chunk_weight[:, None] > 0), 0).to(tl.float32) + offset = (row * q_head_num + head[:, None]) * 512 + dim[None, :] + previous = tl.load(output + offset, (head[:, None] < q_head_num) & (previous_weight[:, None] > 0), 0) + current = tl.load(chunk_output + offset, (head[:, None] < q_head_num) & (chunk_weight[:, None] > 0), 0).to( + tl.float32 + ) result = (previous * previous_weight[:, None] + current * chunk_weight[:, None]) / tl.where( denominator > 0, denominator, 1 )[:, None] - tl.store(Output + offset, result, head[:, None] < NUM_HEADS) + tl.store(output + offset, result, head[:, None] < q_head_num) tl.store( - Lse + head * NUM_Q_TOKENS + row, + softmax_lse + head * q_token_num + row, tl.where(denominator > 0, maximum + tl.log(denominator), -float("inf")), - head < NUM_HEADS, + head < q_head_num, ) @triton.jit def _selected_prefill_compact_kernel( - Topk, - CompactPhysical, - Lengths, - PageTable, - CuQ, - packed_rows: tl.constexpr, - topk_stride_s: tl.constexpr, - topk_stride_d: tl.constexpr, - BATCH: tl.constexpr, + topk_mem_indices, + compact_topk_mem_indices, + cache_seqlens, + page_table, + cu_seqlens_q, + kv_cache_capacity: tl.constexpr, + topk_mem_indices_stride_s: tl.constexpr, + topk_mem_indices_stride_d: tl.constexpr, + q_token_num: tl.constexpr, BLOCK: tl.constexpr, ): qid = tl.program_id(0) lane = tl.arange(0, BLOCK) - local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) + local = tl.load( + topk_mem_indices + qid * topk_mem_indices_stride_s + lane * topk_mem_indices_stride_d, lane < 2176, -1 + ).to(tl.int64) valid = (lane < 2176) & (local >= 0) physical = local - valid = valid & (physical >= 0) & (physical < packed_rows) + valid = valid & (physical >= 0) & (physical < kv_cache_capacity) rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 - tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) + tl.store(compact_topk_mem_indices + qid * 2176 + rank, physical, valid) # Every selected slot has a stable contiguous page; invalid tail pages are - # harmless because cache_seqlens limits FA3 to Lengths[qid]. - tl.store(PageTable + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) + # harmless because cache_seqlens limits FA3 to cache_seqlens[qid]. + tl.store(page_table + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) if qid == 0: - tl.store(CuQ + lane, lane.to(tl.int32), lane <= BATCH) - tl.store(Lengths + qid, tl.sum(valid.to(tl.int32), axis=0)) + tl.store(cu_seqlens_q + lane, lane.to(tl.int32), lane <= q_token_num) + tl.store(cache_seqlens + qid, tl.sum(valid.to(tl.int32), axis=0)) @triton.jit def _selected_prefill_gather_kernel( - CompactPhysical, - Lengths, - PrefillMemIndex, - PrefillKV, - PackedKV, - SelectedKV, - CuK, - packed_stride_s: tl.constexpr, - prefill_stride_s: tl.constexpr, - prefill_stride_d: tl.constexpr, - BATCH: tl.constexpr, - BATCH_PAD: tl.constexpr, + compact_topk_mem_indices, + cache_seqlens, + mem_index, + prefill_cache_kv, + packed_kv, + selected_kv, + cu_seqlens_k_new, + packed_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_s: tl.constexpr, + prefill_cache_kv_stride_d: tl.constexpr, + q_token_num: tl.constexpr, + BLOCK_Q_TOKENS: tl.constexpr, ): qid = tl.program_id(0) row = tl.program_id(1) * 4 + tl.arange(0, 4) - length = tl.load(Lengths + qid) + length = tl.load(cache_seqlens + qid) live = row < length - physical = tl.load(CompactPhysical + qid * 2176 + row, live, 0).to(tl.int64) + physical = tl.load(compact_topk_mem_indices + qid * 2176 + row, live, 0).to(tl.int64) current_row = tl.full((4,), -1, tl.int32) - for j in tl.static_range(BATCH_PAD): - current = j < BATCH - current_mem = tl.load(PrefillMemIndex + j, current, 0).to(tl.int64) + for j in tl.static_range(BLOCK_Q_TOKENS): + current = j < q_token_num + current_mem = tl.load(mem_index + j, current, 0).to(tl.int64) current_row = tl.where(live & current & (physical == current_mem), j, current_row) cached = live & (current_row < 0) current = live & (current_row >= 0) dim = tl.arange(0, 512) - packed_row = PackedKV + physical[:, None] * packed_stride_s + packed_row = packed_kv + physical[:, None] * packed_kv_stride_s group = tl.arange(0, 8) raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) @@ -451,20 +461,20 @@ def _selected_prefill_gather_kernel( ).to(tl.float32) cached_value = tl.where(dim[None, :] < 448, fp8, tail) current_value = tl.load( - PrefillKV + current_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + prefill_cache_kv + current_row[:, None] * prefill_cache_kv_stride_s + dim[None, :] * prefill_cache_kv_stride_d, current[:, None], 0.0, ).to(tl.float32) value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) - tl.store(SelectedKV + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) + tl.store(selected_kv + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) if (qid == 0) & (tl.program_id(1) == 0): - x = tl.arange(0, BATCH_PAD) - lengths = tl.load(Lengths + x, x < BATCH, 0) - tl.store(CuK, 0) - tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) + x = tl.arange(0, BLOCK_Q_TOKENS) + lengths = tl.load(cache_seqlens + x, x < q_token_num, 0) + tl.store(cu_seqlens_k_new, 0) + tl.store(cu_seqlens_k_new + 1 + x, tl.cumsum(lengths, axis=0), x < q_token_num) -def can_use_selected_prefill(q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens): +def can_use_selected_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): return ( q.ndim == 3 and 2 <= q.shape[0] <= 8 @@ -475,75 +485,75 @@ def can_use_selected_prefill(q, packed_kv, prefill_mem_index, prefill_kv, topk_m and packed_kv.shape[1:] == (1, 584) and packed_kv.dtype == torch.uint8 and packed_kv.device == q.device - and total_kv_tokens >= 65536 - and prefill_mem_index.dtype == torch.int32 - and prefill_mem_index.device == q.device - and prefill_mem_index.shape == (q.shape[0],) - and prefill_kv.dtype == torch.bfloat16 - and prefill_kv.device == q.device - and prefill_kv.shape == (q.shape[0], 1, 512) + and total_token_num >= 65536 + and mem_index.dtype == torch.int32 + and mem_index.device == q.device + and mem_index.shape == (q.shape[0],) + and prefill_cache_kv.dtype == torch.bfloat16 + and prefill_cache_kv.device == q.device + and prefill_cache_kv.shape == (q.shape[0], 1, 512) and topk_mem_indices.dtype == torch.int32 and topk_mem_indices.device == q.device and topk_mem_indices.shape == (q.shape[0], 2176) ) -def allocate_selected_prefill_workspace(q, batch_capacity=None): - batch = q.shape[0] if batch_capacity is None else batch_capacity +def allocate_selected_prefill_workspace(q, q_token_capacity=None): + q_token_num = q.shape[0] if q_token_capacity is None else q_token_capacity return { - "compact": torch.empty((batch, 2176), dtype=torch.int64, device=q.device), - "lengths": torch.empty(batch, dtype=torch.int32, device=q.device), - "page_table": torch.empty((batch, 2176), dtype=torch.int32, device=q.device), - "cu_q": torch.empty(batch + 1, dtype=torch.int32, device=q.device), - "cu_k": torch.empty(batch + 1, dtype=torch.int32, device=q.device), - "selected_kv": torch.empty((batch * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), + "compact_topk_mem_indices": torch.empty((q_token_num, 2176), dtype=torch.int64, device=q.device), + "cache_seqlens": torch.empty(q_token_num, dtype=torch.int32, device=q.device), + "page_table": torch.empty((q_token_num, 2176), dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.empty(q_token_num + 1, dtype=torch.int32, device=q.device), + "cu_seqlens_k_new": torch.empty(q_token_num + 1, dtype=torch.int32, device=q.device), + "selected_kv": torch.empty((q_token_num * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), } -def allocate_generic_prefill_workspace(q, topk_indices, query_capacity=None): - capacity = q.shape[0] if query_capacity is None else query_capacity +def allocate_generic_prefill_workspace(q, topk_indices, q_token_capacity=None): + capacity = q.shape[0] if q_token_capacity is None else q_token_capacity return { - "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), - "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), + "chunk_topk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), + "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), } @torch.no_grad() def _context_attention_fwd_fp8_selected_prefill( - q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, selected_workspace=None + q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, selected_workspace=None ): from sgl_kernel.flash_attn import flash_attn_with_kvcache - batch = q.shape[0] + q_token_num = q.shape[0] workspace = selected_workspace or allocate_selected_prefill_workspace(q) - batch_pad = triton.next_power_of_2(batch) - _selected_prefill_compact_kernel[(batch,)]( + block_q_tokens = triton.next_power_of_2(q_token_num) + _selected_prefill_compact_kernel[(q_token_num,)]( topk_mem_indices, - workspace["compact"], - workspace["lengths"], + workspace["compact_topk_mem_indices"], + workspace["cache_seqlens"], workspace["page_table"], - workspace["cu_q"], + workspace["cu_seqlens_q"], packed_kv.shape[0], topk_mem_indices.stride(0), topk_mem_indices.stride(1), - BATCH=batch, + q_token_num=q_token_num, BLOCK=4096, num_warps=8, ) - _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( - workspace["compact"], - workspace["lengths"], - prefill_mem_index, - prefill_kv, + _selected_prefill_gather_kernel[(q_token_num, triton.cdiv(2176, 4))]( + workspace["compact_topk_mem_indices"], + workspace["cache_seqlens"], + mem_index, + prefill_cache_kv, packed_kv, workspace["selected_kv"], - workspace["cu_k"], + workspace["cu_seqlens_k_new"], packed_kv.stride(0), - prefill_kv.stride(0), - prefill_kv.stride(2), - BATCH=batch, - BATCH_PAD=batch_pad, + prefill_cache_kv.stride(0), + prefill_cache_kv.stride(2), + q_token_num=q_token_num, + BLOCK_Q_TOKENS=block_q_tokens, num_warps=4, ) return flash_attn_with_kvcache( @@ -552,9 +562,9 @@ def _context_attention_fwd_fp8_selected_prefill( k_cache=None, v_cache=workspace["selected_kv"].view(-1, 1, 1, 512), page_table=workspace["page_table"], - cache_seqlens=workspace["lengths"], - cu_seqlens_q=workspace["cu_q"], - cu_seqlens_k_new=workspace["cu_k"], + cache_seqlens=workspace["cache_seqlens"], + cu_seqlens_q=workspace["cu_seqlens_q"], + cu_seqlens_k_new=workspace["cu_seqlens_k_new"], max_seqlen_q=1, softmax_scale=softmax_scale, causal=False, @@ -566,69 +576,68 @@ def _context_attention_fwd_fp8_selected_prefill( def context_attention_fwd_fp8( q, packed_kv, - prefill_mem_index, - prefill_kv, - kv_starts, + mem_index, + prefill_cache_kv, + q_kv_start_loc, topk_indices, topk_mem_indices, - req_to_token_index, + req_to_token_indexs, b_req_idx, - cu_kv_seq_len, - cu_q_seq_len, - total_kv_tokens, + b1_cu_kv_seq_len, + b1_cu_q_seq_len, + total_token_num, softmax_scale, selected_workspace=None, direct_workspace=None, generic_workspace=None, ): """Run NoPE sparse prefill from request mappings and physical top-k slots.""" - eligibility = (q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, total_kv_tokens) + eligibility = (q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num) if can_use_direct_smallq(*eligibility): return _context_attention_fwd_fp8_direct_smallq( - q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, direct_workspace + q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace ) if can_use_selected_prefill(*eligibility): return _context_attention_fwd_fp8_selected_prefill( - q, packed_kv, prefill_mem_index, prefill_kv, topk_mem_indices, softmax_scale, selected_workspace + q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, selected_workspace ) from sgl_kernel.flash_attn import flash_attn_with_kvcache - num_q_tokens, num_heads, _ = q.shape - num_kv_tokens = total_kv_tokens - gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, num_kv_tokens) + q_token_num, q_head_num, _ = q.shape + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) workspace = generic_workspace or allocate_generic_prefill_workspace(q, topk_indices) - chunk_indices = workspace["chunk_indices"] - chunk_lengths = workspace["chunk_lengths"] + chunk_topk_indices = workspace["chunk_topk_indices"] + chunk_cache_seqlens = workspace["chunk_cache_seqlens"] cu_seqlens_q = workspace["cu_seqlens_q"] # Keep this potentially 2 GiB buffer local to a DSA layer. gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) - output = lse = None - for chunk_start in range(0, num_kv_tokens, gather_chunk_size): - num_chunk_tokens = min(gather_chunk_size, num_kv_tokens - chunk_start) - _gather_prefill_kv_kernel[(num_chunk_tokens // 4 + b_req_idx.numel(),)]( + output = softmax_lse = None + for chunk_start in range(0, total_token_num, gather_chunk_size): + chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) + _gather_prefill_kv_kernel[(chunk_token_num // 4 + b_req_idx.numel(),)]( packed_kv, - prefill_kv, + prefill_cache_kv, gathered_kv, - req_to_token_index, + req_to_token_indexs, b_req_idx, - cu_kv_seq_len, - cu_q_seq_len, + b1_cu_kv_seq_len, + b1_cu_q_seq_len, chunk_start, - num_chunk_tokens, + chunk_token_num, packed_kv.stride(0), - prefill_kv.stride(0), - prefill_kv.stride(2), - REQ_STRIDE=req_to_token_index.stride(0), - TOKEN_STRIDE=req_to_token_index.stride(1), - BATCH=b_req_idx.numel(), + prefill_cache_kv.stride(0), + prefill_cache_kv.stride(2), + req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), + req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), + batch_size=b_req_idx.numel(), BLOCK_ROWS=4, ) - chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) + chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) if ( - num_q_tokens > 8 - and num_heads == 64 - and num_kv_tokens > num_q_tokens - and num_kv_tokens <= KV_GATHER_CHUNK_SIZE + q_token_num > 8 + and q_head_num == 64 + and total_token_num > q_token_num + and total_token_num <= KV_GATHER_CHUNK_SIZE and q.dtype == torch.bfloat16 and q.is_contiguous() and packed_kv.dtype == torch.uint8 @@ -639,42 +648,42 @@ def context_attention_fwd_fp8( ): from sgl_kernel.flash_mla import flash_mla_sparse_fwd - _compact_topk_indices_kernel[(num_q_tokens,)]( + _compact_topk_indices_kernel[(q_token_num,)]( topk_indices, - kv_starts, - chunk_indices, - chunk_lengths, + q_kv_start_loc, + chunk_topk_indices, + chunk_cache_seqlens, chunk_start, - chunk_start + num_chunk_tokens, + chunk_start + chunk_token_num, TOPK=topk_indices.shape[1], BLOCK=triton.next_power_of_2(topk_indices.shape[1]), COMPACT=False, ) return flash_mla_sparse_fwd( q=q, - kv=gathered_kv[:num_chunk_tokens], - indices=chunk_indices.unsqueeze(1), + kv=gathered_kv[:chunk_token_num], + indices=chunk_topk_indices.unsqueeze(1), sm_scale=softmax_scale, d_v=512, )[0] - _compact_topk_indices_kernel[(num_q_tokens,)]( + _compact_topk_indices_kernel[(q_token_num,)]( topk_indices, - kv_starts, - chunk_indices, - chunk_lengths, + q_kv_start_loc, + chunk_topk_indices, + chunk_cache_seqlens, chunk_start, - chunk_start + num_chunk_tokens, + chunk_start + chunk_token_num, TOPK=topk_indices.shape[1], BLOCK=triton.next_power_of_2(topk_indices.shape[1]), ) - chunk_output, chunk_lse, *_ = flash_attn_with_kvcache( + chunk_output, chunk_softmax_lse, *_ = flash_attn_with_kvcache( # only_qv ignores Q/K; these views avoid allocating placeholder tensors. q=q[..., :64], k_cache=chunk_kv[..., :64], qv=q, v_cache=chunk_kv, - page_table=chunk_indices, - cache_seqlens=chunk_lengths, + page_table=chunk_topk_indices, + cache_seqlens=chunk_cache_seqlens, cu_seqlens_q=cu_seqlens_q, max_seqlen_q=1, softmax_scale=softmax_scale, @@ -684,20 +693,20 @@ def context_attention_fwd_fp8( return_softmax_lse=True, ) if chunk_start == 0: - output, lse = chunk_output, chunk_lse + output, softmax_lse = chunk_output, chunk_softmax_lse else: # Keep the online reduction in FP32 across all remaining chunks. output = output.float() - _merge_att_output_kernel[(num_q_tokens, triton.cdiv(num_heads, 4))]( + _merge_att_output_kernel[(q_token_num, triton.cdiv(q_head_num, 4))]( output, - lse, + softmax_lse, chunk_output, - chunk_lse, - chunk_lse.stride(1), - chunk_lse.stride(0), - NUM_HEADS=num_heads, - NUM_Q_TOKENS=num_q_tokens, + chunk_softmax_lse, + chunk_softmax_lse.stride(1), + chunk_softmax_lse.stride(0), + q_head_num=q_head_num, + q_token_num=q_token_num, BLOCK_HEADS=4, ) - del chunk_output, chunk_lse + del chunk_output, chunk_softmax_lse return output.to(q.dtype) diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index dd2f4c0216..1bbf054354 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -7,63 +7,81 @@ @triton.jit -def _get_query_block(CuQLens, BATCH, BLOCK: tl.constexpr): +def _get_query_block(b1_cu_q_seq_len, batch_size, BLOCK: tl.constexpr): block = tl.program_id(0) # Request b owns [cu[b] // BLOCK + b, cu[b + 1] // BLOCK + b + 1). - # The extra block covers unaligned boundaries without padding to max_q_len. - if BATCH == 1: + # The extra block covers unaligned boundaries without padding to max_q_seq_len. + if batch_size == 1: batch = 0 else: batch = tl.full((), 0, tl.int32) # Keep the search tile independent of the scheduler's request count. lanes = tl.arange(0, 128) - for first in range(0, BATCH, 128): + for first in range(0, batch_size, 128): batches = first + lanes - q_ends = tl.load(CuQLens + batches + 1, batches < BATCH, 0) + q_ends = tl.load(b1_cu_q_seq_len + batches + 1, batches < batch_size, 0) block_ends = q_ends // BLOCK + batches + 1 - batch += tl.sum(((block >= block_ends) & (batches < BATCH)).to(tl.int32), 0) - q_start = tl.load(CuQLens + batch) - q_end = tl.load(CuQLens + batch + 1) + batch += tl.sum(((block >= block_ends) & (batches < batch_size)).to(tl.int32), 0) + q_start = tl.load(b1_cu_q_seq_len + batch) + q_end = tl.load(b1_cu_q_seq_len + batch + 1) local_block = block - (q_start // BLOCK + batch) return batch, q_start, q_end, local_block @triton.jit -def _prefill_nsa_params(PositionIds, ReqIdx, CuQ, CuKV, Lengths, QueryReq, KVStarts, BATCH, BLOCK: tl.constexpr): - batch, q_start, q_end, block = _get_query_block(CuQ, BATCH, BLOCK) +def _prefill_nsa_params( + position_ids, + b_req_idx, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, + q_kv_seq_len, + q_req_idx, + q_kv_start_loc, + batch_size, + BLOCK: tl.constexpr, +): + batch, q_start, q_end, block = _get_query_block(b1_cu_q_seq_len, batch_size, BLOCK) rows = q_start + block * BLOCK + tl.arange(0, BLOCK) valid = rows < q_end - tl.store(Lengths + rows, tl.load(PositionIds + rows, valid, 0) + 1, valid) - tl.store(QueryReq + rows, tl.load(ReqIdx + batch), valid) - tl.store(KVStarts + rows, tl.load(CuKV + batch), valid) + tl.store(q_kv_seq_len + rows, tl.load(position_ids + rows, valid, 0) + 1, valid) + tl.store(q_req_idx + rows, tl.load(b_req_idx + batch), valid) + tl.store(q_kv_start_loc + rows, tl.load(b1_cu_kv_seq_len + batch), valid) -def prefill_nsa_params(position_ids, req_idx, cu_q_lens, cu_kv_lens): +def prefill_nsa_params(position_ids, b_req_idx, b1_cu_q_seq_len, b1_cu_kv_seq_len): """Expand existing request metadata to query rows without visiting historical KV.""" - lengths, query_req, kv_starts = [torch.empty_like(position_ids) for _ in range(3)] - _prefill_nsa_params[(position_ids.numel() // 256 + req_idx.numel(),)]( - position_ids, req_idx, cu_q_lens, cu_kv_lens, lengths, query_req, kv_starts, req_idx.numel(), BLOCK=256 + q_kv_seq_len, q_req_idx, q_kv_start_loc = [torch.empty_like(position_ids) for _ in range(3)] + _prefill_nsa_params[(position_ids.numel() // 256 + b_req_idx.numel(),)]( + position_ids, + b_req_idx, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, + q_kv_seq_len, + q_req_idx, + q_kv_start_loc, + b_req_idx.numel(), + BLOCK=256, ) - return lengths, query_req, kv_starts + return q_kv_seq_len, q_req_idx, q_kv_start_loc @triton.jit def _compress_pools( - Raw, - Tail, - Packed, - Ape, - Lengths, - MemIndex, - ReqIdx, - CuQLens, - MtpIndex, - RAW_STRIDE: tl.constexpr, - TAIL_REQ_STRIDE: tl.constexpr, - TAIL_SLOT_STRIDE: tl.constexpr, - PACKED_STRIDE: tl.constexpr, - HOLD_REQ: tl.constexpr, - BATCH, + k_and_gate, + indexer_tail, + indexer_k, + index_kpool_compress_ape, + q_kv_seq_len, + mem_index, + b_req_idx, + b1_cu_q_seq_len, + b_mtp_index, + k_and_gate_stride: tl.constexpr, + indexer_tail_stride_b: tl.constexpr, + indexer_tail_stride_s: tl.constexpr, + indexer_k_stride: tl.constexpr, + HOLD_REQUEST_ID: tl.constexpr, + batch_size, SINGLE_QUERY: tl.constexpr, TAIL_SIZE: tl.constexpr, HAS_MTP_INDEX: tl.constexpr, @@ -71,72 +89,74 @@ def _compress_pools( if SINGLE_QUERY: batch = tl.program_id(0) row = batch - length = tl.load(Lengths + row) + length = tl.load(q_kv_seq_len + row) valid = length > 0 and length % 4 == 0 q_start = row if HAS_MTP_INDEX: - q_start -= tl.load(MtpIndex + row) + q_start -= tl.load(b_mtp_index + row) else: - batch, q_start, q_end, pool_index = _get_query_block(CuQLens, BATCH, 4) - prefix = tl.load(Lengths + q_start, q_start < q_end, 1) - 1 + batch, q_start, q_end, pool_index = _get_query_block(b1_cu_q_seq_len, batch_size, 4) + prefix = tl.load(q_kv_seq_len + q_start, q_start < q_end, 1) - 1 # Only process pool-closing rows, including a pool spanning the old tail. length = (prefix // 4 + pool_index + 1) * 4 row = q_start + length - prefix - 1 valid = row < q_end - req = tl.load(ReqIdx + batch) - if valid and req != HOLD_REQ: + req = tl.load(b_req_idx + batch) + if valid and req != HOLD_REQUEST_ID: pool = tl.arange(0, 4) cols = tl.arange(0, 128) raw_rows = row - 3 + pool from_chunk = raw_rows >= q_start - chunk_ptr = Raw + raw_rows[:, None] * RAW_STRIDE + cols[None, :] + chunk_ptr = k_and_gate + raw_rows[:, None] * k_and_gate_stride + cols[None, :] tail_slots = (length - 4 + pool) % TAIL_SIZE - tail_ptr = Tail + req * TAIL_REQ_STRIDE + tail_slots[:, None] * TAIL_SLOT_STRIDE + cols[None, :] - raw = tl.where( + tail_ptr = ( + indexer_tail + req * indexer_tail_stride_b + tail_slots[:, None] * indexer_tail_stride_s + cols[None, :] + ) + pool_k = tl.where( from_chunk[:, None], tl.load(chunk_ptr, from_chunk[:, None], 0), tl.load(tail_ptr, ~from_chunk[:, None], 0), ).to(tl.float32) - score = tl.where( + pool_scores = tl.where( from_chunk[:, None], tl.load(chunk_ptr + 128, from_chunk[:, None], 0), tl.load(tail_ptr + 128, ~from_chunk[:, None], 0), ).to(tl.float32) - score += tl.load(Ape + pool[:, None] * 128 + cols[None, :]) - score = tl.exp(score - tl.max(score, 0)[None, :]) - weights = score / tl.sum(score, 0)[None, :] - key = tl.sum(raw * weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) + pool_scores += tl.load(index_kpool_compress_ape + pool[:, None] * 128 + cols[None, :]) + pool_scores = tl.exp(pool_scores - tl.max(pool_scores, 0)[None, :]) + pool_weights = pool_scores / tl.sum(pool_scores, 0)[None, :] + pooled_k = tl.sum(pool_k * pool_weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) for step in tl.static_range(7): - key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) - key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) - scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) - key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) + pooled_k = _butterfly_stage(pooled_k, 64 >> step, 1 << step, 1, 128) + pooled_k = (pooled_k * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) + k_scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(pooled_k), 1), 1e-4) / 448.0))) + pooled_k = tl.minimum(tl.maximum(pooled_k / k_scale[:, None], -448.0), 448.0).to(tl.float8e4nv) # A completed pool lives at its closing token, which is a current query row. - loc = tl.load(MemIndex + row).to(tl.int64) - dest = Packed + loc * PACKED_STRIDE - tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) - tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + loc = tl.load(mem_index + row).to(tl.int64) + dest = indexer_k + loc * indexer_k_stride + tl.store(dest + cols, pooled_k.reshape(128).to(tl.uint8, bitcast=True)) + tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(k_scale, 0)) @triton.jit def _save_pool_tails( - Raw, - Tail, - ReqIdx, - CuQLens, - SeqLens, - RAW_STRIDE: tl.constexpr, - TAIL_REQ_STRIDE: tl.constexpr, - TAIL_SLOT_STRIDE: tl.constexpr, - HOLD_REQ: tl.constexpr, + k_and_gate, + indexer_tail, + b_req_idx, + b1_cu_q_seq_len, + b_seq_len, + k_and_gate_stride: tl.constexpr, + indexer_tail_stride_b: tl.constexpr, + indexer_tail_stride_s: tl.constexpr, + HOLD_REQUEST_ID: tl.constexpr, TAIL_SIZE: tl.constexpr, ): batch = tl.program_id(0) - req = tl.load(ReqIdx + batch) - if req != HOLD_REQ: - length = tl.load(SeqLens + batch) - q_start = tl.load(CuQLens + batch) - q_end = tl.load(CuQLens + batch + 1) + req = tl.load(b_req_idx + batch) + if req != HOLD_REQUEST_ID: + length = tl.load(b_seq_len + batch) + q_start = tl.load(b1_cu_q_seq_len + batch) + q_end = tl.load(b1_cu_q_seq_len + batch + 1) slots = tl.arange(0, triton.next_power_of_2(TAIL_SIZE)) cols = tl.arange(0, 256) positions = length - TAIL_SIZE + slots @@ -144,189 +164,212 @@ def _save_pool_tails( # Absolute positions address a short ring, preserving the raw history # needed after rejecting candidates or revisiting draft positions. mask = ((slots < TAIL_SIZE) & (positions >= 0) & (raw_rows >= q_start))[:, None] - raw = tl.load(Raw + raw_rows[:, None] * RAW_STRIDE + cols[None, :], mask, 0) + k_and_gate_block = tl.load(k_and_gate + raw_rows[:, None] * k_and_gate_stride + cols[None, :], mask, 0) tl.store( - Tail + req * TAIL_REQ_STRIDE + (positions % TAIL_SIZE)[:, None] * TAIL_SLOT_STRIDE + cols[None, :], - raw, + indexer_tail + + req * indexer_tail_stride_b + + (positions % TAIL_SIZE)[:, None] * indexer_tail_stride_s + + cols[None, :], + k_and_gate_block, mask, ) def compress_pools( - raw, tail, packed_buffer, ape, lengths, mem_index, req_idx, cu_q_lens, seq_lens, max_q_len, mtp_index=None + k_and_gate, + indexer_tail, + indexer_k, + index_kpool_compress_ape, + q_kv_seq_len, + mem_index, + b_req_idx, + b1_cu_q_seq_len, + b_seq_len, + max_q_seq_len, + b_mtp_index=None, ): - batch = req_idx.numel() - single_query = max_q_len == 1 and lengths.numel() == batch - blocks = batch if single_query else lengths.numel() // 4 + batch + batch_size = b_req_idx.numel() + single_query = max_q_seq_len == 1 and q_kv_seq_len.numel() == batch_size + blocks = batch_size if single_query else q_kv_seq_len.numel() // 4 + batch_size _compress_pools[(blocks,)]( - raw, - tail, - packed_buffer, - ape, - lengths, + k_and_gate, + indexer_tail, + indexer_k, + index_kpool_compress_ape, + q_kv_seq_len, mem_index, - req_idx, - cu_q_lens, - mtp_index, - raw.stride(0), - tail.stride(0), - tail.stride(1), - packed_buffer.stride(0), - tail.shape[0] - 1, - batch, + b_req_idx, + b1_cu_q_seq_len, + b_mtp_index, + k_and_gate.stride(0), + indexer_tail.stride(0), + indexer_tail.stride(1), + indexer_k.stride(0), + indexer_tail.shape[0] - 1, + batch_size, single_query, - tail.shape[1], - mtp_index is not None, + indexer_tail.shape[1], + b_mtp_index is not None, num_warps=4, ) # Complete all boundary pools before replacing the previous chunk's tail. - _save_pool_tails[(req_idx.numel(),)]( - raw, - tail, - req_idx, - cu_q_lens, - seq_lens, - raw.stride(0), - tail.stride(0), - tail.stride(1), - tail.shape[0] - 1, - tail.shape[1], + _save_pool_tails[(b_req_idx.numel(),)]( + k_and_gate, + indexer_tail, + b_req_idx, + b1_cu_q_seq_len, + b_seq_len, + k_and_gate.stride(0), + indexer_tail.stride(0), + indexer_tail.stride(1), + indexer_tail.shape[0] - 1, + indexer_tail.shape[1], num_warps=4, ) @triton.jit def _get_pool_ranges( - Lengths, - CuQLens, - Starts, - Ends, - PoolLengths, - BATCH, - TOKENS, - MAX_POOLS, + q_kv_seq_len, + b1_cu_q_seq_len, + ks, + ke, + q_pool_seq_len, + batch_size, + q_token_num, + max_pools, SINGLE_QUERY: tl.constexpr, BLOCK: tl.constexpr, ): if SINGLE_QUERY: row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) batch = row - valid = row < TOKENS + valid = row < q_token_num else: - batch, q_start, q_end, local_block = _get_query_block(CuQLens, BATCH, BLOCK) + batch, q_start, q_end, local_block = _get_query_block(b1_cu_q_seq_len, batch_size, BLOCK) row = q_start + local_block * BLOCK + tl.arange(0, BLOCK) valid = row < q_end - length = tl.load(Lengths + row, valid, 0) // 4 - start = batch * MAX_POOLS - tl.store(Starts + row, start, valid) - tl.store(Ends + row, start + length, valid) - tl.store(PoolLengths + row, length, valid) + length = tl.load(q_kv_seq_len + row, valid, 0) // 4 + start = batch * max_pools + tl.store(ks + row, start, valid) + tl.store(ke + row, start + length, valid) + tl.store(q_pool_seq_len + row, length, valid) -def get_pool_ranges(lengths, cu_q_lens, max_q_len, max_pools): +def get_pool_ranges(q_kv_seq_len, b1_cu_q_seq_len, max_q_seq_len, max_pools): """Build per-query pool bounds directly from the packed request boundaries.""" - starts, ends, pool_lengths = [torch.empty_like(lengths) for _ in range(3)] - batch, tokens = cu_q_lens.numel() - 1, lengths.numel() - single_query = max_q_len == 1 and tokens == batch + ks, ke, q_pool_seq_len = [torch.empty_like(q_kv_seq_len) for _ in range(3)] + batch_size, q_token_num = b1_cu_q_seq_len.numel() - 1, q_kv_seq_len.numel() + single_query = max_q_seq_len == 1 and q_token_num == batch_size block = 256 - blocks = triton.cdiv(tokens, block) if single_query else tokens // block + batch + blocks = triton.cdiv(q_token_num, block) if single_query else q_token_num // block + batch_size _get_pool_ranges[(blocks,)]( - lengths, cu_q_lens, starts, ends, pool_lengths, batch, tokens, max_pools, single_query, block, num_warps=4 + q_kv_seq_len, + b1_cu_q_seq_len, + ks, + ke, + q_pool_seq_len, + batch_size, + q_token_num, + max_pools, + single_query, + block, + num_warps=4, ) - return starts, ends, pool_lengths + return ks, ke, q_pool_seq_len -@triton.jit(do_not_specialize=["MAX_POOLS_PER_ROW"]) +@triton.jit(do_not_specialize=["max_pools"]) def _gather_pools( - packed_buffer, + indexer_k, req_to_token_indexs, b_req_idx, b_seq_len, - keys, - scales, - packed_stride: tl.constexpr, - stride_req_to_token_m: tl.constexpr, - MAX_POOLS_PER_ROW, + k_fp8, + k_scale, + indexer_k_stride: tl.constexpr, + req_to_token_indexs_stride_b: tl.constexpr, + max_pools, POOLS_PER_BLOCK: tl.constexpr, ): batch_row = tl.program_id(1) pool_indices = tl.program_id(0) * POOLS_PER_BLOCK + tl.arange(0, POOLS_PER_BLOCK) sequence_length = tl.load(b_seq_len + batch_row) cur_req_idx = tl.load(b_req_idx + batch_row) - valid_pools = (pool_indices < MAX_POOLS_PER_ROW) & (pool_indices < sequence_length // 4) + valid_pools = (pool_indices < max_pools) & (pool_indices < sequence_length // 4) # Four tokens form one pool; its key is stored at the last token. token_slots = tl.load( - req_to_token_indexs + cur_req_idx * stride_req_to_token_m + pool_indices * 4 + 3, valid_pools, 0 + req_to_token_indexs + cur_req_idx * req_to_token_indexs_stride_b + pool_indices * 4 + 3, valid_pools, 0 ).to(tl.int64) key_offsets = tl.arange(0, 128) key_bytes = tl.load( - packed_buffer + token_slots[:, None] * packed_stride + key_offsets[None, :], valid_pools[:, None], 0 + indexer_k + token_slots[:, None] * indexer_k_stride + key_offsets[None, :], valid_pools[:, None], 0 ) pool_scales = tl.load( - (packed_buffer + token_slots * packed_stride + 128).to(tl.pointer_type(tl.float32)), valid_pools, 1.0 + (indexer_k + token_slots * indexer_k_stride + 128).to(tl.pointer_type(tl.float32)), valid_pools, 1.0 ) - output_rows = batch_row.to(tl.int64) * MAX_POOLS_PER_ROW + pool_indices + output_rows = batch_row.to(tl.int64) * max_pools + pool_indices tl.store( - keys + output_rows[:, None] * 128 + key_offsets[None, :], + k_fp8 + output_rows[:, None] * 128 + key_offsets[None, :], key_bytes.to(tl.float8e4nv, bitcast=True), - (pool_indices < MAX_POOLS_PER_ROW)[:, None], + (pool_indices < max_pools)[:, None], ) - tl.store(scales + output_rows, pool_scales, pool_indices < MAX_POOLS_PER_ROW) + tl.store(k_scale + output_rows, pool_scales, pool_indices < max_pools) -def gather_pools(packed_buffer, req_to_token_indexs, b_req_idx, b_seq_len, max_pools): +def gather_pools(indexer_k, req_to_token_indexs, b_req_idx, b_seq_len, max_pools): """Gather FP8 keys and FP32 scales for prefill. b_seq_len counts tokens; max_pools is the pool capacity per row. """ - keys = torch.empty( + k_fp8 = torch.empty( (b_req_idx.numel() * max_pools, 128), - device=packed_buffer.device, + device=indexer_k.device, dtype=torch.float8_e4m3fn, ) - scales = torch.empty((keys.shape[0],), device=keys.device, dtype=torch.float32) + k_scale = torch.empty((k_fp8.shape[0],), device=k_fp8.device, dtype=torch.float32) # Use grid X for pools because grid Y is limited to 65535 blocks. _gather_pools[(triton.cdiv(max_pools, 32), b_req_idx.numel())]( - packed_buffer, + indexer_k, req_to_token_indexs, b_req_idx, b_seq_len, - keys, - scales, - packed_buffer.stride(0), + k_fp8, + k_scale, + indexer_k.stride(0), req_to_token_indexs.stride(0), max_pools, POOLS_PER_BLOCK=32, num_warps=4, ) - return keys, scales + return k_fp8, k_scale -@triton.jit(do_not_specialize=["MAX_PAGES_PER_ROW"]) +@triton.jit(do_not_specialize=["max_pages_per_row"]) def _gather_paged_pools( - packed_buffer, + indexer_k, req_to_token_indexs, b_req_idx, - pool_lengths, - pages, + q_pool_seq_len, + indexer_k_pages, block_table, - packed_stride: tl.constexpr, - stride_req_to_token_m: tl.constexpr, - MAX_PAGES_PER_ROW, + indexer_k_stride: tl.constexpr, + req_to_token_indexs_stride_b: tl.constexpr, + max_pages_per_row, REUSE_GROUP_SIZE: tl.constexpr, ): batch_row = tl.program_id(1) # Avoid compiling a separate kernel for each batch size. num_batch_rows = tl.num_programs(1) cur_req_idx = tl.load(b_req_idx + batch_row) - gather_pool_count = tl.load(pool_lengths + batch_row) + gather_pool_count = tl.load(q_pool_seq_len + batch_row) # The first row copies shared pages; other rows only update the page table. owner_row = batch_row if REUSE_GROUP_SIZE > 1: group_start_row = batch_row // REUSE_GROUP_SIZE * REUSE_GROUP_SIZE group_rows = group_start_row + tl.arange(0, REUSE_GROUP_SIZE) group_req_idx = tl.load(b_req_idx + group_rows, group_rows < num_batch_rows, -1) - group_pool_counts = tl.load(pool_lengths + group_rows, group_rows < num_batch_rows, 0) + group_pool_counts = tl.load(q_pool_seq_len + group_rows, group_rows < num_batch_rows, 0) can_share_pages = ( (group_start_row + REUSE_GROUP_SIZE <= num_batch_rows) & (tl.min(group_req_idx, 0) == tl.max(group_req_idx, 0)) @@ -340,112 +383,118 @@ def _gather_paged_pools( for logical_page in range(tl.program_id(0), tl.cdiv(tl.maximum(gather_pool_count, 1), 64), tl.num_programs(0)): if batch_row != owner_row: tl.store( - block_table + batch_row * MAX_PAGES_PER_ROW + logical_page, - owner_row * MAX_PAGES_PER_ROW + logical_page, + block_table + batch_row * max_pages_per_row + logical_page, + owner_row * max_pages_per_row + logical_page, ) else: pool_indices = logical_page * 64 + pool_offsets valid_pools = pool_indices < gather_pool_count token_slots = tl.load( - req_to_token_indexs + cur_req_idx * stride_req_to_token_m + pool_indices * 4 + 3, valid_pools, 0 + req_to_token_indexs + cur_req_idx * req_to_token_indexs_stride_b + pool_indices * 4 + 3, valid_pools, 0 ).to(tl.int64) key_bytes = tl.load( - packed_buffer + token_slots[:, None] * packed_stride + key_offsets[None, :], + indexer_k + token_slots[:, None] * indexer_k_stride + key_offsets[None, :], valid_pools[:, None], 0, ) pool_scales = tl.load( - (packed_buffer + token_slots * packed_stride + 128).to(tl.pointer_type(tl.float32)), + (indexer_k + token_slots * indexer_k_stride + 128).to(tl.pointer_type(tl.float32)), valid_pools, 1.0, ) - output_page_ptr = pages + (batch_row.to(tl.int64) * MAX_PAGES_PER_ROW + logical_page) * (64 * 132) + output_page_ptr = indexer_k_pages + (batch_row.to(tl.int64) * max_pages_per_row + logical_page) * (64 * 132) # Page layout: [64 x 128 key bytes][64 FP32 scales]. tl.store(output_page_ptr + pool_offsets[:, None] * 128 + key_offsets[None, :], key_bytes) tl.store((output_page_ptr + 64 * 128).to(tl.pointer_type(tl.float32)) + pool_offsets, pool_scales) -def gather_paged_pools(packed_buffer, req_to_token_indexs, b_req_idx, pool_lengths, max_pools, reuse_group=4): +def gather_paged_pools(indexer_k, req_to_token_indexs, b_req_idx, q_pool_seq_len, max_pools, reuse_group=4): """Gather decode pages and the page table. - pool_lengths counts valid pools; max_pools is the pool capacity per row. + q_pool_seq_len counts valid pools; max_pools is the pool capacity per row. """ if reuse_group not in (1, 2, 4, 8): raise ValueError("reuse_group must be 1, 2, 4 or 8") batch_size = b_req_idx.numel() max_pages_per_row = triton.cdiv(max_pools, 64) # Reserve each row's capacity so sharing can change on graph replay. - pages = torch.empty((batch_size * max_pages_per_row, 64, 1, 132), device=packed_buffer.device, dtype=torch.uint8) - block_table = torch.arange(batch_size * max_pages_per_row, device=pages.device, dtype=torch.int32).view( + indexer_k_pages = torch.empty( + (batch_size * max_pages_per_row, 64, 1, 132), device=indexer_k.device, dtype=torch.uint8 + ) + block_table = torch.arange(batch_size * max_pages_per_row, device=indexer_k_pages.device, dtype=torch.int32).view( batch_size, max_pages_per_row ) estimated_copy_groups = max(1, triton.cdiv(batch_size, reuse_group)) if max_pools >= 4096 else batch_size page_blocks_per_row = min(max_pages_per_row, triton.cdiv(get_device_sm_count() * 4, estimated_copy_groups)) _gather_paged_pools[(page_blocks_per_row, batch_size)]( - packed_buffer, + indexer_k, req_to_token_indexs, b_req_idx, - pool_lengths, - pages, + q_pool_seq_len, + indexer_k_pages, block_table, - packed_buffer.stride(0), + indexer_k.stride(0), req_to_token_indexs.stride(0), max_pages_per_row, REUSE_GROUP_SIZE=reuse_group if batch_size >= reuse_group else 1, num_warps=4, ) - return pages, block_table + return indexer_k_pages, block_table @triton.jit def _expand_topk( - Groups, - Lengths, - QueryReq, - ReqToToken, - Out, - Relative, - REQ_STRIDE: tl.constexpr, - TOKEN_STRIDE: tl.constexpr, - TOPK: tl.constexpr, + topk_pool_indices, + q_kv_seq_len, + q_req_idx, + req_to_token_indexs, + topk_mem_indices, + topk_indices, + req_to_token_indexs_stride_b: tl.constexpr, + req_to_token_indexs_stride_s: tl.constexpr, + index_topk: tl.constexpr, WIDTH: tl.constexpr, DENSE: tl.constexpr, BLOCK: tl.constexpr, ): row = tl.program_id(0) lane = tl.arange(0, BLOCK) - length = tl.load(Lengths + row) - req = tl.load(QueryReq + row).to(tl.int64) + length = tl.load(q_kv_seq_len + row) + req = tl.load(q_req_idx + row).to(tl.int64) if DENSE: token = lane valid = lane < length else: - closed_tokens = tl.minimum(length // 4 * 4, TOPK) - group = tl.load(Groups + row * (TOPK // 4) + lane // 4, lane < closed_tokens, -1) + closed_tokens = tl.minimum(length // 4 * 4, index_topk) + group = tl.load(topk_pool_indices + row * (index_topk // 4) + lane // 4, lane < closed_tokens, -1) token = tl.where(lane < closed_tokens, group * 4 + lane % 4, length // 4 * 4 + lane - closed_tokens) valid = (lane < closed_tokens + length % 4) & (token >= 0) - mem = tl.load(ReqToToken + req * REQ_STRIDE + token * TOKEN_STRIDE, valid & (lane < WIDTH), -1) - tl.store(Out + row * WIDTH + lane, mem, lane < WIDTH) - tl.store(Relative + row * WIDTH + lane, tl.where(valid, token, -1), lane < WIDTH) + mem = tl.load( + req_to_token_indexs + req * req_to_token_indexs_stride_b + token * req_to_token_indexs_stride_s, + valid & (lane < WIDTH), + -1, + ) + tl.store(topk_mem_indices + row * WIDTH + lane, mem, lane < WIDTH) + tl.store(topk_indices + row * WIDTH + lane, tl.where(valid, token, -1), lane < WIDTH) -def expand_topk(groups, lengths, query_req, req_to_token, topk, dense=False): - width = triton.cdiv(topk + 3, 128) * 128 - out = torch.empty((lengths.numel(), width), dtype=torch.int32, device=lengths.device) - relative = torch.empty_like(out) - _expand_topk[(lengths.numel(),)]( - groups, - lengths, - query_req, - req_to_token, - out, - relative, - req_to_token.stride(0), - req_to_token.stride(1), - topk, +def expand_topk(topk_pool_indices, q_kv_seq_len, q_req_idx, req_to_token_indexs, index_topk, dense=False): + width = triton.cdiv(index_topk + 3, 128) * 128 + topk_mem_indices = torch.empty((q_kv_seq_len.numel(), width), dtype=torch.int32, device=q_kv_seq_len.device) + topk_indices = torch.empty_like(topk_mem_indices) + _expand_topk[(q_kv_seq_len.numel(),)]( + topk_pool_indices, + q_kv_seq_len, + q_req_idx, + req_to_token_indexs, + topk_mem_indices, + topk_indices, + req_to_token_indexs.stride(0), + req_to_token_indexs.stride(1), + index_topk, width, dense, triton.next_power_of_2(width), num_warps=4, ) - return out, relative + return topk_mem_indices, topk_indices diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index 6b791b8e04..624793c8e4 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -114,8 +114,8 @@ def forbidden(*args, **kwargs): cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() - assert state.lengths is infer.b_seq_len - assert state.query_req is infer.b_req_idx + assert state.q_kv_seq_len is infer.b_seq_len + assert state.q_req_idx is infer.b_req_idx assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): @@ -123,7 +123,7 @@ def forbidden(*args, **kwargs): infer.b_seq_len[:3].sub_(100) infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) graph.replay() - assert state.lengths is infer.b_seq_len and state.query_req is infer.b_req_idx + assert state.q_kv_seq_len is infer.b_seq_len and state.q_req_idx is infer.b_req_idx @pytest.mark.parametrize( @@ -216,7 +216,7 @@ def check(output): check(run()) generic = backend._prefill_workspaces[0]["generic"] if generic is not None: - assert set(generic) == {"chunk_indices", "chunk_lengths", "cu_seqlens_q"} + assert set(generic) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = run() diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 9993b2e162..33352b0b03 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -52,32 +52,36 @@ def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=N self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) total, queries = sum(seq_lens), sum(q_lens) capacity = capacity or total + 17 - self.table = torch.full((2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32) + self.req_to_token_indexs = torch.full( + (2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32 + ) slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] - cu_q, cu_kv, query_req, positions = [0], [0], [], [] + cu_q, cu_kv, q_req_idx, positions = [0], [0], [], [] for req, length, count in zip(self.requests, seq_lens, q_lens): - self.table[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] - query_req.extend([req] * count) + self.req_to_token_indexs[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] + q_req_idx.extend([req] * count) positions.extend(range(length - count, length)) cu_q.append(cu_q[-1] + count) cu_kv.append(cu_kv[-1] + length) - self.query_req, self.positions = tensor(query_req).long(), tensor(positions).long() + self.q_req_idx, self.position_ids = tensor(q_req_idx).long(), tensor(positions).long() self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) - self.source = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) - self.packed = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) + self.source_kv = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) + self.packed_kv = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) - write_fp8_e8m0_packed_kv(self.source, self.all_slots, self.packed) - self.prefill = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] + write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) + self.prefill_cache_kv = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] visible = tensor(positions) + 1 - self.topk = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) - self.topk.remainder_(visible[:, None]) - self.topk.masked_fill_(torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1) - self.topk[:, ::19] = -1 - self.topk[:, 0] = visible - 1 - self.topk[:, 1] = visible - 1 - self.physical = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] - self.infer = _make_infer_state( - req_manager=SimpleNamespace(req_to_token_indexs=self.table), + self.topk_indices = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) + self.topk_indices.remainder_(visible[:, None]) + self.topk_indices.masked_fill_( + torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1 + ) + self.topk_indices[:, ::19] = -1 + self.topk_indices[:, 0] = visible - 1 + self.topk_indices[:, 1] = visible - 1 + self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] + self.infer_state = _make_infer_state( + req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), b_req_idx=tensor(self.requests), b_seq_len=tensor(seq_lens), b1_cu_q_seq_len=tensor(cu_q), @@ -87,13 +91,13 @@ def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=N mem_index=torch.empty(queries, device=device, dtype=torch.int32), ) self.backend = backend or _make_workspace_backend(queries) - _, self.state = self.new_state() + _, self.att_state = self.new_state() self.control = AttControl( nsa_prefill=True, nsa_prefill_dict={ - "topk_indices": self.topk, - "topk_mem_indices": self.physical, - "prefill_cache_kv": self.prefill, + "topk_indices": self.topk_indices, + "topk_mem_indices": self.topk_mem_indices, + "prefill_cache_kv": self.prefill_cache_kv, "softmax_scale": softmax_scale, "kv_lora_rank": 512, }, @@ -101,40 +105,44 @@ def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=N def new_state(self, **overrides): if overrides: - infer = _make_infer_state(**vars(self.infer)) + infer = _make_infer_state(**vars(self.infer_state)) vars(infer).update(overrides) else: - infer = self.infer + infer = self.infer_state state = self.backend.create_att_prefill_state(infer) state.init_state() return infer, state def refresh(self): - self.infer.mem_index.copy_(self.table[self.query_req, self.positions]) - self.physical.copy_(self.table[self.query_req[:, None], self.topk.clamp_min(0).long()]) - self.physical.masked_fill_(self.topk < 0, -1) - self.control.nsa_prefill_dict["topk_indices"] = self.topk + self.infer_state.mem_index.copy_(self.req_to_token_indexs[self.q_req_idx, self.position_ids]) + self.topk_mem_indices.copy_( + self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()] + ) + self.topk_mem_indices.masked_fill_(self.topk_indices < 0, -1) + self.control.nsa_prefill_dict["topk_indices"] = self.topk_indices def run(self, state=None): self.refresh() - active = self.state if state is None else state - return active.prefill_att(self.q, self.packed[..., :584], None, self.control) + active = self.att_state if state is None else state + return active.prefill_att(self.q, self.packed_kv[..., :584], None, self.control) def check(self, actual): - values = _dequantize_kv_cache(self.packed) - values[self.infer.mem_index.long()] = self.prefill - physical = self.table[self.query_req[:, None], self.topk.clamp_min(0).long()].masked_fill(self.topk < 0, -1) + values = _dequantize_kv_cache(self.packed_kv) + values[self.infer_state.mem_index.long()] = self.prefill_cache_kv + physical = self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()].masked_fill( + self.topk_indices < 0, -1 + ) expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) def mutate(self): self.q.normal_() - self.prefill.normal_() - self.source.normal_() - write_fp8_e8m0_packed_kv(self.source, self.all_slots, self.packed) + self.prefill_cache_kv.normal_() + self.source_kv.normal_() + write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) for req, length in zip(self.requests, self.seq_lens): - self.table[req, :length].copy_(self.table[req, :length].roll(1)) - self.topk.copy_(self.topk.roll(1, dims=1)) + self.req_to_token_indexs[req, :length].copy_(self.req_to_token_indexs[req, :length].roll(1)) + self.topk_indices.copy_(self.topk_indices.roll(1, dims=1)) @pytest.mark.parametrize("draft", [False, True]) @@ -248,9 +256,9 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - REQ_STRIDE=11, - TOKEN_STRIDE=1, - BATCH=1, + req_to_token_indexs_stride_b=11, + req_to_token_indexs_stride_s=1, + batch_size=1, BLOCK_ROWS=4, ) expected = _dequantize_kv_cache(packed_kv) @@ -313,10 +321,10 @@ def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chun prefix = 4 if cached else 0 case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - case.topk[0].fill_(-1) - case.topk[1, :2] = prefix - case.topk[3].fill_(-1) - case.topk[3, 0] = prefix + case.topk_indices[0].fill_(-1) + case.topk_indices[1, :2] = prefix + case.topk_indices[3].fill_(-1) + case.topk_indices[3, 0] = prefix case.check(case.run()) graph, output = _capture_cuda_graph(case.run) graph.replay() @@ -333,10 +341,10 @@ def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softm torch.manual_seed(53) case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") - case.topk.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) - case.topk[1, 0] = 1016 - case.topk[2, :3] = torch.tensor([2, 3, 4], device="cuda") - case.topk[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) + case.topk_indices.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) + case.topk_indices[1, 0] = 1016 + case.topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") + case.topk_indices[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) case.check(case.run()) graph, output = _capture_cuda_graph(case.run) graph.replay() @@ -367,7 +375,7 @@ def observed_sparse(*args, **kwargs): case.mutate() graph.replay() case.check(output) - case.topk.fill_(-1) + case.topk_indices.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 assert len(observed) == 3 @@ -404,8 +412,8 @@ def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num case = _PrefillCase(seqs, [1] * num_q, num_heads) case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) - backing[:, ::2] = case.topk - case.topk = backing[:, ::2] + backing[:, ::2] = case.topk_indices + case.topk_indices = backing[:, ::2] output = case.run() case.check(output) assert output.is_contiguous() and torch.count_nonzero(output) > 0 @@ -427,7 +435,7 @@ def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): infer, state = case.new_state() case.check(case.run(state)) assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers - fresh_infer, fresh = case.new_state(mem_index=case.infer.mem_index.clone()) + fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) case.check(case.run(fresh)) assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers case.mutate() @@ -444,14 +452,14 @@ def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): for _ in range(10): infer, state = case.new_state() case.check(case.run(state)) - fresh_infer, fresh = case.new_state(mem_index=case.infer.mem_index.clone()) + fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) case.check(case.run(fresh)) assert {name: value.data_ptr() for name, value in workspace.items()} == pointers graph, output = _capture_cuda_graph(case.run) case.mutate() graph.replay() case.check(output) - case.topk.fill_(-1) + case.topk_indices.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 @@ -937,7 +945,7 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba case = _PrefillCase([32768, 16384, 16384], [5, 4, 4], 64, backend=backend) case.check(case.run()) workspace = backend._prefill_workspaces[0]["generic"] - assert set(workspace) == {"chunk_indices", "chunk_lengths", "cu_seqlens_q"} + assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): infer, state = case.new_state() diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py index 76646c5224..1a4ca0f184 100644 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -122,7 +122,7 @@ def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, max_q_seq_len=1, req_manager=SimpleNamespace(req_to_token_indexs=table), ) - att_state = SimpleNamespace(lengths=lengths) + att_state = SimpleNamespace(q_kv_seq_len=lengths) logits_outputs = [] select_topk = indexer.select_topk_indices diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index c382930a82..49b2a08eb6 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -220,15 +220,15 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): if cuda_graph: # Compile before capture, then restore this chunk's input cache state. saved_tail, saved_packed = tail.clone(), packed.clone() - compress_pools(*args, max_q_len=max_q_len) + compress_pools(*args, max_q_seq_len=max_q_len) tail.copy_(saved_tail) packed.copy_(saved_packed) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - compress_pools(*args, max_q_len=max_q_len) + compress_pools(*args, max_q_seq_len=max_q_len) graph.replay() else: - compress_pools(*args, max_q_len=max_q_len) + compress_pools(*args, max_q_seq_len=max_q_len) for req, _, end in chunks: offset = 0 if req == 1 else 9 assert torch.equal(tail[req, : end % 4], source[offset + end // 4 * 4 : offset + end]) @@ -247,7 +247,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) starts = torch.tensor([0, 9], device="cuda", dtype=torch.int32) groups = torch.tensor([[1, 0], [0, -1]], device="cuda", dtype=torch.int32) - out, relative = expand_topk(groups, lengths, req_idx, table, topk=8) + out, relative = expand_topk(groups, lengths, req_idx, table, index_topk=8) for batch, expected in enumerate([[4, 5, 6, 7, 0, 1, 2, 3, 8], list(range(7))]): assert relative[batch, : len(expected)].tolist() == expected assert ( @@ -273,7 +273,7 @@ def test_kpool_decode_cuda_graph_and_padding(): cu_q_lens = torch.arange(5, device="cuda", dtype=torch.int32) def run(): - compress_pools(raw, tail, packed, ape, lengths, mem_index, req_idx, cu_q_lens, lengths, max_q_len=1) + compress_pools(raw, tail, packed, ape, lengths, mem_index, req_idx, cu_q_lens, lengths, max_q_seq_len=1) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -402,7 +402,7 @@ def test_kpool_packed_kv_above_two_gib(): zero = torch.zeros(4, device="cuda", dtype=torch.int32) ragged = torch.tensor([0, 1, 2, loc], device="cuda", dtype=torch.int32) cu_q_lens = torch.tensor([0, 4], device="cuda", dtype=torch.int32) - compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) + compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_seq_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) @@ -611,8 +611,8 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): max_q_seq_len=tokens, ) state = SimpleNamespace( - lengths=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), - query_req=torch.zeros(tokens, device="cuda", dtype=torch.int32), + q_kv_seq_len=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), + q_req_idx=torch.zeros(tokens, device="cuda", dtype=torch.int32), ) indexer = Glm5NextNsaInfer( 0, @@ -637,8 +637,8 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): infer.mem_index = ragged[pos : pos + 1] infer.max_kv_seq_len = pos + 1 infer.b_seq_len.fill_(pos + 1) - state.lengths = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) - state.query_req = infer.b_req_idx + state.q_kv_seq_len = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) + state.q_req_idx = infer.b_req_idx _, decoded = indexer._get_indices(hidden[pos : pos + 1], hidden[pos : pos + 1], infer, state, weights) valid = 2048 + (pos + 1) % 4 assert set(decoded[0, :valid].tolist()) == set(full[pos, :valid].tolist()) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 354521250f..8a1d70e3ac 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -114,7 +114,7 @@ def prefill(values, destination, tail): torch.arange(count + 1, device="cuda", dtype=torch.int32), lengths, 1, - mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) prefill(raw[: start + count], expected, torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16)) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 1e14047c24..e72e4d6bee 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -268,16 +268,16 @@ def test_pd_kv_and_runtime_state_roundtrip( for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): new_indexes = [i for i in range(32) if i not in indexes][:count] compress_pools( - raw=raw, - tail=mem.tail.buffer[0], - packed_buffer=mem.full.get_layer_view(3).indexer_k, - ape=ape, - lengths=tensor(range(length + 1, length + count + 1)), + k_and_gate=raw, + indexer_tail=mem.tail.buffer[0], + indexer_k=mem.full.get_layer_view(3).indexer_k, + index_kpool_compress_ape=ape, + q_kv_seq_len=tensor(range(length + 1, length + count + 1)), mem_index=tensor(new_indexes), - req_idx=tensor([req_idx]), - cu_q_lens=tensor([0, count]), - seq_lens=tensor([length + count]), - max_q_len=count, + b_req_idx=tensor([req_idx]), + b1_cu_q_seq_len=tensor([0, count]), + b_seq_len=tensor([length + count]), + max_q_seq_len=count, ) outputs.append(mem.full.get_layer_view(3).indexer_k[new_indexes[-1]].clone()) assert torch.equal(*outputs) @@ -345,21 +345,21 @@ def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, def prefill(values, destination, tail): count = len(values) compress_pools( - raw=values, - tail=tail, - packed_buffer=destination, - ape=ape, - lengths=torch.arange( + k_and_gate=values, + indexer_tail=tail, + indexer_k=destination, + index_kpool_compress_ape=ape, + q_kv_seq_len=torch.arange( start=1, end=count + 1, device="cuda", dtype=torch.int32, ), mem_index=ragged[:count], - req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), - cu_q_lens=torch.tensor([0, count], device="cuda", dtype=torch.int32), - seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), - max_q_len=count, + b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, count], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([count], device="cuda", dtype=torch.int32), + max_q_seq_len=count, ) prefill(raw[:prefix], packed, source.get_layer_view(0)) @@ -404,17 +404,17 @@ def prefill(values, destination, tail): dtype=torch.int32, ) compress_pools( - raw=raw[start : start + count], - tail=target.get_layer_view(0), - packed_buffer=packed, - ape=ape, - lengths=lengths, + k_and_gate=raw[start : start + count], + indexer_tail=target.get_layer_view(0), + indexer_k=packed, + index_kpool_compress_ape=ape, + q_kv_seq_len=lengths, mem_index=ragged[start : start + count], - req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), - cu_q_lens=torch.arange(count + 1, device="cuda", dtype=torch.int32), - seq_lens=lengths, - max_q_len=1, - mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + b_req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.arange(count + 1, device="cuda", dtype=torch.int32), + b_seq_len=lengths, + max_q_seq_len=1, + b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) prefill(raw[: start + count], expected, torch.zeros_like(source.get_layer_view(0))) From d0c1db87936397874a0886c5518ac64742dd91ba Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:10:41 +0000 Subject: [PATCH 05/10] perf(glm): simplify FP8 prefill workspaces and dispatch --- .../basemodel/attention/nsa/fp8_glm5_next.py | 41 ++--- .../context_flashattention_nopad_fp8.py | 170 ++++++++++-------- .../triton_kernel/fp8_sparse_decode.py | 4 +- unit_tests/models/glm5_next/test_fp8_kv.py | 106 ++++++++--- 4 files changed, 192 insertions(+), 129 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index 129248fff0..e5240742e3 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -10,33 +10,17 @@ class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - self._prefill_workspaces = [{"smallq": None, "generic": None} for _ in range(2)] + self._prefill_workspaces = [{"direct": None, "generic": None} for _ in range(2)] - def get_smallq_workspace(self, infer_state, q): + def get_direct_prefill_workspace(self, infer_state, q): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - allocate_direct_smallq_workspace, - allocate_selected_prefill_workspace, + allocate_direct_prefill_workspace, ) slot = self._prefill_workspaces[infer_state.microbatch_index] - if slot["smallq"] is None: - slot["smallq"] = { - "direct": allocate_direct_smallq_workspace(q[:1]), - "selected": allocate_selected_prefill_workspace(q, q_token_capacity=8), - } - if q.shape[0] == 1: - return slot["smallq"]["direct"] - q_token_num = q.shape[0] - selected = slot["smallq"]["selected"] - topk_size = selected["compact_topk_mem_indices"].shape[1] - return { - "compact_topk_mem_indices": selected["compact_topk_mem_indices"][:q_token_num], - "cache_seqlens": selected["cache_seqlens"][:q_token_num], - "page_table": selected["page_table"][:q_token_num], - "cu_seqlens_q": selected["cu_seqlens_q"][: q_token_num + 1], - "cu_seqlens_k_new": selected["cu_seqlens_k_new"][: q_token_num + 1], - "selected_kv": selected["selected_kv"][: q_token_num * topk_size], - } + if slot["direct"] is None: + slot["direct"] = allocate_direct_prefill_workspace(q) + return slot["direct"] def get_generic_prefill_workspace(self, infer_state, q, topk_indices): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( @@ -79,14 +63,14 @@ class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): def _nsa_prefill_att(self, q, packed_kv, att_control): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( context_attention_fwd_fp8, - can_use_direct_smallq, + can_use_direct_prefill, can_use_selected_prefill, ) nsa_prefill_dict = att_control.nsa_prefill_dict infer_state = self.infer_state topk_mem_indices = nsa_prefill_dict["topk_mem_indices"] - direct_workspace = selected_workspace = generic_workspace = None + direct_workspace = generic_workspace = None if q.shape[1] == 64: eligibility = ( q, @@ -96,11 +80,9 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): topk_mem_indices, infer_state.total_token_num, ) - if can_use_direct_smallq(*eligibility): - direct_workspace = self.backend.get_smallq_workspace(infer_state, q) - elif can_use_selected_prefill(*eligibility): - selected_workspace = self.backend.get_smallq_workspace(infer_state, q) - else: + if can_use_direct_prefill(*eligibility): + direct_workspace = self.backend.get_direct_prefill_workspace(infer_state, q) + elif not can_use_selected_prefill(*eligibility): graph_guard = ( infer_state.is_cuda_graph_capture or self.backend.model.prefill_graph is not None @@ -124,7 +106,6 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, total_token_num=infer_state.total_token_num, softmax_scale=nsa_prefill_dict["softmax_scale"], - selected_workspace=selected_workspace, direct_workspace=direct_workspace, generic_workspace=generic_workspace, ) diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index dcbd691260..6d25dd5a04 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -8,7 +8,7 @@ @triton.jit -def _direct_smallq_packed_kernel( +def _direct_prefill_packed_kernel( q, packed_kv, mem_index, @@ -120,7 +120,7 @@ def _direct_smallq_packed_kernel( @triton.jit -def reduce_direct_smallq_kernel( +def reduce_sparse_attention_kernel( partial_output, partial_softmax_lse, output, @@ -147,7 +147,7 @@ def reduce_direct_smallq_kernel( tl.store(output + (qid * q_head_num + h) * 512 + d, value) -def can_use_direct_smallq(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): +def can_use_direct_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): return ( q.ndim == 3 and 1 <= q.shape[0] <= 8 @@ -166,7 +166,7 @@ def can_use_direct_smallq(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_in ) -def allocate_direct_smallq_workspace(q): +def allocate_direct_prefill_workspace(q): q_token_num, q_head_num, _ = q.shape num_splits = 34 if q_token_num == 1 else 17 return { @@ -180,7 +180,7 @@ def allocate_direct_smallq_workspace(q): @torch.no_grad() -def _context_attention_fwd_fp8_direct_smallq( +def _context_attention_fwd_fp8_direct_prefill( q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace=None ): q_token_num = q.shape[0] @@ -195,12 +195,12 @@ def _context_attention_fwd_fp8_direct_smallq( output = torch.empty(q.shape, dtype=q.dtype, device=q.device) launch_grid = (q_token_num, num_splits) else: - workspace = direct_workspace or allocate_direct_smallq_workspace(q) + workspace = direct_workspace or allocate_direct_prefill_workspace(q) partial_output = workspace["partial_output"] partial_softmax_lse = workspace["partial_softmax_lse"] output = workspace["output"] launch_grid = (q_token_num, num_splits, triton.cdiv(q_head_num, block_heads)) - _direct_smallq_packed_kernel[launch_grid]( + _direct_prefill_packed_kernel[launch_grid]( q, packed_kv, mem_index, @@ -226,7 +226,7 @@ def _context_attention_fwd_fp8_direct_smallq( softmax_scale=softmax_scale, num_warps=8, ) - reduce_direct_smallq_kernel[(q_token_num, q_head_num, 4)]( + reduce_sparse_attention_kernel[(q_token_num, q_head_num, 4)]( partial_output, partial_softmax_lse, output, @@ -385,67 +385,79 @@ def _merge_att_output_kernel( @triton.jit def _selected_prefill_compact_kernel( topk_mem_indices, - compact_topk_mem_indices, + compact_topk_offsets, cache_seqlens, - page_table, - cu_seqlens_q, kv_cache_capacity: tl.constexpr, topk_mem_indices_stride_s: tl.constexpr, topk_mem_indices_stride_d: tl.constexpr, - q_token_num: tl.constexpr, BLOCK: tl.constexpr, ): qid = tl.program_id(0) lane = tl.arange(0, BLOCK) - local = tl.load( + physical = tl.load( topk_mem_indices + qid * topk_mem_indices_stride_s + lane * topk_mem_indices_stride_d, lane < 2176, -1 ).to(tl.int64) - valid = (lane < 2176) & (local >= 0) - physical = local - valid = valid & (physical >= 0) & (physical < kv_cache_capacity) + valid = (lane < 2176) & (physical >= 0) & (physical < kv_cache_capacity) rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 - tl.store(compact_topk_mem_indices + qid * 2176 + rank, physical, valid) - # Every selected slot has a stable contiguous page; invalid tail pages are - # harmless because cache_seqlens limits FA3 to cache_seqlens[qid]. - tl.store(page_table + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) - if qid == 0: - tl.store(cu_seqlens_q + lane, lane.to(tl.int32), lane <= q_token_num) + tl.store(compact_topk_offsets + qid * 2176 + rank, lane, valid) tl.store(cache_seqlens + qid, tl.sum(valid.to(tl.int32), axis=0)) @triton.jit def _selected_prefill_gather_kernel( - compact_topk_mem_indices, + compact_topk_offsets, cache_seqlens, - mem_index, + topk_indices, + topk_mem_indices, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, prefill_cache_kv, packed_kv, selected_kv, - cu_seqlens_k_new, packed_kv_stride_s: tl.constexpr, prefill_cache_kv_stride_s: tl.constexpr, prefill_cache_kv_stride_d: tl.constexpr, - q_token_num: tl.constexpr, - BLOCK_Q_TOKENS: tl.constexpr, + topk_indices_stride_s: tl.constexpr, + topk_indices_stride_d: tl.constexpr, + topk_mem_indices_stride_s: tl.constexpr, + topk_mem_indices_stride_d: tl.constexpr, + batch_size: tl.constexpr, + BLOCK_ROWS: tl.constexpr, ): qid = tl.program_id(0) - row = tl.program_id(1) * 4 + tl.arange(0, 4) + row = tl.program_id(1) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) length = tl.load(cache_seqlens + qid) live = row < length - physical = tl.load(compact_topk_mem_indices + qid * 2176 + row, live, 0).to(tl.int64) - current_row = tl.full((4,), -1, tl.int32) - for j in tl.static_range(BLOCK_Q_TOKENS): - current = j < q_token_num - current_mem = tl.load(mem_index + j, current, 0).to(tl.int64) - current_row = tl.where(live & current & (physical == current_mem), j, current_row) - cached = live & (current_row < 0) - current = live & (current_row >= 0) + offset = tl.load(compact_topk_offsets + qid * 2176 + row, live, 0) + physical = tl.load( + topk_mem_indices + qid * topk_mem_indices_stride_s + offset * topk_mem_indices_stride_d, live, 0 + ).to(tl.int64) + position = tl.load(topk_indices + qid * topk_indices_stride_s + offset * topk_indices_stride_d, live, 0) + if batch_size == 1: + batch = 0 + else: + lower, upper = 0, batch_size + while lower < upper: + middle = (lower + upper) // 2 + q_end = tl.load(b1_cu_q_seq_len + middle + 1) + before = qid < q_end + upper = tl.where(before, middle, upper) + lower = tl.where(before, lower, middle + 1) + batch = lower + q_start = tl.load(b1_cu_q_seq_len + batch) + q_end = tl.load(b1_cu_q_seq_len + batch + 1) + kv_start = tl.load(b1_cu_kv_seq_len + batch) + kv_end = tl.load(b1_cu_kv_seq_len + batch + 1) + cache_len = kv_end - kv_start - (q_end - q_start) + current = live & (position >= cache_len) + cached = live & (position < cache_len) + current_row = q_start + position - cache_len dim = tl.arange(0, 512) packed_row = packed_kv + physical[:, None] * packed_kv_stride_s group = tl.arange(0, 8) raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) - scales = tl.broadcast_to(scales[:, :, None], (4, 8, 64)).reshape(4, 512) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_ROWS, 8, 64)).reshape(BLOCK_ROWS, 512) fp8 = ( tl.load( packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], @@ -467,17 +479,13 @@ def _selected_prefill_gather_kernel( ).to(tl.float32) value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) tl.store(selected_kv + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) - if (qid == 0) & (tl.program_id(1) == 0): - x = tl.arange(0, BLOCK_Q_TOKENS) - lengths = tl.load(cache_seqlens + x, x < q_token_num, 0) - tl.store(cu_seqlens_k_new, 0) - tl.store(cu_seqlens_k_new + 1 + x, tl.cumsum(lengths, axis=0), x < q_token_num) def can_use_selected_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): + # Bound the BF16 scratch space and avoid gathering more KV than the full context. return ( q.ndim == 3 - and 2 <= q.shape[0] <= 8 + and 2 <= q.shape[0] <= 128 and q.shape[1:] == (64, 512) and q.dtype == torch.bfloat16 and q.is_cuda @@ -485,7 +493,7 @@ def can_use_selected_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem and packed_kv.shape[1:] == (1, 584) and packed_kv.dtype == torch.uint8 and packed_kv.device == q.device - and total_token_num >= 65536 + and total_token_num >= max(65536 if q.shape[0] <= 8 else 131072, q.shape[0] * 2048) and mem_index.dtype == torch.int32 and mem_index.device == q.device and mem_index.shape == (q.shape[0],) @@ -498,15 +506,12 @@ def can_use_selected_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem ) -def allocate_selected_prefill_workspace(q, q_token_capacity=None): - q_token_num = q.shape[0] if q_token_capacity is None else q_token_capacity +def allocate_selected_prefill_workspace(q): + q_token_num = q.shape[0] return { - "compact_topk_mem_indices": torch.empty((q_token_num, 2176), dtype=torch.int64, device=q.device), + "compact_topk_offsets": torch.empty((q_token_num, 2176), dtype=torch.int32, device=q.device), "cache_seqlens": torch.empty(q_token_num, dtype=torch.int32, device=q.device), - "page_table": torch.empty((q_token_num, 2176), dtype=torch.int32, device=q.device), - "cu_seqlens_q": torch.empty(q_token_num + 1, dtype=torch.int32, device=q.device), - "cu_seqlens_k_new": torch.empty(q_token_num + 1, dtype=torch.int32, device=q.device), - "selected_kv": torch.empty((q_token_num * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), + "selected_kv": torch.empty((q_token_num, 2176, 1, 512), dtype=q.dtype, device=q.device), } @@ -521,55 +526,62 @@ def allocate_generic_prefill_workspace(q, topk_indices, q_token_capacity=None): @torch.no_grad() def _context_attention_fwd_fp8_selected_prefill( - q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, selected_workspace=None + q, + packed_kv, + prefill_cache_kv, + topk_indices, + topk_mem_indices, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, + softmax_scale, ): from sgl_kernel.flash_attn import flash_attn_with_kvcache q_token_num = q.shape[0] - workspace = selected_workspace or allocate_selected_prefill_workspace(q) - block_q_tokens = triton.next_power_of_2(q_token_num) + workspace = allocate_selected_prefill_workspace(q) _selected_prefill_compact_kernel[(q_token_num,)]( topk_mem_indices, - workspace["compact_topk_mem_indices"], + workspace["compact_topk_offsets"], workspace["cache_seqlens"], - workspace["page_table"], - workspace["cu_seqlens_q"], packed_kv.shape[0], topk_mem_indices.stride(0), topk_mem_indices.stride(1), - q_token_num=q_token_num, BLOCK=4096, num_warps=8, ) _selected_prefill_gather_kernel[(q_token_num, triton.cdiv(2176, 4))]( - workspace["compact_topk_mem_indices"], + workspace["compact_topk_offsets"], workspace["cache_seqlens"], - mem_index, + topk_indices, + topk_mem_indices, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, prefill_cache_kv, packed_kv, workspace["selected_kv"], - workspace["cu_seqlens_k_new"], packed_kv.stride(0), prefill_cache_kv.stride(0), prefill_cache_kv.stride(2), - q_token_num=q_token_num, - BLOCK_Q_TOKENS=block_q_tokens, + topk_indices_stride_s=topk_indices.stride(0), + topk_indices_stride_d=topk_indices.stride(1), + topk_mem_indices_stride_s=topk_mem_indices.stride(0), + topk_mem_indices_stride_d=topk_mem_indices.stride(1), + batch_size=b1_cu_q_seq_len.numel() - 1, + BLOCK_ROWS=4, num_warps=4, ) + qv = q.unsqueeze(1) + selected_kv = workspace["selected_kv"] return flash_attn_with_kvcache( - q=None, - qv=q, - k_cache=None, - v_cache=workspace["selected_kv"].view(-1, 1, 1, 512), - page_table=workspace["page_table"], + q=qv[..., :64], + qv=qv, + k_cache=selected_kv[..., :64], + v_cache=selected_kv, cache_seqlens=workspace["cache_seqlens"], - cu_seqlens_q=workspace["cu_seqlens_q"], - cu_seqlens_k_new=workspace["cu_seqlens_k_new"], - max_seqlen_q=1, softmax_scale=softmax_scale, causal=False, only_qv=True, - ) + ).squeeze(1) @torch.no_grad() @@ -587,19 +599,25 @@ def context_attention_fwd_fp8( b1_cu_q_seq_len, total_token_num, softmax_scale, - selected_workspace=None, direct_workspace=None, generic_workspace=None, ): """Run NoPE sparse prefill from request mappings and physical top-k slots.""" eligibility = (q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num) - if can_use_direct_smallq(*eligibility): - return _context_attention_fwd_fp8_direct_smallq( + if can_use_direct_prefill(*eligibility): + return _context_attention_fwd_fp8_direct_prefill( q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace ) if can_use_selected_prefill(*eligibility): return _context_attention_fwd_fp8_selected_prefill( - q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, selected_workspace + q, + packed_kv, + prefill_cache_kv, + topk_indices, + topk_mem_indices, + b1_cu_q_seq_len, + b1_cu_kv_seq_len, + softmax_scale, ) from sgl_kernel.flash_attn import flash_attn_with_kvcache diff --git a/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py index 7e1907dc60..c6a41eae9a 100644 --- a/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py +++ b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py @@ -117,7 +117,7 @@ def can_use_fp8_sparse_direct_decode(q_nope, packed_kv, topk, kv_lora_rank, boun def fp8_sparse_direct_decode(q_nope, packed_kv, topk, softmax_scale): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - reduce_direct_smallq_kernel, + reduce_sparse_attention_kernel, ) batch = q_nope.shape[0] @@ -150,7 +150,7 @@ def fp8_sparse_direct_decode(q_nope, packed_kv, topk, softmax_scale): num_warps=4 if block_k == 16 else 8, num_stages=1, ) - reduce_direct_smallq_kernel[(batch, heads, 4)]( + reduce_sparse_attention_kernel[(batch, heads, 4)]( partial_o, partial_lse, out, diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 33352b0b03..5e92cd9c4d 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -39,8 +39,10 @@ def _make_infer_state(**attributes): def _make_workspace_backend(batch_max_tokens): backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(batch_max_tokens=batch_max_tokens, prefill_graph=None) - backend._prefill_workspaces = [{"smallq": None, "generic": None} for _ in range(2)] + backend.model = SimpleNamespace( + batch_max_tokens=batch_max_tokens, prefill_graph=None, config={"index_topk": 2048, "index_kpool": 4} + ) + backend._prefill_workspaces = [{"direct": None, "generic": None} for _ in range(2)] return backend @@ -81,6 +83,7 @@ def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=N self.topk_indices[:, 1] = visible - 1 self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] self.infer_state = _make_infer_state( + is_prefill=True, req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), b_req_idx=tensor(self.requests), b_seq_len=tensor(seq_lens), @@ -354,7 +357,7 @@ def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softm case.check(output) -@pytest.mark.parametrize("num_q", [16, 33]) +@pytest.mark.parametrize("num_q", [64, 129]) def test_fp8_h64_largeq_native_prefill_request_table_graph(monkeypatch, num_q): if torch.cuda.get_device_capability()[0] != 9: pytest.skip("Validated FlashMLA sparse prefill target is Hopper") @@ -406,7 +409,7 @@ def unexpected_native(*args, **kwargs): "num_q,num_heads", [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], ) -def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num_heads): +def test_fp8_direct_prefill_dispatch_strided_holes_and_current(num_q, num_heads): torch.manual_seed(53016 + num_q) seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] case = _PrefillCase(seqs, [1] * num_q, num_heads) @@ -425,36 +428,33 @@ def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num case.check(output) -def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): +def test_fp8_h64_direct_prefill_state_workspace_reuse_and_fresh_state(): torch.manual_seed(53150) case = _PrefillCase([65536], [1], 64) case.check(case.run()) - cache = case.backend._prefill_workspaces[0]["smallq"] - pointers = {name: value.data_ptr() for name, value in cache["direct"].items()} + workspace = case.backend._prefill_workspaces[0]["direct"] + pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): infer, state = case.new_state() case.check(case.run(state)) - assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers + assert {name: value.data_ptr() for name, value in workspace.items()} == pointers fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) case.check(case.run(fresh)) - assert {name: value.data_ptr() for name, value in cache["direct"].items()} == pointers + assert {name: value.data_ptr() for name, value in workspace.items()} == pointers case.mutate() case.check(case.run()) -@pytest.mark.parametrize("num_q", [2, 3, 5, 8]) -def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): +@pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) +def test_fp8_h64_selected_prefill_cached_chunks_and_graph(num_q): torch.manual_seed(53200 + num_q) - case = _PrefillCase([32768] * num_q, [1] * num_q, 64) + case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) case.check(case.run()) - workspace = case.backend._prefill_workspaces[0]["smallq"]["selected"] - pointers = {name: value.data_ptr() for name, value in workspace.items()} - for _ in range(10): + for _ in range(2): infer, state = case.new_state() case.check(case.run(state)) fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) case.check(case.run(fresh)) - assert {name: value.data_ptr() for name, value in workspace.items()} == pointers graph, output = _capture_cuda_graph(case.run) case.mutate() graph.replay() @@ -464,6 +464,70 @@ def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): assert torch.count_nonzero(output) == 0 +def test_fp8_selected_prefill_graph_with_empty_requests_and_changing_boundaries(): + case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + replacement = _PrefillCase([65536] * 4, [2, 0, 5, 0], 64) + for name in ( + "q", + "packed_kv", + "prefill_cache_kv", + "req_to_token_indexs", + "q_req_idx", + "position_ids", + "topk_indices", + "topk_mem_indices", + ): + getattr(case, name).copy_(getattr(replacement, name)) + case.infer_state.b1_cu_q_seq_len.copy_(replacement.infer_state.b1_cu_q_seq_len) + case.infer_state.position_ids.copy_(replacement.infer_state.position_ids) + for name in ("q_kv_seq_len", "q_req_idx", "q_kv_start_loc"): + getattr(case.att_state, name).copy_(getattr(replacement.att_state, name)) + graph.replay() + case.check(output) + + +def test_fp8_one_query_prefill_and_decode_read_their_own_kv_sources(): + case = _PrefillCase([65536], [1], 64) + case.q.zero_() + case.prefill_cache_kv.fill_(1) + case.source_kv.fill_(2) + write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + case.topk_indices.fill_(-1) + case.topk_indices[:, 0] = 65535 + case.refresh() + decode_infer = _make_infer_state(**vars(case.infer_state)) + decode_infer.is_prefill = False + decode_infer.is_cuda_graph_capture = True + decode_state = case.backend.create_att_decode_state(decode_infer) + decode_state.init_state() + control = AttControl( + nsa_decode=True, + nsa_decode_dict={ + "topk_mem_indices": case.topk_mem_indices, + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + }, + ) + + def run(): + prefill = case.run() + decode = decode_state.decode_att((case.q, case.q[..., :0]), case.packed_kv[..., :584], None, control) + return prefill, decode + + prefill, decode = run() + torch.testing.assert_close(prefill, torch.ones_like(prefill), rtol=0, atol=0) + torch.testing.assert_close(decode, torch.full_like(decode, 2), rtol=0, atol=0) + graph, (prefill, decode) = _capture_cuda_graph(run) + case.prefill_cache_kv.fill_(3) + case.source_kv.fill_(4) + write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + graph.replay() + torch.testing.assert_close(prefill, torch.full_like(prefill, 3), rtol=0, atol=0) + torch.testing.assert_close(decode, torch.full_like(decode, 4), rtol=0, atol=0) + + def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): if torch.cuda.get_device_capability()[0] != 9: pytest.skip("Validated FlashMLA sparse decode target is Hopper") @@ -941,8 +1005,8 @@ def direct_must_not_run(*args, **kwargs): def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): torch.manual_seed(53450) - backend = _make_workspace_backend(26) - case = _PrefillCase([32768, 16384, 16384], [5, 4, 4], 64, backend=backend) + backend = _make_workspace_backend(66) + case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], 64, backend=backend) case.check(case.run()) workspace = backend._prefill_workspaces[0]["generic"] assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} @@ -952,7 +1016,7 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba case.check(case.run(state)) case.mutate() case.check(case.run()) - fresh = _PrefillCase([16384, 32768, 16384], [3, 5, 5], 64, capacity=131072, backend=backend) + fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], 64, capacity=131072, backend=backend) fresh.check(fresh.run()) assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0]["generic"].items()} == pointers for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): @@ -978,14 +1042,14 @@ def forbidden(): } control = AttControl(nsa_prefill=True, nsa_prefill_dict=params) for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): - monkeypatch.setattr(fp8_prefill, "can_use_direct_smallq", lambda *args, d=direct: d) + monkeypatch.setattr(fp8_prefill, "can_use_direct_prefill", lambda *args, d=direct: d) monkeypatch.setattr(fp8_prefill, "can_use_selected_prefill", lambda *args, s=selected: s) infer = _make_infer_state( mem_index=torch.empty(1, dtype=torch.int32), req_manager=SimpleNamespace(req_to_token_indexs=torch.empty((1, 1), dtype=torch.int32)), ) backend = _make_workspace_backend(1) - backend.get_smallq_workspace = lambda *args: {} + backend.get_direct_prefill_workspace = lambda *args: {} state = Fp8Glm5NextSparsePrefillState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) assert ( state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), None, control).shape[1] == heads From 914e949fd99cff9891f7de8dcf249218711e871b Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 12:25:51 +0000 Subject: [PATCH 06/10] refactor(glm): unify FP8 prefill on FlashAttention --- .../basemodel/attention/nsa/fp8_glm5_next.py | 65 +-- .../context_flashattention_nopad_fp8.py | 549 ++---------------- .../triton_kernel/fp8_sparse_decode.py | 32 +- .../glm5_next/test_direct_nsa_indices.py | 6 +- unit_tests/models/glm5_next/test_fp8_kv.py | 131 ++--- 5 files changed, 130 insertions(+), 653 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index e5240742e3..cb17eb4253 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -10,25 +10,14 @@ class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - self._prefill_workspaces = [{"direct": None, "generic": None} for _ in range(2)] + self._prefill_workspaces = [None, None] - def get_direct_prefill_workspace(self, infer_state, q): + def get_prefill_workspace(self, infer_state, q, topk_indices): from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - allocate_direct_prefill_workspace, + allocate_prefill_workspace, ) - slot = self._prefill_workspaces[infer_state.microbatch_index] - if slot["direct"] is None: - slot["direct"] = allocate_direct_prefill_workspace(q) - return slot["direct"] - - def get_generic_prefill_workspace(self, infer_state, q, topk_indices): - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - allocate_generic_prefill_workspace, - ) - - slot = self._prefill_workspaces[infer_state.microbatch_index] - buffers = slot["generic"] + buffers = self._prefill_workspaces[infer_state.microbatch_index] q_token_num, topk_size = topk_indices.shape if ( buffers is None @@ -36,10 +25,10 @@ def get_generic_prefill_workspace(self, infer_state, q, topk_indices): or buffers["chunk_topk_indices"].shape[1] != topk_size or buffers["chunk_topk_indices"].dtype != topk_indices.dtype ): - buffers = allocate_generic_prefill_workspace( + buffers = allocate_prefill_workspace( q, topk_indices, q_token_capacity=max(self.model.batch_max_tokens or 0, q_token_num) ) - slot["generic"] = buffers + self._prefill_workspaces[infer_state.microbatch_index] = buffers return { "chunk_topk_indices": buffers["chunk_topk_indices"][:q_token_num], "chunk_cache_seqlens": buffers["chunk_cache_seqlens"][:q_token_num], @@ -61,53 +50,31 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): def _nsa_prefill_att(self, q, packed_kv, att_control): - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - context_attention_fwd_fp8, - can_use_direct_prefill, - can_use_selected_prefill, - ) + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8 nsa_prefill_dict = att_control.nsa_prefill_dict infer_state = self.infer_state - topk_mem_indices = nsa_prefill_dict["topk_mem_indices"] - direct_workspace = generic_workspace = None - if q.shape[1] == 64: - eligibility = ( - q, - packed_kv, - infer_state.mem_index, - nsa_prefill_dict["prefill_cache_kv"], - topk_mem_indices, - infer_state.total_token_num, - ) - if can_use_direct_prefill(*eligibility): - direct_workspace = self.backend.get_direct_prefill_workspace(infer_state, q) - elif not can_use_selected_prefill(*eligibility): - graph_guard = ( - infer_state.is_cuda_graph_capture - or self.backend.model.prefill_graph is not None - or hasattr(infer_state, "prefill_cuda_graph_exe_list") - ) - if not graph_guard and not torch.cuda.is_current_stream_capturing(): - generic_workspace = self.backend.get_generic_prefill_workspace( - infer_state, q, nsa_prefill_dict["topk_indices"] - ) + workspace = None + graph_guard = ( + infer_state.is_cuda_graph_capture + or self.backend.model.prefill_graph is not None + or hasattr(infer_state, "prefill_cuda_graph_exe_list") + ) + if not graph_guard and not torch.cuda.is_current_stream_capturing(): + workspace = self.backend.get_prefill_workspace(infer_state, q, nsa_prefill_dict["topk_indices"]) return context_attention_fwd_fp8( q=q, packed_kv=packed_kv, - mem_index=infer_state.mem_index, prefill_cache_kv=nsa_prefill_dict["prefill_cache_kv"], q_kv_start_loc=self.q_kv_start_loc, topk_indices=nsa_prefill_dict["topk_indices"], - topk_mem_indices=topk_mem_indices, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, b_req_idx=infer_state.b_req_idx, b1_cu_kv_seq_len=infer_state.b1_cu_kv_seq_len, b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, total_token_num=infer_state.total_token_num, softmax_scale=nsa_prefill_dict["softmax_scale"], - direct_workspace=direct_workspace, - generic_workspace=generic_workspace, + workspace=workspace, ) diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index 6d25dd5a04..1395bbbbfd 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -7,237 +7,6 @@ KV_GATHER_CHUNK_SIZE = 2097152 -@triton.jit -def _direct_prefill_packed_kernel( - q, - packed_kv, - mem_index, - prefill_cache_kv, - topk_mem_indices, - partial_output, - partial_softmax_lse, - kv_cache_capacity, - q_stride_s: tl.constexpr, - q_stride_h: tl.constexpr, - q_stride_d: tl.constexpr, - packed_kv_stride_s: tl.constexpr, - prefill_cache_kv_stride_s: tl.constexpr, - prefill_cache_kv_stride_d: tl.constexpr, - topk_mem_indices_stride_s: tl.constexpr, - topk_mem_indices_stride_d: tl.constexpr, - q_head_num: tl.constexpr, - BLOCK_HEADS: tl.constexpr, - q_token_num: tl.constexpr, - NUM_SPLITS: tl.constexpr, - KV_BLOCKS_PER_SPLIT: tl.constexpr, - KV_SLOTS_PER_SPLIT: tl.constexpr, - softmax_scale: tl.constexpr, -): - qid = tl.program_id(0) - split = tl.program_id(1) - if q_head_num == 8: - # H8 fits in one fixed 16-lane tile; the head mask excludes inactive lanes. - h = tl.arange(0, 16) - else: - h = tl.program_id(2) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) - d = tl.arange(0, 512) - hm = h < q_head_num - q_block = tl.load( - q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, - hm[:, None], - 0.0, - ).to(tl.bfloat16) - if q_head_num == 8: - acc = tl.zeros((16, 512), tl.float32) - m = tl.full((16,), -float("inf"), tl.float32) - denominator = tl.zeros((16,), tl.float32) - else: - acc = tl.zeros((BLOCK_HEADS, 512), tl.float32) - m = tl.full((BLOCK_HEADS,), -float("inf"), tl.float32) - denominator = tl.zeros((BLOCK_HEADS,), tl.float32) - for b in tl.range(0, KV_BLOCKS_PER_SPLIT, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): - n = tl.arange(0, 64) - slot = split * KV_SLOTS_PER_SPLIT + b * 64 + n - top = tl.load( - topk_mem_indices + qid * topk_mem_indices_stride_s + slot * topk_mem_indices_stride_d, - slot < 2176, - -1, - ).to(tl.int64) - valid = (top >= 0) & (top < kv_cache_capacity) - kv_mem_index = tl.where(valid, top, 0) - current_row = tl.full((64,), -1, tl.int32) - for j in tl.static_range(q_token_num): - current_row = tl.where(valid & (kv_mem_index == tl.load(mem_index + j).to(tl.int64)), j, current_row) - cached = valid & (current_row < 0) - current = valid & (current_row >= 0) - packed_row = packed_kv + kv_mem_index[:, None] * packed_kv_stride_s - fp8 = tl.load( - packed_row.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], - cached[:, None] & (d[None, :] < 448), - 0.0, - ).to(tl.float32) - scales = tl.load( - packed_row + 576 + d[None, :] // 64, - cached[:, None] & (d[None, :] < 448), - 127, - ) - fp8 = fp8 * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) - tail = tl.load( - (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + d[None, :] - 448, - cached[:, None] & (d[None, :] >= 448), - 0.0, - ).to(tl.float32) - cached_v = tl.where(d[None, :] < 448, fp8, tail) - current_v = tl.load( - prefill_cache_kv - + current_row[:, None] * prefill_cache_kv_stride_s - + d[None, :] * prefill_cache_kv_stride_d, - current[:, None], - 0.0, - ).to(tl.float32) - kv = tl.where(current[:, None], current_v, cached_v).to(tl.bfloat16) - score = tl.dot(q_block, tl.trans(kv), out_dtype=tl.float32) * softmax_scale - score = tl.where(valid[None, :], score, -float("inf")) - block_m = tl.max(score, axis=1) - new_m = tl.maximum(m, block_m) - alpha = tl.where(m == -float("inf"), 0.0, tl.exp(m - new_m)) - p = tl.where(valid[None, :], tl.exp(score - new_m[:, None]), 0.0) - denominator = denominator * alpha + tl.sum(p, axis=1) - acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv, out_dtype=tl.float32) - m = new_m - has = denominator > 0.0 - out = tl.where(has[:, None], acc / denominator[:, None], 0.0) - tl.store( - partial_output + ((qid * NUM_SPLITS + split) * q_head_num + h[:, None]) * 512 + d[None, :], - out, - hm[:, None], - ) - tl.store( - partial_softmax_lse + (qid * NUM_SPLITS + split) * q_head_num + h, - tl.where(has, m + tl.log(denominator), -float("inf")), - hm, - ) - - -@triton.jit -def reduce_sparse_attention_kernel( - partial_output, - partial_softmax_lse, - output, - q_head_num: tl.constexpr, - NUM_SPLITS: tl.constexpr, - BLOCK_SPLITS: tl.constexpr, -): - qid = tl.program_id(0) - h = tl.program_id(1) - d = tl.program_id(2) * 128 + tl.arange(0, 128) - s = tl.arange(0, BLOCK_SPLITS) - sm = s < NUM_SPLITS - softmax_lse = tl.load(partial_softmax_lse + (qid * NUM_SPLITS + s) * q_head_num + h, sm, -float("inf")) - maximum = tl.max(softmax_lse, axis=0) - safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) - weight = tl.where(sm & (softmax_lse > -float("inf")), tl.exp(softmax_lse - safe_maximum), 0.0) - denominator = tl.sum(weight, axis=0) - partial = tl.load( - partial_output + ((qid * NUM_SPLITS + s[:, None]) * q_head_num + h) * 512 + d[None, :], - sm[:, None], - 0.0, - ).to(tl.float32) - value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) - tl.store(output + (qid * q_head_num + h) * 512 + d, value) - - -def can_use_direct_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): - return ( - q.ndim == 3 - and 1 <= q.shape[0] <= 8 - and q.dtype == torch.bfloat16 - and (q.shape[1:] == (8, 512) or (q.shape[0] == 1 and q.shape[1:] == (64, 512))) - and packed_kv.ndim == 3 - and packed_kv.shape[1] == 1 - and packed_kv.shape[2] >= 584 - and total_token_num >= 65536 - and mem_index.dtype == torch.int32 - and mem_index.numel() == q.shape[0] - and prefill_cache_kv.dtype == torch.bfloat16 - and prefill_cache_kv.shape == (q.shape[0], 1, 512) - and topk_mem_indices.dtype == torch.int32 - and topk_mem_indices.shape == (q.shape[0], 2176) - ) - - -def allocate_direct_prefill_workspace(q): - q_token_num, q_head_num, _ = q.shape - num_splits = 34 if q_token_num == 1 else 17 - return { - "partial_output": torch.empty( - (q_token_num, num_splits, q_head_num, q.shape[-1]), dtype=torch.float32, device=q.device - ), - "partial_softmax_lse": torch.empty((q_token_num, num_splits, q_head_num), dtype=torch.float32, device=q.device), - # The reducer requires a contiguous [Q, H, 512] destination. - "output": torch.empty(q.shape, dtype=q.dtype, device=q.device), - } - - -@torch.no_grad() -def _context_attention_fwd_fp8_direct_prefill( - q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace=None -): - q_token_num = q.shape[0] - q_head_num = q.shape[1] - # Q1 uses 34 64-key splits; H64 uses two H32 head tiles per split. - num_splits, kv_blocks_per_split, kv_slots_per_split = (34, 1, 64) if q_token_num == 1 else (17, 2, 128) - block_heads = 32 if q_head_num == 64 else 16 - if q_head_num == 8: - # H8 uses local buffers because its direct path has no reusable workspace. - partial_output = torch.empty((q_token_num, num_splits, 8, 512), dtype=torch.float32, device=q.device) - partial_softmax_lse = torch.empty((q_token_num, num_splits, 8), dtype=torch.float32, device=q.device) - output = torch.empty(q.shape, dtype=q.dtype, device=q.device) - launch_grid = (q_token_num, num_splits) - else: - workspace = direct_workspace or allocate_direct_prefill_workspace(q) - partial_output = workspace["partial_output"] - partial_softmax_lse = workspace["partial_softmax_lse"] - output = workspace["output"] - launch_grid = (q_token_num, num_splits, triton.cdiv(q_head_num, block_heads)) - _direct_prefill_packed_kernel[launch_grid]( - q, - packed_kv, - mem_index, - prefill_cache_kv, - topk_mem_indices, - partial_output, - partial_softmax_lse, - packed_kv.shape[0], - q.stride(0), - q.stride(1), - q.stride(2), - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - topk_mem_indices.stride(0), - topk_mem_indices.stride(1), - q_head_num=q_head_num, - BLOCK_HEADS=block_heads, - q_token_num=q_token_num, - NUM_SPLITS=num_splits, - KV_BLOCKS_PER_SPLIT=kv_blocks_per_split, - KV_SLOTS_PER_SPLIT=kv_slots_per_split, - softmax_scale=softmax_scale, - num_warps=8, - ) - reduce_sparse_attention_kernel[(q_token_num, q_head_num, 4)]( - partial_output, - partial_softmax_lse, - output, - q_head_num=q_head_num, - NUM_SPLITS=num_splits, - BLOCK_SPLITS=triton.next_power_of_2(num_splits), - num_warps=4, - ) - return output - - @triton.jit def _gather_prefill_kv_kernel( packed_kv, @@ -318,25 +87,19 @@ def _compact_topk_indices_kernel( chunk_cache_seqlens, chunk_start, chunk_end, + topk_indices_stride_s: tl.constexpr, + topk_indices_stride_d: tl.constexpr, TOPK: tl.constexpr, BLOCK: tl.constexpr, - COMPACT: tl.constexpr = True, ): row = tl.program_id(0) lane = tl.arange(0, BLOCK) - index = tl.load(topk_indices + row * TOPK + lane, lane < TOPK, -1) + index = tl.load(topk_indices + row * topk_indices_stride_s + lane * topk_indices_stride_d, lane < TOPK, -1) kv_index = index + tl.load(q_kv_start_loc + row) valid = (lane < TOPK) & (index >= 0) & (kv_index >= chunk_start) & (kv_index < chunk_end) - if COMPACT: - compact_index = tl.cumsum(valid.to(tl.int32)) - 1 - tl.store(chunk_topk_indices + row * TOPK + compact_index, kv_index - chunk_start, valid) - tl.store(chunk_cache_seqlens + row, tl.sum(valid.to(tl.int32))) - else: - tl.store( - chunk_topk_indices + row * TOPK + lane, - tl.where(valid, kv_index - chunk_start, -1), - lane < TOPK, - ) + compact_index = tl.cumsum(valid.to(tl.int32)) - 1 + tl.store(chunk_topk_indices + row * TOPK + compact_index, kv_index - chunk_start, valid) + tl.store(chunk_cache_seqlens + row, tl.sum(valid.to(tl.int32))) @triton.jit @@ -382,140 +145,7 @@ def _merge_att_output_kernel( ) -@triton.jit -def _selected_prefill_compact_kernel( - topk_mem_indices, - compact_topk_offsets, - cache_seqlens, - kv_cache_capacity: tl.constexpr, - topk_mem_indices_stride_s: tl.constexpr, - topk_mem_indices_stride_d: tl.constexpr, - BLOCK: tl.constexpr, -): - qid = tl.program_id(0) - lane = tl.arange(0, BLOCK) - physical = tl.load( - topk_mem_indices + qid * topk_mem_indices_stride_s + lane * topk_mem_indices_stride_d, lane < 2176, -1 - ).to(tl.int64) - valid = (lane < 2176) & (physical >= 0) & (physical < kv_cache_capacity) - rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 - tl.store(compact_topk_offsets + qid * 2176 + rank, lane, valid) - tl.store(cache_seqlens + qid, tl.sum(valid.to(tl.int32), axis=0)) - - -@triton.jit -def _selected_prefill_gather_kernel( - compact_topk_offsets, - cache_seqlens, - topk_indices, - topk_mem_indices, - b1_cu_q_seq_len, - b1_cu_kv_seq_len, - prefill_cache_kv, - packed_kv, - selected_kv, - packed_kv_stride_s: tl.constexpr, - prefill_cache_kv_stride_s: tl.constexpr, - prefill_cache_kv_stride_d: tl.constexpr, - topk_indices_stride_s: tl.constexpr, - topk_indices_stride_d: tl.constexpr, - topk_mem_indices_stride_s: tl.constexpr, - topk_mem_indices_stride_d: tl.constexpr, - batch_size: tl.constexpr, - BLOCK_ROWS: tl.constexpr, -): - qid = tl.program_id(0) - row = tl.program_id(1) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) - length = tl.load(cache_seqlens + qid) - live = row < length - offset = tl.load(compact_topk_offsets + qid * 2176 + row, live, 0) - physical = tl.load( - topk_mem_indices + qid * topk_mem_indices_stride_s + offset * topk_mem_indices_stride_d, live, 0 - ).to(tl.int64) - position = tl.load(topk_indices + qid * topk_indices_stride_s + offset * topk_indices_stride_d, live, 0) - if batch_size == 1: - batch = 0 - else: - lower, upper = 0, batch_size - while lower < upper: - middle = (lower + upper) // 2 - q_end = tl.load(b1_cu_q_seq_len + middle + 1) - before = qid < q_end - upper = tl.where(before, middle, upper) - lower = tl.where(before, lower, middle + 1) - batch = lower - q_start = tl.load(b1_cu_q_seq_len + batch) - q_end = tl.load(b1_cu_q_seq_len + batch + 1) - kv_start = tl.load(b1_cu_kv_seq_len + batch) - kv_end = tl.load(b1_cu_kv_seq_len + batch + 1) - cache_len = kv_end - kv_start - (q_end - q_start) - current = live & (position >= cache_len) - cached = live & (position < cache_len) - current_row = q_start + position - cache_len - dim = tl.arange(0, 512) - packed_row = packed_kv + physical[:, None] * packed_kv_stride_s - group = tl.arange(0, 8) - raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) - scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) - scales = tl.broadcast_to(scales[:, :, None], (BLOCK_ROWS, 8, 64)).reshape(BLOCK_ROWS, 512) - fp8 = ( - tl.load( - packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], - cached[:, None] & (dim[None, :] < 448), - 0.0, - ).to(tl.float32) - * scales - ) - tail = tl.load( - (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, - cached[:, None] & (dim[None, :] >= 448), - 0.0, - ).to(tl.float32) - cached_value = tl.where(dim[None, :] < 448, fp8, tail) - current_value = tl.load( - prefill_cache_kv + current_row[:, None] * prefill_cache_kv_stride_s + dim[None, :] * prefill_cache_kv_stride_d, - current[:, None], - 0.0, - ).to(tl.float32) - value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) - tl.store(selected_kv + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) - - -def can_use_selected_prefill(q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num): - # Bound the BF16 scratch space and avoid gathering more KV than the full context. - return ( - q.ndim == 3 - and 2 <= q.shape[0] <= 128 - and q.shape[1:] == (64, 512) - and q.dtype == torch.bfloat16 - and q.is_cuda - and q.stride(-1) == 1 - and packed_kv.shape[1:] == (1, 584) - and packed_kv.dtype == torch.uint8 - and packed_kv.device == q.device - and total_token_num >= max(65536 if q.shape[0] <= 8 else 131072, q.shape[0] * 2048) - and mem_index.dtype == torch.int32 - and mem_index.device == q.device - and mem_index.shape == (q.shape[0],) - and prefill_cache_kv.dtype == torch.bfloat16 - and prefill_cache_kv.device == q.device - and prefill_cache_kv.shape == (q.shape[0], 1, 512) - and topk_mem_indices.dtype == torch.int32 - and topk_mem_indices.device == q.device - and topk_mem_indices.shape == (q.shape[0], 2176) - ) - - -def allocate_selected_prefill_workspace(q): - q_token_num = q.shape[0] - return { - "compact_topk_offsets": torch.empty((q_token_num, 2176), dtype=torch.int32, device=q.device), - "cache_seqlens": torch.empty(q_token_num, dtype=torch.int32, device=q.device), - "selected_kv": torch.empty((q_token_num, 2176, 1, 512), dtype=q.dtype, device=q.device), - } - - -def allocate_generic_prefill_workspace(q, topk_indices, q_token_capacity=None): +def allocate_prefill_workspace(q, topk_indices, q_token_capacity=None): capacity = q.shape[0] if q_token_capacity is None else q_token_capacity return { "chunk_topk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), @@ -524,166 +154,61 @@ def allocate_generic_prefill_workspace(q, topk_indices, q_token_capacity=None): } -@torch.no_grad() -def _context_attention_fwd_fp8_selected_prefill( - q, - packed_kv, - prefill_cache_kv, - topk_indices, - topk_mem_indices, - b1_cu_q_seq_len, - b1_cu_kv_seq_len, - softmax_scale, -): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - q_token_num = q.shape[0] - workspace = allocate_selected_prefill_workspace(q) - _selected_prefill_compact_kernel[(q_token_num,)]( - topk_mem_indices, - workspace["compact_topk_offsets"], - workspace["cache_seqlens"], - packed_kv.shape[0], - topk_mem_indices.stride(0), - topk_mem_indices.stride(1), - BLOCK=4096, - num_warps=8, - ) - _selected_prefill_gather_kernel[(q_token_num, triton.cdiv(2176, 4))]( - workspace["compact_topk_offsets"], - workspace["cache_seqlens"], - topk_indices, - topk_mem_indices, - b1_cu_q_seq_len, - b1_cu_kv_seq_len, - prefill_cache_kv, - packed_kv, - workspace["selected_kv"], - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - topk_indices_stride_s=topk_indices.stride(0), - topk_indices_stride_d=topk_indices.stride(1), - topk_mem_indices_stride_s=topk_mem_indices.stride(0), - topk_mem_indices_stride_d=topk_mem_indices.stride(1), - batch_size=b1_cu_q_seq_len.numel() - 1, - BLOCK_ROWS=4, - num_warps=4, - ) - qv = q.unsqueeze(1) - selected_kv = workspace["selected_kv"] - return flash_attn_with_kvcache( - q=qv[..., :64], - qv=qv, - k_cache=selected_kv[..., :64], - v_cache=selected_kv, - cache_seqlens=workspace["cache_seqlens"], - softmax_scale=softmax_scale, - causal=False, - only_qv=True, - ).squeeze(1) - - @torch.no_grad() def context_attention_fwd_fp8( q, packed_kv, - mem_index, prefill_cache_kv, q_kv_start_loc, topk_indices, - topk_mem_indices, req_to_token_indexs, b_req_idx, b1_cu_kv_seq_len, b1_cu_q_seq_len, total_token_num, softmax_scale, - direct_workspace=None, - generic_workspace=None, + workspace=None, ): - """Run NoPE sparse prefill from request mappings and physical top-k slots.""" - eligibility = (q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, total_token_num) - if can_use_direct_prefill(*eligibility): - return _context_attention_fwd_fp8_direct_prefill( - q, packed_kv, mem_index, prefill_cache_kv, topk_mem_indices, softmax_scale, direct_workspace - ) - if can_use_selected_prefill(*eligibility): - return _context_attention_fwd_fp8_selected_prefill( - q, - packed_kv, - prefill_cache_kv, - topk_indices, - topk_mem_indices, - b1_cu_q_seq_len, - b1_cu_kv_seq_len, - softmax_scale, - ) + """Run GLM sparse prefill with FlashAttention and request-table KV gathering.""" + q_token_num, q_head_num, _ = q.shape from sgl_kernel.flash_attn import flash_attn_with_kvcache - q_token_num, q_head_num, _ = q.shape gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) - workspace = generic_workspace or allocate_generic_prefill_workspace(q, topk_indices) + workspace = workspace or allocate_prefill_workspace(q, topk_indices) chunk_topk_indices = workspace["chunk_topk_indices"] chunk_cache_seqlens = workspace["chunk_cache_seqlens"] cu_seqlens_q = workspace["cu_seqlens_q"] # Keep this potentially 2 GiB buffer local to a DSA layer. - gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + has_history = total_token_num > q_token_num + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) if has_history else None + if not has_history and prefill_cache_kv.stride(-1) != 1: + prefill_cache_kv = prefill_cache_kv.contiguous() output = softmax_lse = None for chunk_start in range(0, total_token_num, gather_chunk_size): chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) - _gather_prefill_kv_kernel[(chunk_token_num // 4 + b_req_idx.numel(),)]( - packed_kv, - prefill_cache_kv, - gathered_kv, - req_to_token_indexs, - b_req_idx, - b1_cu_kv_seq_len, - b1_cu_q_seq_len, - chunk_start, - chunk_token_num, - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), - req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), - batch_size=b_req_idx.numel(), - BLOCK_ROWS=4, - ) - chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) - if ( - q_token_num > 8 - and q_head_num == 64 - and total_token_num > q_token_num - and total_token_num <= KV_GATHER_CHUNK_SIZE - and q.dtype == torch.bfloat16 - and q.is_contiguous() - and packed_kv.dtype == torch.uint8 - and packed_kv.shape[1:] == (1, 584) - and topk_indices.dtype == torch.int32 - and topk_indices.shape[1] == 2176 - and topk_indices.is_contiguous() - ): - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - _compact_topk_indices_kernel[(q_token_num,)]( - topk_indices, - q_kv_start_loc, - chunk_topk_indices, - chunk_cache_seqlens, + if has_history: + _gather_prefill_kv_kernel[(chunk_token_num // 4 + b_req_idx.numel(),)]( + packed_kv, + prefill_cache_kv, + gathered_kv, + req_to_token_indexs, + b_req_idx, + b1_cu_kv_seq_len, + b1_cu_q_seq_len, chunk_start, - chunk_start + chunk_token_num, - TOPK=topk_indices.shape[1], - BLOCK=triton.next_power_of_2(topk_indices.shape[1]), - COMPACT=False, + chunk_token_num, + packed_kv.stride(0), + prefill_cache_kv.stride(0), + prefill_cache_kv.stride(2), + req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), + req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), + batch_size=b_req_idx.numel(), + BLOCK_ROWS=4, ) - return flash_mla_sparse_fwd( - q=q, - kv=gathered_kv[:chunk_token_num], - indices=chunk_topk_indices.unsqueeze(1), - sm_scale=softmax_scale, - d_v=512, - )[0] + chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) + else: + # Current BF16 KV already has the complete request layout. + chunk_kv = prefill_cache_kv[chunk_start : chunk_start + chunk_token_num].view(-1, 1, 1, 512) _compact_topk_indices_kernel[(q_token_num,)]( topk_indices, q_kv_start_loc, @@ -691,6 +216,8 @@ def context_attention_fwd_fp8( chunk_cache_seqlens, chunk_start, chunk_start + chunk_token_num, + topk_indices.stride(0), + topk_indices.stride(1), TOPK=topk_indices.shape[1], BLOCK=triton.next_power_of_2(topk_indices.shape[1]), ) diff --git a/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py index c6a41eae9a..62aea7b20c 100644 --- a/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py +++ b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py @@ -85,6 +85,34 @@ def _fp8_sparse_direct_decode_kernel( ) +@triton.jit +def reduce_sparse_attention_kernel( + partial_output, + partial_softmax_lse, + output, + q_head_num: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, +): + qid = tl.program_id(0) + h = tl.program_id(1) + d = tl.program_id(2) * 128 + tl.arange(0, 128) + s = tl.arange(0, BLOCK_SPLITS) + sm = s < NUM_SPLITS + softmax_lse = tl.load(partial_softmax_lse + (qid * NUM_SPLITS + s) * q_head_num + h, sm, -float("inf")) + maximum = tl.max(softmax_lse, axis=0) + safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weight = tl.where(sm & (softmax_lse > -float("inf")), tl.exp(softmax_lse - safe_maximum), 0.0) + denominator = tl.sum(weight, axis=0) + partial = tl.load( + partial_output + ((qid * NUM_SPLITS + s[:, None]) * q_head_num + h) * 512 + d[None, :], + sm[:, None], + 0.0, + ).to(tl.float32) + value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) + tl.store(output + (qid * q_head_num + h) * 512 + d, value) + + def can_use_fp8_sparse_direct_decode(q_nope, packed_kv, topk, kv_lora_rank, bound, is_cuda_graph_capture=False): """Validate the GLM packed format and choose the direct decode implementation.""" # The producer guarantees valid indices occupy at most the first 2051 slots. @@ -116,10 +144,6 @@ def can_use_fp8_sparse_direct_decode(q_nope, packed_kv, topk, kv_lora_rank, boun def fp8_sparse_direct_decode(q_nope, packed_kv, topk, softmax_scale): - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - reduce_sparse_attention_kernel, - ) - batch = q_nope.shape[0] heads = q_nope.shape[1] block_h = 16 if batch * heads <= 64 else min(32, max(16, triton.next_power_of_2(heads))) diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index 624793c8e4..cc45742bfc 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -214,9 +214,9 @@ def check(output): torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) check(run()) - generic = backend._prefill_workspaces[0]["generic"] - if generic is not None: - assert set(generic) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} + workspace = backend._prefill_workspaces[0] + if workspace is not None: + assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = run() diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 5e92cd9c4d..161cb5b20b 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -42,7 +42,7 @@ def _make_workspace_backend(batch_max_tokens): backend.model = SimpleNamespace( batch_max_tokens=batch_max_tokens, prefill_graph=None, config={"index_topk": 2048, "index_kpool": 4} ) - backend._prefill_workspaces = [{"direct": None, "generic": None} for _ in range(2)] + backend._prefill_workspaces = [None, None] return backend @@ -337,6 +337,28 @@ def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chun case.check(output) +@pytest.mark.parametrize("num_heads", [8, 64]) +@pytest.mark.parametrize("chunk_size", [3, 32768]) +def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) + case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] + case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv + + def check_kv_storage(*args, **kwargs): + assert kwargs["v_cache"].untyped_storage().data_ptr() == case.prefill_cache_kv.untyped_storage().data_ptr() + return flash_attn_with_kvcache(*args, **kwargs) + + monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", check_kv_storage) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) + + @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): @@ -358,22 +380,25 @@ def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softm @pytest.mark.parametrize("num_q", [64, 129]) -def test_fp8_h64_largeq_native_prefill_request_table_graph(monkeypatch, num_q): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse prefill target is Hopper") - from sgl_kernel.flash_mla import flash_mla_sparse_fwd +@pytest.mark.parametrize("num_heads", [8, 64]) +def test_fp8_prefill_uses_flashattention_and_request_table_graph(monkeypatch, num_q, num_heads): + from sgl_kernel.flash_attn import flash_attn_with_kvcache torch.manual_seed(53100 + num_q) - case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], 64, capacity=131072) + case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], num_heads, capacity=131072) observed = [] - def observed_sparse(*args, **kwargs): - observed.append((kwargs["q"].shape, kwargs["kv"].shape, kwargs["indices"].shape)) - return flash_mla_sparse_fwd(*args, **kwargs) + def observed_flashattention(*args, **kwargs): + observed.append(kwargs["qv"].shape) + return flash_attn_with_kvcache(*args, **kwargs) - monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", observed_sparse) + def unexpected_flashmla(*args, **kwargs): + raise AssertionError("FP8 prefill must use its single FlashAttention backend") + + monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", observed_flashattention) + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_flashmla) case.check(case.run()) - assert observed == [((num_q, 64, 512), (65536, 1, 512), (num_q, 1, 2176))] + assert observed == [(num_q, num_heads, 512)] graph, output = _capture_cuda_graph(case.run) case.mutate() graph.replay() @@ -381,35 +406,13 @@ def observed_sparse(*args, **kwargs): case.topk_indices.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 - assert len(observed) == 3 - - -def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): - import sgl_kernel.flash_attn as flash_attn - - torch.manual_seed(53180) - case = _PrefillCase([65536], [16], 8, capacity=65536) - observed = [] - original = flash_attn.flash_attn_with_kvcache - - def observed_flashattention(*args, **kwargs): - observed.append(True) - return original(*args, **kwargs) - - def unexpected_native(*args, **kwargs): - raise AssertionError("H8 large-Q prefill entered the native H64 sparse path") - - monkeypatch.setattr(flash_attn, "flash_attn_with_kvcache", observed_flashattention) - monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_native) - case.check(case.run()) - assert observed @pytest.mark.parametrize( "num_q,num_heads", [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], ) -def test_fp8_direct_prefill_dispatch_strided_holes_and_current(num_q, num_heads): +def test_fp8_cached_prefill_strided_holes_and_current(num_q, num_heads): torch.manual_seed(53016 + num_q) seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] case = _PrefillCase(seqs, [1] * num_q, num_heads) @@ -428,25 +431,8 @@ def test_fp8_direct_prefill_dispatch_strided_holes_and_current(num_q, num_heads) case.check(output) -def test_fp8_h64_direct_prefill_state_workspace_reuse_and_fresh_state(): - torch.manual_seed(53150) - case = _PrefillCase([65536], [1], 64) - case.check(case.run()) - workspace = case.backend._prefill_workspaces[0]["direct"] - pointers = {name: value.data_ptr() for name, value in workspace.items()} - for _ in range(10): - infer, state = case.new_state() - case.check(case.run(state)) - assert {name: value.data_ptr() for name, value in workspace.items()} == pointers - fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) - case.check(case.run(fresh)) - assert {name: value.data_ptr() for name, value in workspace.items()} == pointers - case.mutate() - case.check(case.run()) - - @pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) -def test_fp8_h64_selected_prefill_cached_chunks_and_graph(num_q): +def test_fp8_cached_prefill_multi_query_and_graph(num_q): torch.manual_seed(53200 + num_q) case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) case.check(case.run()) @@ -464,7 +450,7 @@ def test_fp8_h64_selected_prefill_cached_chunks_and_graph(num_q): assert torch.count_nonzero(output) == 0 -def test_fp8_selected_prefill_graph_with_empty_requests_and_changing_boundaries(): +def test_fp8_prefill_graph_with_empty_requests_and_changing_boundaries(): case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) case.check(case.run()) graph, output = _capture_cuda_graph(case.run) @@ -1003,12 +989,13 @@ def direct_must_not_run(*args, **kwargs): torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) -def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): +@pytest.mark.parametrize("num_heads", [8, 64]) +def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(num_heads): torch.manual_seed(53450) backend = _make_workspace_backend(66) - case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], 64, backend=backend) + case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) case.check(case.run()) - workspace = backend._prefill_workspaces[0]["generic"] + workspace = backend._prefill_workspaces[0] assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): @@ -1016,9 +1003,9 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba case.check(case.run(state)) case.mutate() case.check(case.run()) - fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], 64, capacity=131072, backend=backend) + fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], num_heads, capacity=131072, backend=backend) fresh.check(fresh.run()) - assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0]["generic"].items()} == pointers + assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0].items()} == pointers for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): infer, state = case.new_state(**overrides) case.check(case.run(state)) @@ -1028,34 +1015,6 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba case.check(output) -def test_fp8_generic_prefill_workspace_graph_query_is_generic_h64_only(monkeypatch): - def forbidden(): - raise AssertionError("Non-generic prefill queried CUDA graph capture state") - - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", forbidden) - monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", lambda **kwargs: kwargs["q"]) - params = { - "prefill_cache_kv": torch.empty((1, 1, 512), dtype=torch.bfloat16), - "topk_indices": torch.empty((1, 2176), dtype=torch.int32), - "topk_mem_indices": torch.empty((1, 2176), dtype=torch.int32), - "softmax_scale": 0.0625, - } - control = AttControl(nsa_prefill=True, nsa_prefill_dict=params) - for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): - monkeypatch.setattr(fp8_prefill, "can_use_direct_prefill", lambda *args, d=direct: d) - monkeypatch.setattr(fp8_prefill, "can_use_selected_prefill", lambda *args, s=selected: s) - infer = _make_infer_state( - mem_index=torch.empty(1, dtype=torch.int32), - req_manager=SimpleNamespace(req_to_token_indexs=torch.empty((1, 1), dtype=torch.int32)), - ) - backend = _make_workspace_backend(1) - backend.get_direct_prefill_workspace = lambda *args: {} - state = Fp8Glm5NextSparsePrefillState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - assert ( - state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), None, control).shape[1] == heads - ) - - def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state From e40c932c286a3268e80da6a8e63d4d096fa486f6 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:01:29 +0000 Subject: [PATCH 07/10] perf(glm): streamline FP8 prefill KV gathering --- .../context_flashattention_nopad_fp8.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py index f96c97ef3f..f574d7adbb 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -59,19 +59,22 @@ def _gather_prefill_kv_kernel( cached[:, None] & (dim[None, :] < 448), 0.0, ).to(tl.float32) - scales = tl.load(packed_row + 576 + dim[None, :] // 64, cached[:, None] & (dim[None, :] < 448), 127) - values = quantized * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + group = tl.arange(0, 8) + exponent = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) + scales = (exponent.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_ROWS, 8, 64)).reshape(BLOCK_ROWS, 512) + values = quantized * scales tail = tl.load( (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, cached[:, None] & (dim[None, :] >= 448), 0, ).to(tl.float32) - values = tl.where(dim[None, :] < 448, values, tail) + values = tl.where(dim[None, :] < 448, values, tail).to(tl.bfloat16) current_values = tl.load( prefill_cache_kv + prefill_row[:, None] * prefill_cache_kv_stride_s + dim[None, :] * prefill_cache_kv_stride_d, current[:, None], 0, - ).to(tl.float32) + ) tl.store( output + row[:, None] * 512 + dim[None, :], tl.where(current[:, None], current_values, values), @@ -187,7 +190,7 @@ def context_attention_fwd_fp8( for chunk_start in range(0, total_token_num, gather_chunk_size): chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) if has_history: - _gather_prefill_kv_kernel[(chunk_token_num // 4 + b_req_idx.numel(),)]( + _gather_prefill_kv_kernel[(chunk_token_num // 2 + b_req_idx.numel(),)]( packed_kv, prefill_cache_kv, gathered_kv, @@ -203,7 +206,8 @@ def context_attention_fwd_fp8( req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), batch_size=b_req_idx.numel(), - BLOCK_ROWS=4, + BLOCK_ROWS=2, + num_warps=2, ) chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) else: From 4e62741ed673ded747969cad6b35c5d0f6e92b9c Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:12:14 +0000 Subject: [PATCH 08/10] refactor(glm): move FP8 prefill into attention backend --- .../basemodel/attention/nsa/fp8_glm5_next.py | 132 +++++++++++++++--- ...hattention_nopad_fp8.py => fp8_prefill.py} | 118 ---------------- .../glm5_next/test_compact_topk_modes.py | 2 +- .../glm5_next/test_direct_nsa_indices.py | 4 +- .../glm5_next/test_fp8_compile_cache.py | 2 +- unit_tests/models/glm5_next/test_fp8_kv.py | 10 +- 6 files changed, 122 insertions(+), 146 deletions(-) rename lightllm/models/glm5_next/triton_kernel/{context_flashattention_nopad_fp8.py => fp8_prefill.py} (56%) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index cb17eb4253..5f9e2170f6 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -2,21 +2,31 @@ import weakref import torch +import triton from .glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState from ..base_att import BaseAttBackend +# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). +KV_GATHER_CHUNK_SIZE = 2097152 + + +def allocate_prefill_workspace(q, topk_indices, q_token_capacity=None): + capacity = q.shape[0] if q_token_capacity is None else q_token_capacity + return { + "chunk_topk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), + "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), + } + + class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) self._prefill_workspaces = [None, None] def get_prefill_workspace(self, infer_state, q, topk_indices): - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( - allocate_prefill_workspace, - ) - buffers = self._prefill_workspaces[infer_state.microbatch_index] q_token_num, topk_size = topk_indices.shape if ( @@ -49,8 +59,14 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): + @torch.no_grad() def _nsa_prefill_att(self, q, packed_kv, att_control): - from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8 + from sgl_kernel.flash_attn import flash_attn_with_kvcache + from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( + _gather_prefill_kv_kernel, + _compact_topk_indices_kernel, + _merge_att_output_kernel, + ) nsa_prefill_dict = att_control.nsa_prefill_dict infer_state = self.infer_state @@ -62,20 +78,98 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): ) if not graph_guard and not torch.cuda.is_current_stream_capturing(): workspace = self.backend.get_prefill_workspace(infer_state, q, nsa_prefill_dict["topk_indices"]) - return context_attention_fwd_fp8( - q=q, - packed_kv=packed_kv, - prefill_cache_kv=nsa_prefill_dict["prefill_cache_kv"], - q_kv_start_loc=self.q_kv_start_loc, - topk_indices=nsa_prefill_dict["topk_indices"], - req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, - b_req_idx=infer_state.b_req_idx, - b1_cu_kv_seq_len=infer_state.b1_cu_kv_seq_len, - b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, - total_token_num=infer_state.total_token_num, - softmax_scale=nsa_prefill_dict["softmax_scale"], - workspace=workspace, - ) + q_token_num, q_head_num, _ = q.shape + prefill_cache_kv = nsa_prefill_dict["prefill_cache_kv"] + topk_indices = nsa_prefill_dict["topk_indices"] + softmax_scale = nsa_prefill_dict["softmax_scale"] + req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + b_req_idx = infer_state.b_req_idx + b1_cu_kv_seq_len = infer_state.b1_cu_kv_seq_len + b1_cu_q_seq_len = infer_state.b1_cu_q_seq_len + total_token_num = infer_state.total_token_num + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) + workspace = workspace or allocate_prefill_workspace(q, topk_indices) + chunk_topk_indices = workspace["chunk_topk_indices"] + chunk_cache_seqlens = workspace["chunk_cache_seqlens"] + cu_seqlens_q = workspace["cu_seqlens_q"] + # Keep this potentially 2 GiB buffer local to a DSA layer. + has_history = total_token_num > q_token_num + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) if has_history else None + if not has_history and prefill_cache_kv.stride(-1) != 1: + prefill_cache_kv = prefill_cache_kv.contiguous() + output = softmax_lse = None + for chunk_start in range(0, total_token_num, gather_chunk_size): + chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) + if has_history: + _gather_prefill_kv_kernel[(chunk_token_num // 2 + b_req_idx.numel(),)]( + packed_kv, + prefill_cache_kv, + gathered_kv, + req_to_token_indexs, + b_req_idx, + b1_cu_kv_seq_len, + b1_cu_q_seq_len, + chunk_start, + chunk_token_num, + packed_kv.stride(0), + prefill_cache_kv.stride(0), + prefill_cache_kv.stride(2), + req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), + req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), + batch_size=b_req_idx.numel(), + BLOCK_ROWS=2, + num_warps=2, + ) + chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) + else: + # Current BF16 KV already has the complete request layout. + chunk_kv = prefill_cache_kv[chunk_start : chunk_start + chunk_token_num].view(-1, 1, 1, 512) + _compact_topk_indices_kernel[(q_token_num,)]( + topk_indices, + self.q_kv_start_loc, + chunk_topk_indices, + chunk_cache_seqlens, + chunk_start, + chunk_start + chunk_token_num, + topk_indices.stride(0), + topk_indices.stride(1), + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + ) + chunk_output, chunk_softmax_lse, *_ = flash_attn_with_kvcache( + # only_qv ignores Q/K; these views avoid allocating placeholder tensors. + q=q[..., :64], + k_cache=chunk_kv[..., :64], + qv=q, + v_cache=chunk_kv, + page_table=chunk_topk_indices, + cache_seqlens=chunk_cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + num_splits=1, + return_softmax_lse=True, + ) + if chunk_start == 0: + output, softmax_lse = chunk_output, chunk_softmax_lse + else: + # Keep the online reduction in FP32 across all remaining chunks. + output = output.float() + _merge_att_output_kernel[(q_token_num, triton.cdiv(q_head_num, 4))]( + output, + softmax_lse, + chunk_output, + chunk_softmax_lse, + chunk_softmax_lse.stride(1), + chunk_softmax_lse.stride(0), + q_head_num=q_head_num, + q_token_num=q_token_num, + BLOCK_HEADS=4, + ) + del chunk_output, chunk_softmax_lse + return output.to(q.dtype) @dataclasses.dataclass diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/fp8_prefill.py similarity index 56% rename from lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py rename to lightllm/models/glm5_next/triton_kernel/fp8_prefill.py index f574d7adbb..befde8e67f 100644 --- a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py +++ b/lightllm/models/glm5_next/triton_kernel/fp8_prefill.py @@ -1,12 +1,7 @@ -import torch import triton import triton.language as tl -# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). -KV_GATHER_CHUNK_SIZE = 2097152 - - @triton.jit(do_not_specialize=["chunk_token_num"]) def _gather_prefill_kv_kernel( packed_kv, @@ -146,116 +141,3 @@ def _merge_att_output_kernel( tl.where(denominator > 0, maximum + tl.log(denominator), -float("inf")), head < q_head_num, ) - - -def allocate_prefill_workspace(q, topk_indices, q_token_capacity=None): - capacity = q.shape[0] if q_token_capacity is None else q_token_capacity - return { - "chunk_topk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), - "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), - "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), - } - - -@torch.no_grad() -def context_attention_fwd_fp8( - q, - packed_kv, - prefill_cache_kv, - q_kv_start_loc, - topk_indices, - req_to_token_indexs, - b_req_idx, - b1_cu_kv_seq_len, - b1_cu_q_seq_len, - total_token_num, - softmax_scale, - workspace=None, -): - """Run GLM sparse prefill with FlashAttention and request-table KV gathering.""" - q_token_num, q_head_num, _ = q.shape - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) - workspace = workspace or allocate_prefill_workspace(q, topk_indices) - chunk_topk_indices = workspace["chunk_topk_indices"] - chunk_cache_seqlens = workspace["chunk_cache_seqlens"] - cu_seqlens_q = workspace["cu_seqlens_q"] - # Keep this potentially 2 GiB buffer local to a DSA layer. - has_history = total_token_num > q_token_num - gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) if has_history else None - if not has_history and prefill_cache_kv.stride(-1) != 1: - prefill_cache_kv = prefill_cache_kv.contiguous() - output = softmax_lse = None - for chunk_start in range(0, total_token_num, gather_chunk_size): - chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) - if has_history: - _gather_prefill_kv_kernel[(chunk_token_num // 2 + b_req_idx.numel(),)]( - packed_kv, - prefill_cache_kv, - gathered_kv, - req_to_token_indexs, - b_req_idx, - b1_cu_kv_seq_len, - b1_cu_q_seq_len, - chunk_start, - chunk_token_num, - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), - req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), - batch_size=b_req_idx.numel(), - BLOCK_ROWS=2, - num_warps=2, - ) - chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) - else: - # Current BF16 KV already has the complete request layout. - chunk_kv = prefill_cache_kv[chunk_start : chunk_start + chunk_token_num].view(-1, 1, 1, 512) - _compact_topk_indices_kernel[(q_token_num,)]( - topk_indices, - q_kv_start_loc, - chunk_topk_indices, - chunk_cache_seqlens, - chunk_start, - chunk_start + chunk_token_num, - topk_indices.stride(0), - topk_indices.stride(1), - TOPK=topk_indices.shape[1], - BLOCK=triton.next_power_of_2(topk_indices.shape[1]), - ) - chunk_output, chunk_softmax_lse, *_ = flash_attn_with_kvcache( - # only_qv ignores Q/K; these views avoid allocating placeholder tensors. - q=q[..., :64], - k_cache=chunk_kv[..., :64], - qv=q, - v_cache=chunk_kv, - page_table=chunk_topk_indices, - cache_seqlens=chunk_cache_seqlens, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=1, - softmax_scale=softmax_scale, - causal=False, - only_qv=True, - num_splits=1, - return_softmax_lse=True, - ) - if chunk_start == 0: - output, softmax_lse = chunk_output, chunk_softmax_lse - else: - # Keep the online reduction in FP32 across all remaining chunks. - output = output.float() - _merge_att_output_kernel[(q_token_num, triton.cdiv(q_head_num, 4))]( - output, - softmax_lse, - chunk_output, - chunk_softmax_lse, - chunk_softmax_lse.stride(1), - chunk_softmax_lse.stride(0), - q_head_num=q_head_num, - q_token_num=q_token_num, - BLOCK_HEADS=4, - ) - del chunk_output, chunk_softmax_lse - return output.to(q.dtype) diff --git a/unit_tests/models/glm5_next/test_compact_topk_modes.py b/unit_tests/models/glm5_next/test_compact_topk_modes.py index 7f82bea6e3..973ef62c07 100644 --- a/unit_tests/models/glm5_next/test_compact_topk_modes.py +++ b/unit_tests/models/glm5_next/test_compact_topk_modes.py @@ -2,7 +2,7 @@ import torch from triton.runtime.jit import JITFunction -from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import _compact_topk_indices_kernel +from lightllm.models.glm5_next.triton_kernel.fp8_prefill import _compact_topk_indices_kernel pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index cc45742bfc..0891c42827 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -15,7 +15,7 @@ ) from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params @@ -139,7 +139,7 @@ def forbidden(*args, **kwargs): ], ) def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): - monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) torch.manual_seed(1685) batch = min(queries, 2) diff --git a/unit_tests/models/glm5_next/test_fp8_compile_cache.py b/unit_tests/models/glm5_next/test_fp8_compile_cache.py index e2ece2270c..6a07b19811 100644 --- a/unit_tests/models/glm5_next/test_fp8_compile_cache.py +++ b/unit_tests/models/glm5_next/test_fp8_compile_cache.py @@ -2,7 +2,7 @@ import torch from lightllm.models.glm5_next.triton_kernel import ( - context_flashattention_nopad_fp8 as fp8, + fp8_prefill as fp8, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 161cb5b20b..ce67f15656 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -22,8 +22,8 @@ ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill -from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( _gather_prefill_kv_kernel, ) @@ -319,7 +319,7 @@ def check(actual): @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("chunk_size", [3, 7, 32768]) def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): - monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) torch.manual_seed(53) prefix = 4 if cached else 0 case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) @@ -342,7 +342,7 @@ def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chun def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): from sgl_kernel.flash_attn import flash_attn_with_kvcache - monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv @@ -362,7 +362,7 @@ def check_kv_storage(*args, **kwargs): @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): - monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", 129) + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 129) torch.manual_seed(53) case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") From 8cd61293d4d714be90c636753035c14579253a1f Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:25:15 +0000 Subject: [PATCH 09/10] refactor(glm): allocate prefill workspaces through the backend --- .../basemodel/attention/nsa/fp8_glm5_next.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index 5f9e2170f6..3f8136e1c3 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -12,20 +12,21 @@ KV_GATHER_CHUNK_SIZE = 2097152 -def allocate_prefill_workspace(q, topk_indices, q_token_capacity=None): - capacity = q.shape[0] if q_token_capacity is None else q_token_capacity - return { - "chunk_topk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), - "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), - "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), - } - - class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) self._prefill_workspaces = [None, None] + def allocate_prefill_workspace(self, q, topk_indices, q_token_capacity=None): + capacity = q.shape[0] if q_token_capacity is None else q_token_capacity + return { + "chunk_topk_indices": torch.empty( + (capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device + ), + "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), + } + def get_prefill_workspace(self, infer_state, q, topk_indices): buffers = self._prefill_workspaces[infer_state.microbatch_index] q_token_num, topk_size = topk_indices.shape @@ -35,7 +36,7 @@ def get_prefill_workspace(self, infer_state, q, topk_indices): or buffers["chunk_topk_indices"].shape[1] != topk_size or buffers["chunk_topk_indices"].dtype != topk_indices.dtype ): - buffers = allocate_prefill_workspace( + buffers = self.allocate_prefill_workspace( q, topk_indices, q_token_capacity=max(self.model.batch_max_tokens or 0, q_token_num) ) self._prefill_workspaces[infer_state.microbatch_index] = buffers @@ -88,7 +89,7 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): b1_cu_q_seq_len = infer_state.b1_cu_q_seq_len total_token_num = infer_state.total_token_num gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) - workspace = workspace or allocate_prefill_workspace(q, topk_indices) + workspace = workspace or self.backend.allocate_prefill_workspace(q, topk_indices) chunk_topk_indices = workspace["chunk_topk_indices"] chunk_cache_seqlens = workspace["chunk_cache_seqlens"] cu_seqlens_q = workspace["cu_seqlens_q"] From 645114d95b857bb79bf1592844c89a7e751ad0dd Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 10 Nov 2026 15:05:47 +0000 Subject: [PATCH 10/10] perf(glm): reserve reusable prefill workspaces before KV sizing --- .../basemodel/attention/nsa/fp8_glm5_next.py | 81 ++++++------ lightllm/models/glm5_next/model.py | 10 ++ lightllm/models/glm5_next_mtp/model.py | 8 ++ .../glm5_next/test_direct_nsa_indices.py | 31 ++++- unit_tests/models/glm5_next/test_fp8_kv.py | 26 +++- .../glm5_next/test_prefill_workspace.py | 116 ++++++++++++++++++ 6 files changed, 226 insertions(+), 46 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_prefill_workspace.py diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py index 3f8136e1c3..261c2a16c1 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py @@ -8,40 +8,56 @@ from ..base_att import BaseAttBackend -# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). +# Cap each microbatch's BF16 KV workspace at 2 GiB (2097152 × 512 × 2 bytes). KV_GATHER_CHUNK_SIZE = 2097152 class Fp8Glm5NextSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - self._prefill_workspaces = [None, None] - - def allocate_prefill_workspace(self, q, topk_indices, q_token_capacity=None): - capacity = q.shape[0] if q_token_capacity is None else q_token_capacity - return { - "chunk_topk_indices": torch.empty( - (capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device - ), - "chunk_cache_seqlens": torch.empty(capacity, dtype=torch.int32, device=q.device), - "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), - } - - def get_prefill_workspace(self, infer_state, q, topk_indices): + self._prefill_workspaces = self.allocate_prefill_workspaces(model) + + @classmethod + def allocate_prefill_workspaces(cls, model): + regions = 1 + int(model.args.enable_prefill_microbatch_overlap) + q_token_capacity = model.batch_max_tokens + topk_size = triton.cdiv(model.config["index_topk"] + 3, 128) * 128 + kv_token_capacity = min(KV_GATHER_CHUNK_SIZE, model.max_req_num * model.max_seq_length + q_token_capacity) + if model.args.run_mode == "decode": + kv_token_capacity = min(kv_token_capacity, q_token_capacity) + prefix = f"glm5_fp8_prefill_{model.att_state_role}" + gathered_kv = cls.get_gpu_workspace_buffer( + prefix + "_kv", regions * kv_token_capacity * 512, torch.bfloat16 + ).view(regions, kv_token_capacity, 1, 512) + chunk_topk_indices = cls.get_gpu_workspace_buffer( + prefix + "_topk", regions * q_token_capacity * topk_size, torch.int32 + ).view(regions, q_token_capacity, topk_size) + chunk_cache_seqlens = cls.get_gpu_workspace_buffer( + prefix + "_seqlens", regions * q_token_capacity, torch.int32 + ).view(regions, q_token_capacity) + cu_seqlens_q = cls.get_gpu_workspace_buffer( + prefix + "_cu_q", regions * (q_token_capacity + 1), torch.int32 + ).view(regions, q_token_capacity + 1) + cu_seqlens_q.copy_(torch.arange(q_token_capacity + 1, dtype=torch.int32, device=cu_seqlens_q.device)) + return [ + { + "gathered_kv": gathered_kv[index], + "chunk_topk_indices": chunk_topk_indices[index], + "chunk_cache_seqlens": chunk_cache_seqlens[index], + "cu_seqlens_q": cu_seqlens_q[index], + } + for index in range(regions) + ] + + def get_prefill_workspace(self, infer_state, topk_indices): buffers = self._prefill_workspaces[infer_state.microbatch_index] q_token_num, topk_size = topk_indices.shape - if ( - buffers is None - or buffers["chunk_topk_indices"].shape[0] < q_token_num - or buffers["chunk_topk_indices"].shape[1] != topk_size - or buffers["chunk_topk_indices"].dtype != topk_indices.dtype - ): - buffers = self.allocate_prefill_workspace( - q, topk_indices, q_token_capacity=max(self.model.batch_max_tokens or 0, q_token_num) - ) - self._prefill_workspaces[infer_state.microbatch_index] = buffers + assert q_token_num <= buffers["chunk_cache_seqlens"].numel() + assert topk_size <= buffers["chunk_topk_indices"].shape[1] + chunk_topk_indices = buffers["chunk_topk_indices"].view(-1) return { - "chunk_topk_indices": buffers["chunk_topk_indices"][:q_token_num], + "gathered_kv": buffers["gathered_kv"], + "chunk_topk_indices": chunk_topk_indices[: q_token_num * topk_size].view(q_token_num, topk_size), "chunk_cache_seqlens": buffers["chunk_cache_seqlens"][:q_token_num], "cu_seqlens_q": buffers["cu_seqlens_q"][: q_token_num + 1], } @@ -71,14 +87,6 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): nsa_prefill_dict = att_control.nsa_prefill_dict infer_state = self.infer_state - workspace = None - graph_guard = ( - infer_state.is_cuda_graph_capture - or self.backend.model.prefill_graph is not None - or hasattr(infer_state, "prefill_cuda_graph_exe_list") - ) - if not graph_guard and not torch.cuda.is_current_stream_capturing(): - workspace = self.backend.get_prefill_workspace(infer_state, q, nsa_prefill_dict["topk_indices"]) q_token_num, q_head_num, _ = q.shape prefill_cache_kv = nsa_prefill_dict["prefill_cache_kv"] topk_indices = nsa_prefill_dict["topk_indices"] @@ -88,14 +96,13 @@ def _nsa_prefill_att(self, q, packed_kv, att_control): b1_cu_kv_seq_len = infer_state.b1_cu_kv_seq_len b1_cu_q_seq_len = infer_state.b1_cu_q_seq_len total_token_num = infer_state.total_token_num - gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, total_token_num) - workspace = workspace or self.backend.allocate_prefill_workspace(q, topk_indices) + workspace = self.backend.get_prefill_workspace(infer_state, topk_indices) + gathered_kv = workspace["gathered_kv"] + gather_chunk_size = min(gathered_kv.shape[0], total_token_num) chunk_topk_indices = workspace["chunk_topk_indices"] chunk_cache_seqlens = workspace["chunk_cache_seqlens"] cu_seqlens_q = workspace["cu_seqlens_q"] - # Keep this potentially 2 GiB buffer local to a DSA layer. has_history = total_token_num > q_token_num - gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) if has_history else None if not has_history and prefill_cache_kv.stride(-1) != 1: prefill_cache_kv = prefill_cache_kv.contiguous() output = softmax_lse = None diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 5b5235c76e..e15efe0143 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -178,6 +178,16 @@ def create_full_att_states( ) return [state] + def prepare_attention_states(self) -> None: + super().prepare_attention_states() + if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( + self.att_state_role + ): + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + + # Reserve fixed scratch space before the shared KV capacity is profiled. + Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _init_att_backend(self): from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index a59c6f1ab1..6c5f80a414 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -153,6 +153,14 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config.update(num_hidden_layers=1, n_layer=1, mhc=False) return config + def prepare_attention_states(self) -> None: + super().prepare_attention_states() + if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( + self.att_state_role + ): + # Reserve fixed scratch space before the shared KV capacity is profiled. + Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _init_att_backend(self): backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index 0891c42827..d0e6dcb9bd 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -105,7 +105,19 @@ def forbidden(*args, **kwargs): raise AssertionError("GLM must not generate or copy a full ragged index") monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) - model = type("Model", (), {"config": {"index_kpool": 4, "index_topk": 2048}})() + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + model = type( + "Model", + (), + { + "config": {"index_kpool": 4, "index_topk": 2048}, + "batch_max_tokens": 4, + "max_req_num": 1, + "max_seq_length": 32, + "att_state_role": "decode_test", + "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="decode"), + }, + )() backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) assert not hasattr(backend, "ragged_mem_buffers") infer = InferStateInfo() @@ -183,7 +195,20 @@ def refresh_physical(): infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) infer.total_token_num = total_kv infer.req_manager = SimpleNamespace(req_to_token_indexs=table) - model = type("Model", (), {"batch_max_tokens": queries, "prefill_graph": None})() + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + model = type( + "Model", + (), + { + "batch_max_tokens": queries, + "prefill_graph": None, + "max_req_num": batch, + "max_seq_length": max(seqs), + "att_state_role": "indices_test", + "config": {"index_topk": 2048}, + "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="normal"), + }, + )() backend = Fp8Glm5NextSparseAttBackend(model) state = backend.create_att_prefill_state(infer) state.init_state() @@ -216,7 +241,7 @@ def check(output): check(run()) workspace = backend._prefill_workspaces[0] if workspace is not None: - assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} + assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = run() diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index ce67f15656..a6a5c7c9f2 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -31,18 +31,29 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.fixture(autouse=True) +def current_cuda_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + def _make_infer_state(**attributes): state = InferStateInfo() vars(state).update(attributes) return state -def _make_workspace_backend(batch_max_tokens): +def _make_workspace_backend(batch_max_tokens, overlap=False): backend = object.__new__(Fp8Glm5NextSparseAttBackend) backend.model = SimpleNamespace( - batch_max_tokens=batch_max_tokens, prefill_graph=None, config={"index_topk": 2048, "index_kpool": 4} + batch_max_tokens=batch_max_tokens, + prefill_graph=None, + config={"index_topk": 2048, "index_kpool": 4}, + max_req_num=1, + max_seq_length=fp8_glm5_next.KV_GATHER_CHUNK_SIZE, + att_state_role="unit_test", + args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode="normal"), ) - backend._prefill_workspaces = [None, None] + backend._prefill_workspaces = backend.allocate_prefill_workspaces(backend.model) return backend @@ -156,10 +167,13 @@ def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) - model.args = SimpleNamespace(llm_kv_type=kv_type) + model.args = SimpleNamespace(llm_kv_type=kv_type, enable_prefill_microbatch_overlap=False, run_mode="normal") model.config = {"index_topk": 2048, "index_kpool": 4} model.graph_max_batch_size = 4 model.max_seq_length = 128 + model.max_req_num = 2 + model.batch_max_tokens = 4 + model.mtp_previous_draft_models = [] model._init_att_backend() infer_state = InferStateInfo() prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) @@ -990,13 +1004,13 @@ def direct_must_not_run(*args, **kwargs): @pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(num_heads): +def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_replay(num_heads): torch.manual_seed(53450) backend = _make_workspace_backend(66) case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) case.check(case.run()) workspace = backend._prefill_workspaces[0] - assert set(workspace) == {"chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} + assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} pointers = {name: value.data_ptr() for name, value in workspace.items()} for _ in range(10): infer, state = case.new_state() diff --git a/unit_tests/models/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py new file mode 100644 index 0000000000..0d3a85d965 --- /dev/null +++ b/unit_tests/models/glm5_next/test_prefill_workspace.py @@ -0,0 +1,116 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from unit_tests.models.glm5_next.test_fp8_kv import _PrefillCase, _make_workspace_backend, _capture_cuda_graph + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def current_cuda_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("run_mode", ["normal", "decode"]) +def test_prefill_workspace_is_reserved_once_and_reused_at_runtime(monkeypatch, overlap, run_mode): + model = SimpleNamespace( + args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode=run_mode), + batch_max_tokens=8, + max_req_num=2, + max_seq_length=13, + config={"index_topk": 2048}, + att_state_role=f"workspace_test_{run_mode}_{overlap}", + ) + backend_class = fp8_glm5_next.Fp8Glm5NextSparseAttBackend + reserved = backend_class.allocate_prefill_workspaces(model) + pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in reserved] + backend = object.__new__(backend_class) + backend.__init__(model) + assert [ + {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces + ] == pointers + assert len(reserved) == 1 + int(overlap) + capacity = 8 if run_mode == "decode" else 34 + assert reserved[0]["gathered_kv"].shape == (capacity, 1, 512) + assert reserved[0]["chunk_topk_indices"].shape == (8, 2176) + + def forbidden(*args, **kwargs): + raise AssertionError("Runtime must reuse the workspace reserved before KV capacity planning") + + monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) + monkeypatch.setattr(backend, "get_gpu_workspace_buffer", forbidden) + for slot in range(1 + int(overlap)): + for q_tokens in (1, 8, 3): + topk = torch.empty((q_tokens, 2176), dtype=torch.int32, device="cuda") + views = backend.get_prefill_workspace(SimpleNamespace(microbatch_index=slot), topk) + assert {key: value.data_ptr() for key, value in views.items()} == pointers[slot] + torch.testing.assert_close( + views["cu_seqlens_q"], torch.arange(q_tokens + 1, dtype=torch.int32, device="cuda") + ) + + +def test_prefill_microbatches_use_independent_scratch_during_graph_replay(monkeypatch): + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 7) + backend = _make_workspace_backend(8, overlap=True) + first = _PrefillCase([15], [3], 8, backend=backend) + second = _PrefillCase([17], [3], 8, backend=backend) + second.infer_state.microbatch_index = 1 + first.check(first.run()) + second.check(second.run()) + pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces] + assert all(pointers[0][key] != pointers[1][key] for key in pointers[0]) + + def forbidden(*args, **kwargs): + raise AssertionError("Attention must not allocate a new prefill workspace") + + monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) + first_graph, first_output = _capture_cuda_graph(first.run) + second_graph, second_output = _capture_cuda_graph(second.run) + first_stream, second_stream = torch.cuda.Stream(), torch.cuda.Stream() + first.mutate() + second.mutate() + for _ in range(3): + first_stream.wait_stream(torch.cuda.current_stream()) + second_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(first_stream): + first_graph.replay() + with torch.cuda.stream(second_stream): + second_graph.replay() + torch.cuda.current_stream().wait_stream(first_stream) + torch.cuda.current_stream().wait_stream(second_stream) + first.check(first_output) + second.check(second_output) + assert [ + {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces + ] == pointers + + +@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_glm_reserves_prefill_workspace_during_resource_preparation(monkeypatch, model_class, kv_type): + model = object.__new__(model_class) + model.args = SimpleNamespace(llm_kv_type=kv_type) + model.mtp_previous_draft_models = [] + model.att_state_manager = SimpleNamespace(get_full_attention_states=lambda role: [object()]) + calls = [] + + def prepare(self): + assert self is model + calls.append("prepare_states") + + def reserve(cls, prepared_model): + assert prepared_model is model + calls.append("reserve_workspace") + + monkeypatch.setattr(TpPartBaseModel, "prepare_attention_states", prepare) + monkeypatch.setattr(fp8_glm5_next.Fp8Glm5NextSparseAttBackend, "allocate_prefill_workspaces", classmethod(reserve)) + model.prepare_attention_states() + assert calls == (["prepare_states", "reserve_workspace"] if kv_type == "fp8kv_dsa" else ["prepare_states"])