vllm-omni/.claude/skills/add-diffusion-model/references/transformer-adaptation.md

Transformer Adaptation Reference

Adapting a Diffusers Transformer to vLLM-Omni

Step-by-step Checklist

1. Copy the transformer class from diffusers source
2. Remove all mixin classes — inherit only from nn.Module
3. Replace attention dispatch with vllm_omni.diffusion.attention.layer.Attention
4. Replace logger with vllm.logger.init_logger
5. Add od_config: OmniDiffusionConfig | None = None to __init__
6. Remove training-only code (gradient checkpointing, dropout)
7. Add load_weights() method for weight loading from safetensors
8. Add class-level attributes for acceleration features

Mixin Removal

Remove these diffusers mixins (and their imports):

# Remove all of these:
from diffusers.models.modeling_utils import ModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.attention_processor import AttentionModuleMixin
from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin

# Replace:
class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
# With:
class MyTransformer(nn.Module):

Also remove @register_to_config decorators from __init__.

Attention Replacement

The vLLM-Omni Attention layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.

QKV tensor shape must be [batch, seq_len, num_heads, head_dim].

Self-Attention Pattern

from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata

class SelfAttentionBlock(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.to_q = nn.Linear(dim, dim)
        self.to_k = nn.Linear(dim, dim)
        self.to_v = nn.Linear(dim, dim)
        self.to_out = nn.Linear(dim, dim)

        self.attn = Attention(
            num_heads=num_heads,
            head_size=self.head_dim,
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=num_heads,
            role="self",
        )

    def forward(self, x, attn_mask=None):
        B, S, _ = x.shape
        q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
        k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
        v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)

        attn_metadata = AttentionMetadata(attn_mask=attn_mask)
        out = self.attn(q, k, v, attn_metadata=attn_metadata)
        out = out.reshape(B, S, -1)
        return self.to_out(out)

Declaring the attention role. Always pass a role string to Attention(...). It is the lookup key users target with --diffusion-attention-config.per_role.<role>.backend to swap kernels per site without touching model code.

Convention	When to use
role="self"	Q/K/V come from the same hidden state
role="cross"	K/V come from a separate encoder_hidden_states; pair with skip_sequence_parallel=True if K/V is replicated across SP ranks

For multi-modal sites that don't fit self / cross, use a dot-namespaced role and provide role_category so user config can fall back to the generic category:

self.audio_to_video_attn = Attention(
    num_heads=num_heads, head_size=self.head_dim,
    softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False,
    role="mymodel.audio_to_video",
    role_category="cross",
)

Fused QKV with TP (Advanced)

For tensor parallelism, use vLLM's parallel linear layers:

from vllm.model_executor.layers.linear import (
    QKVParallelLinear, RowParallelLinear
)

class TPSelfAttention(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.to_qkv = QKVParallelLinear(
            hidden_size=dim,
            head_size=self.head_dim,
            total_num_heads=num_heads,
            total_num_kv_heads=num_heads,
        )
        self.to_out = RowParallelLinear(dim, dim)

        self.attn = Attention(
            num_heads=num_heads,
            head_size=self.head_dim,
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=num_heads,
            role="self",
        )

Logger Replacement

# Replace:
from diffusers.utils import logging
logger = logging.get_logger(__name__)

# With:
from vllm.logger import init_logger
logger = init_logger(__name__)

Custom Layers from vLLM-Omni

Available utility layers:

from vllm.model_executor.layers.layernorm import RMSNorm
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm

Config Support

from vllm_omni.diffusion.data import OmniDiffusionConfig

class MyTransformer(nn.Module):
    def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
        super().__init__()
        self.od_config = od_config
        self.parallel_config = od_config.parallel_config if od_config else None
        # ... build layers

The transformer config values come from model_index.json → config.json in the transformer subfolder. The pipeline uses get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass) to filter config keys to match the __init__ signature.

Weight Loading

The load_weights method receives an iterable of (name, tensor) from safetensors files, with the prefix (e.g., "transformer.") already stripped by the loader.

from vllm.model_executor.model_loader.weight_utils import default_weight_loader

class MyTransformer(nn.Module):
    def load_weights(self, weights):
        params = dict(self.named_parameters())
        loaded = set()
        for name, tensor in weights:
            # Optional: remap names from diffusers to vllm-omni naming
            # e.g., "ff.net.0.proj" -> "ff.net_0.proj"

            if name in params:
                param = params[name]
                if hasattr(param, "weight_loader"):
                    param.weight_loader(param, tensor)
                else:
                    default_weight_loader(param, tensor)
                loaded.add(name)
        return loaded

QKV Fusion in load_weights

If you fused separate Q/K/V into a QKVParallelLinear, you need to map diffusers' separate weight names:

stacked_params_mapping = [
    ("to_qkv", "to_q", "q"),
    ("to_qkv", "to_k", "k"),
    ("to_qkv", "to_v", "v"),
]

def load_weights(self, weights):
    params = dict(self.named_parameters())
    loaded = set()
    for name, tensor in weights:
        for fused_name, orig_name, shard_id in stacked_params_mapping:
            if orig_name in name:
                name = name.replace(orig_name, fused_name)
                param = params[name]
                param.weight_loader(param, tensor, shard_id)
                loaded.add(name)
                break
        else:
            # Normal loading
            ...
    return loaded

Class-Level Attributes for Features

class MyTransformer(nn.Module):
    # torch.compile: list block class names that repeat and can be compiled
    _repeated_blocks = ["MyTransformerBlock"]

    # CPU offload: attribute name of the nn.ModuleList containing blocks
    _layerwise_offload_blocks_attrs = ["blocks"]

    # LoRA: mapping of fused param names to original param names
    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}

    # Sequence parallelism plan (advanced — add after basic impl works)
    _sp_plan = {
        "blocks.0": SequenceParallelInput(split_dim=1),
        "proj_out": SequenceParallelOutput(gather_dim=1),
    }