[RFC] Multimodal encoder TP — in-tree implementation vs. upstream import

[zhaochenyang20]

Splitting this off from #89 / the encoder-TP thread in #188 — @ischencheng asked the right question above about how encoder TP fits into the post-refactor architecture (encoders living under SimpleScheduler). #89 was closed when we folded it into the broader refactor, but the work itself didn't go away; it deserves its own tracking issue so the design conversation has a home.

Why encoder TP matters more than text-side TP

The text encoder story is essentially "tokenizer + embed lookup" — bytes in, ids out, near-zero GPU footprint. Multimodal encoders are real models: the Qwen3-Omni audio encoder and image encoder are both on the order of hundreds of millions of parameters with activation peaks that scale with input length. On the current Qwen3-Omni speech pipeline, the audio + image encoders co-located on the thinker GPU are the main reason long-video inputs OOM (see the discussion in #327 and the --encoder-mem-reserve workaround in #339). The activation footprint of the encoder grows with input length in a way the thinker AR loop simply doesn't.

So encoder TP isn't a marginal optimization. For any pipeline serving long video or high-resolution multi-image inputs, it's a production requirement. We should pick it back up as a first-class workstream, not let it die in the cracks of the v1 refactor.

Two paths

Path A — build encoder TP inside sglang-omni. Wrap the encoder modules under models/<name>/components/ with a TP layer, manage the process group from inside SimpleScheduler (or a dedicated EncoderScheduler if SimpleScheduler stays minimal), use NCCL all-reduce inside encoder forward, share the rest of the stage machinery with the non-TP path. Self-contained, lands fast. Cost: we own a parallel TP implementation for every encoder we ever add, and we duplicate work that sglang main has already done for its own VLM serving.

Path B — import encoders from sglang main and inherit native TP. sglang upstream already has the multimodal towers (Qwen2-VL, Qwen3-VL, Qwen3-Omni vision/audio encoders) registered in its model registry, and those implementations are wired into sglang's TP infrastructure — ColumnParallelLinear, RowParallelLinear, parallel state, etc. — directly. If sglang-omni imports those encoder definitions from sglang main instead of carrying its own copies under models/<name>/components/image_encoder.py, encoder TP comes for free, and we don't maintain a second implementation. This is also consistent with the broader positioning of sglang-omni as the out-of-tree omni framework: anything that's already core sglang (model definitions with native parallelism, scheduler primitives, KV cache) we re-use rather than reimplement.

Why I lean toward Path B as the long-term direction

Three reasons.

First, maintenance cost. The encoder ecosystem actually lives upstream — every new VLM that lands in sglang main gets a vision tower added there. Tracking upstream is much cheaper than re-wrapping each one in sglang-omni.

Second, correctness. sglang main's encoders are exercised by upstream VLM serving CI, real users, real workloads. Our reimplementation wouldn't have that surface area; subtly divergent attention masking or RoPE in the encoder would surface much later in production.

Third, this directly serves #188's line-count goal. Every encoder definition we delete from models/<name>/components/ and replace with from sglang.srt.models.X import VisionEncoder is a contribution to the ~33K → ~10K reduction.

What Path B actually needs (the hard part)

Path B sounds simple — "just import from upstream" — but the design surface is non-trivial:

1. Interface contract with SimpleScheduler. sglang main's encoder modules are designed to be invoked from inside a sglang Scheduler/ModelRunner that owns a TP process group. SimpleScheduler currently doesn't own a process group — it's a single-process inbox→fn→outbox loop. Three options to settle: either SimpleScheduler grows TP-awareness, or encoder stages route through a different scheduler type that does, or we wrap each encoder in a minimal sglang TpModelWorker and let sglang manage the parallelism while SimpleScheduler just dispatches batches into it. The third option is probably the cleanest — it keeps SimpleScheduler genuinely simple and doesn't fork sglang's TP execution path.

2. Stable upstream import surface. Same coupling concern as the OmniScheduler ↔ sglang Scheduler discussion in SGLang-Omni Refactoring Proposal #188 — the more sglang internals we touch, the more we break on upgrade. Encoder modules are probably more stable than scheduler internals, but we should still treat them as a versioned contract surface (which import paths we depend on, which class signatures we promise to handle), not a free dependency.

3. Encoders not yet upstreamed. Some encoders we use are in sglang-omni but not yet in sglang main (custom Fish audio tokenizer, future Boson audio encoder). For those, the choice is upstream them first or live with Path A locally for them. The framework needs to support both modes side by side: encoder loaded from sglang main with native TP, or encoder loaded from sglang-omni components/ without it. The pipeline config should not have to know which mode a stage is in.

4. Memory accounting under TP. sglang main's TP encoders assume ownership of an allocator budget on each TP rank. Our co-location story (encoder + thinker on the same GPU, [Feat] Expose encoder mem reserve as --encoder-mem-reserve CLI flag #339) collides with this directly. The total_gpu_mem_fraction / mem_fraction_static design from Ratish's comment in SGLang-Omni Refactoring Proposal #188 has to extend to encoder stages cleanly, including the case where the encoder is itself sharded across N GPUs.

[zhaochenyang20]

SGLang Native TP for Scheduler is supported. We could use SGLang Scheduler's abstraction for Encoder and add TP for it. from @JingwenGu0829

[zhaochenyang20]

https://njpxzlfqwrqq.jp.larksuite.com/wiki/VTcEwGVEuihm6xkquEdjc424pNb

@ischencheng

[ischencheng]

As discussed offline , @Hayden727 and I'll take this.

[Hayden727]

Here is the implementation direction I propose for the next step. We will strictly follow Path B from the issue: reuse SGLang main’s encoder/model definitions and TP infrastructure, while keeping SGLang-Omni responsible for pipeline orchestration, payload adaptation, and stage lifecycle management.

Key Changes

Add a Generic EncoderScheduler

Introduce a reusable EncoderScheduler to replace the current assumption that encoder stages must be backed by SimpleScheduler.

• Keep SimpleScheduler simple and avoid adding TP-specific logic to it.
• Make EncoderScheduler implement the v1 scheduler interface: inbox, outbox, start, stop, and abort.
• On the leader rank, receive requests from the pipeline inbox, perform batching/cache handling, run encoder forward, and emit downstream payloads.
• On follower ranks, do not receive pipeline requests directly. Instead, follow the leader through the TP group and enter encoder forward in the same order.

Add a Generic Encoder Backend/Adapter Layer

Introduce a reusable backend and adapter abstraction for encoder execution.

• EncoderBackend: supports local and sglang implementations, exposing a unified forward(model_inputs) -> dict interface.
• EncoderAdapter: owns payload/model conversion, including build_request, apply_result, batch_inputs, and split_outputs.
• Each encoder registers its adapter/backend spec.
• The scheduler remains model-agnostic and does not contain model-specific branches.

Use SGLang Main for Encoder Execution

The sglang backend should use encoder/model definitions from SGLang main.

• Load upstream encoder or model wrappers through a centralized import contract.
• Rely on SGLang main for TP layers, weight partitioning, parallel state, and NCCL collectives.
• If an encoder does not have an upstream SGLang spec, tp_size > 1 should fail clearly.
• For tp_size == 1, local fallback remains available.

Reuse v1 Stage-Level TP

Use the TP process model already present in sglang_omni_v1.

• Continue using StageConfig.tp_size and gpu: list[int].
• Reuse StageGroup and StageProcessSpec for leader/follower process management.
• Encoder factories should accept the already-injected TP arguments: gpu_id, tp_rank, tp_size, and nccl_port.
• Do not introduce a separate follower-process system inside the encoder executor or backend.

cc @ischencheng @zhaochenyang20

[ischencheng]

For human:
https://github.com/ischencheng/sglang-omni/blob/encoder-tp-plan-b-rfc/docs/developer_reference/encoder_tp_path_b_design_lean.md

For AI:
https://github.com/ischencheng/sglang-omni/blob/encoder-tp-plan-b-rfc/docs/developer_reference/encoder_tp_path_b_design.md

Moved to here

[zhaochenyang20]

@ischencheng — we chatted offline about how the encoder gets imported and partially loaded, and you had a pretty clear picture of how to handle the weight scope (avoiding the full-model load on every encoder process so co-located encoder + thinker doesn't OOM).

I think that part really needs to land in the RFC explicitly — it's load-bearing for the whole long-video OOM story this design is supposed to fix. Right now the worker just calls get_model(...) and the question of "what weights actually get loaded under encoder_only=True" isn't answered anywhere.

Could you add a short "Weight loading scope" section next to SGLangEncoderWorker? Just write down what you walked me through — the partial-load approach on our side, why we're not blocking on an upstream PR, and the note that we'll propose an EncoderModelLoader helper to upstream after Phase 1. Important enough that I'd rather it be a first-class section than an Open Question.

[zhaochenyang20]

Naming: SGLangEncoderWorker should be SGLangEncoderRunner (and create_image_encoder_executor → ..._runner) to match #188's vocabulary. Worker / Executor are exactly the terms #188 is collapsing — let's not reintroduce them on the encoder path.

[zhaochenyang20]

One more ask: the RFC jumps straight into architecture without a clear motivation section. Could you add a short one up front?

The core point to make explicit: this is solving the long-sequence OOM problem. Encoder weights are small (~2.5 GB combined for Qwen3-Omni vision + audio); the issue is activation memory scaling with input length. A 1-minute video can push encoder activations to 30+ GB on a single GPU, which is what's pushing the thinker OOM today and forcing the --encoder-mem-reserve workaround in #339. TP shards the activation, not the weights — that's the real win.

Without this framing, reviewers (and our future selves) will read the RFC and ask "why not just give the encoder a bigger GPU" or "why not DP". A two-paragraph motivation section that names the activation-vs-weight distinction and points at #327 / #339 as the concrete pain saves that round-trip.

@ischencheng

[zhaochenyang20]

Small naming nit on the new code: in EncoderScheduler and SGLangEncoderWorker, can we use entry_rank / non_entry_rank consistently instead of leader / follower?

Two reasons:

First, "leader/follower" carries Raft / Kafka connotations — leader election, failover, consensus. None of which apply here. Rank 0 is hardcoded, not elected; followers don't take over on failure. The asymmetry is just who owns external IO (ZMQ recv, relay reader, outbox emit), which "entry rank" describes precisely without the misleading vocabulary.

Second, the RFC already uses is_entry_rank in places, and OmniScheduler uses the same term (scheduling/omni_scheduler.py:393-403). Standardizing on it across the new scheduler keeps the v1 vocabulary self-consistent.

Stage-level single/leader/follower is fine to leave alone — that's existing v1 surface area and not worth churning. Just the new EncoderScheduler / Worker code.

[zhaochenyang20]

@ischencheng — the depth here is great, but I think the RFC could be roughly half its current length without losing the load-bearing parts.

A lot of what's in it reads as implementation rather than design: full Python sketches, per-sub-step unit tests with assertion-level detail, per-key justification in _ENCODER_PROTECTED_KEYS. Those are better as code comments or a test-plan tracking issue linked from the Phase 0 PR. The RFC's job is the decisions and contracts; everything else competes for the reviewer's attention.

Worth settling on a leaner convention now — easier to write, easier to review, easier for the next person to find the key calls six months from now.