Add Liquid Foundation Model (LFM2) (#16890)
This commit is contained in:
Piotr Mazurek
@@ -12,6 +12,7 @@ from sglang.srt.configs.jet_vlm import JetVLMConfig
|
||||
from sglang.srt.configs.kimi_linear import KimiLinearConfig
|
||||
from sglang.srt.configs.kimi_vl import KimiVLConfig
|
||||
from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
|
||||
from sglang.srt.configs.lfm2 import Lfm2Config
|
||||
from sglang.srt.configs.longcat_flash import LongcatFlashConfig
|
||||
from sglang.srt.configs.nano_nemotron_vl import NemotronH_Nano_VL_V2_Config
|
||||
from sglang.srt.configs.nemotron_h import NemotronHConfig
|
||||
@@ -42,6 +43,7 @@ __all__ = [
|
||||
"DotsVLMConfig",
|
||||
"DotsOCRConfig",
|
||||
"FalconH1Config",
|
||||
"Lfm2Config",
|
||||
"NemotronHConfig",
|
||||
"NemotronH_Nano_VL_V2_Config",
|
||||
"JetNemotronConfig",
|
||||
|
||||
102
python/sglang/srt/configs/lfm2.py
Normal file
102
python/sglang/srt/configs/lfm2.py
Normal file
@@ -0,0 +1,102 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2024 Liquid AI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""LFM2 (Liquid Foundation Model 2) configuration"""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from transformers import CONFIG_MAPPING
|
||||
from transformers import Lfm2Config as HFLfm2Config
|
||||
from transformers.utils import logging
|
||||
|
||||
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class Lfm2Config(HFLfm2Config):
|
||||
"""
|
||||
SGLang configuration for LFM2 models.
|
||||
|
||||
Extends HuggingFace's Lfm2Config with hybrid model properties needed by SGLang.
|
||||
LFM2 uses a hybrid architecture mixing full attention and ShortConv layers.
|
||||
"""
|
||||
|
||||
@property
|
||||
def full_attention_layer_ids(self) -> List[int]:
|
||||
"""Return indices of attention layers for KV cache."""
|
||||
return [i for i, lt in enumerate(self.layer_types) if lt == "full_attention"]
|
||||
|
||||
@property
|
||||
def linear_layer_ids(self) -> List[int]:
|
||||
"""Return indices of conv layers for conv state cache."""
|
||||
return [
|
||||
i for i, lt in enumerate(self.layer_types) if lt in ("conv", "short_conv")
|
||||
]
|
||||
|
||||
@property
|
||||
def mamba_chunk_size(self) -> int:
|
||||
"""Return chunk size for Mamba2 backend. LFM2 doesn't use chunking, return 1."""
|
||||
return 1
|
||||
|
||||
@property
|
||||
def mamba2_cache_params(self) -> Optional[Mamba2CacheParams]:
|
||||
"""
|
||||
Get cache params for HybridReqToTokenPool initialization.
|
||||
|
||||
LFM2 uses ShortConv layers with a small fixed-size cache (kernel_size - 1).
|
||||
Unlike full Mamba2 models, LFM2 only uses the conv state, not SSM temporal state.
|
||||
"""
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
||||
|
||||
conv_layer_ids = self.linear_layer_ids
|
||||
if not conv_layer_ids:
|
||||
return None
|
||||
|
||||
hidden_size = self.hidden_size
|
||||
# conv_L_cache in config is kernel_size (e.g., 3)
|
||||
conv_kernel = int(self.conv_L_cache)
|
||||
L_cache = conv_kernel - 1 # actual cache size (e.g., 2 for kernel=3)
|
||||
|
||||
# get_attention_tp_size() requires initialization, default to 1 if not available
|
||||
try:
|
||||
tp_size = get_attention_tp_size()
|
||||
except (AssertionError, RuntimeError):
|
||||
tp_size = 1
|
||||
|
||||
# For ShortConv layers, we use a simplified Mamba2StateShape
|
||||
# LFM2 doesn't use SSM state (state_size=0), only conv state
|
||||
shape = Mamba2StateShape.create(
|
||||
tp_world_size=tp_size,
|
||||
intermediate_size=hidden_size,
|
||||
n_groups=1, # ShortConv doesn't use grouping
|
||||
num_heads=1, # ShortConv is not multi-head
|
||||
head_dim=hidden_size, # Conv operates on full hidden dim
|
||||
state_size=0, # No SSM temporal state for ShortConv
|
||||
conv_kernel=conv_kernel,
|
||||
)
|
||||
|
||||
# Uses default mamba2_state_dtype() which reads SGLANG_MAMBA_CONV_DTYPE env var
|
||||
# (defaults to bfloat16). Set SGLANG_MAMBA_CONV_DTYPE=float16 for fp16 inference.
|
||||
return Mamba2CacheParams(
|
||||
shape=shape,
|
||||
layers=conv_layer_ids,
|
||||
)
|
||||
|
||||
|
||||
# Override HuggingFace's Lfm2Config with our extended version
|
||||
# Cannot use .register() because lfm2 is already registered by transformers
|
||||
# Directly modify the internal _extra_content dict instead
|
||||
CONFIG_MAPPING._extra_content["lfm2"] = Lfm2Config
|
||||
logger.info("Registered SGLang Lfm2Config to override HuggingFace's version")
|
||||
@@ -10,7 +10,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, etc."""
|
||||
"""Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, LFM2, etc."""
|
||||
|
||||
import os
|
||||
from abc import ABC
|
||||
@@ -41,16 +41,19 @@ class Mamba2StateDType:
|
||||
temporal: torch.dtype
|
||||
|
||||
|
||||
CONV_DTYPE = torch.bfloat16
|
||||
|
||||
|
||||
def mamba2_state_dtype() -> Mamba2StateDType:
|
||||
dtype_map = {
|
||||
"float32": torch.float32,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float16": torch.float16,
|
||||
}
|
||||
ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]]
|
||||
return Mamba2StateDType(conv=CONV_DTYPE, temporal=ssm_dtype)
|
||||
conv_dtype = dtype_map.get(
|
||||
os.environ.get("SGLANG_MAMBA_CONV_DTYPE", "bfloat16"), torch.bfloat16
|
||||
)
|
||||
ssm_dtype = dtype_map.get(
|
||||
os.environ.get("SGLANG_MAMBA_SSM_DTYPE", "float32"), torch.float32
|
||||
)
|
||||
return Mamba2StateDType(conv=conv_dtype, temporal=ssm_dtype)
|
||||
|
||||
|
||||
@dataclass(kw_only=True, frozen=True)
|
||||
|
||||
@@ -19,6 +19,7 @@ from sglang.srt.function_call.glm47_moe_detector import Glm47MoeDetector
|
||||
from sglang.srt.function_call.gpt_oss_detector import GptOssDetector
|
||||
from sglang.srt.function_call.internlm_detector import InternlmDetector
|
||||
from sglang.srt.function_call.kimik2_detector import KimiK2Detector
|
||||
from sglang.srt.function_call.lfm2_detector import Lfm2Detector
|
||||
from sglang.srt.function_call.llama32_detector import Llama32Detector
|
||||
from sglang.srt.function_call.mimo_detector import MiMoDetector
|
||||
from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector
|
||||
@@ -51,6 +52,7 @@ class FunctionCallParser:
|
||||
"glm47": Glm47MoeDetector,
|
||||
"gpt-oss": GptOssDetector,
|
||||
"kimi_k2": KimiK2Detector,
|
||||
"lfm2": Lfm2Detector,
|
||||
"llama3": Llama32Detector,
|
||||
"mimo": MiMoDetector,
|
||||
"mistral": MistralDetector,
|
||||
|
||||
387
python/sglang/srt/function_call/lfm2_detector.py
Normal file
387
python/sglang/srt/function_call/lfm2_detector.py
Normal file
@@ -0,0 +1,387 @@
|
||||
"""
|
||||
Detector for LFM2 (Liquid Foundation Model 2) function call format.
|
||||
|
||||
Format Structure (Pythonic style):
|
||||
```
|
||||
<|tool_call_start|>[function_name(arg1="value1", arg2="value2")]<|tool_call_end|>
|
||||
```
|
||||
|
||||
Multiple tool calls:
|
||||
```
|
||||
<|tool_call_start|>[func1(arg="val"), func2(arg="val")]<|tool_call_end|>
|
||||
```
|
||||
|
||||
Also supports JSON format:
|
||||
```
|
||||
<|tool_call_start|>[{"name": "func_name", "arguments": {...}}]<|tool_call_end|>
|
||||
```
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import Tool
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.function_call.base_format_detector import BaseFormatDetector
|
||||
from sglang.srt.function_call.core_types import (
|
||||
StreamingParseResult,
|
||||
StructureInfo,
|
||||
ToolCallItem,
|
||||
_GetInfoFunc,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Lfm2Detector(BaseFormatDetector):
|
||||
"""
|
||||
Detector for LFM2 (Liquid Foundation Model 2) function call format.
|
||||
|
||||
Supports both Pythonic and JSON formats:
|
||||
|
||||
Pythonic:
|
||||
```
|
||||
<|tool_call_start|>[calculator(expression="5 * 7")]<|tool_call_end|>
|
||||
```
|
||||
|
||||
JSON:
|
||||
```
|
||||
<|tool_call_start|>[{"name": "calculator", "arguments": {"expression": "5 * 7"}}]<|tool_call_end|>
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Initializes the detector with necessary state variables.
|
||||
"""
|
||||
super().__init__()
|
||||
self.bot_token = "<|tool_call_start|>"
|
||||
self.eot_token = "<|tool_call_end|>"
|
||||
self.tool_call_separator = ""
|
||||
|
||||
def has_tool_call(self, text: str) -> bool:
|
||||
"""Check if the text contains an LFM2 format tool call."""
|
||||
return self.bot_token in text
|
||||
|
||||
def _get_parameter_value(self, val: ast.AST) -> Any:
|
||||
"""
|
||||
Extract Python literal value from AST node.
|
||||
|
||||
Handles constants, dicts, and lists recursively.
|
||||
Reuses pattern from PythonicDetector.
|
||||
"""
|
||||
if isinstance(val, ast.Constant):
|
||||
return val.value
|
||||
elif isinstance(val, ast.Dict):
|
||||
return {
|
||||
self._get_parameter_value(k): self._get_parameter_value(v)
|
||||
for k, v in zip(val.keys, val.values)
|
||||
if k is not None # Handle {**kwargs} case where key is None
|
||||
}
|
||||
elif isinstance(val, ast.List):
|
||||
return [self._get_parameter_value(v) for v in val.elts]
|
||||
elif isinstance(val, ast.Tuple):
|
||||
return tuple(self._get_parameter_value(v) for v in val.elts)
|
||||
elif isinstance(val, ast.Name):
|
||||
# Handle True, False, None as names in older Python
|
||||
if val.id == "True":
|
||||
return True
|
||||
elif val.id == "False":
|
||||
return False
|
||||
elif val.id == "None":
|
||||
return None
|
||||
else:
|
||||
raise ValueError(f"Unsupported name reference: {val.id}")
|
||||
elif isinstance(val, ast.UnaryOp) and isinstance(val.op, ast.USub):
|
||||
# Handle negative numbers like -5
|
||||
inner = self._get_parameter_value(val.operand)
|
||||
if isinstance(inner, (int, float)):
|
||||
return -inner
|
||||
raise ValueError(f"Cannot negate non-numeric value: {inner}")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Tool call arguments must be literals, got: {type(val).__name__}"
|
||||
)
|
||||
|
||||
def _parse_pythonic_call(
|
||||
self, call: ast.Call, call_index: int, tool_indices: Dict[str, int]
|
||||
) -> Optional[ToolCallItem]:
|
||||
"""
|
||||
Parse a single AST Call node into a ToolCallItem.
|
||||
|
||||
Args:
|
||||
call: AST Call node representing a function call
|
||||
call_index: Index of this call in the list of calls
|
||||
tool_indices: Mapping of tool names to their indices
|
||||
|
||||
Returns:
|
||||
ToolCallItem if successful, None if the call should be skipped
|
||||
"""
|
||||
if not isinstance(call.func, ast.Name):
|
||||
logger.warning(
|
||||
f"Tool call function must be a simple name, got: {type(call.func).__name__}"
|
||||
)
|
||||
return None
|
||||
|
||||
function_name = call.func.id
|
||||
|
||||
# Validate that the function exists in the tools
|
||||
if function_name not in tool_indices:
|
||||
logger.warning(
|
||||
f"Model attempted to call undefined function: {function_name}"
|
||||
)
|
||||
if not envs.SGLANG_FORWARD_UNKNOWN_TOOLS.get():
|
||||
return None # Skip unknown tools (default legacy behavior)
|
||||
|
||||
# Parse arguments
|
||||
arguments = {}
|
||||
for keyword in call.keywords:
|
||||
if keyword.arg is None:
|
||||
# **kwargs unpacking - skip for now
|
||||
logger.warning("Tool call with **kwargs unpacking is not supported")
|
||||
continue
|
||||
try:
|
||||
arguments[keyword.arg] = self._get_parameter_value(keyword.value)
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to parse argument {keyword.arg}: {e}")
|
||||
return None
|
||||
|
||||
return ToolCallItem(
|
||||
tool_index=call_index, # Use the call index in the response, not tool position
|
||||
name=function_name,
|
||||
parameters=json.dumps(arguments, ensure_ascii=False),
|
||||
)
|
||||
|
||||
def _parse_pythonic_content(
|
||||
self, content: str, tools: List[Tool]
|
||||
) -> Tuple[List[ToolCallItem], str]:
|
||||
"""
|
||||
Parse Pythonic format tool calls using AST.
|
||||
|
||||
Args:
|
||||
content: The content between tool call tags (without the tags)
|
||||
tools: List of available tools
|
||||
|
||||
Returns:
|
||||
Tuple of (list of parsed calls, error message if any)
|
||||
"""
|
||||
content = content.strip()
|
||||
tool_indices = self._get_tool_indices(tools)
|
||||
|
||||
try:
|
||||
module = ast.parse(content)
|
||||
parsed = getattr(module.body[0], "value", None) if module.body else None
|
||||
|
||||
if parsed is None:
|
||||
return [], "Empty or invalid Python expression"
|
||||
|
||||
# Handle both single call and list of calls
|
||||
if isinstance(parsed, ast.List):
|
||||
call_nodes = parsed.elts
|
||||
elif isinstance(parsed, ast.Call):
|
||||
call_nodes = [parsed]
|
||||
else:
|
||||
return (
|
||||
[],
|
||||
f"Expected function call or list, got: {type(parsed).__name__}",
|
||||
)
|
||||
|
||||
# Validate all elements are calls
|
||||
if not all(isinstance(e, ast.Call) for e in call_nodes):
|
||||
return [], "Not all elements in list are function calls"
|
||||
|
||||
calls = []
|
||||
for call_index, call in enumerate(call_nodes):
|
||||
item = self._parse_pythonic_call(call, call_index, tool_indices)
|
||||
if item is not None:
|
||||
calls.append(item)
|
||||
|
||||
return calls, ""
|
||||
|
||||
except SyntaxError as e:
|
||||
return [], f"Python syntax error: {e}"
|
||||
except Exception as e:
|
||||
logger.exception("Unexpected error in pythonic tool call parsing")
|
||||
return [], f"Unexpected error: {e}"
|
||||
|
||||
def _parse_json_content(
|
||||
self, content: str, tools: List[Tool]
|
||||
) -> Tuple[List[ToolCallItem], str]:
|
||||
"""
|
||||
Parse JSON format tool calls.
|
||||
|
||||
Uses parse_base_json from BaseFormatDetector for consistent handling
|
||||
of SGLANG_FORWARD_UNKNOWN_TOOLS and tool validation.
|
||||
|
||||
Args:
|
||||
content: The content between tool call tags (without the tags)
|
||||
tools: List of available tools
|
||||
|
||||
Returns:
|
||||
Tuple of (list of parsed calls, error message if any)
|
||||
"""
|
||||
content = content.strip()
|
||||
|
||||
try:
|
||||
parsed = json.loads(content)
|
||||
# parse_base_json handles list/dict normalization, tool validation,
|
||||
# and SGLANG_FORWARD_UNKNOWN_TOOLS consistently with other detectors
|
||||
calls = self.parse_base_json(parsed, tools)
|
||||
return calls, ""
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
return [], f"JSON parse error: {e}"
|
||||
|
||||
def _parse_tool_calls_content(
|
||||
self, content: str, tools: List[Tool]
|
||||
) -> List[ToolCallItem]:
|
||||
"""
|
||||
Parse the content between tool call tags.
|
||||
Handles both JSON and Pythonic formats.
|
||||
"""
|
||||
content = content.strip()
|
||||
|
||||
# First, try JSON format (faster check)
|
||||
if content.startswith("[{") or content.startswith("{"):
|
||||
calls, error = self._parse_json_content(content, tools)
|
||||
if calls:
|
||||
return calls
|
||||
# If JSON parsing failed but it looked like JSON, log the error
|
||||
if error:
|
||||
logger.debug(f"JSON parsing failed: {error}, trying Pythonic format")
|
||||
|
||||
# Try Pythonic format
|
||||
calls, error = self._parse_pythonic_content(content, tools)
|
||||
if calls:
|
||||
return calls
|
||||
|
||||
if error:
|
||||
logger.warning(f"Failed to parse tool calls: {error}")
|
||||
|
||||
return []
|
||||
|
||||
def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
|
||||
"""
|
||||
One-time parsing: Detects and parses tool calls in the provided text.
|
||||
"""
|
||||
idx = text.find(self.bot_token)
|
||||
normal_text = text[:idx].strip() if idx != -1 else text
|
||||
|
||||
if self.bot_token not in text:
|
||||
return StreamingParseResult(normal_text=normal_text, calls=[])
|
||||
|
||||
# Find all <|tool_call_start|>...<|tool_call_end|> blocks
|
||||
pattern = rf"{re.escape(self.bot_token)}(.*?){re.escape(self.eot_token)}"
|
||||
match_result_list = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
calls = []
|
||||
for match_result in match_result_list:
|
||||
parsed_calls = self._parse_tool_calls_content(match_result, tools)
|
||||
calls.extend(parsed_calls)
|
||||
|
||||
return StreamingParseResult(normal_text=normal_text, calls=calls)
|
||||
|
||||
def _strip_special_tokens(self, text: str) -> str:
|
||||
"""Remove special tokens from text."""
|
||||
return text.replace(self.bot_token, "").replace(self.eot_token, "")
|
||||
|
||||
def parse_streaming_increment(
|
||||
self, new_text: str, tools: List[Tool]
|
||||
) -> StreamingParseResult:
|
||||
"""
|
||||
Streaming incremental parsing for LFM2 tool calls.
|
||||
|
||||
This implementation properly handles Pythonic format by:
|
||||
1. Buffering until we see complete <|tool_call_start|>[...]<|tool_call_end|>
|
||||
2. Emitting normal text before tool calls immediately
|
||||
3. Parsing complete tool call blocks using detect_and_parse
|
||||
|
||||
Based on PythonicDetector streaming logic.
|
||||
"""
|
||||
self._buffer += new_text
|
||||
|
||||
# Check for partial bot_token at the end
|
||||
partial_bot = self._ends_with_partial_token(self._buffer, self.bot_token)
|
||||
partial_eot = self._ends_with_partial_token(self._buffer, self.eot_token)
|
||||
|
||||
# Find bot_token position
|
||||
bot_pos = self._buffer.find(self.bot_token)
|
||||
|
||||
if bot_pos == -1:
|
||||
# No tool call start found
|
||||
if partial_bot:
|
||||
# Might be partial bot_token, hold back that part
|
||||
safe_text = self._buffer[:-partial_bot]
|
||||
self._buffer = self._buffer[-partial_bot:]
|
||||
return StreamingParseResult(normal_text=safe_text)
|
||||
else:
|
||||
# No tool call, emit all as normal text
|
||||
normal_text = self._strip_special_tokens(self._buffer)
|
||||
self._buffer = ""
|
||||
return StreamingParseResult(normal_text=normal_text)
|
||||
|
||||
# We have bot_token - extract any normal text before it
|
||||
normal_text_before = self._buffer[:bot_pos] if bot_pos > 0 else ""
|
||||
|
||||
# Look for the end token
|
||||
eot_pos = self._buffer.find(self.eot_token, bot_pos + len(self.bot_token))
|
||||
|
||||
if eot_pos == -1:
|
||||
# No end token yet - check if we might have a partial one
|
||||
if partial_eot:
|
||||
# Hold back the partial token, but we need to keep buffering
|
||||
# Just emit any normal text before the tool call
|
||||
if normal_text_before:
|
||||
self._buffer = self._buffer[bot_pos:]
|
||||
return StreamingParseResult(normal_text=normal_text_before)
|
||||
# Keep buffering
|
||||
return StreamingParseResult(normal_text="")
|
||||
|
||||
# No end token and no partial - keep buffering but emit normal text
|
||||
if normal_text_before:
|
||||
self._buffer = self._buffer[bot_pos:]
|
||||
return StreamingParseResult(normal_text=normal_text_before)
|
||||
|
||||
# Just keep buffering
|
||||
return StreamingParseResult(normal_text="")
|
||||
|
||||
# We have a complete tool call block
|
||||
tool_call_block = self._buffer[bot_pos : eot_pos + len(self.eot_token)]
|
||||
remaining = self._buffer[eot_pos + len(self.eot_token) :]
|
||||
|
||||
# Parse the complete block
|
||||
result = self.detect_and_parse(tool_call_block, tools)
|
||||
|
||||
# Update buffer with remaining text
|
||||
self._buffer = remaining
|
||||
|
||||
# Add any normal text before the tool call
|
||||
if normal_text_before:
|
||||
result.normal_text = normal_text_before + (result.normal_text or "")
|
||||
|
||||
return result
|
||||
|
||||
def supports_structural_tag(self) -> bool:
|
||||
"""
|
||||
Return False because LFM2 uses Pythonic format which is not JSON-compatible.
|
||||
|
||||
structural_tag only supports JSON-compatible content between begin and end,
|
||||
so it cannot parse Pythonic function call syntax like `func(arg="val")`.
|
||||
"""
|
||||
return False
|
||||
|
||||
def structure_info(self) -> _GetInfoFunc:
|
||||
"""
|
||||
Return structure info for constrained generation.
|
||||
|
||||
Note: This is provided for completeness but won't be used since
|
||||
supports_structural_tag() returns False.
|
||||
"""
|
||||
return lambda name: StructureInfo(
|
||||
begin="<|tool_call_start|>[" + name + "(",
|
||||
end=")]<|tool_call_end|>",
|
||||
trigger="<|tool_call_start|>",
|
||||
)
|
||||
@@ -35,6 +35,7 @@ from sglang.srt.configs import (
|
||||
JetNemotronConfig,
|
||||
JetVLMConfig,
|
||||
KimiLinearConfig,
|
||||
Lfm2Config,
|
||||
NemotronH_Nano_VL_V2_Config,
|
||||
NemotronHConfig,
|
||||
Qwen3NextConfig,
|
||||
@@ -1491,7 +1492,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
pattern = getattr(config, "mtp_hybrid_override_pattern", None)
|
||||
if pattern is not None and "M" not in pattern:
|
||||
return None
|
||||
if isinstance(config, FalconH1Config | NemotronHConfig):
|
||||
if isinstance(config, FalconH1Config | NemotronHConfig | Lfm2Config):
|
||||
return config
|
||||
if isinstance(config, NemotronH_Nano_VL_V2_Config):
|
||||
return config.llm_config
|
||||
|
||||
566
python/sglang/srt/models/lfm2.py
Normal file
566
python/sglang/srt/models/lfm2.py
Normal file
@@ -0,0 +1,566 @@
|
||||
"""
|
||||
LFM2 (Liquid Foundation Model 2) implementation for SGLang.
|
||||
|
||||
This is a hybrid architecture with both attention and short conv layers.
|
||||
- Attention layers use standard KV cache (RadixAttention)
|
||||
- Conv layers use MambaPool for state caching (via HybridReqToTokenPool)
|
||||
|
||||
The model uses a gated 1D causal convolution (kernel=3) instead of attention
|
||||
in some layers, providing linear memory complexity for those layers.
|
||||
|
||||
Uses optimized causal_conv1d kernels from the mamba package for fast inference.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Iterable, Optional, Set, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from sglang.srt.configs.lfm2 import Lfm2Config
|
||||
from sglang.srt.distributed import get_pp_group
|
||||
from sglang.srt.layers.attention.mamba.causal_conv1d import (
|
||||
causal_conv1d_fn,
|
||||
causal_conv1d_update,
|
||||
)
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.layers.radix_attention import RadixAttention
|
||||
from sglang.srt.layers.rotary_embedding import get_rope
|
||||
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.utils import add_prefix, make_layers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# We don't use it, we keep it for reference. If we run sglang.srt.layers.layernorm.RMSNorm
|
||||
# kernel the difference in logprobs slightly increases, but to an acceptable degree
|
||||
# class Lfm2RMSNorm(nn.Module):
|
||||
# """LFM2-specific RMSNorm: weight * x (not (1 + weight) * x like Gemma)."""
|
||||
|
||||
# def __init__(self, hidden_size: int, eps: float = 1e-6):
|
||||
# super().__init__()
|
||||
# self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
# self.variance_epsilon = eps
|
||||
|
||||
# def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
# input_dtype = hidden_states.dtype
|
||||
# hidden_states = hidden_states.to(torch.float32)
|
||||
# variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
# hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
# return (self.weight * hidden_states).to(input_dtype)
|
||||
|
||||
|
||||
class Lfm2MLP(nn.Module):
|
||||
"""MLP with SwiGLU activation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
intermediate_size = config.intermediate_size
|
||||
|
||||
if config.block_auto_adjust_ff_dim:
|
||||
intermediate_size = int(2 * intermediate_size / 3)
|
||||
if config.block_ffn_dim_multiplier is not None:
|
||||
intermediate_size = int(
|
||||
config.block_ffn_dim_multiplier * intermediate_size
|
||||
)
|
||||
intermediate_size = config.block_multiple_of * (
|
||||
(intermediate_size + config.block_multiple_of - 1)
|
||||
// config.block_multiple_of
|
||||
)
|
||||
|
||||
self.w1 = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
intermediate_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("w1", prefix),
|
||||
)
|
||||
self.w3 = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
intermediate_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("w3", prefix),
|
||||
)
|
||||
self.w2 = RowParallelLinear(
|
||||
intermediate_size,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("w2", prefix),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
gate, _ = self.w1(x)
|
||||
up, _ = self.w3(x)
|
||||
out, _ = self.w2(F.silu(gate) * up)
|
||||
return out
|
||||
|
||||
|
||||
class Lfm2Attention(nn.Module):
|
||||
"""Grouped-query attention with RoPE and Q/K layernorm."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
layer_id: int,
|
||||
attn_layer_id: int,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.total_num_heads = config.num_attention_heads
|
||||
self.total_num_kv_heads = config.num_key_value_heads
|
||||
self.head_dim = getattr(config, "head_dim", None) or (
|
||||
self.hidden_size // self.total_num_heads
|
||||
)
|
||||
self.scaling = self.head_dim**-0.5
|
||||
|
||||
rope_parameters = getattr(config, "rope_parameters", None)
|
||||
if rope_parameters is not None and "rope_theta" in rope_parameters:
|
||||
rope_theta = rope_parameters["rope_theta"]
|
||||
else:
|
||||
rope_theta = getattr(config, "rope_theta", 10000)
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
head_size=self.head_dim,
|
||||
rotary_dim=self.head_dim,
|
||||
max_position=getattr(config, "max_position_embeddings", 8192),
|
||||
rope_scaling=getattr(config, "rope_scaling", None),
|
||||
base=rope_theta,
|
||||
is_neox_style=True,
|
||||
dtype=torch.get_default_dtype(),
|
||||
)
|
||||
|
||||
self.qkv_proj = QKVParallelLinear(
|
||||
self.hidden_size,
|
||||
self.head_dim,
|
||||
self.total_num_heads,
|
||||
self.total_num_kv_heads,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("qkv_proj", prefix),
|
||||
)
|
||||
self.out_proj = RowParallelLinear(
|
||||
self.total_num_heads * self.head_dim,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("out_proj", prefix),
|
||||
)
|
||||
|
||||
self.q_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
|
||||
self.k_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
|
||||
|
||||
self.num_local_q_heads = self.qkv_proj.num_heads
|
||||
self.num_local_kv_heads = self.qkv_proj.num_kv_heads
|
||||
|
||||
self.attn = RadixAttention(
|
||||
num_heads=self.num_local_q_heads,
|
||||
head_dim=self.head_dim,
|
||||
scaling=self.scaling,
|
||||
num_kv_heads=self.num_local_kv_heads,
|
||||
layer_id=layer_id,
|
||||
prefix=add_prefix("attn", prefix),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
T = hidden_states.shape[0]
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
|
||||
q_size = self.num_local_q_heads * self.head_dim
|
||||
kv_size = self.num_local_kv_heads * self.head_dim
|
||||
q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
q = q.reshape(T, self.num_local_q_heads, self.head_dim)
|
||||
k = k.reshape(T, self.num_local_kv_heads, self.head_dim)
|
||||
|
||||
q = self.q_layernorm(q.reshape(-1, self.head_dim)).reshape(
|
||||
T, self.num_local_q_heads, self.head_dim
|
||||
)
|
||||
k = self.k_layernorm(k.reshape(-1, self.head_dim)).reshape(
|
||||
T, self.num_local_kv_heads, self.head_dim
|
||||
)
|
||||
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
|
||||
attn_out = self.attn(q.reshape(T, -1), k.reshape(T, -1), v, forward_batch)
|
||||
out, _ = self.out_proj(attn_out)
|
||||
return out
|
||||
|
||||
|
||||
class Lfm2ShortConv(nn.Module):
|
||||
"""
|
||||
Gated short convolution layer using optimized causal_conv1d kernels.
|
||||
|
||||
Architecture: in_proj -> split(B, C, x) -> Bx -> conv1d -> C*conv_out -> out_proj
|
||||
- Uses double gating: B (before conv) and C (after conv)
|
||||
- Fixed-size cache: stores last (kernel_size - 1) tokens
|
||||
- Uses causal_conv1d_fn for prefill and causal_conv1d_update for decode
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
layer_idx: int,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.layer_idx = layer_idx
|
||||
self.conv_kernel = int(config.conv_L_cache)
|
||||
self.L_cache = self.conv_kernel - 1
|
||||
self.use_bias = bool(config.conv_bias)
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
self.in_proj = nn.Linear(
|
||||
config.hidden_size, 3 * config.hidden_size, bias=self.use_bias
|
||||
)
|
||||
self.out_proj = nn.Linear(
|
||||
config.hidden_size, config.hidden_size, bias=self.use_bias
|
||||
)
|
||||
|
||||
# Conv weights stored in format matching causal_conv1d: (hidden_size, kernel_size)
|
||||
# Weight loading will handle conversion from HF's (hidden_size, 1, kernel_size)
|
||||
self.conv_weight = nn.Parameter(
|
||||
torch.empty(config.hidden_size, self.conv_kernel)
|
||||
)
|
||||
if self.use_bias:
|
||||
self.conv_bias = nn.Parameter(torch.empty(config.hidden_size))
|
||||
else:
|
||||
self.register_parameter("conv_bias", None)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
if forward_batch.forward_mode.is_idle():
|
||||
return hidden_states
|
||||
|
||||
layer_cache = forward_batch.req_to_token_pool.mamba2_layer_cache(self.layer_idx)
|
||||
conv_state = layer_cache.conv[0]
|
||||
req_pool_indices = forward_batch.req_pool_indices
|
||||
|
||||
# Project and split into gates: B (pre-conv), C (post-conv), x (input)
|
||||
proj = self.in_proj(hidden_states)
|
||||
B_gate, C_gate, x = proj.chunk(3, dim=-1)
|
||||
Bx = B_gate * x
|
||||
|
||||
if forward_batch.forward_mode.is_decode():
|
||||
# Decode: single token per request, use optimized update kernel
|
||||
conv_out = causal_conv1d_update(
|
||||
Bx,
|
||||
conv_state,
|
||||
self.conv_weight,
|
||||
self.conv_bias,
|
||||
activation=None,
|
||||
conv_state_indices=req_pool_indices.to(torch.int32),
|
||||
)
|
||||
else:
|
||||
# Prefill: multiple tokens, use varlen kernel
|
||||
T = hidden_states.shape[0]
|
||||
Bx_t = Bx.transpose(0, 1).contiguous()
|
||||
|
||||
# Build query_start_loc: [0, cumsum(seq_lens)...]
|
||||
extend_start_loc = forward_batch.extend_start_loc
|
||||
if extend_start_loc is not None and len(extend_start_loc) > 1:
|
||||
query_start_loc = torch.cat(
|
||||
[
|
||||
extend_start_loc,
|
||||
torch.tensor(
|
||||
[T], dtype=torch.int32, device=hidden_states.device
|
||||
),
|
||||
]
|
||||
)
|
||||
cache_indices = req_pool_indices.to(torch.int32)
|
||||
else:
|
||||
query_start_loc = torch.tensor(
|
||||
[0, T], dtype=torch.int32, device=hidden_states.device
|
||||
)
|
||||
cache_indices = req_pool_indices[:1].to(torch.int32)
|
||||
|
||||
conv_out = causal_conv1d_fn(
|
||||
Bx_t,
|
||||
self.conv_weight,
|
||||
self.conv_bias,
|
||||
query_start_loc=query_start_loc,
|
||||
cache_indices=cache_indices,
|
||||
has_initial_state=None,
|
||||
conv_states=conv_state,
|
||||
activation=None,
|
||||
).transpose(0, 1)
|
||||
|
||||
return self.out_proj(C_gate * conv_out)
|
||||
|
||||
|
||||
class Lfm2DecoderLayer(nn.Module):
|
||||
"""Decoder layer - either attention or conv based on config."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
layer_id: int,
|
||||
attn_layer_id: int,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.layer_type = config.layer_types[layer_id]
|
||||
self.is_attention_layer = self.layer_type == "full_attention"
|
||||
|
||||
self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
||||
self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
||||
|
||||
if self.is_attention_layer:
|
||||
self.self_attn = Lfm2Attention(
|
||||
config=config,
|
||||
layer_id=layer_id,
|
||||
attn_layer_id=attn_layer_id,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("self_attn", prefix),
|
||||
)
|
||||
else:
|
||||
self.conv = Lfm2ShortConv(
|
||||
config=config,
|
||||
layer_idx=layer_id,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("conv", prefix),
|
||||
)
|
||||
|
||||
self.feed_forward = Lfm2MLP(
|
||||
config=config,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("feed_forward", prefix),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
layer_id: int,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
residual: Optional[torch.Tensor],
|
||||
forward_batch: ForwardBatch,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
residual = hidden_states
|
||||
normed = self.operator_norm(hidden_states)
|
||||
|
||||
if self.is_attention_layer:
|
||||
hidden_states = self.self_attn(positions, normed, forward_batch)
|
||||
else:
|
||||
hidden_states = self.conv(normed, forward_batch)
|
||||
|
||||
hidden_states = hidden_states + residual
|
||||
hidden_states = hidden_states + self.feed_forward(
|
||||
self.ffn_norm(hidden_states)
|
||||
)
|
||||
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
class Lfm2Model(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
org_num_embeddings=config.vocab_size,
|
||||
prefix=add_prefix("embed_tokens", prefix),
|
||||
)
|
||||
|
||||
# Compute attention layer IDs for KV cache
|
||||
attn_layer_ids = []
|
||||
attn_count = 0
|
||||
for layer_type in config.layer_types:
|
||||
if layer_type == "full_attention":
|
||||
attn_layer_ids.append(attn_count)
|
||||
attn_count += 1
|
||||
else:
|
||||
attn_layer_ids.append(-1)
|
||||
|
||||
self.num_attention_layers = attn_count
|
||||
|
||||
def get_layer(idx: int, prefix: str, **kwargs):
|
||||
return Lfm2DecoderLayer(
|
||||
config=config,
|
||||
layer_id=idx,
|
||||
attn_layer_id=attn_layer_ids[idx],
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
self.layers = make_layers(
|
||||
config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
|
||||
)
|
||||
self.embedding_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
hidden_states = (
|
||||
inputs_embeds if inputs_embeds is not None else self.embed_tokens(input_ids)
|
||||
)
|
||||
|
||||
residual = None
|
||||
for i in range(len(self.layers)):
|
||||
hidden_states, residual = self.layers[i](
|
||||
layer_id=i,
|
||||
positions=positions,
|
||||
hidden_states=hidden_states,
|
||||
residual=residual,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
|
||||
return self.embedding_norm(hidden_states)
|
||||
|
||||
|
||||
class Lfm2ForCausalLM(nn.Module):
|
||||
"""LFM2 for causal language modeling with hybrid attention/conv architecture."""
|
||||
|
||||
fall_back_to_pt_during_load = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Lfm2Config,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.pp_group = get_pp_group()
|
||||
assert self.pp_group.is_first_rank and self.pp_group.is_last_rank
|
||||
|
||||
self.quant_config = quant_config
|
||||
self.model = Lfm2Model(config, quant_config, prefix=add_prefix("model", prefix))
|
||||
self.lm_head = ParallelLMHead(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
org_num_embeddings=config.vocab_size,
|
||||
prefix=add_prefix("lm_head", prefix),
|
||||
)
|
||||
self.logits_processor = LogitsProcessor(config)
|
||||
self.num_attention_layers = self.model.num_attention_layers
|
||||
|
||||
def get_num_kv_cache_layers(self) -> int:
|
||||
return self.num_attention_layers
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
):
|
||||
hidden_states = self.model(input_ids, positions, forward_batch, inputs_embeds)
|
||||
return self.logits_processor(
|
||||
input_ids, hidden_states, self.lm_head, forward_batch
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False
|
||||
) -> Set[str]:
|
||||
stacked_params_mapping = [
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: Set[str] = set()
|
||||
embed_tokens_weight = None
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
|
||||
if "embed_tokens.weight" in name:
|
||||
embed_tokens_weight = loaded_weight
|
||||
|
||||
# Handle conv.weight -> conv_weight conversion for ShortConv layers
|
||||
# HF shape: (hidden_size, 1, kernel_size) -> squeeze to (hidden_size, kernel_size)
|
||||
if ".conv.weight" in name:
|
||||
name = name.replace(".conv.weight", ".conv_weight")
|
||||
# Squeeze out the middle dimension: (D, 1, K) -> (D, K)
|
||||
loaded_weight = loaded_weight.squeeze(1)
|
||||
|
||||
# Handle conv.bias -> conv_bias conversion
|
||||
if ".conv.bias" in name:
|
||||
name = name.replace(".conv.bias", ".conv_bias")
|
||||
|
||||
# Handle QKV stacking
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
break
|
||||
if name not in params_dict:
|
||||
break
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader")
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(name)
|
||||
break
|
||||
else:
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
if name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
# Handle tied lm_head weight
|
||||
if "lm_head.weight" not in loaded_params and "lm_head.weight" in params_dict:
|
||||
if embed_tokens_weight is not None:
|
||||
param = params_dict["lm_head.weight"]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, embed_tokens_weight)
|
||||
loaded_params.add("lm_head.weight")
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
EntryClass = [Lfm2ForCausalLM]
|
||||
@@ -1592,6 +1592,33 @@ class ServerArgs:
|
||||
)
|
||||
self.disable_radix_cache = True
|
||||
self.disable_overlap_schedule = False
|
||||
elif model_arch in ["Lfm2ForCausalLM"]:
|
||||
assert (
|
||||
not self.enable_mamba_extra_buffer()
|
||||
), f"mamba extra_buffer is not supported for {model_arch} model"
|
||||
if not self.disable_radix_cache:
|
||||
logger.warning(
|
||||
"Disabling overlap schedule since mamba no_buffer is not compatible with "
|
||||
"overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary"
|
||||
)
|
||||
self.disable_overlap_schedule = True
|
||||
if is_sm100_supported():
|
||||
if self.attention_backend is None:
|
||||
self.attention_backend = "flashinfer"
|
||||
logger.info(
|
||||
f"Use flashinfer as attention backend on sm100 for {model_arch}"
|
||||
)
|
||||
if self.attention_backend == "trtllm_mha":
|
||||
logger.warning(
|
||||
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
|
||||
"Try to use --attention-backend flashinfer if radix cache is necessary."
|
||||
)
|
||||
self.disable_radix_cache = True
|
||||
self.disable_overlap_schedule = False
|
||||
assert self.attention_backend != "triton", (
|
||||
f"{model_arch} does not support triton attention backend, "
|
||||
"as the first layer might not be an attention layer"
|
||||
)
|
||||
|
||||
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
|
||||
self.disable_overlap_schedule = True
|
||||
|
||||
Reference in New Issue
Block a user