fix: break attention state ownership cycles

lightllm/common/basemodel/attention/base_att.py

@@ -1,4 +1,5 @@
 import threading
+import weakref
 
 import torch
 from abc import ABC, abstractmethod
@@ -55,10 +56,10 @@ def get_gpu_workspace_buffer(key_name: str, workspace_size: int, dtype: torch.dt
                 BaseAttBackend._workspace_buffers[buffer_key] = workspace_buffer
             return workspace_buffer
 
-    def create_att_prefill_state(self) -> "BasePrefillAttState":
+    def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "BasePrefillAttState":
         raise NotImplementedError("not impl")
 
-    def create_att_decode_state(self) -> "BaseDecodeAttState":
+    def create_att_decode_state(self, infer_state: "InferStateInfo") -> "BaseDecodeAttState":
         raise NotImplementedError("not impl")
 
     def uses_dynamic_spec_verify_layout(self) -> bool:
@@ -115,15 +116,45 @@ class AttControl:
 
 
 @dataclass
-class BasePrefillAttState(ABC):
+class BaseAttState(ABC):
+    """Attention metadata with a weak reference to its owning InferStateInfo."""
 
     backend: BaseAttBackend = None
-    infer_state: "InferStateInfo" = None
+    # InferStateInfo 已强引用本状态，这里反向使用弱引用，避免形成强引用环。
+    # 否则状态及其 Tensor 需等待循环 GC 回收，可能延迟显存复用、增加显存峰值。
+    # InferStateInfo 的生命周期由推理调用方或 CUDA Graph 缓存管理，不由 attention state 延长。
+    _infer_state_weak_ref: Optional[weakref.ReferenceType["InferStateInfo"]] = None
+
+    @property
+    def infer_state(self) -> Optional["InferStateInfo"]:
+        if self._infer_state_weak_ref is None:
+            return None
+        infer_state = self._infer_state_weak_ref()
+        if infer_state is None:
+            raise ReferenceError("attention state's InferStateInfo has been released")
+        return infer_state
+
+    def check_no_infer_state(self) -> None:
+        """Reject InferStateInfo stored directly in any instance attribute.
+
+        This shallow check does not inspect containers or callback closures.
+        """
+        # InferStateInfo imports attention states, so import it only when checking.
+        from lightllm.common.basemodel.infer_struct import InferStateInfo
+
+        for name, value in vars(self).items():
+            if isinstance(value, InferStateInfo):
+                raise RuntimeError(
+                    f"{type(self).__name__}.{name} must not retain InferStateInfo; use _infer_state_weak_ref instead"
+                )
 
     @abstractmethod
-    def init_state(self):
+    def init_state(self) -> None:
         pass
 
+
+@dataclass
+class BasePrefillAttState(BaseAttState):
     @abstractmethod
     def prefill_att(
         self,
@@ -137,15 +168,9 @@ def prefill_att(
 
 
 @dataclass
-class BaseDecodeAttState(ABC):
-    backend: BaseAttBackend = None
-    infer_state: "InferStateInfo" = None
-
-    @abstractmethod
-    def init_state(self):
-        pass
-
+class BaseDecodeAttState(BaseAttState):
     def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"):
+        """Copy tensor metadata into the buffers owned by the captured graph."""
         for attr_name, attr_value in vars(new_state).items():
             if isinstance(attr_value, torch.Tensor):
                 attr_ = getattr(self, attr_name, None)

lightllm/common/basemodel/attention/fa3/fp.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 import triton
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
@@ -76,10 +77,10 @@ def get_page_table_view(self, att_batch_size, max_kv_len, microbatch_index):
         )
 
     def create_att_prefill_state(self, infer_state) -> "Fa3PrefillAttState":
-        return Fa3PrefillAttState(backend=self, infer_state=infer_state)
+        return Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "Fa3DecodeAttState":
-        return Fa3DecodeAttState(backend=self, infer_state=infer_state)
+        return Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/fa3/fp8.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 from ..base_att import AttControl
 from lightllm.utils.sgl_utils import flash_attn_with_kvcache
@@ -24,10 +25,10 @@ def _init_infer_page_size(self):
         )
 
     def create_att_prefill_state(self, infer_state) -> "Fp8Fa3PrefillAttState":
-        return Fp8Fa3PrefillAttState(backend=self, infer_state=infer_state)
+        return Fp8Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "Fp8Fa3DecodeAttState":
-        return Fp8Fa3DecodeAttState(backend=self, infer_state=infer_state)
+        return Fp8Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/fa3/mla.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 import triton
 from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl
@@ -13,10 +14,10 @@
 
 class MlaFa3AttBackend(Fa3AttBackend):
     def create_att_prefill_state(self, infer_state) -> "MlaFa3PrefillAttState":
-        return MlaFa3PrefillAttState(backend=self, infer_state=infer_state)
+        return MlaFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "MlaFa3DecodeAttState":
-        return MlaFa3DecodeAttState(backend=self, infer_state=infer_state)
+        return MlaFa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/fa3/neo.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import inspect
 import warnings
 
@@ -40,7 +41,7 @@
 
 class NeoFa3AttBackend(Fa3AttBackend):
     def create_att_prefill_state(self, infer_state) -> "NeoFa3PrefillAttState":
-        return NeoFa3PrefillAttState(backend=self, infer_state=infer_state)
+        return NeoFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/flashinfer/fp.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 import triton
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
@@ -43,10 +44,10 @@ def _init_infer_page_size(self):
         )
 
     def create_att_prefill_state(self, infer_state) -> "FlashInferPrefillAttState":
