Implement BGE-M3 Sparse Embeddings in SGLang (#10869)

Co-authored-by: Christian Bahls <christian.bahls@planet-ai.de>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Christian Bahls
2025-10-22 22:46:16 +02:00
committed by GitHub
parent 5dccf69713
commit 164302c7df
6 changed files with 187 additions and 8 deletions

View File

@@ -77,6 +77,7 @@ from sglang.srt.model_loader.utils import (
DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
0.8 # Reserve 20% GPU memory headroom for ModelOpt calibration
)
from sglang.srt.environ import envs
from sglang.srt.model_loader.weight_utils import (
download_safetensors_index_file_from_hf,
download_weights_from_hf,
@@ -244,10 +245,19 @@ def _initialize_model(
quant_config = _get_quantization_config(
model_config, load_config, packed_modules_mapping
)
return model_class(
config=model_config.hf_config,
quant_config=quant_config,
)
# Build kwargs conditionally
kwargs = {
"config": model_config.hf_config,
"quant_config": quant_config,
}
# Only add sparse head kwargs if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
kwargs["sparse_head"] = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.value
kwargs["model_path"] = model_config.model_path
return model_class(**kwargs)
class BaseModelLoader(ABC):