minimax_h3 ¶

def __init__(
    self,
    time_embed_dim: int,
    hidden_size: int,
    quant_config: QuantizationConfig | None = None,
    prefix: str = "",
    apply_silu: bool = True,
) -> None:
    super().__init__()
    self.hidden_size = hidden_size
    self.apply_silu = apply_silu
    self.linear = ReplicatedLinear(
        time_embed_dim,
        6 * hidden_size * MINIMAX_H3_MODALITY_NUM,
        bias=True,
        quant_config=quant_config,
        prefix=f"{prefix}.linear",
    )

def __init__(
    self,
    hidden_size: int,
    time_embed_dim: int,
    eps: float,
    quant_config: QuantizationConfig | None = None,
    prefix: str = "",
    apply_silu: bool = True,
) -> None:
    super().__init__()
    self.norm = nn.RMSNorm(hidden_size, eps=eps)
    self.apply_silu = apply_silu
    self.linear = ReplicatedLinear(
        time_embed_dim,
        2 * hidden_size,
        bias=True,
        quant_config=quant_config,
        prefix=f"{prefix}.linear",
    )

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    attention_head_dim: int,
    qk_norm_eps: float,
    supported_attention_backends: tuple[AttentionBackendEnum, ...],
    quant_config: QuantizationConfig | None,
    prefix: str,
    fuse_qknorm_rope: bool = False,
    fa4_packed_varlen: bool = False,
    exact_rope: bool = False,
) -> None:
    super().__init__()
    self.num_attention_heads = num_attention_heads
    self.attention_head_dim = attention_head_dim
    inner_dim = num_attention_heads * attention_head_dim
    self.to_q = ReplicatedLinear(
        hidden_size,
        inner_dim,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.to_q",
    )
    self.to_k = ReplicatedLinear(
        hidden_size,
        inner_dim,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.to_k",
    )
    self.to_v = ReplicatedLinear(
        hidden_size,
        inner_dim,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.to_v",
    )
    self.norm_q = nn.RMSNorm(attention_head_dim, eps=qk_norm_eps)
    self.norm_k = nn.RMSNorm(attention_head_dim, eps=qk_norm_eps)
    self.to_out = ReplicatedLinear(
        inner_dim,
        hidden_size,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.to_out",
    )
    self.fuse_qknorm_rope = fuse_qknorm_rope
    self.exact_rope = exact_rope
    # VSA carries a learned gate on its pooled-compression branch. The H3
    # checkpoint has no such weight, so the loader zero-initializes it
    # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled
    # until finetuned. Built only when VSA-H3 actually resolves, keeping
    # the FLASH/SDPA paths and their state_dict untouched.
    resolved_backend = get_attn_backend(attention_head_dim,
                                        get_compute_dtype(),
                                        supported_attention_backends=supported_attention_backends)
    use_vsa = resolved_backend.get_name() == "VIDEO_SPARSE_ATTN_H3"
    attention_cls = DistributedAttention_VSA if use_vsa else DistributedAttention
    self.distributed_attention = attention_cls(
        num_heads=num_attention_heads,
        head_size=attention_head_dim,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=prefix,
        fa4_packed_varlen=fa4_packed_varlen,
    )
    # Opt-in inference route: VSA-H3 selection on the block-sparse FP4
    # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path.
    self._layer_idx = layer_idx_from_prefix(prefix, default=-1)
    self._vsa_fp4 = use_vsa and vsa_fp4_requested()
    self._vsa_tile_first = use_vsa and envs.FASTVIDEO_H3_VSA_TILE_FIRST.get()
    self.to_gate_compress: ReplicatedLinear | None = None
    # None = unchecked; the first forward tests the loaded weight once and
    # skips the gate branch entirely while it is structurally zero.
    self._gate_compress_active: bool | None = None
    if use_vsa:
        self.to_gate_compress = ReplicatedLinear(
            hidden_size,
            inner_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.to_gate_compress",
        )

def __init__(
    self,
    hidden_size: int,
    ffn_dim: int,
    quant_config: QuantizationConfig | None = None,
    prefix: str = "",
    fuse_swiglu: bool = False,
    exact_swiglu: bool = False,
) -> None:
    super().__init__()
    self.fc_in = ReplicatedLinear(
        hidden_size,
        2 * ffn_dim,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.fc_in",
    )
    self.fc_out = ReplicatedLinear(
        ffn_dim,
        hidden_size,
        bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.fc_out",
    )
    self.fuse_swiglu = fuse_swiglu
    self.exact_swiglu = exact_swiglu
    self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance(
        self.fc_out.quant_method, MXFP8QuantizeMethod)
    # Inference-only token chunking: the 2 * ffn_dim intermediate is ~5.3x the block input
    # (4.5 GiB at 78k tokens), so chunks bound the activation peak on 24-32 GB GPUs.
    self.chunk_tokens = envs.FASTVIDEO_H3_FFN_CHUNK_TOKENS.get()

