[Feature] Stronger transformers modeling backend with TP, PP, MoE, VLMs, and torch compile

[adarshxs]

Motivation

Adds a generic modeling backend that uses HF transformers models directly via AutoModel.from_config(), enabling any model with a tp_plan, pp_plan and custom attention support to run on SGLang without a dedicated model implementation.

Modifications

Architecture

Mixin-based design that composes capabilities:
• TransformersBase - core class: meta-device init, recursive module replacement (Linear -> TP, RMSNorm -> fused), attention instance creation, PP support, weight loading via
AutoWeightsLoader
• CausalMixin - LM head + logits processor
• EmbeddingMixin - pooling for embedding models
• MoEMixin - auto-detects experts modules and replaces with TransformersFusedMoE (custom-op backed, fused kernels, EPLB recording)
• MultiModalMixin - vision/audio/video encoder dispatch, M-RoPE, token_type_ids propagation

Concrete classes cover all combinations: TransformersForCausalLM, TransformersMoEForCausalLM, TransformersMultiModalForCausalLM, TransformersMultiModalMoEForCausalLM, plus embedding
and sequence classification variants.

Testing

Conducted the following tests. Infrastructure = (1xH100, 4xA100):

Model	Config	Status
Qwen3-0.6B	TP=1, TP=2, TP=1 + Torch Compile, TP=2 + Torch Compile	PASS
Qwen3-30B-A3B (MoE, 128 experts)	TP=1, TP=2	PASS
Gemma3-4B-IT (VLM)	TP=4, TP=4 + Torch Compile, text+image	PASS
Qwen3-VL-2B (VLM)	TP=1, TP=1 + Torch Compile, text	PASS

#	Category	Model	Architecture	Config
1	Non-MoE LLM Tests	Qwen3-0.6B	TransformersForCausalLM	TP=1
2	Non-MoE LLM Tests	Qwen3-0.6B	TransformersForCausalLM	TP=2
3	Non-MoE LLM Tests	Qwen3-0.6B	TransformersForCausalLM	TP=1+compile
4	Non-MoE LLM Tests	Qwen3-0.6B	TransformersForCausalLM	TP=2+compile
5	Non-MoE LLM Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	Single GPU
6	Non-MoE LLM Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	TP=2
7	Non-MoE LLM Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	TP=4
8	MoE LLM Tests	Qwen3-30B-A3B (128 experts)	TransformersMoEForCausalLM	TP=1
9	MoE LLM Tests	Qwen3-30B-A3B (128 experts)	TransformersMoEForCausalLM	TP=2
10	MoE LLM Tests	Qwen3-30B-A3B (128 experts)	TransformersMoEForCausalLM	EP=2 (tp=2, ep=2)
11	MoE LLM Tests	Qwen3-30B-A3B (128 experts)	TransformersMoEForCausalLM	TP=4
12	MoE LLM Tests	Qwen3-30B-A3B (128 experts)	TransformersMoEForCausalLM	EP=4 (tp=4, ep=4)
13	VLM Tests	Qwen2.5-VL-7B-Instruct	TransformersMultiModalForCausalLM	Single GPU, text+image
14	VLM Tests	Qwen3-VL-2B	TransformersMultiModalForCausalLM	TP=1, text
15	VLM Tests	Qwen3-VL-2B	TransformersMultiModalForCausalLM	TP=1+compile, text
16	VLM Tests	Gemma3-4B-IT	TransformersForCausalLM	Single GPU, text
17	VLM Tests	Gemma3-4B-IT	TransformersMultiModalForCausalLM	Single GPU, text
18	VLM Tests	Gemma3-4B-IT	TransformersMultiModalForCausalLM	TP=4
19	VLM Tests	Gemma3-4B-IT	TransformersMultiModalForCausalLM	TP=4+compile
20	VLM Tests	Gemma3-4B-IT	TransformersMultiModalForCausalLM	text+image
21	VLM Tests	Gemma3-4B-IT	TransformersMultiModalForCausalLM	Single GPU, image
22	Pipeline Parallel Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	PP=2
23	Pipeline Parallel Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	PP=4
24	Data Parallel Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	DP=2
25	Data Parallel Tests	Qwen2.5-7B-Instruct	TransformersForCausalLM	DP=4

Checklist

• Format your code according to the Format code with pre-commit.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @adarshxs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances SGLang's model compatibility by introducing a generic backend that directly integrates with Hugging Face Transformers models. The new architecture supports advanced features like tensor and pipeline parallelism, Mixture-of-Experts, and multimodal capabilities, allowing for broader model deployment and more efficient execution, particularly for complex models like VLMs and MoE architectures. This change aims to reduce the effort required to onboard new models and improve performance through optimized operations and torch.compile integration.

