From 1b97ff11881173d9926c67cb9ace9457f5e57954 Mon Sep 17 00:00:00 2001 From: Ding Zuhao Date: Sun, 27 Sep 2026 16:48:08 +0800 Subject: [PATCH] [Model] Support SenseNova-U1-A3B-MoT (MoE) (#8172) Signed-off-by: Ding Zuhao --- docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 2 + recipes/SenseNova/SenseNova-U1.md | 120 ++++++++++ .../sensenova_u1/test_sensenova_u1_moe.py | 214 ++++++++++++++++++ .../test_sensenova_u1_moe_cuda.py | 158 +++++++++++++ .../sensenova_u1/pipeline_sensenova_u1.py | 105 +++++++-- .../sensenova_u1/sensenova_u1_transformer.py | 155 +++++++++++-- .../transformers_utils/configs/__init__.py | 2 + .../configs/sensenova_u1.py | 96 +++++++- 9 files changed, 807 insertions(+), 46 deletions(-) create mode 100644 tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe.py create mode 100644 tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 96d333819..3ea7fe2f1 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -101,6 +101,7 @@ th { | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | | `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | ✅︎ | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) | +| `SenseNovaU1Pipeline` | SenseNova-U1-A3B (MoE) | `sensenova/SenseNova-U1-A3B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.md#a3b-moe-variant) | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 198e61845..cad0713fc 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -164,6 +164,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **SenseNova-U1-A3B** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | @@ -179,6 +180,7 @@ The following tables show which models support each feature: > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks. +> 8. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident. ### VideoGen diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md index 8e9091d63..a2f2bd090 100644 --- a/recipes/SenseNova/SenseNova-U1.md +++ b/recipes/SenseNova/SenseNova-U1.md @@ -34,6 +34,126 @@ text-to-text chat. - HuggingFace model page: [SenseNova/SenseNova-U1-8B-MoT](https://huggingface.co/SenseNova/SenseNova-U1-8B-MoT) +## A3B (MoE) variant + +[sensenova/SenseNova-U1-A3B-MoT](https://huggingface.co/sensenova/SenseNova-U1-A3B-MoT) +runs on this same recipe and the same `SenseNovaU1Pipeline` — the checkpoint keeps +`model_type: neo_chat`, so it resolves without `--model-class-name`. Only the `--model` +id changes. + +The backbone is Qwen3-MoE instead of dense Qwen3, and both MoT branches are sparse: +`mlp` routes over 128 experts and `mlp_mot_gen` over 32, top-8 each, with per-expert +width `moe_intermediate_size=768`. `SenseNovaU1Config` picks `SenseNovaU1MoELLMConfig` +whenever `llm_config` carries `num_experts`, so the dense 8B path is untouched. The +gen-path knobs (`gen_num_experts`, `gen_num_experts_per_tok`, `gen_moe_intermediate_size`) +fall back to their understanding-path counterparts when absent. + +Experts run through vLLM's `FusedMoE`, whose `moe_forward` op resolves its layer through +`vllm.forward_context` rather than the diffusion runner's own context, so +`SenseNovaU1ForCausalLM.forward` enters that context when the model has MoE layers. + +### Memory + +The weights dominate, and every expert stays resident even though only the top-8 of each +tower run per token -- the "A3B" in the name is the ~4B activated, not the 38.7B stored. +Counted from the checkpoint's safetensors headers: + +| Component | Params | BF16 | +| --- | --- | --- | +| Understanding-path experts (128/layer) | 28.99 B | 54.0 GiB | +| Generation-path experts (32/layer) | 7.25 B | 13.5 GiB | +| Attention (dual tower) | 1.81 B | 3.4 GiB | +| Embedding + LM head | 0.62 B | 1.2 GiB | +| Routers, norms, vision and flow-matching heads | 0.06 B | 0.1 GiB | +| Total | 38.74 B | 72.2 GiB | + +At 1024x1024 the 72.9 GiB peak is those 72.2 GiB of weights plus about 0.7 GiB of +activations and workspace. The 2048x2048 edit, the largest activation case below, peaks at +74.6 GiB, 2.4 GiB above the weights. Text-to-image only routes through the generation tower, +leaving the 54 GiB of understanding-path experts idle; that tower is reached by the text +prefix, think mode, `img2text` and `text2text`. One 96 GB card fits TP=1, and TP=2 halves +the per-GPU footprint. + +### Measured (1x H20 96GB, BF16) + +Fixed case used for the stage split and the torch.profiler kernel table: + +- Prompt: `Close portrait of an elderly woman by a farmhouse window, warm natural light.` +- 1024x1024, 50 Euler steps, think off, seed 42, `cfg_scale` 4.0, `cfg_norm` none, + `timestep_shift` 3.0, `t_eps` 0.02 +- Python 3.12, vLLM 0.30.0, torch 2.13.0+cu130 +- Peak GPU memory is the reserved high-water mark the runner records (`peak_memory_mb`) + +| | Time | Memory | +| --- | --- | --- | +| Weights after load | ~13 s | 72.2 GiB | +| Pipeline `forward` | 7.68 s | 72.9 GiB peak | + +`--enable-diffusion-pipeline-profiler` on that request: + +| Stage | Time | Share | +| --- | --- | --- | +| Text prefix (`_t2i_prefix_forward`, understanding tower) | 0.06 s | 0.7% | +| Denoising loop (`_run_denoising_loop`, generation tower, 50 steps) | 7.56 s | 98.5% | +| Pipeline `forward` | 7.68 s | 100% | + +There is no VAE. The loop ends in unpatchify / denorm / PIL, which is inside the denoise +row and is a few tens of milliseconds at 1024x1024. + +A warm 8-step generate at the same resolution, wrapped in `torch.profiler` after one +unprofiled warmup request, spent 1.19 s of self-CUDA time as: + +| Kernel family | Self CUDA | Share | What it is | +| --- | --- | --- | --- | +| `fused_moe_kernel` | 576 ms | 48% | top-8 expert GEMMs on the 32-expert generation tower | +| FlashAttention `flash_fwd_kernel` | 222 ms | 19% | dual-tower attention | +| `nvjet` GEMMs | ~233 ms | 20% | `qkv_proj` / `o_proj` (and the small dense leftovers) | +| `qk_norm_rope_kernel` | 58 ms | 5% | fused RMSNorm + RoPE | +| MoE routing / `act_and_mul` / elementwise | ~40 ms | 3% | top-k, SwiGLU, residual | + +So a denoise step is about half FusedMoE and about two-fifths attention + dense GEMM. +`--cfg-parallel-size 2` on the same 8-step t2i request (2-branch `cond`+`uncond`) +is 0.95 s vs 1.48 s on one GPU (~1.56x). Img2img CFG-P=2 also completes for both +branch counts at 8-step 1024x1024: 2-way (`img_cfg_scale=1`, `cond`+`img_cond`) +1.09 s; 3-way (`cfg_scale=4`, `img_cfg_scale=2`, `cond`+`img_cond`+`uncond`) +1.74 s (3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT and +layerwise CPU offload do not take effect on A3B; see +[diffusion features](../../docs/user_guide/diffusion_features.md) note 8. + +Other configurations on the same machine (second run, warm kernels): + +| Case | Total | Peak GPU memory | +| --- | --- | --- | +| text2img 1024x1024, 50 steps, think on | 9.16 s | 73.0 GiB | +| text2img 1024x1024, 50 steps, think off, TP=2 | 5.64 s | 36.7 GiB per GPU | +| img2img 2048x2048 output, 25 steps, think off | 20.93 s | 74.6 GiB | + +Think-on spends 1.37 s in `_generate_think` (215 tokens). At the same seed TP=2 differs +from TP=1 only numerically (MAE 2.97/255, identical composition) from the changed +reduction order. + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model sensenova/SenseNova-U1-A3B-MoT \ + --prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \ + --width 1024 --height 1024 \ + --seed 42 --num-inference-steps 50 --cfg-scale 4.0 \ + --extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \ + --output sensenova_a3b_t2i.png +``` + +`img2text` and `text2text` go through the server, as for the dense checkpoint: + +```bash +vllm serve sensenova/SenseNova-U1-A3B-MoT --omni --port 8091 + +python examples/online_serving/sensenova_u1/openai_chat_client.py \ + -s http://127.0.0.1:8091 -m img2text -i input.png -p "Describe this image." +``` + +All four modalities were verified on A3B: text2img (with and without think), img2img +editing, img2text and text2text. + ## Hardware Support ## GPU diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe.py new file mode 100644 index 000000000..1f35c9abc --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU tests for SenseNova-U1-A3B MoE config, layer dispatch and weight loading. + +FusedMoE numerics live in ``test_sensenova_u1_moe_cuda.py`` so the ready +``core_model and cpu`` lane and the ``core_model and cuda`` model lane collect +disjoint items. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline +from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( + _is_moe, + _is_sparse_und_layer, +) +from vllm_omni.transformers_utils.configs.sensenova_u1 import ( + SenseNovaU1Config, + SenseNovaU1LLMConfig, + SenseNovaU1MoELLMConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_A3B_LLM = { + "architectures": ["Qwen3MoeForCausalLM"], + "model_type": "qwen3_moe", + "hidden_size": 2048, + "intermediate_size": 6144, + "moe_intermediate_size": 768, + "num_experts": 128, + "gen_num_experts": 32, + "num_experts_per_tok": 8, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "num_hidden_layers": 48, + "head_dim": 128, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "vocab_size": 151936, + "max_position_embeddings": 262144, + "max_position_embeddings_hw": 10000, + "rope_theta": 10000000, + "rope_theta_hw": 10000.0, + "decoder_sparse_step": 1, + "mlp_only_layers": [], + "norm_topk_prob": True, +} + + +def _tiny_moe_llm(**overrides) -> SenseNovaU1MoELLMConfig: + kwargs = dict( + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=8, + num_experts_per_tok=2, + moe_intermediate_size=16, + # A3B sets this explicitly; the Qwen3MoeConfig default is False. + norm_topk_prob=True, + ) + kwargs.update(overrides) + return SenseNovaU1MoELLMConfig(**kwargs) + + +def _tiny_dense_llm() -> SenseNovaU1LLMConfig: + return SenseNovaU1LLMConfig(hidden_size=32, num_hidden_layers=1, num_attention_heads=4) + + +def test_dense_checkpoint_still_uses_qwen3_config(): + cfg = SenseNovaU1Config(llm_config={"hidden_size": 64, "num_hidden_layers": 1, "num_attention_heads": 2}) + assert isinstance(cfg.llm_config, SenseNovaU1LLMConfig) + assert not _is_moe(cfg.llm_config) + + +def test_a3b_dict_selects_moe_config_and_gen_defaults(): + cfg = SenseNovaU1Config(llm_config=_A3B_LLM) + assert isinstance(cfg.llm_config, SenseNovaU1MoELLMConfig) + llm = cfg.llm_config + assert llm.num_experts == 128 + assert llm.gen_num_experts == 32 + assert llm.gen_num_experts_per_tok == 8 + assert llm.gen_moe_intermediate_size == 768 + assert llm.moe_intermediate_size == 768 + assert len(llm.layer_types) == 48 + assert llm.layer_types == ["full_attention"] * 48 + assert llm.rope_theta_hw == 10000.0 + + +def test_gen_moe_knobs_fall_back_to_und_path(): + llm = dict(_A3B_LLM) + llm.pop("gen_num_experts") + cfg = SenseNovaU1Config(llm_config=llm) + assert cfg.llm_config.gen_num_experts == 128 + assert cfg.llm_config.gen_num_experts_per_tok == 8 + assert cfg.llm_config.gen_moe_intermediate_size == 768 + + +def test_sparse_und_layer_matches_qwen3_moe_step(): + config = _tiny_moe_llm(num_hidden_layers=3, mlp_only_layers=[0], decoder_sparse_step=2) + assert not _is_sparse_und_layer(config, 0) # mlp_only + assert _is_sparse_und_layer(config, 1) # (1+1) % 2 == 0 + assert not _is_sparse_und_layer(config, 2) + assert not _is_sparse_und_layer(_tiny_dense_llm(), 0) + + +class _Routed(nn.Module): + def __init__(self): + super().__init__() + self.w13_weight = nn.Parameter(torch.zeros(2, 8, 4)) + self.w2_weight = nn.Parameter(torch.zeros(2, 4, 8)) + + +class _Experts(nn.Module): + def __init__(self): + super().__init__() + self.routed_experts = _Routed() + + +class _Layer(nn.Module): + def __init__(self): + super().__init__() + self.mlp = nn.Module() + self.mlp.gate_up_proj = nn.Linear(4, 16, bias=False) + self.mlp.experts = _Experts() + self.mlp_mot_gen = nn.Module() + self.mlp_mot_gen.experts = _Experts() + + +class _MoeStub(SenseNovaU1Pipeline): + """Only the parameter tree matters here, so skip the real __init__.""" + + def __init__(self, llm_cfg: SenseNovaU1LLMConfig | SenseNovaU1MoELLMConfig): + nn.Module.__init__(self) + self.llm_cfg = llm_cfg + self.language_model = nn.Module() + self.language_model.model = nn.Module() + self.language_model.model.layers = nn.ModuleList([_Layer()]) + + +def test_load_weights_routes_expert_gate_proj_to_fused_w13(): + """A3B shards are ``experts.{i}.gate_proj``; they must not hit dense gate_up.""" + pipe = _MoeStub(_tiny_moe_llm(num_experts=2, gen_num_experts=2, num_experts_per_tok=1)) + layer = pipe.language_model.model.layers[0] + with torch.no_grad(): + layer.mlp.gate_up_proj.weight.zero_() + layer.mlp.experts.routed_experts.w13_weight.zero_() + + calls: list[tuple] = [] + + def _loader(param, loaded, name, shard_id=None, expert_id=None, return_success=True): + calls.append((name, shard_id, expert_id, tuple(loaded.shape))) + return True + + layer.mlp.experts.routed_experts.w13_weight.weight_loader = _loader + layer.mlp_mot_gen.experts.routed_experts.w13_weight.weight_loader = _loader + + def _stacked_loader(param, loaded_weight, shard_id): + rows = loaded_weight.shape[0] + if shard_id == 0: + param.data[:rows].copy_(loaded_weight) + elif shard_id == 1: + param.data[rows:].copy_(loaded_weight) + else: + raise AssertionError(shard_id) + + layer.mlp.gate_up_proj.weight.weight_loader = _stacked_loader + + hidden, inter = 4, 8 + loaded = pipe.load_weights( + [ + ( + "language_model.model.layers.0.mlp.experts.0.gate_proj.weight", + torch.ones(inter, hidden), + ), + ( + "language_model.model.layers.0.mlp.experts.0.up_proj.weight", + torch.full((inter, hidden), 2.0), + ), + ( + "language_model.model.layers.0.mlp.gate_proj.weight", + torch.full((inter, hidden), 3.0), + ), + ] + ) + + assert torch.equal(layer.mlp.gate_up_proj.weight[:inter], torch.full((inter, hidden), 3.0)) + assert torch.equal(layer.mlp.gate_up_proj.weight[inter:], torch.zeros(inter, hidden)) + assert any(c[1] == "w1" and c[2] == 0 for c in calls) + assert any(c[1] == "w3" and c[2] == 0 for c in calls) + assert any(name.endswith("mlp.experts.routed_experts.w13_weight") for name, *_ in calls) + assert "language_model.model.layers.0.mlp.experts.routed_experts.w13_weight" in loaded + assert "language_model.model.layers.0.mlp.gate_up_proj.weight" in loaded + + +def test_load_weights_skips_stacked_mapping_on_expert_names(): + """Without the skip, ``.gate_proj`` would rewrite expert keys to gate_up_proj.""" + pipe = _MoeStub(_tiny_dense_llm()) + layer = pipe.language_model.model.layers[0] + with torch.no_grad(): + layer.mlp.gate_up_proj.weight.fill_(1.0) + + pipe.load_weights( + [ + ( + "language_model.model.layers.0.mlp.experts.0.gate_proj.weight", + torch.full((8, 4), 9.0), + ), + ] + ) + assert torch.equal(layer.mlp.gate_up_proj.weight, torch.ones(16, 4)) diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py new file mode 100644 index 000000000..44c8899ad --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA numerics for SenseNova-U1-A3B FusedMoE vs. the official Qwen3-MoE loop. + +Collected by the ready ``Diffusion · Model Test`` job +(``tests/diffusion/models/ -m 'core_model and cuda'``), not by the CPU simple lane. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline +from vllm_omni.transformers_utils.configs.sensenova_u1 import SenseNovaU1MoELLMConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +def _tiny_moe_llm(**overrides) -> SenseNovaU1MoELLMConfig: + kwargs = dict( + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=8, + num_experts_per_tok=2, + moe_intermediate_size=16, + norm_topk_prob=True, + ) + kwargs.update(overrides) + return SenseNovaU1MoELLMConfig(**kwargs) + + +def _reference_qwen3_moe( + x: torch.Tensor, + router_logits: torch.Tensor, + gate_w: torch.Tensor, + up_w: torch.Tensor, + down_w: torch.Tensor, + top_k: int, +) -> torch.Tensor: + """Per-expert loop matching the official Qwen3MoeSparseMoeBlock used by SenseNova-U1. + + Softmax over the router logits in fp32, top-k, renormalize, then each selected + expert's SwiGLU weighted by its routing probability. Accumulates in fp32. + """ + probs = torch.softmax(router_logits.float(), dim=-1) + top_p, top_i = torch.topk(probs, top_k, dim=-1) + top_p = top_p / top_p.sum(dim=-1, keepdim=True) + out = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + for expert in range(gate_w.shape[0]): + token_idx, slot_idx = torch.where(top_i == expert) + if token_idx.numel() == 0: + continue + xe = x[token_idx].float() + h = F.silu(xe @ gate_w[expert].float().t()) * (xe @ up_w[expert].float().t()) + out.index_add_(0, token_idx, (h @ down_w[expert].float().t()) * top_p[token_idx, slot_idx, None]) + return out + + +@pytest.fixture +def cuda_vllm_env(monkeypatch): + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + + monkeypatch.setenv("MASTER_ADDR", "localhost") + monkeypatch.setenv("MASTER_PORT", "29617") + vllm_config = VllmConfig(device_config=DeviceConfig(device="cuda")) + with set_current_vllm_config(vllm_config): + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") + initialize_model_parallel() + # The modular FusedMoE kernel draws scratch buffers from here; the + # diffusion worker initializes it at startup. + init_workspace_manager(torch.device("cuda")) + try: + yield vllm_config + finally: + reset_workspace_manager() + cleanup_dist_env_and_memory() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for FusedMoE") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_sparse_moe_block_matches_reference_after_checkpoint_loading(cuda_vllm_env): + """Checkpoint-format expert shards, loaded through the production loader, must + reproduce the reference Qwen3-MoE output. Covers the w13 gate/up order, the + expert-id mapping and the routing math in one pass.""" + import vllm.forward_context as vllm_forward_context + from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase + + from vllm_omni.diffusion.layers.fused_moe import FusedMoE + from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import SenseNovaU1SparseMoeBlock + + num_experts, top_k, hidden, inter, num_tokens = 16, 4, 128, 64, 32 + torch.manual_seed(0) + # Checkpoint layout: experts.{e}.gate_proj / up_proj are [inter, hidden], down_proj is [hidden, inter]. + gate_w = (torch.randn(num_experts, inter, hidden) / hidden**0.5).to("cuda", torch.bfloat16) + up_w = (torch.randn(num_experts, inter, hidden) / hidden**0.5).to("cuda", torch.bfloat16) + down_w = (torch.randn(num_experts, hidden, inter) / inter**0.5).to("cuda", torch.bfloat16) + router_w = (torch.randn(num_experts, hidden) / hidden**0.5).to("cuda", torch.bfloat16) + + config = _tiny_moe_llm( + hidden_size=hidden, num_experts=num_experts, num_experts_per_tok=top_k, moe_intermediate_size=inter + ) + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + with torch.device("cuda"): + block = SenseNovaU1SparseMoeBlock( + config, + num_experts=num_experts, + num_experts_per_tok=top_k, + moe_intermediate_size=inter, + prefix="model.layers.0.mlp", + ) + finally: + torch.set_default_dtype(default_dtype) + + params = dict(block.named_parameters()) + mapping = FusedMoE.make_expert_params_mapping( + block, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) + loaded: set[str] = set() + for expert in range(num_experts): + for proj, weights in (("gate_proj", gate_w), ("up_proj", up_w), ("down_proj", down_w)): + assert SenseNovaU1Pipeline._load_expert_weight( + f"experts.{expert}.{proj}.weight", weights[expert], params, mapping, loaded + ) + with torch.no_grad(): + params["gate.weight"].copy_(router_w) + for module in block.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, QuantizeMethodBase): + quant_method.process_weights_after_loading(module) + + x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), vllm_forward_context.set_forward_context(None, cuda_vllm_env): + router_logits, _ = block.gate(x) + actual = block(x).float() + + expected = _reference_qwen3_moe(x, router_logits, gate_w, up_w, down_w, top_k) + rel_err = ((actual - expected).norm() / expected.norm()).item() + # Measured 0.5% on H20 in BF16; the slack covers kernel tiling differences across GPUs. + assert rel_err < 2e-2, f"FusedMoE deviates from the reference: relative error {rel_err:.4f}" + + # The tolerance must be able to fail: swapping gate and up breaks SwiGLU. + swapped = _reference_qwen3_moe(x, router_logits, up_w, gate_w, down_w, top_k) + assert ((actual - swapped).norm() / swapped.norm()).item() > 0.1 diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 2be1a9c7a..79d42d455 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -34,6 +34,7 @@ from vllm.logger import init_logger from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.layers.fused_moe import FusedMoE from vllm_omni.diffusion.lora.loader import ( LoraLoaderMixin, _apply_diffusers_lora_alpha_scaling, @@ -576,7 +577,14 @@ class SenseNovaU1Pipeline( ] self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler + profiler_targets=[ + "_t2i_prefix_forward", + "_it2i_prefix_forward", + "_generate_think", + "_generate_text", + "_run_denoising_loop", + ], + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, ) # ----------------------------------------------------------------------- @@ -1650,29 +1658,43 @@ class SenseNovaU1Pipeline( applied += 1 return applied + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + num_experts = max( + getattr(self.llm_cfg, "num_experts", 0) or 0, + getattr(self.llm_cfg, "gen_num_experts", 0) or 0, + ) + if num_experts <= 1: + return [] + return FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = self.stacked_params_mapping + expert_params_mapping = self.get_expert_mapping() params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: - loaded = False - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - stacked_name = name.replace(weight_name, param_name) - if stacked_name not in params_dict: - break - param = params_dict[stacked_name] - weight_loader = getattr(param, "weight_loader", None) - if weight_loader is not None: - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(stacked_name) - loaded = True - break + # Expert FFN shards must skip the stacked mapping: ``.gate_proj`` is + # a substring of ``experts.{i}.gate_proj``, which would fuse them + # into a dense gate_up_proj instead of reaching the FusedMoE loader. + is_expert_weight = ".experts." in name - if loaded: + if not is_expert_weight and self._load_stacked_weight( + name, loaded_weight, params_dict, stacked_params_mapping, loaded_params + ): + continue + + if is_expert_weight: + # A shard that matches no mapping belongs to an expert this rank + # does not own, so dropping it here is expected under TP/EP. + self._load_expert_weight(name, loaded_weight, params_dict, expert_params_mapping, loaded_params) continue if name not in params_dict: @@ -1690,3 +1712,54 @@ class SenseNovaU1Pipeline( loaded_params.add(name) return loaded_params + + @staticmethod + def _load_stacked_weight( + name: str, + loaded_weight: torch.Tensor, + params_dict: dict[str, torch.nn.Parameter], + stacked_params_mapping: list[tuple[str, str, str | int]], + loaded_params: set[str], + ) -> bool: + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + stacked_name = name.replace(weight_name, param_name) + if stacked_name not in params_dict: + return False + param = params_dict[stacked_name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is not None: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(stacked_name) + return True + return False + + @staticmethod + def _load_expert_weight( + name: str, + loaded_weight: torch.Tensor, + params_dict: dict[str, torch.nn.Parameter], + expert_params_mapping: list[tuple[str, str, int, str]], + loaded_params: set[str], + ) -> bool: + for param_name, weight_name, expert_id, shard_id in expert_params_mapping: + if weight_name not in name: + continue + stacked_name = name.replace(weight_name, param_name) + param = params_dict.get(stacked_name) + weight_loader = getattr(param, "weight_loader", None) if param is not None else None + if weight_loader is None: + continue + success = weight_loader( + param, + loaded_weight, + stacked_name, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(stacked_name) + return True + return False diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py index ab69d3baa..267206e2f 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py +++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py @@ -12,17 +12,22 @@ Ported from the sensenova_u1 package with vllm tensor-parallel support: from __future__ import annotations import copy +from contextlib import contextmanager from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F +import vllm.forward_context as vllm_forward_context from cache_dit import ForwardPattern from transformers.cache_utils import DynamicCache +from vllm.config import get_current_vllm_config, get_current_vllm_config_or_none +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -34,6 +39,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, SensenovaCachedAdapter +from vllm_omni.diffusion.layers.fused_moe import FusedMoE logger = init_logger(__name__) @@ -230,6 +236,95 @@ class SenseNovaU1MLP(nn.Module): return x +def _is_moe(config) -> bool: + num_experts = getattr(config, "num_experts", None) + return isinstance(num_experts, int) and num_experts > 1 + + +def _is_sparse_und_layer(config, layer_idx: int) -> bool: + if not _is_moe(config): + return False + mlp_only_layers = list(getattr(config, "mlp_only_layers", None) or []) + decoder_sparse_step = int(getattr(config, "decoder_sparse_step", 1) or 1) + return layer_idx not in mlp_only_layers and (layer_idx + 1) % decoder_sparse_step == 0 + + +class SenseNovaU1SparseMoeBlock(nn.Module): + def __init__( + self, + config, + *, + num_experts: int, + num_experts_per_tok: int, + moe_intermediate_size: int, + quant_config=None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + if self.tp_size > num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than the number of experts {num_experts}." + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + num_experts, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate", + ) + enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel + self.experts = FusedMoE( + num_experts=num_experts, + top_k=num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=False, + num_redundant_experts=0, + pcp_size=None if enable_expert_parallel else 1, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + return final_hidden_states.view(orig_shape) + + +def _build_mlp(config, layer_idx: int, *, gen_path: bool, quant_config=None, prefix: str = ""): + if gen_path and _is_moe(config): + return SenseNovaU1SparseMoeBlock( + config, + num_experts=config.gen_num_experts, + num_experts_per_tok=config.gen_num_experts_per_tok, + moe_intermediate_size=config.gen_moe_intermediate_size, + quant_config=quant_config, + prefix=prefix, + ) + if not gen_path and _is_sparse_und_layer(config, layer_idx): + return SenseNovaU1SparseMoeBlock( + config, + num_experts=config.num_experts, + num_experts_per_tok=config.num_experts_per_tok, + moe_intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + prefix=prefix, + ) + return SenseNovaU1MLP( + config.hidden_size, + config.intermediate_size, + config.hidden_act, + quant_config=quant_config, + prefix=prefix, + ) + + # --------------------------------------------------------------------------- # Attention with MoT # --------------------------------------------------------------------------- @@ -557,19 +652,9 @@ class SenseNovaU1DecoderLayer(nn.Module): self.self_attn = SenseNovaU1Attention( config, layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn" ) - self.mlp = SenseNovaU1MLP( - config.hidden_size, - config.intermediate_size, - config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - self.mlp_mot_gen = SenseNovaU1MLP( - config.hidden_size, - config.intermediate_size, - config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp_mot_gen", + self.mlp = _build_mlp(config, layer_idx, gen_path=False, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.mlp_mot_gen = _build_mlp( + config, layer_idx, gen_path=True, quant_config=quant_config, prefix=f"{prefix}.mlp_mot_gen" ) self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.input_layernorm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -761,6 +846,31 @@ class SenseNovaU1ForCausalLM(nn.Module): # LogitsProcessor handles the TP all-gather of vocab-sharded ParallelLMHead # outputs so callers see full-vocab logits regardless of tp_size. self.logits_processor = LogitsProcessor(config.vocab_size) + self.has_moe = any( + isinstance(layer.mlp, SenseNovaU1SparseMoeBlock) or isinstance(layer.mlp_mot_gen, SenseNovaU1SparseMoeBlock) + for layer in self.model.layers + ) + + @contextmanager + def _vllm_forward_context(self): + """Enter vLLM's own forward context for the A3B MoE layers. + + ``moe_forward`` resolves its layer through ``vllm.forward_context``, + which the diffusion runner's context does not populate -- that one is + ``vllm_omni.diffusion.forward_context``. Dense checkpoints have no + FusedMoE and stay on the runner's context alone. + """ + if not self.has_moe or vllm_forward_context.is_forward_context_available(): + yield + return + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + # Warmup paths can run outside the runner's context; let the MoE op + # raise on its own rather than masking it with a second failure. + yield + return + with vllm_forward_context.set_forward_context(None, vllm_config): + yield def forward( self, @@ -787,15 +897,16 @@ class SenseNovaU1ForCausalLM(nn.Module): inputs_embeds=self.model.embed_tokens(input_ids), ) - outputs = self.model( - input_ids=input_ids, - indexes=indexes, - attention_mask=attention_mask, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - **kwargs, - ) + with self._vllm_forward_context(): + outputs = self.model( + input_ids=input_ids, + indexes=indexes, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) logits = self.logits_processor(self.lm_head, outputs.last_hidden_state) if compute_logits else None return SenseNovaU1CausalLMOutput( logits=logits, diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 63e38b01e..c8a5c19f5 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "AuKConfig": "vllm_omni.transformers_utils.configs.auk", "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", + "SenseNovaU1MoELLMConfig": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ @@ -63,6 +64,7 @@ __all__ = [ "Qwen3VLMoeVisionConfig", "WhisperEncoderConfig", "SenseNovaU1Config", + "SenseNovaU1MoELLMConfig", ] diff --git a/vllm_omni/transformers_utils/configs/sensenova_u1.py b/vllm_omni/transformers_utils/configs/sensenova_u1.py index 81e8b3406..bae3aeef5 100644 --- a/vllm_omni/transformers_utils/configs/sensenova_u1.py +++ b/vllm_omni/transformers_utils/configs/sensenova_u1.py @@ -1,12 +1,57 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """HuggingFace-style configuration classes for SenseNova-U1.""" -from transformers import AutoConfig, PretrainedConfig, Qwen3Config +from __future__ import annotations + +from transformers import AutoConfig, PretrainedConfig, Qwen3Config, Qwen3MoeConfig _SENSENOVA_U1_LLM_HIDDEN_DEFAULT = 4096 +def _restore_legacy_rope_theta(config) -> None: + """Expose the v4 rope attribute expected by the vendored model code.""" + if hasattr(config, "rope_theta"): + return + rope_parameters = getattr(config, "rope_parameters", None) or {} + config.rope_theta = float(rope_parameters.get("rope_theta", 10000.0)) + + +def _backfill_layer_types(config) -> None: + """Qwen3MoeConfig does not always populate ``layer_types``; attention reads it.""" + existing = getattr(config, "layer_types", None) + if existing and len(existing) == config.num_hidden_layers: + return + use_swa = bool(getattr(config, "use_sliding_window", False)) and getattr(config, "sliding_window", None) is not None + max_window_layers = int(getattr(config, "max_window_layers", 0) or 0) + config.layer_types = [ + "sliding_attention" if (use_swa and i >= max_window_layers) else "full_attention" + for i in range(config.num_hidden_layers) + ] + + +def _is_moe_llm_config(llm_config) -> bool: + if isinstance(llm_config, dict): + model_type = llm_config.get("model_type", "") + archs = llm_config.get("architectures") or [] + num_experts = llm_config.get("num_experts", 0) or 0 + else: + model_type = getattr(llm_config, "model_type", "") + archs = getattr(llm_config, "architectures", None) or [] + num_experts = getattr(llm_config, "num_experts", 0) or 0 + + if isinstance(model_type, str) and "moe" in model_type.lower(): + return True + for arch in archs: + arch_str = str(arch) + if "Moe" in arch_str or "MoE" in arch_str: + return True + try: + return int(num_experts) > 1 + except (TypeError, ValueError): + return False + + # Adapted from: https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_chat.py class SenseNovaU1LLMConfig(Qwen3Config): """Qwen3-based LLM backbone config with 3D RoPE extensions.""" @@ -26,6 +71,42 @@ class SenseNovaU1LLMConfig(Qwen3Config): super().__init__(**kwargs) +class SenseNovaU1MoELLMConfig(Qwen3MoeConfig): + """Qwen3-MoE LLM backbone config for SenseNova-U1-A3B.""" + + model_type = "sensenova_u1_moe_llm" + + def __init__( + self, + rope_theta_hw: float = 10000.0, + max_position_embeddings_hw: int = 10000, + gen_num_experts: int | None = None, + gen_num_experts_per_tok: int | None = None, + gen_moe_intermediate_size: int | None = None, + **kwargs, + ): + super().__init__(**kwargs) + _restore_legacy_rope_theta(self) + self.rope_theta_hw = rope_theta_hw + self.max_position_embeddings_hw = max_position_embeddings_hw + self.gen_num_experts = int(gen_num_experts) if gen_num_experts is not None else int(self.num_experts) + self.gen_num_experts_per_tok = ( + int(gen_num_experts_per_tok) if gen_num_experts_per_tok is not None else int(self.num_experts_per_tok) + ) + self.gen_moe_intermediate_size = ( + int(gen_moe_intermediate_size) if gen_moe_intermediate_size is not None else int(self.moe_intermediate_size) + ) + _backfill_layer_types(self) + + +def _build_llm_config(llm_config): + if isinstance(llm_config, dict): + if _is_moe_llm_config(llm_config): + return SenseNovaU1MoELLMConfig(**llm_config) + return SenseNovaU1LLMConfig(**llm_config) + return llm_config + + # Adapted from https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_vit.py#L10 class SenseNovaU1VisionConfig(PretrainedConfig): """Vision embedding config (2D RoPE + conv patch embed, no transformer).""" @@ -37,8 +118,8 @@ class SenseNovaU1VisionConfig(PretrainedConfig): num_channels: int = 3, patch_size: int = 16, hidden_size: int = 1024, - llm_hidden_size: list[int] | None = None, - downsample_ratio: list[float] | None = None, + llm_hidden_size: list[int] | int | None = None, + downsample_ratio: list[float] | float | None = None, rope_theta_vision: float = 10000.0, max_position_embeddings_vision: int = 10000, **kwargs, @@ -67,7 +148,7 @@ class SenseNovaU1Config(PretrainedConfig): def __init__( self, - llm_config: dict | SenseNovaU1LLMConfig | None = None, + llm_config: dict | SenseNovaU1LLMConfig | SenseNovaU1MoELLMConfig | None = None, vision_config: dict | SenseNovaU1VisionConfig | None = None, downsample_ratio: float = 0.5, template: str = "neo1_0", @@ -90,9 +171,7 @@ class SenseNovaU1Config(PretrainedConfig): t_eps: float = 0.02, **kwargs, ): - if isinstance(llm_config, dict): - llm_config = SenseNovaU1LLMConfig(**llm_config) - self.llm_config = llm_config or SenseNovaU1LLMConfig() + self.llm_config = _build_llm_config(llm_config) if llm_config is not None else SenseNovaU1LLMConfig() if isinstance(vision_config, dict): vision_config = SenseNovaU1VisionConfig(**vision_config) @@ -125,5 +204,6 @@ AutoConfig.register("sensenova_u1", SenseNovaU1Config) __all__ = [ "SenseNovaU1Config", "SenseNovaU1LLMConfig", + "SenseNovaU1MoELLMConfig", "SenseNovaU1VisionConfig", ]