def __init__(self, rope_freq_dim: int, rope_theta: float) -> None:
    super().__init__()
    inv_freq = 1.0 / (rope_theta**(torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) /
                                   (2 * rope_freq_dim)))
    self.register_buffer("inv_freq", inv_freq, persistent=False)

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    attention_head_dim: int,
    ffn_dim: int,
    num_layers: int,
    norm_eps: float,
    qk_norm_eps: float,
    final_norm_eps: float,
    supported_attention_backends: tuple[AttentionBackendEnum, ...],
    quant_config: QuantizationConfig | None,
    prefix: str,
) -> None:
    super().__init__()
    self.refiner_blocks = nn.ModuleList([
        MiniMaxH3TokenRefinerBlock(
            hidden_size,
            num_attention_heads,
            attention_head_dim,
            ffn_dim,
            norm_eps,
            qk_norm_eps,
            supported_attention_backends,
            quant_config,
            prefix=f"{prefix}.refiner_blocks.{index}",
        ) for index in range(num_layers)
    ])
    self.final_norm = nn.RMSNorm(hidden_size, eps=final_norm_eps)

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    attention_head_dim: int,
    ffn_dim: int,
    norm_eps: float,
    qk_norm_eps: float,
    supported_attention_backends: tuple[AttentionBackendEnum, ...],
    quant_config: QuantizationConfig | None,
    prefix: str,
) -> None:
    super().__init__()
    self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
    self.attn = MiniMaxH3Attention(
        hidden_size,
        num_attention_heads,
        attention_head_dim,
        qk_norm_eps,
        supported_attention_backends,
        quant_config,
        prefix=f"{prefix}.attn",
    )
    self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
    self.ff = MiniMaxH3FeedForward(
        hidden_size,
        ffn_dim,
        quant_config=quant_config,
        prefix=f"{prefix}.ff",
    )