-        return FlashInferPrefillAttState(backend=self, infer_state=infer_state)
+        return FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "FlashInferDecodeAttState":
-        return FlashInferDecodeAttState(backend=self, infer_state=infer_state)
+        return FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass
@@ -151,7 +152,7 @@ class FlashInferDecodeAttState(BaseDecodeAttState):
     decode_wrapper: object = None
 
     def _should_init_decode_wrapper(self) -> bool:
-        return should_init_decode_wrapper(self.backend.model, self.infer_state)
+        return should_init_decode_wrapper(self.infer_state)
 
     def init_state(self):
         import flashinfer

lightllm/common/basemodel/attention/flashinfer/fp8.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 from ..base_att import AttControl
 from .fp import FlashInferAttBackend, FlashInferPrefillAttState, FlashInferDecodeAttState
@@ -28,10 +29,10 @@ def _init_infer_page_size(self):
         )
 
     def create_att_prefill_state(self, infer_state) -> "Fp8FlashInferPrefillAttState":
-        return Fp8FlashInferPrefillAttState(backend=self, infer_state=infer_state)
+        return Fp8FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "Fp8FlashInferDecodeAttState":
-        return Fp8FlashInferDecodeAttState(backend=self, infer_state=infer_state)
+        return Fp8FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/flashinfer/mla.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 import triton
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
@@ -58,10 +59,10 @@ def _init_infer_page_size(self):
         )
 
     def create_att_prefill_state(self, infer_state) -> "MlaFlashInferPrefillAttState":
-        return MlaFlashInferPrefillAttState(backend=self, infer_state=infer_state)
+        return MlaFlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "MlaFlashInferDecodeAttState":
-        return MlaFlashInferDecodeAttState(backend=self, infer_state=infer_state)
+        return MlaFlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass
@@ -136,7 +137,7 @@ class MlaFlashInferDecodeAttState(BaseDecodeAttState):
     decode_wrapper: object = None
 
     def _should_init_decode_wrapper(self) -> bool:
-        return should_init_decode_wrapper(self.backend.model, self.infer_state)
+        return should_init_decode_wrapper(self.infer_state)
 
     def init_state(self):
         import flashinfer

lightllm/common/basemodel/attention/flashinfer/utils.py

