vllm_omni.diffusion.models.seedvr2.nadit ¶

dim = int(dim)

emb_dim = int(emb_dim)

layers = list(layers)

input_split_sizes: tuple[int, ...]

network_exchange_required: bool

num_recv_rows: int

output_split_sizes: tuple[int, ...]

recv_to_dst_indices: Tensor

send_indices: Tensor

global_windows: int

joint_cu_seqlens: Tensor

joint_len: int

joint_lengths: Tensor

joint_order: Tensor

layout_key: WindowLayoutKey

local_windows: int

max_joint_len: int

num_video_tokens: int

text_len: int

txt_src: Tensor

vid_src: Tensor

video_cu_seqlens: Tensor

window_shapes: Tensor

txt = txt

vid = vid

all = factory(dims.vid)

shared_weights = shared_weights

txt = None if vid_only else factory(dims.txt)

vid = factory(dims.vid)

vid_only = vid_only

vid_sample = vid_sample

axis_dim = self.rotary_dim // self.num_axes

mm = True

num_axes = int(num_axes)

rot_dim = self.axis_dim * self.num_axes

rotary_dim = int(rotary_dim)

ada = MMModule(
    lambda d: AdaSingle(
        int(d), emb_dim, layers=["attn", "mlp"]
    ),
    dims,
    shared_weights=shared_weights,
    vid_only=is_last_layer,
)

attn = NaSwinAttention(
    vid_dim=vid_dim,
    txt_dim=txt_dim,
    heads=heads,
    head_dim=head_dim,
    qk_bias=qk_bias,
    qk_norm_eps=norm_eps,
    rope_dim=rope_dim,
    shared_weights=shared_weights,
    use_varlen_kernel=use_varlen_kernel,
)

attn_norm = MMModule(
    lambda d: RMSNorm(
        int(d), eps=norm_eps, elementwise_affine=False
    ),
    dims,
    shared_weights=shared_weights,
)

is_last_layer = bool(is_last_layer)

mlp = MMModule(
    lambda d: SwiGLUMLP(int(d), expand_ratio),
    dims,
    shared_weights=shared_weights,
    vid_only=is_last_layer,
)

mlp_norm = MMModule(
    lambda d: RMSNorm(
        int(d), eps=norm_eps, elementwise_affine=False
    ),
    dims,
    shared_weights=shared_weights,
    vid_only=is_last_layer,
)

attention = Attention(
    num_heads=heads,
    head_size=head_dim,
    causal=False,
    softmax_scale=self.softmax_scale,
    role="self",
    skip_sequence_parallel=True,
)

attention_backend_name = (
    backend.get_name() if backend is not None else None
)

attention_path = (
    "packed_varlen"
    if self.use_varlen_kernel
    else "grouped_sdpa"
)

attention_stats = {
    "packed_varlen_calls": 0,
    "grouped_sdpa_calls": 0,
    "no_local_windows_calls": 0,
}

head_dim = int(head_dim)

heads = int(heads)

inner_dim = int(inner_dim)

norm_k = MMModule(
    lambda d: RMSNorm(
        int(d), eps=qk_norm_eps, elementwise_affine=True
    ),
    MMArg(head_dim, head_dim),
    shared_weights=shared_weights,
)

norm_q = MMModule(
    lambda d: RMSNorm(
        int(d), eps=qk_norm_eps, elementwise_affine=True
    ),
    MMArg(head_dim, head_dim),
    shared_weights=shared_weights,
)

proj_out = MMModule(
    lambda d: nn.Linear(inner_dim, int(d)),
    dims,
    shared_weights=shared_weights,
)

proj_qkv = MMModule(
    lambda d: nn.Linear(
        int(d), 3 * inner_dim, bias=qk_bias
    ),
    dims,
    shared_weights=shared_weights,
)

rope = NaMMRotaryEmbedding3d(
    rotary_dim=rope_dim, num_axes=3
)

softmax_scale = 1.0 / math.sqrt(head_dim)

use_varlen_kernel = requested_varlen and supports_varlen

varlen_fallback_reason: str | None = None

out_scale = nn.Parameter(torch.ones(dim))

out_shift = nn.Parameter(torch.zeros(dim))

eps = float(eps)

weight = nn.Parameter(torch.ones(dim))

global_token_ids: Tensor

layout_key: WindowLayoutKey

num_local_tokens: int

num_local_windows: int

rank: int

video_cu_seqlens: Tensor

window_ids: Tensor

window_shapes: Tensor

blocks = nn.ModuleList(
    [
        NaMMSRTransformerBlock(
            vid_dim=vid_dim,
            txt_dim=txt_dim,
            emb_dim=emb_dim,
            heads=heads,
            head_dim=head_dim,
            expand_ratio=expand_ratio,
            norm_eps=norm_eps,
            qk_bias=False,
            mlp_type=mlp_type,
            shared_weights=not index < mm_layers,
            rope_dim=rope_dim,
            is_last_layer=index == num_layers - 1,
            use_varlen_kernel=use_varlen_kernel,
        )
        for index in range(num_layers)
    ]
)

emb_in = TimeEmbedding(
    sinusoidal_dim=256,
    hidden_dim=max(vid_dim, txt_dim),
    output_dim=emb_dim,
)

num_layers = int(num_layers)

patch_size = tuple(int(v) for v in patch_size)

txt_dim = int(txt_dim)

txt_in = (
    nn.Linear(txt_in_dim, txt_dim)
    if txt_in_dim != txt_dim
    else nn.Identity()
)

vid_dim = int(vid_dim)

vid_in = _NaPatchIn(
    in_channels=vid_in_channels,
    patch_size=self.patch_size,
    dim=vid_dim,
)

vid_out = _NaPatchOut(
    out_channels=vid_out_channels,
    patch_size=self.patch_size,
    dim=vid_dim,
)

vid_out_ada = OutAda(vid_dim)

vid_out_norm = (
    RMSNorm(vid_dim, eps=norm_eps, elementwise_affine=True)
    if vid_out_norm
    else None
)

window = tuple(int(v) for v in window)

window_method = tuple(window_method)

group = group

manager = WindowLayoutManager(
    token_grid,
    group=group,
    world_size=self.world_size,
    rank=self.rank,
    window=window,
    methods=methods,
    num_layers=num_layers,
    planner_version=planner_version,
)

num_layers = int(num_layers)

rank = int(rank)

stats: dict[str, int] = {
    "layout_transitions": 0,
    "network_transitions": 0,
    "local_reorders": 0,
    "remote_video_rows": 0,
    "text_all_reduces": 0,
    "fused_text_mean_calls": 0,
}

text_len = int(text_len)

transitions: int

world_size = int(world_size)

proj_in = nn.Linear(dim, hidden, bias=False)

proj_in_gate = nn.Linear(dim, hidden, bias=False)

proj_out = nn.Linear(hidden, dim, bias=False)

act = nn.SiLU()

proj_hid = nn.Linear(hidden_dim, hidden_dim)

proj_in = nn.Linear(sinusoidal_dim, hidden_dim)

proj_out = nn.Linear(hidden_dim, output_dim)

sinusoidal_dim = int(sinusoidal_dim)

ideal_lower_bound: int

layout_key: WindowLayoutKey

max_rank_tokens: int

max_window_tokens: int = 0

owner_rank: Tensor

planner_version: int

rank_token_counts: tuple[int, ...]

rank_window_counts: tuple[int, ...]

rank_window_ids: tuple[Tensor, ...]

world_size: int

key: WindowLayoutKey

num_tokens: int

num_windows: int

window_offsets: Tensor

window_shapes: Tensor

window_token_ids: Tensor

geometry_fingerprint: str

geometry_version: int

token_grid: tuple[int, int, int]

window_method: str

window_shape: tuple[int, int, int]

cache = shared_plan_cache() if cache is None else cache

device: device

group = group

methods = tuple(methods)

num_layers = int(num_layers)

planner_version = int(planner_version)

rank = int(rank)

token_grid = (
    int(token_grid[0]),
    int(token_grid[1]),
    int(token_grid[2]),
)

transitions = 0

window = tuple(window)

world_size = int(world_size)

hits = 0

misses = 0

dst_key: WindowLayoutKey

input_split_sizes: tuple[int, ...]

network_exchange_required: bool

num_recv_rows: int

num_send_rows: int

output_split_sizes: tuple[int, ...]

planner_version: int = PLANNER_VERSION

rank: int

recv_to_dst_indices: Tensor

send_indices: Tensor

src_key: WindowLayoutKey

world_size: int