def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None:
    super().__init__(config, hf_config)
    arch = config.arch_config
    self.enabled_fusions = _enabled_minimax_h3_fusions()
    if self.enabled_fusions:
        if HAVE_TRITON:
            logger.info(
                "MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled forwards "
                "fall back to eager; torch.compile captures opaque custom-op boundaries).",
                ",".join(sorted(self.enabled_fusions)))
        else:
            logger.warning(
                "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; "
                "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions)))
    self.enabled_exact_kernels = _enabled_minimax_h3_exact_kernels()
    if self.enabled_exact_kernels:
        if exact_kernels.HAVE_TRITON:
            logger.info(
                "MiniMax H3 bit-exact kernels enabled: %s (CUDA BF16 inference-only; other inputs and any "
                "op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keep their existing path).",
                ",".join(sorted(self.enabled_exact_kernels)))
        else:
            logger.warning(
                "FASTVIDEO_MINIMAX_H3_EXACT_KERNELS requested %s but Triton is unavailable; "
                "every forward stays on the eager path.", ",".join(sorted(self.enabled_exact_kernels)))
    sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1
    if arch.num_attention_heads % sp_world_size:
        raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by "
                         f"sequence parallel size ({sp_world_size}).")

    self.hidden_size = arch.hidden_size
    self.num_attention_heads = arch.num_attention_heads
    self.num_channels_latents = arch.in_channels
    self.patch_size = tuple(int(value) for value in arch.patch_size)
    video_patch_dim = arch.in_channels * math.prod(arch.patch_size)

    self.proj_in = ReplicatedLinear(
        video_patch_dim,
        arch.hidden_size,
        bias=True,
        quant_config=config.quant_config,
        prefix=f"{config.prefix}.proj_in",
    )
    self.audio_proj_in = ReplicatedLinear(
        arch.audio_in_channels,
        arch.hidden_size,
        bias=True,
        quant_config=config.quant_config,
        prefix=f"{config.prefix}.audio_proj_in",
    )
    self.context_embedder = ReplicatedLinear(
        arch.text_dim,
        arch.hidden_size,
        bias=True,
        quant_config=config.quant_config,
        prefix=f"{config.prefix}.context_embedder",
    )
    self.time_proj = Timesteps(
        num_channels=arch.freq_dim,
        flip_sin_to_cos=True,
        downscale_freq_shift=0,
    )
    self.time_embedder = MLP(
        arch.freq_dim,
        arch.time_embed_hidden_dim,
        arch.time_embed_dim,
        act_type="silu",
        quant_config=config.quant_config,
        prefix=f"{config.prefix}.time_embedder",
    )
    self.adaln_rank: int | None = arch.adaln_rank
    if self.adaln_rank is not None and config.uniform_parameter_dtype:
        raise ValueError(
            "Rank-reduced AdaLN checkpoints (adaln_rank set) cannot be trained: "
            "uniform_parameter_dtype needs one dtype for every trainable "
            "parameter, but factorized AdaLN weights are pinned to FP16 "
            "(BF16 reconstructs them ~1.7x worse). Fine-tune the full-rank "
            "checkpoint instead, then re-fit the basis with "
            "scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py.")
    adaln_dim = self.adaln_rank or arch.time_embed_dim
    self.adaln_basis = ReplicatedLinear(
        arch.time_embed_dim,
        self.adaln_rank,
        bias=False,
        quant_config=None,
        prefix=f"{config.prefix}.adaln_basis",
    ) if self.adaln_rank else None

    self.rope = MiniMaxH3RotaryPosEmbed(arch.rope_freq_dim, arch.rope_theta)
    # per-generation caches for loop-invariant work (see _rotary_for /
    # _refined_text); plain attrs, never in state_dict
    self._rope_cache: tuple | None = None
    self._text_cache: tuple | None = None
    self.token_refiner = MiniMaxH3TokenRefiner(
        arch.hidden_size,
        arch.num_attention_heads,
        arch.attention_head_dim,
        arch.ffn_dim,
        arch.num_refiner_layers,
        arch.norm_eps,
        arch.qk_norm_eps,
        arch.final_norm_eps,
        # The refiner attends over the text stream only; the packed-sequence
        # VSA backend must never be selected for it.
        tuple(backend for backend in self.supported_attention_backends
              if backend != AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3),
        config.quant_config,
        prefix=f"{config.prefix}.token_refiner",
    )
    self.transformer_blocks = nn.ModuleList([
        MiniMaxH3TransformerBlock(
            arch.hidden_size,
            arch.num_attention_heads,
            arch.attention_head_dim,
            arch.ffn_dim,
            adaln_dim,
            arch.norm_eps,
            arch.qk_norm_eps,
            self.supported_attention_backends,
            config.quant_config,
            prefix=f"{config.prefix}.transformer_blocks.{index}",
            adaln_apply_silu=self.adaln_rank is None,
            fuse_modulate="modulate" in self.enabled_fusions,
            fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions,
            fuse_swiglu="swiglu" in self.enabled_fusions,
            fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(),
            exact_kernels=self.enabled_exact_kernels,
        ) for index in range(arch.num_layers)
    ])
    self.norm_out = MiniMaxH3AdaLayerNormOut(
        arch.hidden_size,
        adaln_dim,
        arch.final_norm_eps,
        quant_config=config.quant_config,
        prefix=f"{config.prefix}.norm_out",
        apply_silu=self.adaln_rank is None,
    )
    # Parallel Decoding Distillation (PDD) students widen both output
    # projections to ``pdd_steps`` heads, one per interval of a fixed fine
    # time grid. The sampler fuses one block of heads per forward
    # (``fuse_pdd_block``), so the fused output keeps the released width.
    self.pdd_steps: int | None = arch.pdd_steps
    self.proj_out = self._output_projection(arch.hidden_size, video_patch_dim, config, "proj_out")
    self.audio_proj_out = self._output_projection(arch.hidden_size, arch.audio_in_channels, config,
                                                  "audio_proj_out")
    self.__post_init__()

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    attention_head_dim: int,
    ffn_dim: int,
    time_embed_dim: int,
    norm_eps: float,
    qk_norm_eps: float,
    supported_attention_backends: tuple[AttentionBackendEnum, ...],
    quant_config: QuantizationConfig | None,
    prefix: str,
    adaln_apply_silu: bool = True,
    fuse_modulate: bool = False,
    fuse_qknorm_rope: bool = False,
    fuse_swiglu: bool = False,
    fa4_packed_varlen: bool = False,
    exact_kernels: frozenset[str] = frozenset(),
) -> None:
    super().__init__()
    self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
    self.attn = MiniMaxH3Attention(
        hidden_size,
        num_attention_heads,
        attention_head_dim,
        qk_norm_eps,
        supported_attention_backends,
        quant_config,
        prefix=f"{prefix}.attn",
        fuse_qknorm_rope=fuse_qknorm_rope,
        fa4_packed_varlen=fa4_packed_varlen,
        exact_rope="rope" in exact_kernels,
    )
    self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
    self.ff = MiniMaxH3FeedForward(
        hidden_size,
        ffn_dim,
        quant_config=quant_config,
        prefix=f"{prefix}.ff",
        fuse_swiglu=fuse_swiglu,
        exact_swiglu="swiglu" in exact_kernels,
    )
    self.adaln_proj = MiniMaxH3AdaLayerNormModulation(
        time_embed_dim,
        hidden_size,
        quant_config=quant_config,
        prefix=f"{prefix}.adaln_proj",
        apply_silu=adaln_apply_silu,
    )
    self.fuse_modulate = fuse_modulate
    self.exact_modulate = "modulate" in exact_kernels