[Feature] Add --disable-normalize-embedding server arg

python/sglang/srt/arg_groups/fields/model.py

@@ -118,6 +118,15 @@ class Model(msgspec.Struct):
         bool,
         "Whether to use a CausalLM as an embedding model.",
     ] = False
+    disable_normalize_embedding: A[
+        bool,
+        "Disable the L2 normalization applied by the embedding Pooler after pooling. "
+        "By default, models that construct their Pooler with normalize=True L2-normalize "
+        "the pooled embedding; this flag turns that off at serve time without modifying "
+        "the model code. Models that already pool without normalization are unaffected, "
+        "and so are models with a custom pooler that does not go through Pooler (for "
+        "example the cross-encoder, BERT-pooler and vision-pooler paths).",
+    ] = False
     revision: A[
         Optional[str],
         "The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version.",

python/sglang/srt/layers/pooler.py

@@ -12,6 +12,7 @@
 from transformers import PretrainedConfig
 
 from sglang.srt.layers.activation import get_cross_encoder_activation_function
+from sglang.srt.runtime_context import get_model, is_config_published
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -165,6 +166,20 @@ def score_and_pool(
     )
 
 
+def _disable_normalize_embedding() -> bool:
+    """Whether ``--disable-normalize-embedding`` is set.
+
+    ``Pooler`` is constructible outside a published server context (offline /
+    standalone model construction), while config bags fail closed until
+    ``publish`` has run. Absent a published model namespace there is no flag to
+    honor, so fall back to the default of keeping normalization. Production
+    model init runs after ``publish``, so serving always sees the real value.
+    """
+    if not is_config_published("model"):
+        return False
+    return get_model().disable_normalize_embedding
+
+
 class Pooler(nn.Module):
     """A layer that pools specific information from hidden states.
     This layer does the following:
@@ -179,6 +194,8 @@ class Pooler(nn.Module):
     def __init__(self, pooling_type: PoolingType, normalize: bool):
         super().__init__()
         self.pooling_type = pooling_type
+        if normalize and _disable_normalize_embedding():
+            normalize = False
         self.normalize = normalize
 
     def forward(

python/sglang/srt/runtime_context.py

@@ -1069,6 +1069,20 @@ def set_server_args(self, server_args: ServerArgs) -> None:
         self._overrides_log = []
         self._publish_role = None
 
+    def is_config_published(self, name: str) -> bool:
+        """Whether ``name``'s config bag has been projected.
+
+        For the rare reader that must tolerate an unpublished context (a layer
+        constructible offline, outside any published server context) and has a
+        well-defined pre-publish default. Regular reads must keep using
+        ``config_bag`` and fail closed. Deliberately does not run the role
+        namespace check: this answers "is it there", not "may I read it", so a
+        role violation still raises from the subsequent ``config_bag`` call
+        rather than being reported as unpublished.
+        """
+        bags = self._config_bags
+        return bool(bags) and name in bags
+
     def config_bag(self, name: str) -> _ConfigBag:
         """Return the top-level config namespace bag (``device`` / ``model`` /
         ``exec`` / ``schedule`` / ``memory`` / ``spec`` / ``lora`` / ``mm`` /
@@ -1374,6 +1388,15 @@ def get_device() -> _ConfigBag:
     return _CONTEXT.config_bag("device")
 
 
+def is_config_published(name: str) -> bool:
+    """Whether the ``name`` config namespace has been projected by ``publish``.
+
+    Only for readers that must stay constructible without a published context;
+    everything else reads the bag directly and fails closed.
+    """
+    return _CONTEXT.is_config_published(name)
+
+
 def get_model() -> _ConfigBag:
     return _CONTEXT.config_bag("model")
 

test/registered/unit/layers/test_pooler_score_and_pool.py

@@ -16,6 +16,7 @@
     PoolingType,
     score_and_pool,
 )
+from sglang.srt.runtime_context import get_context, reset_context
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -187,5 +188,68 @@ def test_empty_delimiter_indices(self):
         self.assertEqual(out.embeddings[0].shape, (0, self.num_labels))
 
 
+class TestDisableNormalizeEmbedding(CustomTestCase):
+    """--disable-normalize-embedding turns off the Pooler's L2 normalization.
+
+    ``Pooler.__init__`` reads the resolved flag from the ``model`` config bag,
+    so these tests publish a context via ``override_server_args`` rather than
+    faking the accessor. Construction against an *unpublished* context stays
+    supported and keeps normalization — see
+    ``test_unpublished_context_keeps_normalization``.
+    """
+
+    def setUp(self):
+        torch.manual_seed(42)
+        self.hidden = torch.randn(4, 8)
+        self.fb = _make_forward_batch(extend_seq_lens=[2, 2])
+
+    def _pooled(self, pooler):
+        return pooler(self.hidden, self.fb).embeddings
+
+    def test_flag_disables_normalization(self):
+        with get_context().override_server_args(disable_normalize_embedding=True):
+            pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+            self.assertFalse(pooler.normalize)
+            pooled = self._pooled(pooler)
+
+        # Un-normalized rows keep the raw hidden-state norms.
+        expected = self.hidden[torch.tensor([1, 3])]
+        torch.testing.assert_close(pooled, expected)
+
+    def test_default_keeps_normalization(self):
+        with get_context().override_server_args(disable_normalize_embedding=False):
+            pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+            self.assertTrue(pooler.normalize)
+            pooled = self._pooled(pooler)
+
+        torch.testing.assert_close(
+            pooled.norm(p=2, dim=-1), torch.ones(pooled.shape[0])
+        )
+
+    def test_flag_does_not_force_enable_normalization(self):
+        """The override is one-way: normalize=False models stay un-normalized."""
+        with get_context().override_server_args(disable_normalize_embedding=True):
+            self.assertFalse(
+                Pooler(pooling_type=PoolingType.LAST, normalize=False).normalize
+            )
+        with get_context().override_server_args(disable_normalize_embedding=False):
+            self.assertFalse(
+                Pooler(pooling_type=PoolingType.LAST, normalize=False).normalize
+            )
+
+    def test_unpublished_context_keeps_normalization(self):
+        """Pooler stays constructible without a published server context.
+
+        Config bags fail closed before ``publish``, so reading the flag
+        unconditionally would make ``Pooler(normalize=True)`` raise in any
+        standalone/offline construction path. With no published model
+        namespace there is no override to honor, so normalization is kept.
+        """
+        self.addCleanup(reset_context)
+        reset_context()
+        pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+        self.assertTrue(pooler.normalize)
+
+
 if __name__ == "__main__":
     unittest.main()