Highlights

• Generic Transformers Backend: Introduced a new generic modeling backend that directly utilizes Hugging Face AutoModel.from_config() to support a wider range of models without requiring dedicated SGLang implementations.
• Mixin-Based Architecture: Implemented a modular, mixin-based design (TransformersBase, CausalMixin, EmbeddingMixin, MoEMixin, MultiModalMixin) to compose capabilities like tensor parallelism (TP), pipeline parallelism (PP), Mixture-of-Experts (MoE), and Vision-Language Models (VLMs).
• MoE and Multimodal Support: Added robust support for MoE models by replacing expert modules with TransformersFusedMoE and enhanced multimodal processing through a new TransformersAutoMultimodalProcessor that leverages Hugging Face's native processing for tokenization and feature extraction, including M-RoPE position handling.
• Torch Compile Integration: Integrated torch.compile with compatibility checks, specifically disabling it for models with dynamic RoPE scaling to ensure stable graph compilation.
• Improved Weight Loading and Module Replacement: Refactored weight loading with a new AutoWeightsLoader for greater flexibility and introduced recursive module replacement for linear layers (TP) and RMSNorm layers (fused) within the Transformers backend.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/srt/configs/model_config.py
  • Updated multimodal disabling logic to be conditional on the model implementation, allowing multimodal features for the new Transformers backend.
  • Expanded multimodal model detection to include models with vision or audio sub-configurations.
• python/sglang/srt/managers/io_struct.py
  • Added an optional token_type_ids field to TokenizedGenerateReqInput to support multimodal models that require it.
• python/sglang/srt/managers/multimodal_processor.py
  • Introduced TransformersAutoMultimodalProcessor to handle generic multimodal processing for the new Transformers backend.
• python/sglang/srt/managers/scheduler.py
  • Propagated token_type_ids from received requests to the tokenized object.
  • Made the pad_input_ids_func call conditional, allowing it to be None for certain backends.
• python/sglang/srt/managers/tokenizer_manager.py
  • Extracted and flattened token_type_ids from multimodal inputs for propagation.
• python/sglang/srt/model_executor/model_runner.py
  • Disabled piecewise CUDA graph for the Transformers backend, recommending torch.compile instead.
  • Added a check to disable torch.compile if the Transformers backend model is not compatible (e.g., due to dynamic RoPE scaling).
• python/sglang/srt/model_loader/utils.py
  • Added helper functions _is_moe_model, _is_sequence_classification_model, and _get_transformers_backend_arch to dynamically determine model architecture types.
  • Modified resolve_transformers_arch to dynamically select the appropriate Transformers backend class based on model characteristics (e.g., MoE, multimodal, causal LM, embedding, sequence classification).
  • Improved compatibility checks and warning messages during model class resolution.
• python/sglang/srt/models/transformers.py
  • Refactored the TransformersForCausalLM class into a flexible mixin-based hierarchy (TransformersBase, CausalMixin, EmbeddingMixin, MoEMixin, MultiModalMixin) to support various model types and features.
  • Implemented TransformersBase for core functionalities including meta-device initialization, recursive module replacement (Linear, RMSNorm), attention instance creation, and pipeline parallel support.
  • Introduced TransformersFusedMoE to wrap SGLang's native MoE implementation, handling expert parallelism and weight loading for MoE models.
  • Added can_enable_torch_compile function to determine torch.compile compatibility, specifically addressing dynamic RoPE scaling.
  • Implemented pipeline parallelism logic, including PPMissingLayer to handle skipped layers on non-participating ranks.
  • Updated sglang_flash_attention_forward to align with the new architecture and parameters.
  • Expanded the EntryClass list to include all new mixin combinations for causal LMs, embedding models, and sequence classification models with optional MoE and multimodal support.
• python/sglang/srt/models/utils.py
  • Added the __or__ operator to WeightsMapper for easier combination of weight mapping rules.
  • Introduced AutoWeightsLoader to provide a more robust and configurable mechanism for loading model weights, including options to skip or ignore specific parameters.
• python/sglang/srt/multimodal/processors/transformers_auto.py
  • Added TransformersAutoMultimodalProcessor to generically process multimodal inputs by directly applying Hugging Face processors, handling token expansion, M-RoPE positions, and token_type_ids.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a generic modeling backend for SGLang that leverages Hugging Face transformers models directly. It supports Tensor Parallelism (TP), Pipeline Parallelism (PP), Mixture of Experts (MoE), and Vision-Language Models (VLMs) through a mixin-based architecture. Key components include TransformersBase for core logic like meta-device initialization and module replacement, and specialized mixins for causal, embedding, and multimodal capabilities. The implementation also integrates torch.compile support and provides a generic multimodal processor. Overall, the design is modular and significantly expands SGLang's model compatibility.

