SGLang-Omni Refactoring Proposal

[zhaochenyang20]

Current Status

SGLang-Omni already has solid support for foundational omni models:

• Fish Audio S2 Pro: A TTS model merged into the main branch, with CUDA Graph support (PR [optimize] Add CUDA Graph Support for S2 Pro Model #153) and streaming inference.
• Qwen3-Omni: End-to-end multimodal understanding + speech synthesis. The Thinker (multimodal understanding) is merged into the main branch; the Talker (speech synthesis) is integrated in PR [WIP] Talker integration + optimizatioon #155 (talker_jing branch) and needs to be rebased onto main.

Leveraging SGLang's kernel-level optimizations, we have further applied torch.compile and CUDA Graph capture to many smaller components, yielding significant performance improvements.

Core Issues

The codebase exhibits clear signs of vibe coding — rapid iteration has accumulated substantial redundancy and over-abstraction. Specifically:

Excessive Abstraction Layers

A single request must pass through the following layers from HTTP API to torch.forward:

HTTP API → Client → Coordinator → Stage → Worker → Executor → Engine → Scheduler → ModelRunner → torch.forward

Among these, Stage → Worker → Executor → Engine are four consecutive layers whose responsibilities heavily overlap: each one essentially "receives a request, delegates execution, and returns results." Based on experience from the SGLang main repo, we should be able to go directly from Stage → Engine, eliminating the Worker and Executor layers entirely.

[We defined the Torch version first, then implemented the SGLang backend. As we will detach the torch implementation, we can make it rather concise.]

[TTS/ASR ModelRunner should be combined? Qwen-TTS using a diffusion-like ModelRunner. We should take into consideration. Adarsh's implementation of Qwen-TTS by Torch backend #76 . We need a dedicated RFC on what SGLang ModelRunner could support.]

[Should we remove the Torch backend or still have it?]

[(Least priority) Pipeline: trying to place multiple stages on a single GPU. Modify SGLang RunTime for this could work. Also considering the stage of fusing to cutting off communication.]

Conflicting Same-Named Concepts Create Cognitive Overhead

Concept	Instances	Location	Justified?
Scheduler	2	engines/omni/scheduler.py vs engines/ar/sglang_backend/scheduler/	Needs unified naming
ModelRunner	3	engines/omni/model_runner.py, engines/omni/runtime/sglang_ar.py, engines/ar/sglang_backend/model_runner.py	Needs unified naming
Message	2	client/types.py vs fishaudio_s2_pro/fish_speech/conversation.py	Justified — different domain concepts

The Connection instances across different relay backends (nixl vs nccl) are justified, as they describe connection semantics for different transport protocols with naturally isolated namespaces. The real issues to address are the naming collisions for Scheduler and ModelRunner.

Too Many Adaptation Layers, Bloated Glue Code

• EngineExecutor uses getattr to dynamically delegate prepare_stream, stream, start, and stop, essentially translating the Engine interface into the Executor interface — the two are nearly isomorphic. Any dynamic attribute delegation is a signal of bad design, indicating unclear interface definitions.
• Engine(ABC) and Executor(ABC) have highly similar method signatures (add_request, get_result, abort), differing only in data types (StagePayload vs Any).
• FusedExecutor chains multiple Executor instances, but in practice, stage fusion in the pipeline config already handles this at compile time.

Code Redundancy

Redundancy	Lines	Description
config/qwen3_omni.py	~378	Entirely duplicates models/qwen3_omni/hf_config.py, with zero references across the codebase
Relay Put/Get pattern	~400	Polling + timeout + callback pattern repeated across 4 relay implementations (shm, nccl, nixl, mooncake)
Test dummy classes	~200	Duplicate definitions across test_scheduler.py, test_scheduler_streaming.py, test_factory.py
_inject_multimodal_embeds	~150	Highly similar mask-based embedding injection logic in sglang_ar.py and s2pro_sglang_ar.py

Large File Concentration

Total codebase is ~33K lines (.py), with 15 files exceeding 500 lines. The largest, text2semantic/modeling.py, reaches 1,484 lines. sglang_ar.py (1,040 lines) consolidates BatchPlanner, ResourceManager, OutputProcessor, IterationController, and ModelRunner into a single file — far too many responsibilities.

Refactoring Goals

1. Eliminate at least two layers of excessive abstraction, reducing the request path from 8–10 layers to 6.
2. Substantially reduce total code volume (target: from ~33K lines to under ~10K lines).
3. Maintain existing functionality and performance — the refactoring must not introduce regressions.
4. Lower the cost of onboarding new models — adding a new model should only require implementing the model definition + pipeline config, without needing to understand this many middleware layers.

[Ratish1]

This comment is resolved.

Suggestions by Ratish

Hey @zhaochenyang20 , I have read through your proposal and verified it against the codebase. Some of my thoughts below:

I agree with the direction of this proposal. The current architecture feels heavier than it needs to be on the execution path, and there is real cleanup value here around adapter bloat, naming drift, and orchestration being concentrated in too few places.

Current Architecture:
Today the hot path is long: HTTP/API -> client -> coordinator -> stage -> worker -> executor -> engine -> scheduler -> model runner -> forward. That layering gives flexibility, but it also makes the system harder to understand and reason about than it should be.

A few things stood out from reviewing the code:

• There is at least some clearly removable duplication, including a duplicate Qwen config.
• The naming collision around Scheduler and ModelRunner is real and creates unnecessary cognitive overhead.
• The Engine/Executor boundary is not clearly defined. Important behaviors like streaming, lifecycle hooks, and feedback wiring are handled through runtime capability checks rather than an explicit interface.
• Some orchestration-heavy files are carrying too many responsibilities at once, especially the SGLang AR runtime and talker executor paths.

We should rather stage this process and create sub issues first under this?. Maybe, we can follow this order below:

• remove confirmed dead duplication first
• clean up overloaded names
• make execution contracts more explicit
• split orchestration-heavy files
• then evaluate how much of the current adapter stack can naturally collapse once responsibilities are clearer.

So overall, I support the motivation behind this issue. we ahould just present it as a staged architectural simplification of the current system based on this.

Response by chenyang

Thanks for the feedback — really appreciate the thoughtful input. @Ratish1

I've had in-depth discussions with the person who originally designed the core primitives. Many of the current abstractions exist because we initially built on a Torch inference backend, and then later integrated the SGLang model runner as a second backend. A lot of the abstraction layers were introduced at the Torch level. Now that we have the SGLang backend, several of those layers can be removed entirely. And objectively, we need to retire the Torch backend — if anyone wants a Torch baseline, the right comparison target is HuggingFace Transformers directly.

I'm fully committed to this refactoring effort, even though this is my first time leading something at this scale. I believe that compared to shipping one or two individual features (like the TTS accuracy benchmark or TTS torch.compile work), system design is a much bigger and more impactful undertaking. Opportunities to do system design at this scope are rare, and I don't take it lightly. I'd love for you — and anyone else reading this — to get involved.

That said, here's the current plan:

1. First, finish the accuracy/performance benchmarks for TTS and Qwen3-Omni, then integrate all four test suites into CI. These benchmarks are our safety net during refactoring — we cannot allow regressions in correctness or performance.
2. Then, begin the refactoring itself.

On the refactoring approach: I'd like to prioritize design over implementation. Once the design is solid, Claude can produce high-quality code. As for how to arrive at good design — in my experience, the first step is a thorough code walkthrough to understand how the current abstractions actually work. My Awesome-ML-SYS-Tutorial repo has many examples of this pattern. Huapeng is currently leading the code walkthrough effort and will share his findings soon. @PopSoda2002

[Kare0638]

Very clear proposal ! I hope that I can participate in this great work with my best efforts. I am new to here and I will get more familiar with the project as soon as possibe during the first benchmarks work.

[zhaochenyang20]

We move the design to https://njpxzlfqwrqq.jp.larksuite.com/wiki/VTcEwGVEuihm6xkquEdjc424pNb?from=from_copylink

[adityavaid]

This comment is resolved:

Comments by adityavaid

Amazing proposal overall.
One suggestion though, let’s also measure success by new model bring-up effort, not just LOC reduction.

Something like:

• how many files need touching to add a model
• how many framework interfaces someone must implement
• how much custom test scaffolding is needed
• time to get first end-to-end request working

If we can make those numbers consistently smaller release-over-release, we’ll improve maintainability and onboarding in a much more meaningful way

Response by Jingwen

Hi @adityavaid When designing this framework the most important thing I considered is actually extensibility.

Prioritized modualizing modelrunner and scheduler which are the key in all types of omni models we need to adapt. I will create a more detailed report explaining this in this couple of days.

It's more than great if you found there's specific scenarios where we can design it to better adapt new model. Comment anytime or dm if you have any specific ideas!

[haojin2]

This comment is resolved:

Comments by Hao

Overall great proposal but a few things might need clarification:

• For Fish S2 Pro model, looks like some post-forward steps are scoped within forward() function, while Qwen3 uses SGL's sampler. I imagine there might also be other models that has some post-forward actions mixed into the forward function. How do we plan to take of such cases in the base ModelRunner? My intuition is to make the "shared" function as no-op by default to avoid un-expected fallback behaviors
• For OmniScheduler it says to reuse Sglang's scheduler, is it copying the same code or inheriting the class?
• By removing "Worker" we might also get rid of the current entrance to stage-level parallelism (which is not used by any of the 2 existing models), but if we do need stage-level parallelism in the future, what are the plans for still supporting that?

Discussion by adityavaid

Hi @JingwenGu0829, Thanks!
So preface of this is that I tried to review the codebase from a “what if I add a new model today” angle.

A few data points stood out:

• There are still multiple Scheduler/ModelRunner paths in code, which makes the extension path unclear for new contributors.
  This is already touched upon in the thread, having separate engines/ar/sglang_backend and engines/omni path. From the gist of the thread, I get that we may be going with only keeping one in future. So for future model addition , is this gonna be enforced in any way ?

• Stream/feedback wiring is still mostly implicit (capability checks/internal wiring), which is powerful but harder to discover when onboarding.

From a bring-up perspective, I think we can consider :

• Adding a fail-fast onboarding checks (registry + factory path validation at startup or during build)

• ( purely a proposal) Adding a Versioned onboarding contract (ModelAdapterSpec), Think of this as a template every new model must follow.
  Example: “You must provide these files, these functions, and these tests.” “Versioned” means if the template changes later (v1 -> v2), we can support both clearly.

• Make stream/feedback capabilities explicit in interfaces . Instead of special-case wiring per model, model declares capabilities up front. Example: supports_feedback_loop=true, supports_hidden_capture=false

Again these are just suggestions based on the issues I faced while trying to add a new model.

Response by Jingwen

Hi @haojin2 Thanks! These are really good questions.

1. Yes, your intuition aligns with the design plan right now. Note that the current implementation of both Fish and Qwen are bad. But in the future at lease qwen's talker and code predictor will be put into one modelrunner to enable sglang's native CUDA graph.

So in the plan of refractor, no-op defaults are the right approach. The base ModelRunner just calls tp_worker.forward_batch_generation() and lets the model do whatever it wants inside forward(). The hooks (prepare_forward/post_forward) are opt-in. If a model handles everything internally(like a vanilla LLM), the callbacks in extract_output_fn just read the buffers. If a model needs external post-processing (Qwen3 talker, fish, future more complex model), pre and post_forward does it.

2. In the current situation we can neither copy nor inherit the base sglang scheduler. Since the native sgalng scheduler is super heavy with 300+ attributes we don't have in sgalng-omni when init. So, aligning with sglang-diffusion, the plan is to use composition instead of inheritance, omniScheduler is a standalone class that uses SGLang's PrefillManager and DecodeManager (internal components of SGLang's Scheduler) for batch selection, but owns its own event loop and request lifecycle.

However, I do think the sglang main repo needs a refractor to get a clean basic scheduler class for main, omni and diffusion scheduler to inherit on. But this is not really possible since it's sglang main repo's long-term design planning.

3. We are not removing the Worker abstraction in sglang native framework. Our ModelWorker wraps SGLang's full TpModelWorker → ModelRunner stack. The 'worker' removed in the proposal is purely another different layer, but with the same 'worker' name. It's non-sglang-native, and it handles async dispatch layer between Stage and Executor that handled request queuing, result routing, and stream management. It's a complete redundant layer and will be removed. However the worker abstraction for things like TP is sglang-native and will be kept for sure.

[ischencheng]

Discussion by Chen Cheng for Encoder TP

Hi @zhaochenyang20, I'm interested in implementing encoder TP (originally #89). I see the refactoring design puts encoders under SimpleScheduler. Do you have any plans for how TP would fit into the new architecture? I'd like to contribute this, want to make sure it aligns with your direction before I start.

This is moved to #375.