Phase A step 1: kill 2 syncs/forward in model LVM heads + EOS cache

Cherry-picked from upstream PR UCSB-AI#3 (namezhenzhang):

* lvm_value_utils.get_eos_token_ids: cache result on req._lvm_eos_token_ids;
  was re-walking sets every candidate filter call.
* qwen2_lvm.py / qwen3_lvm.py / qwen2_5_vl_lvm.py: drop two
  forward_batch.{extend_seq_lens,extend_prefix_lens}.tolist() calls per
  LVM forward by carrying tree_value_cached_prefix_lens in the
  TreeValueSpecInput. With ~1300 LVM forwards in the paper run, that
  removes ~2600 GPU->CPU syncs.
* tree_value_spec.py: vectorize the candidate self-attention diagonal
  using numpy fancy indexing instead of a per-token Python loop.

No algorithm change; same forward outputs.

Co-Authored-By: Zhen Zhang <namezhenzhang@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 3 people

sglang-LenVM/python/sglang/srt/lvm/lvm_value_utils.py

@@ -10,6 +10,10 @@ def get_eos_token_ids(req: Any) -> Set[int]:
     - `req.eos_token_ids`: a Set[int] from `ModelConfig.hf_eos_token_id` (preferred)
     - `req.tokenizer.eos_token_id`: tokenizer-defined EOS id
     """
+    cached = getattr(req, "_lvm_eos_token_ids", None)
+    if cached is not None:
+        return set(cached)
+
     eos: Set[int] = set()
 
     eos_ids = getattr(req, "eos_token_ids", None)
@@ -32,6 +36,10 @@ def get_eos_token_ids(req: Any) -> Set[int]:
         except Exception as exc:
             raise ValueError(f"Invalid LenVM tokenizer eos_token_id: {tok_eos!r}") from exc
 
+    try:
+        setattr(req, "_lvm_eos_token_ids", tuple(sorted(eos)))
+    except Exception:
+        pass
     return eos
 
 

sglang-LenVM/python/sglang/srt/lvm/tree_value_spec.py

@@ -250,8 +250,8 @@ def build_tree_value_custom_mask_and_positions(
         # Candidate rows: attend all prefix tokens (0..L-1) and itself.
         cand_row_start = L - P
         m[cand_row_start : cand_row_start + N, :L] = True
-        for j in range(N):
-            m[cand_row_start + j, L + j] = True
+        cand_arange = np.arange(N)
+        m[cand_row_start + cand_arange, L + cand_arange] = True
 
         mask_off += q_len * k_len
 

sglang-LenVM/python/sglang/srt/models/qwen2_5_vl_lvm.py

@@ -77,22 +77,24 @@ def forward(
         )
 
         if prefix_lens is not None and cand_lens is not None:
-            extend_lens = forward_batch.extend_seq_lens.tolist()
+            # Same upstream-PR-#3 sync removal as qwen2_lvm.py.
             cached_prefix_lens = (
-                forward_batch.extend_prefix_lens.tolist()
-                if forward_batch.extend_prefix_lens is not None
-                else [0] * len(extend_lens)
+                getattr(spec, "tree_value_cached_prefix_lens", None)
+                or [0] * len(prefix_lens)
             )
 
             out: list[torch.Tensor] = []
             offset = 0
-            for i, ext_len in enumerate(extend_lens):
+            for i, (prefix_len, cand_len, cached_prefix_len) in enumerate(
+                zip(prefix_lens, cand_lens, cached_prefix_lens)
+            ):
+                prefix_len = int(prefix_len)
+                cand_len = int(cand_len)
+                cached_prefix_len = int(cached_prefix_len)
+                ext_len = max(prefix_len - cached_prefix_len, 0) + cand_len
                 vals_i = token_values[offset : offset + ext_len]
                 offset += ext_len
 
-                prefix_len = int(prefix_lens[i])
-                cand_len = int(cand_lens[i])
-                cached_prefix_len = int(cached_prefix_lens[i])
                 cand_offset = max(prefix_len - cached_prefix_len, 0)
                 out.append(vals_i[cand_offset : cand_offset + cand_len])
             return EmbeddingPoolerOutput(embeddings=out)

sglang-LenVM/python/sglang/srt/models/qwen2_lvm.py

@@ -86,22 +86,26 @@ def forward(
         prefix_lens = getattr(spec, "tree_value_prefix_lens", None) if spec is not None else None
         cand_lens = getattr(spec, "tree_value_candidate_lens", None) if spec is not None else None
         if prefix_lens is not None and cand_lens is not None:
-            extend_lens = forward_batch.extend_seq_lens.tolist()
+            # Avoid forward_batch.{extend_seq_lens,extend_prefix_lens}.tolist()
+            # syncs on every LVM forward by carrying these lengths in the
+            # TreeValueSpecInput. Sourced from upstream PR #3.
             cached_prefix_lens = (
-                forward_batch.extend_prefix_lens.tolist()
-                if forward_batch.extend_prefix_lens is not None
-                else [0] * len(extend_lens)
+                getattr(spec, "tree_value_cached_prefix_lens", None)
+                or [0] * len(prefix_lens)
             )
 
             out: list[torch.Tensor] = []
             offset = 0
-            for i, ext_len in enumerate(extend_lens):
+            for i, (prefix_len, cand_len, cached_prefix_len) in enumerate(
+                zip(prefix_lens, cand_lens, cached_prefix_lens)
+            ):
+                L = int(prefix_len)
+                N = int(cand_len)
+                P = int(cached_prefix_len)
+                ext_len = max(L - P, 0) + N
                 vals_i = token_values[offset : offset + ext_len]
                 offset += ext_len
 
-                L = int(prefix_lens[i])
-                N = int(cand_lens[i])
-                P = int(cached_prefix_lens[i])
                 cand_offset = max(L - P, 0)
                 out.append(vals_i[cand_offset : cand_offset + N])
             return EmbeddingPoolerOutput(embeddings=out)

sglang-LenVM/python/sglang/srt/models/qwen3_lvm.py

@@ -83,22 +83,24 @@ def forward(
         prefix_lens = getattr(spec, "tree_value_prefix_lens", None) if spec is not None else None
         cand_lens = getattr(spec, "tree_value_candidate_lens", None) if spec is not None else None
         if prefix_lens is not None and cand_lens is not None:
-            extend_lens = forward_batch.extend_seq_lens.tolist()
+            # Same upstream-PR-#3 sync removal as qwen2_lvm.py.
             cached_prefix_lens = (
-                forward_batch.extend_prefix_lens.tolist()
-                if forward_batch.extend_prefix_lens is not None
-                else [0] * len(extend_lens)
+                getattr(spec, "tree_value_cached_prefix_lens", None)
+                or [0] * len(prefix_lens)
             )
 
             out: list[torch.Tensor] = []
             offset = 0
-            for i, ext_len in enumerate(extend_lens):
+            for i, (prefix_len, cand_len, cached_prefix_len) in enumerate(
+                zip(prefix_lens, cand_lens, cached_prefix_lens)
+            ):
+                L = int(prefix_len)
+                N = int(cand_len)
+                P = int(cached_prefix_len)
+                ext_len = max(L - P, 0) + N
                 vals_i = token_values[offset : offset + ext_len]
                 offset += ext_len
 
-                L = int(prefix_lens[i])
-                N = int(cand_lens[i])
-                P = int(cached_prefix_lens[i])
                 cand_offset = max(L - P, 0)
                 out.append(vals_i[cand_offset : cand_offset + N])
             return EmbeddingPoolerOutput(embeddings=out)