Refactor DSPy-driven optimization flow with coordinator-based orchestration by sandlbn · Pull Request #48 · IntelLabs/Xe-Forge

sandlbn · 2026-07-20T18:50:33Z

This PR refactors the DSPy integration to make optimization flow more modular, template-driven, and agentic.

add a new CoordinatorAgent and coordinator tools to orchestrate optimization steps dynamically
simplify and refactor analyzer, planner, ReAct, and optimizer agent flows
move prompt construction toward reusable rendered template instructions
add device-specific prompt helpers and new prompt templates for analysis, planning, optimization, autotuning, and coordinator behavior
update pipeline wiring to support the new coordinator-driven strategy
clean up package exports and related config/model references

Copilot

Pull request overview

This PR refactors the DSPy-driven kernel optimization flow to be template-driven and adds a new coordinator-based orchestration strategy, while consolidating shared agent utilities and updating pipeline wiring/config to support the new strategy.

Changes:

Introduces CoordinatorAgent + coordinator tools to drive dynamic analyze/apply/benchmark/profile loops via dspy.ReActV2.
Moves signature instruction content into reusable Jinja2 templates and injects them via append_instructions() across analyzer/planner/optimizer agents.
Updates pipeline strategy wiring, shared utilities, and result metadata (token usage), and bumps the project version/dependency.

Reviewed changes

Copilot reviewed 26 out of 26 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
src/xe_forge/prompts/templates/sycl_optimization_signature.md.j2	Adds SYCL/CUTLASS optimization instruction template.
src/xe_forge/prompts/templates/sycl_algorithmic_signature.md.j2	Adds SYCL algorithmic optimization instruction template.
src/xe_forge/prompts/templates/planning_signature.md.j2	Adds planning/stage-ordering instruction template.
src/xe_forge/prompts/templates/optimization_signature.md.j2	Adds templated optimizer-stage guidance (device-conditional).
src/xe_forge/prompts/templates/optimization_react_signature.md.j2	Adds templated ReAct optimizer guidance (device-conditional).
src/xe_forge/prompts/templates/coordinator_signature.md.j2	Adds coordinator orchestration instruction template.
src/xe_forge/prompts/templates/autotune_signature.md.j2	Adds autotune instruction template (device-conditional).
src/xe_forge/prompts/templates/analysis_signature.md.j2	Adds analysis instruction template including injected issue-category blocks.
src/xe_forge/prompts/templates/algorithmic_signature.md.j2	Adds algorithmic optimization instruction template.
src/xe_forge/prompts/device_prompts.py	Adds template rendering + device addendum helpers for signatures.
src/xe_forge/prompts/init.py	Adds Jinja2 environment and `render_signature_instructions()` export.
src/xe_forge/planner.py	Refactors planner to inject template-based planning instructions.
src/xe_forge/pipeline.py	Wires coordinator strategy, template-based instructions, token usage tracking, and helper refactors.
src/xe_forge/models.py	Adds `token_usage` field to `OptimizationResult`.
src/xe_forge/config.py	Expands agent `strategy` comment to include `coordinator`.
src/xe_forge/cli.py	Updates LiteLLM HTTP clients configuration for DSPy setup.
src/xe_forge/agents/utils.py	Introduces shared utilities (`extract_gemm_dims`, `verify_sycl`, `SUCCESS_MESSAGE`).
src/xe_forge/agents/react_agent.py	Refactors ReAct agent to use shared utilities, template instructions, and `ReActV2`-style history handling.
src/xe_forge/agents/optimizer_agent.py	Refactors optimizer agent to use shared utilities and injected template instructions.
src/xe_forge/agents/cover.py	Adjusts CoVeR fallback signature and extraction predictor behavior.
src/xe_forge/agents/coordinator.py	Adds CoordinatorAgent implementation using `dspy.ReActV2` and coordinator tools.
src/xe_forge/agents/coordinator_tools.py	Adds CoordinatorState and tool closures for analyze/retrieve/apply/benchmark/profile/status.
src/xe_forge/agents/analyzer_agent.py	Injects template-based analysis guidance and device/DSL naming helpers.
src/xe_forge/agents/init.py	Exports CoordinatorAgent from the agents package.
src/xe_forge/init.py	Switches to lazy import for `XeForgePipeline`.
pyproject.toml	Bumps version and updates DSPy dependency/source configuration.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

