fix(tree-of-thought): avoid Context::fork() to fix output corruption on Vulkan/Windows

[zatchbell1311-wq]

Summary

Context::fork() causes the model to produce incoherent/garbage output immediately after forking, even with a single fork (num_branches=1) and no concurrency. Plain chat completions without fork work correctly with the same model.

Root cause

Traced to CopyD2D (via ggml_backend_tensor_get/set with non-zero offsets) silently returning garbage on the Vulkan backend. Filed as issue #418.

Fix

Each branch now creates a fresh Context::new() and replays prior turns as text (system + user/assistant history) instead of relying on KV-cache sharing via fork().

Verified coherent output with num_branches=1 and num_branches=2 (8 concurrent leaves) at 128 and 512 tokens/step on Qwen3-0.6B (Vulkan/RTX 2050).

Tradeoff

Loses the KV-cache sharing fork() was meant to provide. The underlying fork() bug (likely in runtime/src/context/snapshot.rs, GPU page copy) remains unfixed and may warrant a separate investigation.

Summary by CodeRabbit

• Refactor
  • Improved the tree-of-thought generation flow for more consistent multi-step outputs.
  • Changed how parallel branches are prepared and replayed so each branch now runs with a cleaner, isolated context.

[coderabbitai]

Walkthrough

The tree-of-thought inferlet's parallel generation logic was refactored to replace Context::fork() with fresh Context instances created per branch at each level. Instead of chaining forked contexts, each level replays prior conversation history (propose/execute outputs) as user/assistant messages before generating new content.

Changes

Tree-of-Thought Context Refactor

Layer / File(s)	Summary
Replace context fork with per-branch context replay inferlets/tree-of-thought/src/lib.rs	Level 1 generates propose text using a fresh propose_ctx with an inlined sys_prompt; Level 2 builds a new execute_ctx replaying the propose exchange before sending EXECUTE_PROMPT; Level 3 builds a new reflect_ctx replaying propose/execute exchanges before sending REFLECT_PROMPT. Added clone() calls carry prompt/text state across nested async closures, and print labels were updated accordingly.

Sequence Diagram(s)

sequenceDiagram
  participant Level1 as Level1Task
  participant Level2 as Level2Task
  participant Level3 as Level3Task
  participant Model

  Level1->>Model: load model, propose_ctx.generate(propose prompt)
  Model-->>Level1: propose text
  Level1->>Level2: pass propose prompt + propose text
  Level2->>Model: execute_ctx replay(user: propose, assistant: propose text, user: EXECUTE_PROMPT)
  Model-->>Level2: execute text
  Level2->>Level3: pass propose/execute prompts + texts
  Level3->>Model: reflect_ctx replay(user: propose, assistant: propose text, user: execute, assistant: execute text, user: REFLECT_PROMPT)
  Model-->>Level3: reflect text

Loading

Possibly related issues

• ctx.fork() produces incoherent output on Vulkan/Windows — CopyD2D in aux_server.cpp round-trips through ggml_backend_tensor_get which doesn't support partial tensor reads on ggml-Vulkan #418: Describes fork-induced corruption issues, and this PR's replacement of Context::fork() with fresh per-branch Context::new() plus conversation replay directly addresses that workaround.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly states the main fix: replacing Context::fork() to address Vulkan/Windows output corruption.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

inferlets/tree-of-thought/src/lib.rs (1)

79-84: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Batch replayed turns into one prefill.

system/user/assistant/cue only buffer tokens, and generate() drains that buffer on the first step. Flushing after every replayed message adds several avoidable forward passes per node; build the full replay, call cue(), then generate.

♻️ Proposed refactor

 propose_ctx.system(sys_prompt);
-propose_ctx.flush().await?;
 let propose_prompt = format!("{}{}", PROPOSE_PROMPT_TEMPLATE, question_);
 propose_ctx.user(&propose_prompt);
-propose_ctx.flush().await?;
 propose_ctx.cue();

 // ...

 execute_ctx.system(sys_prompt);
-execute_ctx.flush().await?;
 execute_ctx.user(&propose_prompt_);
-execute_ctx.flush().await?;
 execute_ctx.assistant(&propose_text_);
-execute_ctx.flush().await?;
 execute_ctx.user(EXECUTE_PROMPT);
-execute_ctx.flush().await?;
 execute_ctx.cue();

 // ...

 reflect_ctx.system(sys_prompt);
-reflect_ctx.flush().await?;
 reflect_ctx.user(&propose_prompt_3);
-reflect_ctx.flush().await?;
 reflect_ctx.assistant(&propose_text_3);
-reflect_ctx.flush().await?;
 reflect_ctx.user(EXECUTE_PROMPT);
-reflect_ctx.flush().await?;
 reflect_ctx.assistant(&execute_text_3);
-reflect_ctx.flush().await?;
 reflect_ctx.user(REFLECT_PROMPT);
-reflect_ctx.flush().await?;
 reflect_ctx.cue();

Also applies to: 102-110, 129-141

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@inferlets/tree-of-thought/src/lib.rs` around lines 79 - 84, The replay in
`propose_ctx` is doing separate flushes after `system` and `user`, which causes
extra forward passes instead of one prefill. Update the replay flow in the
`propose_ctx` setup to buffer all replayed turns first, then call `cue()` once
before `generate()`, and apply the same batching approach to the other replay
blocks referenced by `propose_ctx` so the context is built in one pass.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@inferlets/tree-of-thought/src/lib.rs`:
- Around line 79-84: The replay in `propose_ctx` is doing separate flushes after
`system` and `user`, which causes extra forward passes instead of one prefill.
Update the replay flow in the `propose_ctx` setup to buffer all replayed turns
first, then call `cue()` once before `generate()`, and apply the same batching
approach to the other replay blocks referenced by `propose_ctx` so the context
is built in one pass.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b9e1b117-3918-4870-a71e-3ea0bc2c6d2c

📥 Commits

Reviewing files that changed from the base of the PR and between 317abc1 and 83c4d97.

📒 Files selected for processing (1)

• inferlets/tree-of-thought/src/lib.rs