Skip to content

vllm_omni.diffusion.models.minimax_h3.encoder

MiniMax H3 Qwen3-VL layer-50 text/vision encoder.

The encoder is reimplemented on top of vLLM-style tensor-parallel building blocks (column/row/vocab-parallel linears bound to the MiniMax H3 encoder process group) instead of the single-GPU transformers backbone. This makes the large Qwen3-VL text model TP-shardable across text_encoder_tp_size ranks, which removes the rank-0 memory hotspot that dominates no-offload peak memory (the retained 50-layer encoder is ~51.5 GB in BF16).

The computation contract is unchanged from the HF reference path:

  • plain BF16 weights, only the first MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER decoder layers are kept (the checkpoint consumes the unnormalized hidden state after decoder layer 49);
  • cuDNN SDP is enabled during encode;
  • the multimodal backbone runs with an all-ones attention mask and mm_token_type_ids derived from the image/video token ids;
  • DeepStack visual features are injected at the first len(deepstack_visual_indexes) decoder layers;
  • the output is hidden_states[50] of shape [seq, 5120].

Tensor parallelism follows vLLM's TP semantics: vocab-parallel embeddings and row-parallel projections all-reduce so the hidden state stays fully replicated on every encoder rank, while column-parallel projections (QKV / MLP up) keep a local shard. After the final layer every encoder rank therefore holds the complete [seq, 5120] hidden state.

MINIMAX_H3_QWEN3VL_HIDDEN_DIM module-attribute 

MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120

MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER module-attribute 

MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50

logger module-attribute 

logger = init_logger(__name__)

MiniMaxH3Qwen3VLEncoder

Bases: Module

Frozen, TP-capable Qwen3-VL backbone returning layer-50 states.

encoder_group is a GroupCoordinator over the encoder tensor-parallel ranks (by default the first text_encoder_tp_size DiT ranks). Only ranks inside the group load weights; other ranks construct a parameter-free stub that never runs the forward.

device_target instance-attribute 

device_target = device

encoder_group instance-attribute 

encoder_group = encoder_group

image_token_id instance-attribute 

image_token_id = int(config.image_token_id)

is_loaded property 

is_loaded: bool

quant_config instance-attribute 

quant_config = quant_config

text_model instance-attribute 

text_model = MiniMaxH3Qwen3VLTextModel(
    encoder_group,
    config.text_config,
    MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
    dtype,
    quant_config=quant_config,
)

tp_size property 

tp_size: int

video_token_id instance-attribute 

video_token_id = int(config.video_token_id)

vision instance-attribute 

vision = MiniMaxH3Qwen3VLVisionModel(config.vision_config)

encode_ids 

encode_ids(
    input_ids: Tensor,
    *,
    pixel_values: Tensor | None = None,
    image_grid_thw: Tensor | None = None,
    pixel_values_videos: Tensor | None = None,
    video_grid_thw: Tensor | None = None,
) -> Tensor

forward 

forward(
    input_ids: Tensor,
    *,
    pixel_values: Tensor | None = None,
    image_grid_thw: Tensor | None = None,
    pixel_values_videos: Tensor | None = None,
    video_grid_thw: Tensor | None = None,
) -> Tensor

Hook-compatible entry point for model-level CPU offloading.

load_to_device 

load_to_device() -> None

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

offload_to_cpu 

offload_to_cpu() -> None

set_omni_component_cache 

set_omni_component_cache(
    cache: BoundedAllocatorCache | None,
) -> None

MiniMaxH3Qwen3VLMergedColumnParallelLinear

Bases: LinearBase

Packed gate/up projection sharded along the output dimension.

group instance-attribute 

group = group

intermediate_size instance-attribute 

intermediate_size = intermediate_size

intermediate_size_per_partition instance-attribute 

intermediate_size_per_partition = (
    intermediate_size // tp_size
)

output_dtype instance-attribute 

output_dtype = dtype

forward 

forward(input_: Tensor) -> Tensor

weight_loader 

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    loaded_shard_id: str | None = None,
) -> None

MiniMaxH3Qwen3VLQKVParallelLinear

Bases: LinearBase

QKV projection with GQA head sharding along the output dimension.

group instance-attribute 

group = group

head_dim instance-attribute 

head_dim = head_dim

hidden_size instance-attribute 

hidden_size = hidden_size

local_num_heads instance-attribute 

local_num_heads = num_heads // tp_size

