110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
# Copyright 2026 Liquid AI. 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-VL (Liquid Foundation Model 2 Vision-Language) configuration"""
|
|
|
|
from typing import List, Optional
|
|
|
|
from transformers import CONFIG_MAPPING
|
|
from transformers import Lfm2VlConfig as HFLfm2VlConfig
|
|
from transformers.utils import logging
|
|
|
|
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
class Lfm2VlConfig(HFLfm2VlConfig):
|
|
"""
|
|
SGLang configuration for LFM2-VL models.
|
|
|
|
Extends HuggingFace's Lfm2VlConfig with hybrid model properties needed by SGLang.
|
|
LFM2-VL combines:
|
|
- SigLip2 vision encoder with NaFlex variable-resolution support
|
|
- LFM2 language model with hybrid attention + short convolution
|
|
- Multimodal projector with pixel unshuffle downsampling
|
|
"""
|
|
|
|
@property
|
|
def full_attention_layer_ids(self) -> List[int]:
|
|
"""Return indices of attention layers for KV cache (from text_config)."""
|
|
return [
|
|
i
|
|
for i, lt in enumerate(self.text_config.layer_types)
|
|
if lt == "full_attention"
|
|
]
|
|
|
|
@property
|
|
def linear_layer_ids(self) -> List[int]:
|
|
"""Return indices of conv layers for conv state cache (from text_config)."""
|
|
return [
|
|
i
|
|
for i, lt in enumerate(self.text_config.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.text_config.hidden_size
|
|
# conv_L_cache in config is kernel_size (e.g., 3)
|
|
conv_kernel = int(self.text_config.conv_L_cache)
|
|
|
|
# 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
|
|
# We pass num_heads=tp_size so divide(tp_size, tp_size)=1 always works.
|
|
# Since state_size=0, the temporal state shape has zero elements anyway.
|
|
shape = Mamba2StateShape.create(
|
|
tp_world_size=tp_size,
|
|
intermediate_size=hidden_size,
|
|
n_groups=1, # ShortConv doesn't use grouping
|
|
num_heads=tp_size, # Ensures divide works; temporal state is empty anyway
|
|
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 Lfm2VlConfig with our extended version
|
|
# Cannot use .register() because lfm2_vl may already be registered by transformers
|
|
# Directly modify the internal _extra_content dict instead
|
|
CONFIG_MAPPING._extra_content["lfm2_vl"] = Lfm2VlConfig
|