Cache _linear_attn_registry_cache with sentinel (#25436)

This commit is contained in:
fzyzcjy
2026-05-16 09:19:21 +08:00
committed by GitHub
parent 6a1b05d706
commit 0071033ff6
@@ -29,7 +29,7 @@ import uuid
from collections import defaultdict from collections import defaultdict
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Callable, List, Optional, Tuple, Union from typing import Any, Callable, List, Optional, Tuple, Union
import torch import torch
import torch.distributed as dist import torch.distributed as dist
@@ -289,6 +289,8 @@ UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data proces
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_UNSET: Any = object()
def resolve_language_model(model: nn.Module) -> nn.Module: def resolve_language_model(model: nn.Module) -> nn.Module:
model_cls_name = model.__class__.__name__ model_cls_name = model.__class__.__name__
@@ -531,6 +533,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# For hisparse (must be set before initialize() so CUDA graph capture can see it) # For hisparse (must be set before initialize() so CUDA graph capture can see it)
self.hisparse_coordinator = None self.hisparse_coordinator = None
self._linear_attn_registry_cache: Any = _UNSET
# Initialize the model runner # Initialize the model runner
self.initialize(pre_model_load_memory) self.initialize(pre_model_load_memory)
self.check_quantized_moe_compatibility() self.check_quantized_moe_compatibility()
@@ -2313,7 +2317,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
return None return None
def _get_linear_attn_registry_result(self): def _get_linear_attn_registry_result(self):
if not hasattr(self, "_linear_attn_registry_cache"): if self._linear_attn_registry_cache is _UNSET:
self._linear_attn_registry_cache = get_linear_attn_config( self._linear_attn_registry_cache = get_linear_attn_config(
self.model_config.hf_config self.model_config.hf_config
) )