@@ -1,16 +1,17 @@
-def should_init_decode_wrapper(model, infer_state) -> bool:
-    graph = getattr(model, "graph", None)
-    if graph is None:
-        # Cuda graph is disabled, so this state owns a normal decode wrapper.
-        return True
-
-    if infer_state.is_cuda_graph:
-        # This is the captured graph state; it must create the wrapper captured by replay.
-        return True
-
-    if not graph.can_run(infer_state.batch_size, infer_state.max_kv_seq_len):
-        # Cuda graph is enabled, but this input falls outside graph limits and runs normally.
-        return True
-
-    # This is a temporary replay state. Its tensors are copied into the captured graph state.
-    return False
+def should_init_decode_wrapper(infer_state) -> bool:
+    match (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay):
+        case (False, False):
+            # eager 推理（包括 graph 不可用时的回退）会直接执行 attention，需要创建 wrapper。
+            return True
+        case (True, False):
+            # capture 需要创建 wrapper，并随捕获状态保存在图缓存中，供后续 replay 使用。
+            return True
+        case (False, True):
+            # replay 的临时状态只准备待复制的 metadata，实际执行使用图缓存中的 wrapper。
+            return False
+        case _:
+            # 两个标记必须是互斥的布尔值，不能同时处于 capture 和 replay 状态。
+            raise ValueError(
+                f"Invalid CUDA Graph flags: capture={infer_state.is_cuda_graph_capture}, "
+                f"replay={infer_state.is_cuda_graph_replay}"
+            )

lightllm/common/basemodel/attention/linear/gdn.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
@@ -99,10 +100,10 @@ def _rearrange_mixed_qkv(self, mixed_qkv, decode=False):
             return query, key, value
 
     def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState":
-        return LinearAttPrefillAttState(backend=self, infer_state=infer_state)
+        return LinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState":
-        return LinearAttDecodeAttState(backend=self, infer_state=infer_state)
+        return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/linear/kda.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import dataclasses
+import weakref
 from typing import TYPE_CHECKING
 
 import torch
@@ -47,10 +48,10 @@ def __init__(self, model: "TpPartBaseModel"):
         self.mtp_step = model.args.mtp_step
 
     def create_att_prefill_state(self, infer_state: "InferStateInfo"):
-        return KDAPrefillAttState(backend=self, infer_state=infer_state)
+        return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state: "InferStateInfo"):
-        return KDADecodeAttState(backend=self, infer_state=infer_state)
+        return KDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def split_qkv(self, mixed_qkv: torch.Tensor):
         return mixed_qkv.split(self.tp_hidden_size, dim=-1)

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

@@ -2,6 +2,7 @@
 # Uses sgl_kernel.flash_mla and sgl_kernel.flash_attn from the sglang kernel library.
 
 import dataclasses
+import weakref
 import torch
 from typing import Tuple, TYPE_CHECKING
 
@@ -22,10 +23,10 @@ def __init__(self, model):
         ]
 
     def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparsePrefillAttState":
-        return NsaFlashMlaSparsePrefillAttState(backend=self, infer_state=infer_state)
+        return NsaFlashMlaSparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparseDecodeAttState":
-        return NsaFlashMlaSparseDecodeAttState(backend=self, infer_state=infer_state)
+        return NsaFlashMlaSparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

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

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 from typing import TYPE_CHECKING, Tuple
 
@@ -19,10 +20,10 @@ def __init__(self, model):
         ]
 
     def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparsePrefillAttState":
-        return NsaFlashMlaFp8SparsePrefillAttState(backend=self, infer_state=infer_state)
+        return NsaFlashMlaFp8SparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparseDecodeAttState":
-        return NsaFlashMlaFp8SparseDecodeAttState(backend=self, infer_state=infer_state)
+        return NsaFlashMlaFp8SparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

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

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 
 import torch
 
@@ -11,10 +12,10 @@
 
 class Glm5NextSparseAttBackend(NsaFlashMlaSparseAttBackend):
     def create_att_prefill_state(self, infer_state):
-        return Glm5NextSparsePrefillState(backend=self, infer_state=infer_state)
+        return Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state):
-        return Glm5NextSparseDecodeState(backend=self, infer_state=infer_state)
+        return Glm5NextSparseDecodeState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass

lightllm/common/basemodel/attention/triton/fp.py

@@ -1,4 +1,5 @@
 import dataclasses
+import weakref
 import torch
 from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
 from typing import Optional
@@ -7,10 +8,10 @@
 
 class TritonAttBackend(BaseAttBackend):
     def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState":
-        return TritonPrefillAttState(backend=self, infer_state=infer_state)
+        return TritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
     def create_att_decode_state(self, infer_state) -> "TritonDecodeAttState":
-        return TritonDecodeAttState(backend=self, infer_state=infer_state)
+        return TritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
 
 
 @dataclasses.dataclass