local_num_kv_heads instance-attribute 

local_num_kv_heads = num_kv_heads // tp_size

num_heads instance-attribute 

num_heads = num_heads

num_kv_heads instance-attribute 

num_kv_heads = num_kv_heads

output_dtype instance-attribute 

output_dtype = dtype

forward 

forward(input_: Tensor) -> Tensor

weight_loader 

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    loaded_shard_id: str | None = None,
) -> None

MiniMaxH3Qwen3VLRMSNorm

Bases: RMSNorm

Qwen3-VL RMSNorm using the common implementation.

The model-specific name keeps checkpoint weight keys stable, while the common RMSNorm dispatches to torch_npu.npu_rms_norm on Ascend. Gamma uses the checkpoint/model dtype, but native fallbacks keep the RMS variance reduction and scaling in float32.

MiniMaxH3Qwen3VLRowParallelLinear

Bases: LinearBase

group instance-attribute 

group = group

input_is_parallel instance-attribute 

input_is_parallel = input_is_parallel

input_size_per_partition instance-attribute 

input_size_per_partition = input_size // tp_size

output_dtype instance-attribute 

output_dtype = dtype

forward 

forward(input_: Tensor) -> Tensor

weight_loader 

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    loaded_shard_id: str | None = None,
) -> None

MiniMaxH3Qwen3VLTextAttention

Bases: Module

config instance-attribute 

config = config

group instance-attribute 

group = group

head_dim instance-attribute 

head_dim = getattr(
    config,
    "head_dim",
    config.hidden_size // config.num_attention_heads,
)

hidden_size instance-attribute 

hidden_size = config.hidden_size

is_causal instance-attribute 

is_causal = True

k_norm instance-attribute 

k_norm = MiniMaxH3Qwen3VLRMSNorm(
    self.head_dim, eps=config.rms_norm_eps, dtype=dtype
)

num_heads instance-attribute 

num_heads = config.num_attention_heads

num_kv_heads instance-attribute 

num_kv_heads = config.num_key_value_heads

o_proj instance-attribute 

o_proj = MiniMaxH3Qwen3VLRowParallelLinear(
    group,
    input_size=self.num_heads * self.head_dim,
    output_size=self.hidden_size,
    dtype=dtype,
    input_is_parallel=True,
    quant_config=quant_config,
    prefix=f"{prefix}.o_proj",
)

q_norm instance-attribute 

q_norm = MiniMaxH3Qwen3VLRMSNorm(
    self.head_dim, eps=config.rms_norm_eps, dtype=dtype
)

qkv_proj instance-attribute 

qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
    group,
    hidden_size=self.hidden_size,
    num_heads=self.num_heads,
    num_kv_heads=self.num_kv_heads,
    head_dim=self.head_dim,
    dtype=dtype,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj",
)

scaling instance-attribute 

scaling = self.head_dim ** -0.5

forward 

forward(
    hidden_states: Tensor,
    position_embeddings: tuple[Tensor, Tensor],
) -> Tensor

MiniMaxH3Qwen3VLTextDecoderLayer

Bases: Module

hidden_size instance-attribute 

hidden_size = config.hidden_size

input_layernorm instance-attribute 

input_layernorm = MiniMaxH3Qwen3VLRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps, dtype=dtype
)

mlp instance-attribute 

mlp = MiniMaxH3Qwen3VLTextMLP(
    group,
    config,
    dtype,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp",
)

post_attention_layernorm instance-attribute 

post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps, dtype=dtype
)

self_attn instance-attribute 

self_attn = MiniMaxH3Qwen3VLTextAttention(
    group,
    config,
    dtype,
    quant_config=quant_config,
    prefix=f"{prefix}.self_attn",
)

forward 

forward(
    hidden_states: Tensor,
    position_embeddings: tuple[Tensor, Tensor],
) -> Tensor

MiniMaxH3Qwen3VLTextMLP

Bases: Module

down_proj instance-attribute 

down_proj = MiniMaxH3Qwen3VLRowParallelLinear(
    group,
    input_size=self.intermediate_size,
    output_size=self.hidden_size,
    dtype=dtype,
    input_is_parallel=True,
    quant_config=quant_config,
    prefix=f"{prefix}.down_proj",
)

gate_up_proj instance-attribute 

gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
    group,
    input_size=self.hidden_size,
    intermediate_size=self.intermediate_size,
    dtype=dtype,
    quant_config=quant_config,
    prefix=f"{prefix}.gate_up_proj",
)

hidden_size instance-attribute 

hidden_size = config.hidden_size

intermediate_size instance-attribute 

intermediate_size = config.intermediate_size

forward 

forward(x: Tensor) -> Tensor

MiniMaxH3Qwen3VLTextModel

Bases: Module

TP-sharded Qwen3-VL text model returning unnormalized layer-50 states.

config instance-attribute 

config = config

embed_tokens instance-attribute 

embed_tokens = MiniMaxH3Qwen3VLVocabParallelEmbedding(
    group,
    num_embeddings=config.vocab_size,
    embedding_dim=config.hidden_size,
    dtype=dtype,
)

group instance-attribute 

group = group

hidden_size instance-attribute 

hidden_size = config.hidden_size

layers instance-attribute 

layers = nn.ModuleList(
    [
        MiniMaxH3Qwen3VLTextDecoderLayer(
            group,
            config,
            dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.layers.{layer_idx}",
        )
        for layer_idx in range(self.num_layers)
    ]
)

num_layers instance-attribute 

num_layers = min(
    int(getattr(config, "num_hidden_layers", 0)),
    selected_layer,
)

rotary_emb instance-attribute 

rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config)

selected_layer instance-attribute 

selected_layer = selected_layer

vocab_size instance-attribute 

vocab_size = config.vocab_size

forward 

forward(
    inputs_embeds: Tensor,
    positions: Tensor,
    *,
    visual_pos_masks: Tensor | None = None,
    deepstack_visual_embeds: list[Tensor] | None = None,
) -> Tensor

MiniMaxH3Qwen3VLTextRotaryEmbedding

Bases: Module

Qwen3-VL text M-RoPE (interleaved sections). Same math as the reference.

base instance-attribute 

base = float(config.rope_parameters['rope_theta'])

config instance-attribute 

config = config

dim instance-attribute 

dim = int(
    getattr(config, "head_dim", None)
    or config.hidden_size // config.num_attention_heads
)

mrope_section instance-attribute 

mrope_section = config.rope_parameters.get(
    "mrope_section", [24, 20, 20]
)

forward 

forward(
    x: Tensor, position_ids: Tensor
) -> tuple[Tensor, Tensor]

MiniMaxH3Qwen3VLVisionAttention

Bases: Module

dim instance-attribute 

dim = config.hidden_size

head_dim instance-attribute 

head_dim = self.dim // self.num_heads

num_heads instance-attribute 

num_heads = config.num_heads

num_key_value_groups instance-attribute 

num_key_value_groups = 1

proj instance-attribute 

proj = nn.Linear(self.dim, self.dim)

qkv instance-attribute 

qkv = nn.Linear(self.dim, self.dim * 3, bias=True)

scaling instance-attribute 

scaling = self.head_dim ** -0.5

forward 

forward(
    hidden_states: Tensor,
    cu_seqlens: Tensor,
    position_embeddings: tuple[Tensor, Tensor]
    | None = None,
) -> Tensor

MiniMaxH3Qwen3VLVisionBlock

Bases: Module

attn instance-attribute 

attn = MiniMaxH3Qwen3VLVisionAttention(config=config)

mlp instance-attribute 

mlp = MiniMaxH3Qwen3VLVisionMLP(config=config)

norm1 instance-attribute 

norm1 = nn.LayerNorm(config.hidden_size, eps=1e-06)

norm2 instance-attribute 

norm2 = nn.LayerNorm(config.hidden_size, eps=1e-06)

forward 

forward(
    hidden_states: Tensor,
    cu_seqlens: Tensor,
    position_embeddings: tuple[Tensor, Tensor],
) -> Tensor

MiniMaxH3Qwen3VLVisionMLP

Bases: Module

act_fn instance-attribute 

act_fn = nn.GELU(approximate='tanh')

hidden_size instance-attribute 

hidden_size = config.hidden_size

intermediate_size instance-attribute 

intermediate_size = config.intermediate_size

linear_fc1 instance-attribute 

linear_fc1 = nn.Linear(
    self.hidden_size, self.intermediate_size, bias=True
)

linear_fc2 instance-attribute 

linear_fc2 = nn.Linear(
    self.intermediate_size, self.hidden_size, bias=True
)

forward 

forward(hidden_state: Tensor) -> Tensor

