Implementation:SqueezeAILab ETS Tree Expand

Knowledge Sources	ETS
Domains	Tree_Search, Inference_Time_Compute
Last Updated	2026-02-14 02:00 GMT

Overview

Concrete tool for expanding a tree node by forking SGLang state and generating scored reasoning steps, provided by the Tree class in rebase.py.

Description

The Tree.expand() method takes a parent node and a width count, forks the parent's SGLang state into that many copies, then for each fork generates a reasoning step and scores it with the PRM. The generation and scoring behavior varies by policy_model_type:

llemma: Generates up to max_step_tokens tokens stopping at "ки", scores via logits at token ID 8094
mistral: Same stop token, scores via logits at token IDs [648, 387, 12902] (plus_tag, minus_tag, step_tag)
llama: Uses assistant role wrapper, stops at "\n\n", scores via logits at positions [10, 12, -3]

Usage

Called by Tree.select_and_expand() for each node with a positive expansion width. Not called directly by user code.

Code Reference

Source Location

Repository: ETS
File: rebase.py
Lines: 160-183

Signature

def expand(self, node, wid):
    """
    Fork node state and generate + score new reasoning steps.

    Args:
        node (TreeNode): Parent node to expand
        wid (int): Number of children to create (fork width)

    Side effects:
        - Appends (fork, parent) pairs to self.running_list
        - Appends forks to self.history_list
        - Increments global_num_model_calls
    """

Import

# Internal method of Tree class in rebase.py
# Uses SGLang primitives:
from sglang import gen, assistant

I/O Contract

Inputs

Name	Type	Required	Description
node	TreeNode	Yes	Parent node whose state will be forked
wid	int	Yes	Number of child forks to create

Outputs

Name	Type	Description
running_list entries	list[tuple]	(fork_state, parent_node) pairs appended to self.running_list
history_list entries	list	Fork states appended to self.history_list for token counting

Usage Examples

Called by select_and_expand

# From Tree.select_and_expand() at rebase.py:L632-634
for expand_node, width in zip(nodes, widths):
    if width >= 1:
        self.expand(expand_node, width)

# After expansion, insert new nodes:
running_list = tree.get_running_list()
for state, parent in running_list:
    tree.insert(state, parent)

Internal Fork Logic (llemma model)

# Conceptual view of what expand() does for llemma models
state = node.get_state()
forks = state.fork(wid)        # Fork KV cache into wid copies
for fork in forks:
    fork.set_score_backend(self.reward_backend)  # Connect to PRM
    fork += gen("step", max_step_tokens, stop="ки", temperature=T)  # Generate step
    fork += gen("score", max_tokens=0, forward_only=True, logits_require_id=8094)  # Score step

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment