Qwen3 causal LM text encoder for FastVideo diffusion models (e.g. Flux2 Klein).

`fastvideo.models.encoders.qwen3.LocalAttention`

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

`fastvideo.models.encoders.qwen3.LocalAttention.forward`

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
k: torch.Tensor: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
v: torch.Tensor: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor: torch.Tensor: Output tensor after local attention

`fastvideo.models.encoders.qwen3.BaseEncoderOutput`

BaseEncoderOutput(last_hidden_state: torch.FloatTensor | None = None, pooler_output: torch.FloatTensor | None = None, hidden_states: tuple[torch.FloatTensor, ...] | None = None, attentions: tuple[torch.FloatTensor, ...] | None = None, attention_mask: torch.Tensor | None = None)

`fastvideo.models.encoders.qwen3.Qwen3TextConfig`

Qwen3TextConfig(arch_config: TextEncoderArchConfig = Qwen3TextArchConfig(), prefix: str = 'qwen3', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, is_chat_model: bool = True, treat_empty_as_dot: bool = False, *, _resolved_attention_backend: AttentionBackendEnum | None = None, chat_template_enable_thinking: bool = False)

Top-level config for Qwen3 text encoder.

`fastvideo.models.encoders.qwen3.get_tp_world_size`

get_tp_world_size() -> int

Return world size for the tensor model parallel group.

`fastvideo.models.encoders.qwen3.SiluAndMul`

SiluAndMul()

An activation function for SwiGLU.

The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.

Shapes

x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) return: (num_tokens, d) or (batch_size, seq_len, d)

`fastvideo.models.encoders.qwen3.SiluAndMul.forward_native`

forward_native(x: torch.Tensor) -> torch.Tensor

PyTorch-native implementation equivalent to forward().

`fastvideo.models.encoders.qwen3.RMSNorm`

RMSNorm(hidden_size: int, eps: float = 1e-06, dtype: torch.dtype = torch.float32, var_hidden_size: int | None = None, has_weight: bool = True)

Root mean square normalization.

Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467

`fastvideo.models.encoders.qwen3.RMSNorm.forward_native`

forward_native(x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

PyTorch-native implementation equivalent to forward().

`fastvideo.models.encoders.qwen3.RMSNorm.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.qwen3.MergedColumnParallelLinear`

MergedColumnParallelLinear(input_size: int, output_sizes: list[int], bias: bool = True, gather_output: bool = False, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Packed linear layers with column parallelism.

Similar to ColumnParallelLinear, but the weight matrix is concatenated along the output dimension. When the weight matrix is loaded, the different partitions are sharded separately.

Parameters

input_size: int: input dimension of the linear layer.
output_sizes: list[int]: list of output dimensions of the linear layer.
bias: bool: If true, add bias.
gather_output: bool: If true, call all-gather on output and make the output available to all GPUs, otherwise, every GPU will have its own output.
skip_bias_add: bool: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.encoders.qwen3.MergedColumnParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None) -> None

`fastvideo.models.encoders.qwen3.MergedColumnParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None) -> None

`fastvideo.models.encoders.qwen3.QKVParallelLinear`

QKVParallelLinear(hidden_size: int, head_size: int, total_num_heads: int, total_num_kv_heads: int | None = None, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Linear layers for the attention's QKV transformation.

Linear layers for the linear transformation of the query, key, and value vectors in the attention layer. The weight matrix is concatenated along the output dimension. The layer is parallelized along the head dimension. When the number of key/value heads is smaller than the number of query heads (e.g., multi-query/grouped-query attention), the key/value head may be replicated while the query heads are partitioned.

Parameters

hidden_size: int: input hidden state size of the transformer.
head_size: int: size of each attention head.
total_num_heads: int: total number of attention query heads.
total_num_kv_heads: int | None: total number of attention key/value heads. If None, assume total_num_kv_heads = total_num_heads.
bias: bool: If true, add bias.
skip_bias_add: bool: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.encoders.qwen3.QKVParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None)

`fastvideo.models.encoders.qwen3.QKVParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None)

`fastvideo.models.encoders.qwen3.RowParallelLinear`

RowParallelLinear(input_size: int, output_size: int, bias: bool = True, input_is_parallel: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, reduce_results: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = '')

Linear layer with row parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension as: - - | A_1 | | . | A = | . | X = [X_1, ..., X_p] | . | | A_p | - - Arguments: input_size: first dimension of matrix A. output_size: second dimension of matrix A. bias: If true, add bias. Note that bias is not parallelized. input_is_parallel: If true, we assume that the input is already split across the GPUs and we do not split again. skip_bias_add: This was added to enable performance optimization where bias can be fused with other element-wise operations. We skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure.

`fastvideo.models.encoders.qwen3.RowParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.qwen3.RowParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.qwen3.RowParallelLinear.forward`

forward(input_) -> tuple[torch.Tensor, Parameter | None]

`fastvideo.models.encoders.qwen3.RowParallelLinear.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.qwen3.QuantizationConfig`

QuantizationConfig()

Base class for quantization configs.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_name`

get_name() -> QuantizationMethods

Name of the quantization method.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_supported_act_dtypes`

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_min_capability`

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_config_filenames`

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

`fastvideo.models.encoders.qwen3.QuantizationConfig.from_config`

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

`fastvideo.models.encoders.qwen3.QuantizationConfig.override_quantization_method`

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_from_keys`

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_from_keys_or`

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_quant_method`

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module: The layer for the quant method.
prefix: str: The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

`fastvideo.models.encoders.qwen3.QuantizationConfig.validate_runtime`

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

`fastvideo.models.encoders.qwen3.QuantizationConfig.get_cache_scale`

get_cache_scale(name: str) -> str | None

`fastvideo.models.encoders.qwen3.get_rope`

get_rope(head_size: int, rotary_dim: int, max_position: int, base: int | float, is_neox_style: bool = True, rope_scaling: dict[str, Any] | None = None, dtype: torch.dtype | None = None, partial_rotary_factor: float = 1.0) -> RotaryEmbedding

`fastvideo.models.encoders.qwen3.VocabParallelEmbedding`

VocabParallelEmbedding(num_embeddings: int, embedding_dim: int, params_dtype: torch.dtype | None = None, org_num_embeddings: int | None = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = '')

Embedding parallelized in the vocabulary dimension.

Adapted from torch.nn.Embedding, note that we pad the vocabulary size to make sure it is divisible by the number of model parallel GPUs.

In order to support various loading methods, we ensure that LoRA-added embeddings are always at the end of TP-sharded tensors. In other words, we shard base embeddings and LoRA embeddings separately (both padded), and place them in the same tensor. In this example, we will have the original vocab size = 1010, added vocab size = 16 and padding to 64. Therefore, the total vocab size with padding will be 1088 (because we first pad 1010 to 1024, add 16, and then pad to 1088). Therefore, the tensor format looks like the following: TP1, rank 0 (no sharding): |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 |

TP2, rank 0: |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | TP2, rank 1: |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 |

Parameters

num_embeddings: int: vocabulary size.
embedding_dim: int: size of hidden state.
params_dtype: torch.dtype | None: type of the parameters.
org_num_embeddings: int | None: original vocabulary size (without LoRA).
padding_size: int: padding size for the vocabulary.
quant_config: QuantizationConfig | None: quant config for the layer
prefix: str: full name of the layer in the state dict

`fastvideo.models.encoders.qwen3.VocabParallelEmbedding.get_sharded_to_full_mapping`

get_sharded_to_full_mapping() -> list[int] | None

Get a mapping that can be used to reindex the gathered logits for sampling.

During sampling, we gather logits from all ranks. The relationship of index->token_id will follow the same format as outlined in the class docstring. However, after the gather, we want to reindex the final logits tensor to map index->token_id one-to-one (the index is always equal the token_id it corresponds to). The indices returned by this method allow us to do that.

`fastvideo.models.encoders.qwen3.VocabParallelEmbedding.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.qwen3.VocabParallelEmbedding.forward`

forward(input_)

`fastvideo.models.encoders.qwen3.VocabParallelEmbedding.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.qwen3.TextEncoder`

TextEncoder(config: TextEncoderConfig)

Base for native encoders with a model-specific forward output contract.

`fastvideo.models.encoders.qwen3.TextEncoder.forward`

forward(*args: Any, **kwargs: Any) -> TextEncoderOutputT

`fastvideo.models.encoders.qwen3.default_weight_loader`

default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None

Default weight loader.

`fastvideo.models.encoders.qwen3.maybe_remap_kv_scale_name`

maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None

Remap the name of FP8 k/v_scale parameters.

This function handles the remapping of FP8 k/v_scale parameter names. It detects if the given name ends with a suffix and attempts to remap it to the expected name format in the model. If the remapped name is not found in the params_dict, a warning is printed and None is returned.

Parameters

name: str: The original loaded checkpoint parameter name.
params_dict: dict: Dictionary containing the model's named parameters.

Returns

str: str | None: The remapped parameter name if successful, or the original name if no remapping is needed.
None: str | None: If the remapped name is not found in params_dict.

`fastvideo.models.encoders.qwen3.Qwen3MLP`

Qwen3MLP(hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = '')

Qwen3 MLP with SwiGLU activation and tensor parallelism.

Source

class Qwen3MLP(nn.Module):
    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: QuantizationConfig | None = None,
        bias: bool = False,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        return x

`fastvideo.models.encoders.qwen3.Qwen3MLP.forward`

forward(x: torch.Tensor) -> torch.Tensor

Source

def forward(self, x: torch.Tensor) -> torch.Tensor:
    x, _ = self.gate_up_proj(x)
    x = self.act_fn(x)
    x, _ = self.down_proj(x)
    return x

`fastvideo.models.encoders.qwen3.Qwen3Attention`

Qwen3Attention(config: Qwen3TextConfig, hidden_size: int, num_heads: int, num_kv_heads: int, rope_theta: float = 1000000.0, rope_scaling: dict[str, Any] | None = None, max_position_embeddings: int = 40960, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = '')

Qwen3 attention with QK-Norm and tensor parallelism.

Key difference from LLaMA: RMSNorm is applied to Q and K before attention.

Source

class Qwen3Attention(nn.Module):
    """Qwen3 attention with QK-Norm and tensor parallelism.

    Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
    """

    def __init__(
        self,
        config: Qwen3TextConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 1000000.0,
        rope_scaling: dict[str, Any] | None = None,
        max_position_embeddings: int = 40960,
        quant_config: QuantizationConfig | None = None,
        bias: bool = False,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tp_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            assert self.total_num_kv_heads % tp_size == 0
        else:
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.total_num_heads)
        self.rotary_dim = self.head_dim
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.rotary_dim,
            max_position=max_position_embeddings,
            base=int(rope_theta),
            rope_scaling=rope_scaling,
            is_neox_style=True,
        )

        self.attn = LocalAttention(
            self.num_heads,
            self.head_dim,
            self.num_kv_heads,
            softmax_scale=self.scaling,
            causal=True,
            supported_attention_backends=config._supported_attention_backends,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        batch_size, seq_len = q.shape[0], q.shape[1]
        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

        q = self.q_norm(q)
        k = self.k_norm(k)

        q = q.reshape(batch_size, seq_len, -1)
        k = k.reshape(batch_size, seq_len, -1)

        q, k = self.rotary_emb(positions, q, k)

        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

        if attention_mask is None:
            attn_output = self.attn(q, k, v)
        else:
            q_sdpa = q.transpose(1, 2)
            k_sdpa = k.transpose(1, 2)
            v_sdpa = v.transpose(1, 2)
            causal_mask = torch.ones(
                seq_len,
                seq_len,
                device=q.device,
                dtype=torch.bool,
            ).tril()
            key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
            attn_mask = causal_mask[None, None, :, :] & key_mask[:, None, None, :]
            attn_output = torch.nn.functional.scaled_dot_product_attention(
                q_sdpa,
                k_sdpa,
                v_sdpa,
                attn_mask=attn_mask,
                dropout_p=0.0,
                is_causal=False,
                scale=self.scaling,
                enable_gqa=self.num_heads != self.num_kv_heads,
            ).transpose(1, 2)

        attn_output = attn_output.reshape(batch_size, seq_len, -1)

        output, _ = self.o_proj(attn_output)
        return output

`fastvideo.models.encoders.qwen3.Qwen3Attention.forward`

forward(positions: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor

Source

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

    batch_size, seq_len = q.shape[0], q.shape[1]
    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
    v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

    q = self.q_norm(q)
    k = self.k_norm(k)

    q = q.reshape(batch_size, seq_len, -1)
    k = k.reshape(batch_size, seq_len, -1)

    q, k = self.rotary_emb(positions, q, k)

    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

    if attention_mask is None:
        attn_output = self.attn(q, k, v)
    else:
        q_sdpa = q.transpose(1, 2)
        k_sdpa = k.transpose(1, 2)
        v_sdpa = v.transpose(1, 2)
        causal_mask = torch.ones(
            seq_len,
            seq_len,
            device=q.device,
            dtype=torch.bool,
        ).tril()
        key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
        attn_mask = causal_mask[None, None, :, :] & key_mask[:, None, None, :]
        attn_output = torch.nn.functional.scaled_dot_product_attention(
            q_sdpa,
            k_sdpa,
            v_sdpa,
            attn_mask=attn_mask,
            dropout_p=0.0,
            is_causal=False,
            scale=self.scaling,
            enable_gqa=self.num_heads != self.num_kv_heads,
        ).transpose(1, 2)

    attn_output = attn_output.reshape(batch_size, seq_len, -1)

    output, _ = self.o_proj(attn_output)
    return output

`fastvideo.models.encoders.qwen3.Qwen3DecoderLayer`

Qwen3DecoderLayer(config: Qwen3TextConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')

Qwen3 transformer decoder layer.

Source

class Qwen3DecoderLayer(nn.Module):
    """Qwen3 transformer decoder layer."""

    def __init__(
        self,
        config: Qwen3TextConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 1000000.0)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings", 40960)
        attention_bias = getattr(config, "attention_bias", False)

        self.self_attn = Qwen3Attention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            bias=attention_bias,
            prefix=f"{prefix}.self_attn",
        )
        self.mlp = Qwen3MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            bias=getattr(config, "mlp_bias", False),
            prefix=f"{prefix}.mlp",
        )
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)

        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            attention_mask=attention_mask,
        )

        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual

`fastvideo.models.encoders.qwen3.Qwen3DecoderLayer.forward`

forward(positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, attention_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

Source

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    residual: torch.Tensor | None,
    attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    if residual is None:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
    else:
        hidden_states, residual = self.input_layernorm(hidden_states, residual)

    hidden_states = self.self_attn(
        positions=positions,
        hidden_states=hidden_states,
        attention_mask=attention_mask,
    )

    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
    hidden_states = self.mlp(hidden_states)
    return hidden_states, residual

`fastvideo.models.encoders.qwen3.Qwen3ForCausalLM`

Qwen3ForCausalLM(config: Qwen3TextConfig)

Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).

Features: - Tensor parallelism support - FlashAttention/SDPA support via LocalAttention - QK-Norm for better training stability - output_hidden_states for Klein (layers 9, 18, 27)

Source

class Qwen3ForCausalLM(TextEncoder):
    """Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).

    Features:
    - Tensor parallelism support
    - FlashAttention/SDPA support via LocalAttention
    - QK-Norm for better training stability
    - output_hidden_states for Klein (layers 9, 18, 27)
    """

    supports_hf_from_pretrained = True

    def __init__(self, config: Qwen3TextConfig) -> None:
        super().__init__(config)

        self.config = config
        self.quant_config = getattr(config, "quant_config", None)

        if getattr(config, "lora_config", None) is not None:
            max_loras = getattr(config.lora_config, "max_loras", 1)
            lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1)
            lora_vocab = lora_vocab_size * max_loras
        else:
            lora_vocab = 0
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
            quant_config=self.quant_config,
        )

        self.layers = nn.ModuleList([
            Qwen3DecoderLayer(
                config=config,
                quant_config=self.quant_config,
                prefix=f"{config.prefix}.layers.{i}",
            ) for i in range(config.num_hidden_layers)
        ])

        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    @classmethod
    def from_pretrained_local(
        cls,
        model_path: str,
        model_config: Qwen3TextConfig,
        dtype: torch.dtype,
        device: torch.device,
    ) -> nn.Module:
        from transformers import AutoModel

        # FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
        # materializing an unused LM head and full-vocabulary logits, including
        # for checkpoints whose metadata names Qwen3ForCausalLM.
        return AutoModel.from_pretrained(
            model_path,
            local_files_only=True,
            dtype=dtype,
            low_cpu_mem_usage=True,
        ).eval().to(device)

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs: Any,
    ) -> BaseEncoderOutput:
        output_hidden_states = (output_hidden_states
                                if output_hidden_states is not None else self.config.output_hidden_states)

        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            assert input_ids is not None
            hidden_states = self.get_input_embeddings(input_ids)

        residual = None

        if position_ids is None:
            # Expand to [batch_size, seq_len]: the rotary layer flattens
            # positions to ``num_tokens`` and reshapes q/k to
            # ``(num_tokens, -1, head_dim)``. A bare [1, seq_len] only matches
            # ``num_tokens`` when batch_size == 1; for batched inputs it folds
            # the batch dim into the head dim and misaligns RoPE. Expanding to
            # ``batch_size * seq_len`` tokens keeps the layout correct.
            position_ids = torch.arange(
                0,
                hidden_states.shape[1],
                device=hidden_states.device,
                dtype=torch.long,
            ).unsqueeze(0).expand(hidden_states.shape[0], -1)

        all_hidden_states: tuple[Any, ...] | None = (() if output_hidden_states else None)

        for layer in self.layers:
            if all_hidden_states is not None:
                all_hidden_states += ((hidden_states, ) if residual is None else (hidden_states + residual, ))
            hidden_states, residual = layer(
                position_ids,
                hidden_states,
                residual,
                attention_mask=attention_mask,
            )

        hidden_states, _ = self.norm(hidden_states, residual)

        if all_hidden_states is not None:
            all_hidden_states += (hidden_states, )

        return BaseEncoderOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        stacked_params_mapping = self.config.arch_config.stacked_params_mapping
        # A fused destination is initialized either by one already-fused tensor
        # or after every split source projection has been loaded. Include
        # auxiliary quantization parameters (for example scale_weight) rather
        # than limiting completeness checks to weight/bias tensors.
        expected_stacked_shards = {(name, shard_id)
                                   for name in params_dict
                                   for param_name, _, shard_id in stacked_params_mapping if param_name in name}
        fused_param_names = {name for name, _ in expected_stacked_shards}
        loaded_stacked_shards: set[tuple[str, str | int]] = set()
        loaded_fused_params: set[str] = set()

        for name, loaded_weight in weights:
            if name.startswith("model."):
                name = name[6:]

            if "rotary_emb.inv_freq" in name:
                continue
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
                continue

            if "scale" in name:
                kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
                if kv_scale_name is None:
                    continue
                name = kv_scale_name

            matched_stacked_param = False
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                matched_stacked_param = True
                target_name = name.replace(weight_name, param_name)

                if target_name.endswith(".bias") and target_name not in params_dict:
                    break

                if target_name not in params_dict:
                    break

                param = params_dict[target_name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                loaded_params.add(target_name)
                loaded_stacked_shards.add((target_name, shard_id))
                break

            if matched_stacked_param:
                continue

            if name.endswith(".bias") and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict[name]
            if name in fused_param_names and name.endswith(".scale_weight"):
                # Merged scale loaders interpret a missing shard id as shard 0.
                # An exact fused key is already a complete vector, so copy it
                # atomically and retain the default loader's shape validation.
                default_weight_loader(param, loaded_weight)
            else:
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
            if name in fused_param_names:
                loaded_fused_params.add(name)

        required_split_shards = {(name, shard_id)
                                 for name, shard_id in expected_stacked_shards if name not in loaded_fused_params}
        missing_stacked_shards = required_split_shards - loaded_stacked_shards
        if missing_stacked_shards:
            formatted_missing = ", ".join(f"{name}[{shard_id}]" for name, shard_id in sorted(
                missing_stacked_shards,
                key=lambda item: (item[0], str(item[1])),
            ))
            raise ValueError("Missing required stacked checkpoint shards: "
                             f"{formatted_missing}")

        return loaded_params

`fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.from_pretrained_local`

from_pretrained_local(model_path: str, model_config: Qwen3TextConfig, dtype: torch.dtype, device: torch.device) -> nn.Module

Source

@classmethod
def from_pretrained_local(
    cls,
    model_path: str,
    model_config: Qwen3TextConfig,
    dtype: torch.dtype,
    device: torch.device,
) -> nn.Module:
    from transformers import AutoModel

    # FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
    # materializing an unused LM head and full-vocabulary logits, including
    # for checkpoints whose metadata names Qwen3ForCausalLM.
    return AutoModel.from_pretrained(
        model_path,
        local_files_only=True,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).eval().to(device)

`fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.get_input_embeddings`

get_input_embeddings(input_ids: torch.Tensor) -> torch.Tensor

Source

def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
    return self.embed_tokens(input_ids)

`fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.forward`

forward(input_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs: Any) -> BaseEncoderOutput

Source

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs: Any,
) -> BaseEncoderOutput:
    output_hidden_states = (output_hidden_states
                            if output_hidden_states is not None else self.config.output_hidden_states)

    if inputs_embeds is not None:
        hidden_states = inputs_embeds
    else:
        assert input_ids is not None
        hidden_states = self.get_input_embeddings(input_ids)

    residual = None

    if position_ids is None:
        # Expand to [batch_size, seq_len]: the rotary layer flattens
        # positions to ``num_tokens`` and reshapes q/k to
        # ``(num_tokens, -1, head_dim)``. A bare [1, seq_len] only matches
        # ``num_tokens`` when batch_size == 1; for batched inputs it folds
        # the batch dim into the head dim and misaligns RoPE. Expanding to
        # ``batch_size * seq_len`` tokens keeps the layout correct.
        position_ids = torch.arange(
            0,
            hidden_states.shape[1],
            device=hidden_states.device,
            dtype=torch.long,
        ).unsqueeze(0).expand(hidden_states.shape[0], -1)

    all_hidden_states: tuple[Any, ...] | None = (() if output_hidden_states else None)

    for layer in self.layers:
        if all_hidden_states is not None:
            all_hidden_states += ((hidden_states, ) if residual is None else (hidden_states + residual, ))
        hidden_states, residual = layer(
            position_ids,
            hidden_states,
            residual,
            attention_mask=attention_mask,
        )

    hidden_states, _ = self.norm(hidden_states, residual)

    if all_hidden_states is not None:
        all_hidden_states += (hidden_states, )

    return BaseEncoderOutput(
        last_hidden_state=hidden_states,
        hidden_states=all_hidden_states,
    )

`fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.load_weights`

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

Source

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    stacked_params_mapping = self.config.arch_config.stacked_params_mapping
    # A fused destination is initialized either by one already-fused tensor
    # or after every split source projection has been loaded. Include
    # auxiliary quantization parameters (for example scale_weight) rather
    # than limiting completeness checks to weight/bias tensors.
    expected_stacked_shards = {(name, shard_id)
                               for name in params_dict
                               for param_name, _, shard_id in stacked_params_mapping if param_name in name}
    fused_param_names = {name for name, _ in expected_stacked_shards}
    loaded_stacked_shards: set[tuple[str, str | int]] = set()
    loaded_fused_params: set[str] = set()

    for name, loaded_weight in weights:
        if name.startswith("model."):
            name = name[6:]

        if "rotary_emb.inv_freq" in name:
            continue
        if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
            continue

        if "scale" in name:
            kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
            if kv_scale_name is None:
                continue
            name = kv_scale_name

        matched_stacked_param = False
        for param_name, weight_name, shard_id in stacked_params_mapping:
            if weight_name not in name:
                continue
            matched_stacked_param = True
            target_name = name.replace(weight_name, param_name)

            if target_name.endswith(".bias") and target_name not in params_dict:
                break

            if target_name not in params_dict:
                break

            param = params_dict[target_name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            loaded_params.add(target_name)
            loaded_stacked_shards.add((target_name, shard_id))
            break

        if matched_stacked_param:
            continue

        if name.endswith(".bias") and name not in params_dict:
            continue

        if name not in params_dict:
            continue

        param = params_dict[name]
        if name in fused_param_names and name.endswith(".scale_weight"):
            # Merged scale loaders interpret a missing shard id as shard 0.
            # An exact fused key is already a complete vector, so copy it
            # atomically and retain the default loader's shape validation.
            default_weight_loader(param, loaded_weight)
        else:
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
        if name in fused_param_names:
            loaded_fused_params.add(name)

    required_split_shards = {(name, shard_id)
                             for name, shard_id in expected_stacked_shards if name not in loaded_fused_params}
    missing_stacked_shards = required_split_shards - loaded_stacked_shards
    if missing_stacked_shards:
        formatted_missing = ", ".join(f"{name}[{shard_id}]" for name, shard_id in sorted(
            missing_stacked_shards,
            key=lambda item: (item[0], str(item[1])),
        ))
        raise ValueError("Missing required stacked checkpoint shards: "
                         f"{formatted_missing}")

    return loaded_params