MiniMaxH3Qwen3VLVisionModel

Bases: Module

Qwen3-VL vision tower (patch embed + blocks + merger + DeepStack).

blocks instance-attribute 

blocks = nn.ModuleList(
    [
        MiniMaxH3Qwen3VLVisionBlock(config=config)
        for _ in range(config.depth)
    ]
)

config instance-attribute 

config = config

deepstack_merger_list instance-attribute 

deepstack_merger_list = nn.ModuleList(
    [
        MiniMaxH3Qwen3VLVisionPatchMerger(
            config=config, use_postshuffle_norm=True
        )
        for _ in range(
            len(config.deepstack_visual_indexes)
        )
    ]
)

deepstack_visual_indexes instance-attribute 

deepstack_visual_indexes = config.deepstack_visual_indexes

merger instance-attribute 

merger = MiniMaxH3Qwen3VLVisionPatchMerger(
    config=config, use_postshuffle_norm=False
)

num_grid_per_side instance-attribute 

num_grid_per_side = int(config.num_position_embeddings**0.5)

patch_embed instance-attribute 

patch_embed = MiniMaxH3Qwen3VLVisionPatchEmbed(
    config=config
)

patch_size instance-attribute 

patch_size = config.patch_size

pos_embed instance-attribute 

pos_embed = nn.Embedding(
    config.num_position_embeddings, config.hidden_size
)

rotary_pos_emb instance-attribute 

rotary_pos_emb = MiniMaxH3Qwen3VLVisionRotaryEmbedding(
    head_dim // 2
)

spatial_merge_size instance-attribute 

spatial_merge_size = config.spatial_merge_size

spatial_merge_unit instance-attribute 

spatial_merge_unit = (
    self.spatial_merge_size * self.spatial_merge_size
)

fast_pos_embed_interpolate 

fast_pos_embed_interpolate(grid_thw: Tensor) -> Tensor

forward 

forward(
    hidden_states: Tensor, grid_thw: Tensor
) -> tuple[Tensor, list[Tensor]]

rot_pos_emb 

rot_pos_emb(grid_thw: Tensor) -> Tensor

MiniMaxH3Qwen3VLVisionPatchEmbed

Bases: Module

embed_dim instance-attribute 

embed_dim = config.hidden_size

in_channels instance-attribute 

in_channels = config.in_channels

patch_size instance-attribute 

patch_size = config.patch_size

proj instance-attribute 

proj = nn.Conv3d(
    self.in_channels,
    self.embed_dim,
    kernel_size=kernel_size,
    stride=kernel_size,
    bias=True,
)

temporal_patch_size instance-attribute 

temporal_patch_size = config.temporal_patch_size

forward 

forward(hidden_states: Tensor) -> Tensor

MiniMaxH3Qwen3VLVisionPatchMerger

Bases: Module

act_fn instance-attribute 

act_fn = nn.GELU()

hidden_size instance-attribute 

hidden_size = (
    config.hidden_size * config.spatial_merge_size**2
)

linear_fc1 instance-attribute 

linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)

linear_fc2 instance-attribute 

linear_fc2 = nn.Linear(
    self.hidden_size, config.out_hidden_size
)

norm instance-attribute 

norm = nn.LayerNorm(
    self.hidden_size
    if use_postshuffle_norm
    else config.hidden_size,
    eps=1e-06,
)

use_postshuffle_norm instance-attribute 

use_postshuffle_norm = use_postshuffle_norm

forward 

forward(x: Tensor) -> Tensor

MiniMaxH3Qwen3VLVisionRotaryEmbedding

Bases: Module

dim instance-attribute 

dim = dim

theta instance-attribute 

theta = theta

forward 

forward(seqlen: int) -> Tensor

MiniMaxH3Qwen3VLVocabParallelEmbedding

Bases: Module

embedding_dim instance-attribute 

embedding_dim = embedding_dim

group instance-attribute 

group = group

num_embeddings instance-attribute 

num_embeddings = num_embeddings

num_embeddings_per_partition instance-attribute 

num_embeddings_per_partition = num_embeddings // tp_size

weight instance-attribute 

weight = nn.Parameter(
    torch.empty(
        self.num_embeddings_per_partition,
        embedding_dim,
        dtype=dtype,
    )
)

forward 

forward(input_: Tensor) -> Tensor

weight_loader 

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    loaded_shard_id: str | None = None,
) -> None