[gemini-code-assist]

Calling dict(self.named_parameters()) inside load_weights can be expensive, especially since this method is called for every weight shard during model loading. It is better to cache this dictionary in __init__ or lazily upon the first call to load_weights.

References

1. Avoid redundant computations in performance-critical paths like weight loading. ^(link)

[gemini-code-assist]

The device is hardcoded to "cuda" when materializing parameters from the meta device. This will cause failures on other supported hardware backends (e.g., xpu, npu, cpu). Consider using the current default device or passing the target device to the model constructor.

[gemini-code-assist]

The logic next(v for v in output.values() if torch.is_tensor(v)) is non-deterministic and risky if the dictionary contains multiple tensors (e.g., both pooler_output and hidden_states). It is safer to look for specific known keys used by Hugging Face multimodal encoders.

[gemini-code-assist]

itertools.groupby only groups consecutive items with the same key. If the weights iterable is not sorted by name, this will result in multiple calls to _load_module for the same submodule with partial weight sets, which is inefficient and may break submodules that expect all weights at once. Ensure the weights are sorted by name before grouping.

Suggested change

	for prefix, group in itertools.groupby(weights_by_parts, key=lambda x: x[0][0]):
	for prefix, group in itertools.groupby(sorted(weights_by_parts, key=lambda x: x[0][0]), key=lambda x: x[0][0]):

[yuan-luo]

Could you specify the usage commands for these new features and paste the accuracy result?

[adarshxs]

Usage Instructions

Basic Usage (Transformers Backend)

Use the transformers modeling backend explicitly:

python -m sglang.launch_server \
  --model-path Qwen/Qwen3-4B \
  --model-impl transformers \
  --host 0.0.0.0 --port 30000

or

sglang serve \
  --model-path Qwen/Qwen3-4B \
  --model-impl transformers

Supported --model-impl Values

Value	Behavior
auto (default)	Use SGLang native implementation if available, fall back to transformers
sglang	SGLang native implementation (error if not available)
transformers	HF transformers backend

Accuracy Eval

Model: Qwen/Qwen3-4B
Hardware: 1x H100 NVL

Benchmark	SGLang Native	Transformers Backend
MGSM-EN (n=250)	0.916	0.940
MMLU (n=250)	0.720	0.744

[yuan-luo]

/tag-and-rerun-ci

[yuan-luo]

/tag-and-rerun-ci

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

@JustinTong0323 Could you please take another look? LGTM now.

[JustinTong0323]

PR Review Summary

Reviewed with specialized agents focusing on: code quality, silent failures, test coverage, and code simplification. Overall this is a well-structured feature addition with a clean mixin-based architecture — but there are several issues worth addressing before merge.

Critical: 3 | High: 5 | Simplification suggestions: 5

See inline comments for details on each issue.

Test Coverage

The PR adds ~2074 lines of implementation but only 43 lines of tests (a single GSM8K eval). Key untested areas include:

• _is_moe_model / _get_transformers_backend_arch routing logic (pure functions, trivially unit-testable)
• AutoWeightsLoader weight dispatch
• TransformersFusedMoE expert replacement
• MultiModalMixin + TransformersAutoMultimodalProcessor
• Pipeline parallel layer slicing
• Embedding/classification model classes

Simplification Opportunities

1. Consolidate duplicated helpers: _first_attr (in transformers_auto.py) and _getattr_first (in transformers.py) are identical. Same for _uses_mrope / _uses_mrope_positions.
2. Extract tensor-list flattening helper from _to_tensor_output — same logic appears twice in the same method.
3. Simplify replace_rms_norm_class Gemma detection — instantiating a dummy module to check zero weights is fragile; checking class name would be simpler.
4. Split TransformersBase.__init__ (~100+ lines) into focused setup methods for readability.

Strengths

• Clean mixin-based architecture composing TP, PP, MoE, VLM capabilities
• Elegant WeightsMapper composition via __or__ and __init_subclass__
• Comprehensive manual testing across 25 configurations
• Proper hidden_states.clone() in MoE custom op to prevent mutation

[yuan-luo]

/rerun-failed-ci

[adarshxs]

/rerun-failed-ci

[JustinTong0323]

/rerun-failed-ci

[yuan-luo]

please resolve the conflict.

[adarshxs]

/rerun-failed-ci

[adarshxs]

/rerun-failed-ci