sgl-project--sglang/python/sglang/multimodal_gen/test/unit/test_ideogram4.py

import json
import os
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch

import torch
import torch.nn.functional as F
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig

from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
from sglang.multimodal_gen.configs.models.encoders.ideogram import (
    Ideogram4TextEncoderConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
    Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.sample.ideogram import (
    IDEOGRAM4_PRESETS,
    Ideogram4SamplingParams,
)
from sglang.multimodal_gen.registry import _get_config_info, get_model_info
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType, get_module_role
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
    _maybe_shard_bitsandbytes_4bit_quant_state,
)
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
    ModelOptFp4Config,
    ModelOptFp4LinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
    FP8_WEIGHT_DTYPE,
    W8A8_FP8_GEMM_ENV,
    WeightOnlyFP8ColumnParallelLinear,
    WeightOnlyFP8Linear,
    WeightOnlyFP8RowParallelLinear,
    dequantize_rowwise_fp8_weight,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    Qwen3VLTextRotaryEmbedding,
    qwen3_apply_rotary_pos_emb,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
    TransformerLoader,
    _server_args_for_transformer_component,
)
from sglang.multimodal_gen.runtime.loader.fsdp_load import (
    load_model_from_full_model_state_dict,
)
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
from sglang.multimodal_gen.runtime.models.dits.ideogram import (
    Ideogram4RowParallelLinear,
    Ideogram4Transformer2DModel,
)
from sglang.multimodal_gen.runtime.models.encoders.ideogram import (
    IdeogramQwen3VLTextEncoder,
)
from sglang.multimodal_gen.runtime.pipelines.ideogram import (
    _resolve_ideogram4_unconditional_transformer_weights_path,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
    DenoisingContext,
    DenoisingStage,
    DenoisingStepState,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ideogram import (
    IMAGE_POSITION_OFFSET,
    LLM_TOKEN_INDICATOR,
    OUTPUT_IMAGE_INDICATOR,
    Ideogram4DecodingStage,
    Ideogram4DenoisingStage,
    Ideogram4TextEncodingStage,
    make_step_intervals,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
    TextEncodingStage,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.server_args import set_global_server_args


def _reference_qwen3_mrope(position_ids, head_dim, rope_theta, mrope_section):
    batch_size = position_ids.shape[0]
    pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)
    inv_freq = 1.0 / (
        rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
    )
    inv_freq = inv_freq[None, None, :, None].expand(3, batch_size, -1, 1)
    freqs = inv_freq @ pos.unsqueeze(2)
    freqs = freqs.transpose(2, 3)
    freqs_t = freqs[0].clone()
    for axis, offset in ((1, 1), (2, 2)):
        length = mrope_section[axis] * 3
        idx = torch.arange(offset, length, 3, device=freqs_t.device)
        freqs_t[..., idx] = freqs[axis][..., idx]
    emb = torch.cat((freqs_t, freqs_t), dim=-1)
    return emb.cos(), emb.sin()


class DummyTokenizer:
    def apply_chat_template(self, messages, add_generation_prompt, tokenize):
        return messages[0]["content"][0]["text"]

    def __call__(self, text, return_tensors, add_special_tokens):
        values = [int(x) for x in text.split()]
        return {"input_ids": torch.tensor([values], dtype=torch.long)}


class FakeIdeogramTransformer(torch.nn.Module):
    def forward(self, *, x, **kwargs):
        return torch.zeros_like(x)


class FakeIdeogramVAE(torch.nn.Module):
    def decode(self, z):
        return z[:, :3]


class FakeIdeogramPipeline:
    def __init__(self, transformer, unconditional_transformer):
        self.modules = {
            "transformer": transformer,
            "unconditional_transformer": unconditional_transformer,
        }


class FakeBnbQuantState:
    def __init__(
        self,
        absmax,
        shape=None,
        code=None,
        blocksize=None,
        quant_type=None,
        dtype=None,
        offset=None,
        state2=None,
    ):
        self.absmax = absmax
        self.shape = shape
        self.code = code
        self.blocksize = blocksize
        self.quant_type = quant_type
        self.dtype = dtype
        self.offset = offset
        self.state2 = state2
        self.nested = state2 is not None


def _fake_server_args(cfg=None):
    return SimpleNamespace(
        pipeline_config=cfg or Ideogram4PipelineConfig(),
        comfyui_mode=False,
        enable_torch_compile=False,
        enable_breakable_cuda_graph=False,
        attention_backend="torch_sdpa",
        enable_layerwise_nvtx_marker=False,
        model_loaded={"transformer": True},
        model_paths={},
        disable_autocast=False,
        enable_cfg_parallel=False,
        attention_backend_config=None,
    )


def _fake_ideogram_pipeline(transformer, unconditional_transformer):
    return FakeIdeogramPipeline(transformer, unconditional_transformer)


class TestIdeogram4(unittest.TestCase):
    def test_registry_resolves_model_index_class_name(self):
        get_model_info.cache_clear()
        _get_config_info.cache_clear()
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(f"{tmpdir}/model_index.json", "w", encoding="utf-8") as f:
                json.dump(
                    {"_class_name": "Ideogram4Pipeline", "_diffusers_version": "0.0.0"},
                    f,
                )
            for subdir in (
                "scheduler",
                "text_encoder",
                "tokenizer",
                "transformer",
                "unconditional_transformer",
                "vae",
            ):
                os.mkdir(f"{tmpdir}/{subdir}")
            info = get_model_info(tmpdir, backend="sglang")
        self.assertEqual(info.pipeline_cls.__name__, "Ideogram4Pipeline")
        self.assertIs(info.pipeline_config_cls, Ideogram4PipelineConfig)
        self.assertIs(info.sampling_param_cls, Ideogram4SamplingParams)

    def test_registry_resolves_comfy_nvfp4_repo_to_native_pipeline(self):
        get_model_info.cache_clear()
        _get_config_info.cache_clear()

        info = get_model_info("Comfy-Org/Ideogram-4", backend="sglang")

        self.assertEqual(info.pipeline_cls.__name__, "Ideogram4Nvfp4Pipeline")
        self.assertIs(info.pipeline_config_cls, Ideogram4PipelineConfig)
        self.assertIs(info.sampling_param_cls, Ideogram4SamplingParams)

    def test_registry_resolves_official_nf4_repo_to_native_pipeline(self):
        get_model_info.cache_clear()
        _get_config_info.cache_clear()

        with patch(
            "sglang.multimodal_gen.registry.maybe_download_model_index",
            return_value={
                "_class_name": "Ideogram4Pipeline",
                "_diffusers_version": "0.0.0",
            },
        ):
            info = get_model_info("ideogram-ai/ideogram-4-nf4", backend="sglang")

        self.assertEqual(info.pipeline_cls.__name__, "Ideogram4Pipeline")
        self.assertIs(info.pipeline_config_cls, Ideogram4PipelineConfig)
        self.assertIs(info.sampling_param_cls, Ideogram4SamplingParams)

    def test_rowwise_fp8_dequant_uses_output_channel_scale(self):
        weight = torch.tensor(
            [[1.0, 2.0, -3.0], [4.0, -5.0, 6.0]], dtype=FP8_WEIGHT_DTYPE
        )
        scale = torch.tensor([0.5, 2.0], dtype=torch.float32)
        actual = dequantize_rowwise_fp8_weight(weight, scale, torch.float32)
        expected = weight.to(torch.float32) * scale[:, None]
        torch.testing.assert_close(actual, expected)

    def test_shared_qwen3_mrope_matches_ideogram_reference_layout(self):
        position_ids = torch.tensor(
            [
                [[0, 0, 0], [1, 1, 1], [65536, 65536, 65536]],
                [[0, 0, 0], [0, 2, 3], [65536, 65537, 65538]],
            ],
            dtype=torch.long,
        )
        head_dim = 8
        rope_theta = 5_000_000.0
        mrope_section = (2, 1, 1)
        rotary_emb = Qwen3VLTextRotaryEmbedding(
            head_dim=head_dim,
            rope_theta=rope_theta,
            mrope_section=mrope_section,
        )

        cos, sin = rotary_emb(torch.empty((), dtype=torch.float32), position_ids)
        ref_cos, ref_sin = _reference_qwen3_mrope(
            position_ids, head_dim, rope_theta, mrope_section
        )

        torch.testing.assert_close(cos.float(), ref_cos)
        torch.testing.assert_close(sin.float(), ref_sin)

    def test_usp_attention_key_mask_matches_segment_mask_for_valid_tokens(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            torch.manual_seed(0)
            batch_size, seq_len, num_heads, head_dim = 2, 5, 2, 8
            q = torch.randn(batch_size, seq_len, num_heads, head_dim)
            k = torch.randn(batch_size, seq_len, num_heads, head_dim)
            v = torch.randn(batch_size, seq_len, num_heads, head_dim)
            segment_ids = torch.tensor(
                [[-1, -1, 1, 1, 1], [-1, 1, 1, 1, 1]], dtype=torch.long
            )
            position_ids = torch.stack(
                [
                    torch.arange(seq_len).expand(batch_size, -1),
                    torch.arange(seq_len).expand(batch_size, -1) + 1,
                    torch.arange(seq_len).expand(batch_size, -1) + 2,
                ],
                dim=-1,
            )
            rotary_emb = Qwen3VLTextRotaryEmbedding(
                head_dim=head_dim, mrope_section=(2, 1, 1)
            )
            cos, sin = rotary_emb(q, position_ids)
            q, k = qwen3_apply_rotary_pos_emb(q, k, cos.unsqueeze(2), sin.unsqueeze(2))

            full_mask = (
                segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)
            ).unsqueeze(1)
            expected = F.scaled_dot_product_attention(
                q.transpose(1, 2),
                k.transpose(1, 2),
                v.transpose(1, 2),
                attn_mask=full_mask,
            ).transpose(1, 2)

            with (
                patch(
                    "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                    return_value=1,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.attention.layer.get_sequence_parallel_world_size",
                    return_value=1,
                ),
            ):
                attn = USPAttention(
                    num_heads=num_heads,
                    head_size=head_dim,
                    supported_attention_backends={AttentionBackendEnum.TORCH_SDPA},
                )
                with set_forward_context(current_timestep=0, attn_metadata=None):
                    actual = attn(q, k, v, attn_mask=segment_ids > 0)

            valid = segment_ids > 0
            torch.testing.assert_close(actual[valid], expected[valid])
        finally:
            set_global_server_args(prev_args)

    def test_ideogram_preset_guidance_order(self):
        turbo = IDEOGRAM4_PRESETS["V4_TURBO_12"]
        default = IDEOGRAM4_PRESETS["V4_DEFAULT_20"]
        self.assertEqual(turbo["num_steps"], 12)
        self.assertEqual(default["num_steps"], 20)
        self.assertEqual(turbo["guidance_schedule"][0], 3.0)
        self.assertEqual(turbo["guidance_schedule"][-1], 7.0)
        self.assertEqual(tuple(make_step_intervals(2).tolist()), (0.0, 0.5, 1.0))

    def test_ideogram_sampling_params_sync_steps_with_preset(self):
        params = Ideogram4SamplingParams(preset="V4_TURBO_12")
        self.assertEqual(params.num_inference_steps, 12)
        self.assertEqual(params.guidance_scale, 7.0)
        same_steps = Ideogram4SamplingParams(
            preset="V4_TURBO_12", num_inference_steps=12
        )
        self.assertEqual(same_steps.num_inference_steps, 12)
        with self.assertRaisesRegex(ValueError, "derives num_inference_steps"):
            Ideogram4SamplingParams(preset="V4_TURBO_12", num_inference_steps=20)
        same_guidance = Ideogram4SamplingParams(
            preset="V4_TURBO_12", guidance_scale=7.0
        )
        self.assertEqual(same_guidance.guidance_scale, 7.0)
        with self.assertRaisesRegex(ValueError, "guidance_scale cannot be set"):
            Ideogram4SamplingParams(preset="V4_TURBO_12", guidance_scale=6.0)
        with self.assertRaisesRegex(ValueError, "Unknown Ideogram 4 preset"):
            Ideogram4SamplingParams(preset="V4_FAST")

    def test_ideogram_sampling_params_merge_recomputes_preset_fields(self):
        target = Ideogram4SamplingParams()
        user = Ideogram4SamplingParams(
            preset="V4_TURBO_12",
            height=256,
            width=256,
        )

        target._merge_with_user_params(
            user, explicit_fields={"preset", "height", "width"}
        )

        self.assertEqual(target.preset, "V4_TURBO_12")
        self.assertEqual(target.num_inference_steps, 12)
        self.assertEqual(target.guidance_scale, 7.0)
        self.assertEqual(target.height, 256)
        self.assertEqual(target.width, 256)

    def test_unconditional_transformer_uses_denoiser_loader_path(self):
        self.assertIn("unconditional_transformer", TransformerLoader.component_names)
        self.assertEqual(
            get_module_role("unconditional_transformer"), RoleType.DENOISER
        )

        server_args = SimpleNamespace(
            transformer_weights_path="/unused/override.safetensors",
            nunchaku_config={"enabled": True},
            component_transformer_weights_paths={},
        )
        component_args = _server_args_for_transformer_component(
            server_args, "unconditional_transformer"
        )
        self.assertIsNot(component_args, server_args)
        self.assertIsNone(component_args.transformer_weights_path)
        self.assertIsNone(component_args.nunchaku_config)

    def test_transformer_component_uses_per_component_weights_override(self):
        server_args = SimpleNamespace(
            transformer_weights_path=(
                "/ckpt/diffusion_models/ideogram4_nvfp4_mixed.safetensors"
            ),
            nunchaku_config={"enabled": True},
            component_transformer_weights_paths={
                "unconditional_transformer": (
                    "/ckpt/diffusion_models/"
                    "ideogram4_unconditional_nvfp4_mixed.safetensors"
                )
            },
        )

        component_args = _server_args_for_transformer_component(
            server_args,
            "unconditional_transformer",
        )

        self.assertIsNot(component_args, server_args)
        self.assertEqual(
            component_args.transformer_weights_path,
            "/ckpt/diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors",
        )
        self.assertIsNone(component_args.nunchaku_config)

    def test_ideogram_nvfp4_unconditional_transformer_path_uses_sibling_file(self):
        self.assertEqual(
            _resolve_ideogram4_unconditional_transformer_weights_path(
                "/ckpt/diffusion_models/ideogram4_nvfp4_mixed.safetensors"
            ),
            "/ckpt/diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors",
        )
        self.assertIsNone(
            _resolve_ideogram4_unconditional_transformer_weights_path(
                "/ckpt/custom_transformer.safetensors"
            )
        )

    def test_ideogram_denoiser_does_not_request_dtype_cast(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(_fake_server_args())
            transformer = FakeIdeogramTransformer()
            unconditional_transformer = FakeIdeogramTransformer()
            stage = Ideogram4DenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=_fake_ideogram_pipeline(
                    transformer, unconditional_transformer
                ),
            )
            uses = stage.component_uses(_fake_server_args(), "stage")
        finally:
            set_global_server_args(prev_args)
        self.assertEqual(
            [use.component_name for use in uses],
            [
                "transformer",
                "unconditional_transformer",
            ],
        )
        self.assertTrue(all(use.target_dtype is None for use in uses))

    def test_ideogram_stages_inherit_common_stage_bases(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(_fake_server_args())
            text_stage = Ideogram4TextEncodingStage(
                text_encoder=None, tokenizer=DummyTokenizer()
            )
            denoising_stage = Ideogram4DenoisingStage(
                transformer=FakeIdeogramTransformer(),
                unconditional_transformer=FakeIdeogramTransformer(),
            )
        finally:
            set_global_server_args(prev_args)
        self.assertIsInstance(text_stage, TextEncodingStage)
        self.assertIsInstance(denoising_stage, DenoisingStage)

    def test_ideogram_text_encoding_dedup_fingerprint_and_extra_copy(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        cfg = Ideogram4PipelineConfig()
        args = _fake_server_args(cfg)
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(args)
            stage = Ideogram4TextEncodingStage(
                text_encoder=None, tokenizer=DummyTokenizer()
            )
        finally:
            set_global_server_args(prev_args)
        base = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=256,
                width=256,
                num_outputs_per_prompt=1,
            )
        )
        same = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=256,
                width=256,
                num_outputs_per_prompt=1,
            )
        )
        different_height = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=512,
                width=256,
                num_outputs_per_prompt=1,
            )
        )
        different_width = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=256,
                width=512,
                num_outputs_per_prompt=1,
            )
        )
        different_outputs = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=256,
                width=256,
                num_outputs_per_prompt=2,
            )
        )

        base_fingerprint = stage.build_dedup_fingerprint(base, args)
        self.assertEqual(base_fingerprint, stage.build_dedup_fingerprint(same, args))
        self.assertNotEqual(
            base_fingerprint, stage.build_dedup_fingerprint(different_height, args)
        )
        self.assertNotEqual(
            base_fingerprint, stage.build_dedup_fingerprint(different_width, args)
        )
        self.assertNotEqual(
            base_fingerprint, stage.build_dedup_fingerprint(different_outputs, args)
        )

        base.prompt_embeds = [torch.tensor([1.0])]
        base.prompt_embeds_mask = [torch.tensor([True])]
        base.extra["ideogram4"] = {
            "position_ids": torch.tensor([[1]]),
            "metadata": {"grid_h": 16},
        }
        stage.copy_deduplicated_outputs(base, same)

        self.assertIn("ideogram4", same.extra)
        self.assertTrue(
            torch.equal(
                same.extra["ideogram4"]["position_ids"],
                base.extra["ideogram4"]["position_ids"],
            )
        )
        self.assertIsNot(
            same.extra["ideogram4"]["position_ids"],
            base.extra["ideogram4"]["position_ids"],
        )

    def test_ideogram_text_encoding_verifies_custom_outputs(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        args = _fake_server_args()
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(args)
            stage = Ideogram4TextEncodingStage(
                text_encoder=None, tokenizer=DummyTokenizer()
            )
        finally:
            set_global_server_args(prev_args)

        batch = Req(
            sampling_params=Ideogram4SamplingParams(
                prompt="11 12",
                height=256,
                width=256,
                num_outputs_per_prompt=1,
            )
        )
        self.assertTrue(stage.verify_input(batch, args).is_valid())

        batch.do_classifier_free_guidance = True
        batch.negative_prompt = []
        batch.negative_prompt_embeds = []
        batch.prompt_embeds = [torch.zeros(1, 4, 8)]
        batch.prompt_embeds_mask = [torch.ones(1, 4, dtype=torch.bool)]
        batch.extra["ideogram4"] = {"position_ids": torch.zeros(1, 4, 3)}

        self.assertTrue(stage.verify_output(batch, args).is_valid())

    def test_ideogram_denoising_component_names_from_pipeline_modules(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        transformer = FakeIdeogramTransformer()
        unconditional_transformer = FakeIdeogramTransformer()
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(_fake_server_args())
            stage = Ideogram4DenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=_fake_ideogram_pipeline(
                    transformer, unconditional_transformer
                ),
            )
            uses = stage.component_uses(_fake_server_args(), "stage")
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(
            [use.component_name for use in uses],
            ["transformer", "unconditional_transformer"],
        )

    def test_ideogram_attention_backend_is_passed_from_config(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        config = Ideogram4DiTConfig()
        self.assertEqual(
            config.arch_config._supported_attention_backends,
            {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
        )
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(_fake_server_args())
            with patch(
                "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                return_value=1,
            ):
                with torch.device("meta"):
                    model = Ideogram4Transformer2DModel(config, {})
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(
            model.supported_attention_backends,
            config.arch_config._supported_attention_backends,
        )
        self.assertEqual(
            model.layers[0].attention.attn.backend,
            AttentionBackendEnum.TORCH_SDPA,
        )

    def test_ideogram_dit_meta_state_dict_matches_checkpoint_shapes(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with patch(
                "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                return_value=1,
            ):
                with torch.device("meta"):
                    model = Ideogram4Transformer2DModel(Ideogram4DiTConfig(), {})
        finally:
            set_global_server_args(prev_args)
        state = model.state_dict()
        self.assertEqual(len(state), 669)
        self.assertEqual(tuple(state["input_proj.weight"].shape), (4608, 128))
        self.assertEqual(tuple(state["input_proj.weight_scale"].shape), (4608,))
        self.assertEqual(
            tuple(state["layers.0.attention.qkv.weight"].shape), (13824, 4608)
        )
        self.assertEqual(state["layers.0.attention.qkv.weight"].dtype, FP8_WEIGHT_DTYPE)

    def test_ideogram_dit_uses_tp_fp8_linears_when_tp_is_initialized(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with (
                patch(
                    "sglang.multimodal_gen.runtime.models.dits.ideogram.model_parallel_is_initialized",
                    return_value=True,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.models.dits.ideogram.get_tp_world_size",
                    return_value=2,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.linear.get_tp_group",
                    return_value=fake_tp_group,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8.get_tp_group",
                    return_value=fake_tp_group,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                    return_value=1,
                ),
            ):
                with torch.device("meta"):
                    model = Ideogram4Transformer2DModel(Ideogram4DiTConfig(), {})
        finally:
            set_global_server_args(prev_args)

        self.assertIsInstance(model.input_proj, WeightOnlyFP8ColumnParallelLinear)
        self.assertEqual(tuple(model.input_proj.weight.shape), (2304, 128))
        self.assertEqual(
            tuple(model.layers[0].attention.qkv.weight.shape), (6912, 4608)
        )

    def test_ideogram_dit_nvfp4_quant_config_uses_native_fp4_linears(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        quant_config = ModelOptFp4Config(
            is_checkpoint_nvfp4_serialized=True,
            group_size=16,
            exclude_modules=[
                "input_proj",
                "llm_cond_proj",
                "t_embedding.*",
                "adaln_proj",
                "layers.*.adaln_modulation",
                "final_layer.*",
            ],
        )
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with patch(
                "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                return_value=1,
            ):
                with torch.device("meta"):
                    model = Ideogram4Transformer2DModel(
                        Ideogram4DiTConfig(),
                        {},
                        quant_config=quant_config,
                    )
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(model.layers[0].attention.qkv.prefix, "layers.0.attention.qkv")
        self.assertIsInstance(
            model.layers[0].attention.qkv.quant_method,
            ModelOptFp4LinearMethod,
        )
        self.assertIsInstance(model.input_proj.quant_method, UnquantizedLinearMethod)

        state = model.state_dict()
        self.assertEqual(
            tuple(state["layers.0.attention.qkv.weight"].shape),
            (13824, 2304),
        )
        self.assertEqual(state["layers.0.attention.qkv.weight"].dtype, torch.uint8)
        self.assertEqual(
            tuple(state["layers.0.attention.qkv.weight_scale"].shape),
            (13824, 288),
        )
        self.assertEqual(
            state["layers.0.attention.qkv.weight_scale"].dtype,
            FP8_WEIGHT_DTYPE,
        )
        self.assertEqual(
            tuple(state["layers.0.attention.qkv.weight_scale_2"].shape),
            (1,),
        )
        self.assertEqual(
            tuple(state["layers.0.attention.qkv.input_scale"].shape),
            (1,),
        )

    def test_ideogram_dit_tp_nvfp4_uses_megatron_parallel_quant_linears(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
        quant_config = ModelOptFp4Config(
            is_checkpoint_nvfp4_serialized=True,
            group_size=16,
        )
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with (
                patch(
                    "sglang.multimodal_gen.runtime.models.dits.ideogram.model_parallel_is_initialized",
                    return_value=True,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.models.dits.ideogram.get_tp_world_size",
                    return_value=2,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.linear.get_tp_group",
                    return_value=fake_tp_group,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
                    return_value=1,
                ),
            ):
                with torch.device("meta"):
                    model = Ideogram4Transformer2DModel(
                        Ideogram4DiTConfig(),
                        {},
                        quant_config=quant_config,
                    )
        finally:
            set_global_server_args(prev_args)

        self.assertFalse(model.layers[0].attention.qkv.gather_output)
        self.assertEqual(
            tuple(model.layers[0].attention.qkv.weight.shape), (6912, 2304)
        )
        self.assertIsInstance(
            model.layers[0].attention.qkv.quant_method,
            ModelOptFp4LinearMethod,
        )
        self.assertIsInstance(model.layers[0].attention.o, Ideogram4RowParallelLinear)
        self.assertTrue(model.layers[0].attention.o.input_is_parallel)
        self.assertFalse(model.layers[0].feed_forward.w1.gather_output)
        self.assertFalse(model.layers[0].feed_forward.w3.gather_output)
        self.assertIsInstance(
            model.layers[0].feed_forward.w2, Ideogram4RowParallelLinear
        )
        self.assertTrue(model.layers[0].feed_forward.w2.input_is_parallel)

    def test_bitsandbytes_tp_quant_state_uses_local_output_shard(self):
        param = torch.nn.Parameter(
            torch.empty(8, 1, dtype=torch.uint8), requires_grad=False
        )
        param.bnb_full_shape = (4, 8)
        param.bnb_local_shape = (2, 8)
        param.bnb_output_shard_start = 2
        param.bnb_input_shard_start = 0
        quant_state = FakeBnbQuantState(
            absmax=torch.arange(8, dtype=torch.float32),
            shape=torch.Size((4, 8)),
            code=torch.ones(16, dtype=torch.float32),
            blocksize=4,
            quant_type="nf4",
            dtype=torch.bfloat16,
        )

        sharded = _maybe_shard_bitsandbytes_4bit_quant_state(param, quant_state)

        self.assertEqual(sharded.shape, torch.Size((2, 8)))
        torch.testing.assert_close(sharded.absmax, torch.tensor([4.0, 5.0, 6.0, 7.0]))

    def test_assign_load_preserves_bitsandbytes_tp_attrs(self):
        class TinyModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.weight = torch.nn.Parameter(
                    torch.empty(8, 1, dtype=torch.uint8), requires_grad=False
                )
                self.weight.bnb_full_shape = (4, 8)
                self.weight.bnb_local_shape = (2, 8)
                self.weight.bnb_output_shard_start = 2
                self.weight.bnb_input_shard_start = 0

        model = TinyModule()
        load_model_from_full_model_state_dict(
            model,
            iter([("weight", torch.ones(8, 1, dtype=torch.uint8))]),
            torch.device("cpu"),
            param_dtype=None,
            strict=True,
            param_names_mapping=lambda name: (name, None, None),
        )

        self.assertEqual(model.weight.bnb_full_shape, (4, 8))
        self.assertEqual(model.weight.bnb_local_shape, (2, 8))
        self.assertEqual(model.weight.bnb_output_shard_start, 2)
        self.assertEqual(model.weight.bnb_input_shard_start, 0)

    def test_missing_weight_only_fp8_scale_is_fatal(self):
        with torch.device("meta"):
            model = WeightOnlyFP8Linear(3, 2, bias=False)
        weights = iter(
            [
                (
                    "weight",
                    torch.tensor(
                        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
                        dtype=FP8_WEIGHT_DTYPE,
                    ),
                )
            ]
        )
        with self.assertRaisesRegex(ValueError, "Required checkpoint parameter"):
            load_model_from_full_model_state_dict(
                model,
                weights,
                torch.device("cpu"),
                param_dtype=None,
                strict=False,
                param_names_mapping=lambda name: (name, None, None),
            )

    def test_weight_only_fp8_load_accepts_explicit_scale(self):
        with torch.device("meta"):
            model = WeightOnlyFP8Linear(3, 2, bias=False)
        weights = iter(
            [
                (
                    "weight",
                    torch.tensor(
                        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
                        dtype=FP8_WEIGHT_DTYPE,
                    ),
                ),
                ("weight_scale", torch.tensor([0.5, 2.0], dtype=torch.float32)),
            ]
        )
        load_model_from_full_model_state_dict(
            model,
            weights,
            torch.device("cpu"),
            param_dtype=None,
            strict=False,
            param_names_mapping=lambda name: (name, None, None),
        )
        self.assertEqual(model.weight.dtype, FP8_WEIGHT_DTYPE)
        self.assertEqual(model.weight_scale.dtype, torch.float32)

    def test_weight_only_fp8_w8a8_gemm_defaults_to_off(self):
        with patch.dict(os.environ, {W8A8_FP8_GEMM_ENV: "0"}):
            model = WeightOnlyFP8Linear(3, 2, bias=False)

        self.assertFalse(model.enable_fused_w8a8)

    def test_weight_only_fp8_w8a8_gemm_env_opt_in(self):
        with patch.dict(os.environ, {W8A8_FP8_GEMM_ENV: "1"}):
            model = WeightOnlyFP8Linear(3, 2, bias=False)

        self.assertTrue(model.enable_fused_w8a8)

    def test_weight_only_fp8_w8a8_gemm_explicit_flag_overrides_env(self):
        with patch.dict(os.environ, {W8A8_FP8_GEMM_ENV: "1"}):
            model = WeightOnlyFP8Linear(3, 2, bias=False, enable_fused_w8a8=False)

        self.assertFalse(model.enable_fused_w8a8)

    def test_ideogram_text_encoder_post_config_hook_preserves_local_arch(self):
        config = Ideogram4TextEncoderConfig()
        config.arch_config.architectures = ["RemoteQwen3VLTextModel"]
        config.arch_config.ideogram_fp8_weight_only = False
        config.post_diffusers_config_update()
        self.assertEqual(
            config.arch_config.architectures, ["IdeogramQwen3VLTextEncoder"]
        )
        self.assertTrue(config.arch_config.ideogram_fp8_weight_only)
        self.assertFalse(config.arch_config.ideogram_bnb_4bit_weight_only)
        self.assertFalse(config.arch_config.requires_gpu_resident_text_encoder)

    def test_ideogram_text_encoder_post_config_hook_uses_bnb_for_nf4(self):
        config = Ideogram4TextEncoderConfig()
        config.update_model_arch(
            {
                "quantization_config": {
                    "quant_method": "bitsandbytes",
                    "load_in_4bit": True,
                    "bnb_4bit_quant_type": "nf4",
                }
            }
        )

        self.assertEqual(
            config.arch_config.architectures, ["IdeogramQwen3VLTextEncoder"]
        )
        self.assertTrue(config.arch_config.ideogram_bnb_4bit_weight_only)
        self.assertFalse(config.arch_config.ideogram_fp8_weight_only)
        self.assertTrue(config.arch_config.requires_gpu_resident_text_encoder)

    def test_ideogram_text_encoder_swaps_linears_to_weight_only_fp8(self):
        config = Ideogram4TextEncoderConfig()
        config.post_diffusers_config_update()
        config.arch_config.text_config = Qwen3VLTextConfig(
            vocab_size=32,
            hidden_size=16,
            intermediate_size=32,
            num_hidden_layers=1,
            num_attention_heads=2,
            num_key_value_heads=2,
            head_dim=8,
            max_position_embeddings=64,
            pad_token_id=0,
        )
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with (
                patch.dict(os.environ, {W8A8_FP8_GEMM_ENV: "1"}),
                torch.device("meta"),
            ):
                encoder = IdeogramQwen3VLTextEncoder(config)
        finally:
            set_global_server_args(prev_args)
        fp8_linears = [
            module
            for module in encoder.modules()
            if isinstance(module, WeightOnlyFP8Linear)
        ]
        self.assertTrue(fp8_linears)
        self.assertTrue(all(not module.enable_fused_w8a8 for module in fp8_linears))
        self.assertFalse(
            any(isinstance(module, torch.nn.Linear) for module in encoder.modules())
        )

    def test_ideogram_text_encoder_tp_fp8_uses_megatron_parallel_linears(self):
        config = Ideogram4TextEncoderConfig()
        config.post_diffusers_config_update()
        config.arch_config.text_config = Qwen3VLTextConfig(
            vocab_size=32,
            hidden_size=16,
            intermediate_size=32,
            num_hidden_layers=1,
            num_attention_heads=4,
            num_key_value_heads=4,
            head_dim=4,
            max_position_embeddings=64,
            pad_token_id=0,
        )
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(
                SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
            )
            with (
                patch(
                    "sglang.multimodal_gen.runtime.models.encoders.qwen3vl.model_parallel_is_initialized",
                    return_value=True,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.models.encoders.qwen3vl.get_tp_world_size",
                    return_value=2,
                ),
                patch(
                    "sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8.get_tp_group",
                    return_value=fake_tp_group,
                ),
            ):
                with torch.device("meta"):
                    encoder = IdeogramQwen3VLTextEncoder(config)
        finally:
            set_global_server_args(prev_args)

        layer = encoder.language_model.layers[0]
        self.assertEqual(layer.self_attn.num_heads, 2)
        self.assertEqual(layer.self_attn.num_key_value_heads, 2)
        self.assertIsInstance(layer.self_attn.q_proj, WeightOnlyFP8ColumnParallelLinear)
        self.assertFalse(layer.self_attn.q_proj.gather_output)
        self.assertIsInstance(layer.self_attn.o_proj, WeightOnlyFP8RowParallelLinear)
        self.assertTrue(layer.self_attn.o_proj.input_is_parallel)
        self.assertTrue(layer.self_attn.o_proj.reduce_results)
        self.assertIsInstance(layer.mlp.gate_proj, WeightOnlyFP8ColumnParallelLinear)
        self.assertFalse(layer.mlp.gate_proj.gather_output)
        self.assertIsInstance(layer.mlp.up_proj, WeightOnlyFP8ColumnParallelLinear)
        self.assertFalse(layer.mlp.up_proj.gather_output)
        self.assertIsInstance(layer.mlp.down_proj, WeightOnlyFP8RowParallelLinear)
        self.assertTrue(layer.mlp.down_proj.input_is_parallel)
        self.assertTrue(layer.mlp.down_proj.reduce_results)

    def test_denoise_and_decode_shape_check(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        cfg = Ideogram4PipelineConfig()
        args = _fake_server_args(cfg)
        device = get_local_torch_device()
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(args)
            transformer = FakeIdeogramTransformer()
            unconditional_transformer = FakeIdeogramTransformer()
            denoise_stage = Ideogram4DenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=_fake_ideogram_pipeline(
                    transformer, unconditional_transformer
                ),
            )
            decode_stage = Ideogram4DecodingStage(vae=FakeIdeogramVAE())
            batch = Req(
                sampling_params=Ideogram4SamplingParams(
                    prompt="11 12",
                    height=256,
                    width=256,
                    preset="V4_TURBO_12",
                    suppress_logs=True,
                )
            )
            batch.latents = torch.zeros(1, 1, 128, device=device)
            batch.raw_latent_shape = batch.latents.shape
            batch.prompt_embeds = [torch.zeros(1, 2, 8, device=device)]
            batch.extra["ideogram4"] = {
                "max_text_tokens": 1,
                "num_image_tokens": 1,
                "position_ids": torch.zeros(1, 2, 3, dtype=torch.long, device=device),
                "segment_ids": torch.ones(1, 2, dtype=torch.long, device=device),
                "indicator": torch.tensor(
                    [[LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR]],
                    dtype=torch.long,
                    device=device,
                ),
                "grid_h": 1,
                "grid_w": 1,
            }
            denoised = denoise_stage.forward(batch, args)
            self.assertEqual(tuple(denoised.latents.shape), (1, 1, 128))
            decoded = decode_stage.forward(denoised, args)
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(tuple(decoded.output.shape), (1, 3, 2, 2))

    def test_ideogram_bcg_padded_positive_output_is_cropped(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        cfg = Ideogram4PipelineConfig()
        args = _fake_server_args(cfg)
        device = get_local_torch_device()
        prev_args = server_args_module._global_server_args
        try:
            set_global_server_args(args)
            transformer = FakeIdeogramTransformer()
            unconditional_transformer = FakeIdeogramTransformer()
            stage = Ideogram4DenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=_fake_ideogram_pipeline(
                    transformer, unconditional_transformer
                ),
            )
            batch = Req(
                sampling_params=Ideogram4SamplingParams(
                    prompt="11 12",
                    height=256,
                    width=512,
                    preset="V4_TURBO_12",
                    suppress_logs=True,
                )
            )
            batch.prompt_embeds = [torch.zeros(1, 3, 8, device=device)]
            batch.extra["ideogram4"] = {
                "max_text_tokens": 1,
                "num_image_tokens": 2,
                "position_ids": torch.zeros(1, 3, 3, dtype=torch.long, device=device),
                "segment_ids": torch.ones(1, 3, dtype=torch.long, device=device),
                "indicator": torch.tensor(
                    [
                        [
                            LLM_TOKEN_INDICATOR,
                            OUTPUT_IMAGE_INDICATOR,
                            OUTPUT_IMAGE_INDICATOR,
                        ]
                    ],
                    dtype=torch.long,
                    device=device,
                ),
            }
            ctx = DenoisingContext(
                scheduler=None,
                extra_step_kwargs={},
                target_dtype=torch.float32,
                autocast_enabled=False,
                timesteps=torch.tensor([0], device=device),
                num_inference_steps=1,
                num_warmup_steps=0,
                image_kwargs={},
                pos_cond_kwargs={},
                neg_cond_kwargs={},
                latents=torch.zeros(1, 2, 128, device=device),
                boundary_timestep=None,
                z=None,
                reserved_frames_mask=None,
                seq_len=None,
                guidance=torch.ones(1, device=device),
                is_warmup=False,
                extra={
                    "ideogram4_schedule_values": torch.tensor(
                        [1.0, 0.0], device=device
                    ),
                    "ideogram4_schedule_deltas": torch.tensor([1.0], device=device),
                    "ideogram4_guidance_schedule": torch.tensor([1.0], device=device),
                    "ideogram4_text_z_padding": torch.zeros(1, 1, 128, device=device),
                    "ideogram4_attn_mask": torch.ones(
                        1, 3, dtype=torch.bool, device=device
                    ),
                    "ideogram4_attn_mask_meta": None,
                    "ideogram4_neg_position_ids": torch.zeros(
                        1, 2, 3, dtype=torch.long, device=device
                    ),
                    "ideogram4_neg_segment_ids": torch.ones(
                        1, 2, dtype=torch.long, device=device
                    ),
                    "ideogram4_neg_indicator": torch.full(
                        (1, 2),
                        OUTPUT_IMAGE_INDICATOR,
                        dtype=torch.long,
                        device=device,
                    ),
                    "ideogram4_neg_attn_mask": torch.ones(
                        1, 2, dtype=torch.bool, device=device
                    ),
                    "ideogram4_neg_attn_mask_meta": None,
                    "ideogram4_neg_llm_features": torch.zeros(1, 2, 8, device=device),
                },
            )
            step = DenoisingStepState(
                step_index=0,
                t_host=torch.tensor(0),
                t_device=torch.tensor(0, device=device),
                t_int=0,
                current_model=transformer,
                current_guidance_scale=None,
                attn_metadata=None,
            )

            def fake_run(current_model, call_kwargs):
                if current_model is transformer:
                    out = torch.zeros(1, 5, 128, device=device)
                    out[:, 1:3] = 4.0
                    out[:, 3:] = 99.0
                    return out
                return torch.ones(1, 2, 128, device=device)

            with patch.object(stage, "_run_ideogram_transformer", side_effect=fake_run):
                stage._run_denoising_step(ctx, step, batch, args)
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(tuple(ctx.latents.shape), (1, 2, 128))
        self.assertTrue(
            torch.allclose(ctx.latents, torch.full((1, 2, 128), 4.0, device=device))
        )

    def test_text_input_builder_matches_official_layout(self):
        prev_args = None
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            cfg = Ideogram4PipelineConfig()
            args = SimpleNamespace(pipeline_config=cfg, comfyui_mode=False)
            set_global_server_args(args)
            stage = Ideogram4TextEncodingStage(
                text_encoder=None, tokenizer=DummyTokenizer()
            )
            inputs = stage._build_inputs(["11 12 13", "21"], 256, 256, args)
        finally:
            set_global_server_args(prev_args)

        self.assertEqual(inputs["grid_h"], 16)
        self.assertEqual(inputs["grid_w"], 16)
        self.assertEqual(inputs["num_image_tokens"], 256)
        self.assertEqual(inputs["max_text_tokens"], 3)
        self.assertEqual(inputs["token_ids"][0, :3].tolist(), [11, 12, 13])
        self.assertEqual(inputs["token_ids"][1, :2].tolist(), [0, 0])
        self.assertTrue(
            torch.all(inputs["indicator"][0, :3] == LLM_TOKEN_INDICATOR).item()
        )
        self.assertTrue(
            torch.all(inputs["indicator"][0, 3:] == OUTPUT_IMAGE_INDICATOR).item()
        )
        self.assertEqual(inputs["position_ids"][0, 3, 0].item(), IMAGE_POSITION_OFFSET)

    def test_text_input_builder_rejects_unsupported_resolution(self):
        import sglang.multimodal_gen.runtime.server_args as server_args_module

        prev_args = server_args_module._global_server_args
        try:
            cfg = Ideogram4PipelineConfig()
            args = SimpleNamespace(pipeline_config=cfg, comfyui_mode=False)
            set_global_server_args(args)
            stage = Ideogram4TextEncodingStage(
                text_encoder=None, tokenizer=DummyTokenizer()
            )
            with self.assertRaisesRegex(ValueError, "between 256 and 2048"):
                stage._build_inputs(["11"], 128, 256, args)
        finally:
            set_global_server_args(prev_args)


if __name__ == "__main__":
    unittest.main()