Add support for Matryoshka embeddings (#126) (#11142)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

python/sglang/srt/managers/schedule_batch.py

@@ -442,6 +442,7 @@ class Req:
         priority: Optional[int] = None,
         metrics_collector: Optional[SchedulerMetricsCollector] = None,
         extra_key: Optional[str] = None,
+        dimensions: Optional[int] = None,
         http_worker_ipc: Optional[str] = None,
     ):
         # Input and output info
@@ -650,6 +651,9 @@ class Req:
         self.tmp_end_idx: int = -1
         self.metadata_buffer_index: int = -1
 
+        # For Matryoshka embeddings
+        self.dimensions = dimensions
+
     @property
     def seqlen(self):
         return len(self.origin_input_ids) + len(self.output_ids)
@@ -1014,6 +1018,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     encoder_lens_cpu: Optional[List[int]] = None
     encoder_out_cache_loc: Optional[torch.Tensor] = None
 
+    # For matryoshka embeddings
+    dimensions: Optional[list[int]] = None
+
     # For split prefill
     split_index: int = 0
     split_prefill_finished: bool = False
@@ -1177,6 +1184,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         prefix_lens = [len(r.prefix_indices) for r in reqs]
         extend_lens = [r.extend_input_len for r in reqs]
 
+        # For matryoshka embeddings
+        if self.model_config.is_matryoshka and any(
+            r.dimensions is not None for r in reqs
+        ):
+            self.dimensions = [
+                r.dimensions if r.dimensions else self.model_config.hidden_size
+                for r in reqs
+            ]
+
         token_type_ids = [
             r.token_type_ids for r in reqs if r.token_type_ids is not None
         ]
@@ -1765,6 +1781,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             ),
             extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
             is_prefill_only=self.is_prefill_only,
+            dimensions=self.dimensions,
         )
 
     def copy(self):
@@ -1873,5 +1890,8 @@ class ModelWorkerBatch:
     capture_hidden_mode: CaptureHiddenMode = None
     hicache_consumer_index: int = -1
 
+    # For matryoshka embeddings
+    dimensions: Optional[list[int]] = None
+
     # Whether this batch is prefill-only (no token generation needed)
     is_prefill_only: bool = False