The performance speedup we achieved was up to 13.3x, as described in the paper, this must be a typo.

gbenms · 2026-08-03T23:04:27Z

+                candidates.append(result)

-        for attempt in range(best_k):
+            # Fixed stage loop runs only when coordinator strategy is not active.


should be under if-else: missing else for the block of code that does optimization according to a fixed stages should only be active if there is no coordinator (also fits the comment here)- else includes: l379-638 (up to: candidates.append(result))

gbenms · 2026-08-03T23:25:38Z

            if not analysis.detected_issues:
                result.success, result.optimized_code = True, kernel_code
                candidates.append(result)
-                continue


this should be replaced by an else, so the pipeline wouldnt continue to run for an empty issues kernel, ends at: result = max(candidates, key=lambda..

gbenms · 2026-08-03T23:29:36Z

            if not stages_to_apply:
                result.success, result.optimized_code = True, kernel_code
                candidates.append(result)
-                continue


this should be replaced by an else, so the pipeline wouldnt continue to run for an empty stage kernel, ends at: result = max(candidates, key=lambda..

gbenms · 2026-08-03T23:31:26Z

+                    result.total_speedup = speedup
+                candidates.append(result)

-        for attempt in range(best_k):


should update README that top-k was removed

gbenms · 2026-08-03T23:33:50Z


-        for attempt in range(best_k):
+            # Fixed stage loop runs only when coordinator strategy is not active.
+            best_k = max(1, self.config.optimization.best_k) if self.coordinator is None else 0


these two lines should be deleted- best-k not used for optimization anymore.

gbenms · 2026-08-03T23:37:07Z

 from xe_forge.models import OptimizationResult, OptimizationStage
-from xe_forge.pipeline import XeForgePipeline

 __version__ = "0.2.0"


change to 0.3.0 to match pyproject.toml

sandlbn added 2 commits July 13, 2026 22:31

Refactor DSPY

84d78ea

Patch ruff

a42ac3d

sandlbn requested a review from Copilot July 20, 2026 18:50

Copilot started reviewing on behalf of sandlbn July 20, 2026 18:51 View session

Patch ruff

719b6c4

sandlbn self-assigned this Jul 20, 2026

sandlbn requested review from danielfleischer and gbenms July 20, 2026 18:52

Copilot AI reviewed Jul 20, 2026

View reviewed changes

sandlbn and others added 7 commits July 20, 2026 11:57

Potential fix for pull request finding

49ac80d

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

7ff0be2

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Remove hardcoded dspy path

cb8c27c

Potential fix for pull request finding

8a2f7a3

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

872837f

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Fix failing usage tracking

0607011

FA speedups

f1d6def

The performance speedup we achieved was up to 13.3x, as described in the paper, this must be a typo.

gbenms requested changes Aug 3, 2026

View reviewed changes

sandlbn added 2 commits August 3, 2026 17:12

Remove completly best-k

5e5d3f0

Reformat files

a706d9d

gbenms approved these changes Aug 4, 2026

View reviewed changes

Reformat files

40ffbe1

sandlbn merged commit 4dcb508 into main Aug 4, 2026
2 checks passed

sandlbn deleted the mspoczy/refactor branch August 4, 2026 01:52

Navigation Menu

Uh oh!

Refactor DSPy-driven optimization flow with coordinator-based orchestration - #48

Refactor DSPy-driven optimization flow with coordinator-based orchestration#48
sandlbn merged 13 commits into
mainfrom
mspoczy/refactor

sandlbn commented Jul 20, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gbenms Aug 3, 2026

Uh oh!

gbenms Aug 3, 2026

Uh oh!

gbenms Aug 3, 2026

Uh oh!

gbenms Aug 3, 2026

Uh oh!

gbenms Aug 3, 2026

Uh oh!

gbenms Aug 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

sandlbn commented Jul 20, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

gbenms Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants