perf(glm): reuse inference metadata and streamline FP8 prefill

lightllm/common/basemodel/attention/nsa/glm5_next.py

@@ -3,14 +3,10 @@
 
 import torch
 
-from .flashmla_sparse import (
-    NsaFlashMlaSparseAttBackend,
-    NsaFlashMlaSparsePrefillAttState,
-    NsaFlashMlaSparseDecodeAttState,
-)
+from ..base_att import AttControl, BaseAttBackend, BasePrefillAttState, BaseDecodeAttState
 
 
-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))
 
@@ -19,57 +15,86 @@ def create_att_decode_state(self, infer_state):
 
 
 @dataclasses.dataclass
-class Glm5NextSparsePrefillState(NsaFlashMlaSparsePrefillAttState):
+class Glm5NextSparsePrefillState(BasePrefillAttState):
+    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
+
+        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, 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(NsaFlashMlaSparseDecodeAttState):
+class Glm5NextSparseDecodeState(BaseDecodeAttState):
+    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):
-        super().init_state()
-        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))
 
+    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
 
-        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,
         )

lightllm/models/glm5_next/indexer.py

@@ -13,133 +13,159 @@ 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,
-            starts=att_state.ks,
-            ragged=att_state.ragged_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,
+            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,
+            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:
-            return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True)
+        if infer_state.max_kv_seq_len <= self.index_topk:
+            return expand_topk(
+                None,
+                att_state.q_kv_seq_len,
+                att_state.q_req_idx,
+                infer_state.req_manager.req_to_token_indexs,
+                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)
-        return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk)
+            topk_pool_indices = self._get_decode_indices(q_fp8, weights, indexer_k, infer_state, att_state, max_pools)
+        return expand_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_chunk_size(self, max_pools):
         # Account for DeepGEMM padding within the 512 MiB workspace.
-        query_tile = 128 // self.heads
+        query_tile = 128 // self.index_n_heads
         logits_row_bytes = triton.cdiv(max_pools, 256) * 256 * 4
         query_tiles = max(1, 512 * 1024 ** 2 // (query_tile * logits_row_bytes))
         return query_tiles * query_tile
 
-    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)
         chunk_size = self._get_prefill_chunk_size(max_pools)
         import deep_gemm
 
         for start in range(0, q_fp8.shape[0], chunk_size):
             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])
+            self.select_topk_indices(logits, q_pool_seq_len[start:end], topk_pool_indices[start:end])
             # Free before allocating the next chunk.
             del logits
-        return groups
+        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

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