Remove VeOmni runtime dependency

README.md

@@ -92,8 +92,6 @@ conda activate llada-image
 pip install -r requirements.txt
 ```
 
-The published LLaDA2 text encoder uses `veomni.ops.fused_moe_forward`. Install a compatible LLaDA2 / VeOmni runtime before running inference.
-
 ### 2. Run inference
 
 The pipeline accepts a prompt and, for editing, an optional reference image.

text_encoder/fused_moe_ops.py

@@ -0,0 +1,375 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone, inference-only VeOmni v0.1.0 fused-MoE compatibility shim.
+
+This module preserves the ``veomni.ops.fused_moe_forward`` call signature used
+by VeOmni v0.1.0 while removing VeOmni's training, Expert Parallelism (EP), NPU,
+and Seed-kernel dependencies. It is intended for single-device inference only.
+
+The CUDA fast path uses a small Triton grouped-linear kernel. If Triton is not
+available, the tensors are not on CUDA, or ``LLADA_MOE_BACKEND=eager`` is set,
+the implementation falls back to ordinary PyTorch operations.
+
+Replace the original model-code import with, for example,
+``from .fused_moe_v010 import fused_moe_forward``.
+
+Derived from ByteDance-Seed/VeOmni v0.1.0.post1:
+https://github.com/ByteDance-Seed/VeOmni/tree/v0.1.0.post1
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # The eager fallback does not require Triton.
+    triton = None
+    tl = None
+
+
+_SUPPORTED_TRITON_DTYPES = (torch.float16, torch.bfloat16)
+
+
+if triton is not None:
+
+    @triton.jit
+    def _grouped_linear_kernel(
+        input_ptr,
+        weight_ptr,
+        output_ptr,
+        expert_cumsum_ptr,
+        N: tl.constexpr,
+        K: tl.constexpr,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+    ):
+        """Compute per-expert ``input @ weight.T`` for contiguous tensors."""
+        block_m = tl.program_id(axis=0)
+        block_n = tl.program_id(axis=1)
+        expert = tl.program_id(axis=2)
+
+        expert_start = tl.load(expert_cumsum_ptr + expert - 1, mask=expert > 0, other=0)
+        expert_end = tl.load(expert_cumsum_ptr + expert)
+        expert_tokens = expert_end - expert_start
+
+        if block_m * BLOCK_M >= expert_tokens:
+            return
+
+        row_offsets = block_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        col_offsets = block_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        k_offsets = tl.arange(0, BLOCK_K)
+
+        input_ptrs = (
+            input_ptr
+            + (expert_start + row_offsets[:, None]) * K
+            + k_offsets[None, :]
+        )
+        weight_ptrs = (
+            weight_ptr
+            + expert * N * K
+            + col_offsets[None, :] * K
+            + k_offsets[:, None]
+        )
+
+        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+            remaining_k = K - k_block * BLOCK_K
+            inputs = tl.load(
+                input_ptrs,
+                mask=(row_offsets[:, None] < expert_tokens) & (k_offsets[None, :] < remaining_k),
+                other=0.0,
+            )
+            weights = tl.load(
+                weight_ptrs,
+                mask=(col_offsets[None, :] < N) & (k_offsets[:, None] < remaining_k),
+                other=0.0,
+            )
+            accumulator += tl.dot(inputs, weights)
+            input_ptrs += BLOCK_K
+            weight_ptrs += BLOCK_K
+
+        output_ptrs = (
+            output_ptr
+            + (expert_start + row_offsets[:, None]) * N
+            + col_offsets[None, :]
+        )
+        tl.store(
+            output_ptrs,
+            accumulator,
+            mask=(row_offsets[:, None] < expert_tokens) & (col_offsets[None, :] < N),
+        )
+
+
+def _validate_inputs(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> None:
+    if num_experts <= 0:
+        raise ValueError(f"num_experts must be positive, got {num_experts}")
+    if torch.is_grad_enabled():
+        raise RuntimeError(
+            "This standalone fused_moe_forward is inference-only. Call it under "
+            "torch.no_grad() or torch.inference_mode()."
+        )
+    if hidden_states.ndim != 2:
+        raise ValueError(f"hidden_states must have shape [tokens, hidden], got {tuple(hidden_states.shape)}")
+    if routing_weights.ndim != 2 or selected_experts.shape != routing_weights.shape:
+        raise ValueError(
+            "routing_weights and selected_experts must have the same [tokens, top_k] shape, got "
+            f"{tuple(routing_weights.shape)} and {tuple(selected_experts.shape)}"
+        )
+    if routing_weights.shape[1] == 0:
+        raise ValueError("top_k must be positive")
+    if routing_weights.shape[0] != hidden_states.shape[0]:
+        raise ValueError("routing_weights and hidden_states must contain the same number of tokens")
+    if selected_experts.dtype not in (torch.int32, torch.int64):
+        raise TypeError(f"selected_experts must be int32 or int64, got {selected_experts.dtype}")
+    if fc1_1_weight.ndim != 3 or fc1_2_weight.ndim != 3 or fc2_weight.ndim != 3:
+        raise ValueError("expert weights must be rank-3 tensors")
+    if fc1_1_weight.shape != fc1_2_weight.shape:
+        raise ValueError("fc1_1_weight and fc1_2_weight must have identical shapes")
+
+    experts, intermediate_size, hidden_size = fc1_1_weight.shape
+    expected_fc2_shape = (experts, hidden_size, intermediate_size)
+    if experts != num_experts:
+        raise ValueError(f"num_experts={num_experts}, but the weights contain {experts} experts")
+    if hidden_states.shape[1] != hidden_size:
+        raise ValueError(f"hidden size is {hidden_states.shape[1]}, but the weights expect {hidden_size}")
+    if tuple(fc2_weight.shape) != expected_fc2_shape:
+        raise ValueError(f"fc2_weight must have shape {expected_fc2_shape}, got {tuple(fc2_weight.shape)}")
+    if selected_experts.numel():
+        # These scalar checks synchronize CUDA once, before launching harder-to-debug kernels.
+        min_expert = int(selected_experts.min().item())
+        max_expert = int(selected_experts.max().item())
+        if min_expert < 0 or max_expert >= num_experts:
+            raise ValueError(f"selected expert IDs must be in [0, {num_experts}), got [{min_expert}, {max_expert}]")
+
+    devices = {
+        hidden_states.device,
+        routing_weights.device,
+        selected_experts.device,
+        fc1_1_weight.device,
+        fc1_2_weight.device,
+        fc2_weight.device,
+    }
+    if len(devices) != 1:
+        raise ValueError(f"all inputs and weights must be on one device, got {sorted(map(str, devices))}")
+
+
+def _route_tokens(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sort routed token copies by expert and return the inverse permutation."""
+    top_k = selected_experts.shape[1]
+    flat_experts = selected_experts.reshape(-1).to(torch.int64)
+    order = torch.argsort(flat_experts, stable=True)
+    sorted_hidden_states = hidden_states[torch.div(order, top_k, rounding_mode="floor")].contiguous()
+    sorted_routing_weights = routing_weights.reshape(-1)[order].contiguous()
+    tokens_per_expert = torch.bincount(flat_experts, minlength=num_experts)
+    expert_cumsum = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32).contiguous()
+    return sorted_hidden_states, sorted_routing_weights, expert_cumsum, order
+
+
+def _unroute_tokens(
+    sorted_outputs: torch.Tensor,
+    order: torch.Tensor,
+    num_tokens: int,
+    top_k: int,
+) -> torch.Tensor:
+    restored = torch.empty_like(sorted_outputs)
+    restored[order] = sorted_outputs
+    # VeOmni's v0.1.0 gather kernel accumulates the top-k outputs in FP32.
+    return restored.view(num_tokens, top_k, -1).sum(dim=1, dtype=torch.float32).to(sorted_outputs.dtype)
+
+
+def _grouped_linear_triton(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    expert_cumsum: torch.Tensor,
+) -> torch.Tensor:
+    if triton is None:  # pragma: no cover - guarded by the caller
+        raise RuntimeError("Triton is not available")
+    if not inputs.is_contiguous() or not weights.is_contiguous():
+        raise ValueError("the Triton path requires contiguous inputs and expert weights")
+
+    num_experts, output_size, input_size = weights.shape
+    if inputs.shape[1] != input_size:
+        raise ValueError(f"input width is {inputs.shape[1]}, but the weights expect {input_size}")
+
+    output = torch.empty((inputs.shape[0], output_size), dtype=inputs.dtype, device=inputs.device)
+    block_m, block_n, block_k = 128, 128, 32
+    grid = (
+        triton.cdiv(inputs.shape[0], block_m),
+        triton.cdiv(output_size, block_n),
+        num_experts,
+    )
+    with torch.cuda.device(inputs.device):
+        _grouped_linear_kernel[grid](
+            inputs,
+            weights,
+            output,
+            expert_cumsum,
+            N=output_size,
+            K=input_size,
+            BLOCK_M=block_m,
+            BLOCK_N=block_n,
+            BLOCK_K=block_k,
+            num_warps=8,
+            num_stages=3,
+        )
+    return output
+
+
+def _triton_moe_forward(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    sorted_hidden, sorted_routing, expert_cumsum, order = _route_tokens(
+        num_experts, routing_weights, selected_experts, hidden_states
+    )
+    gate = _grouped_linear_triton(sorted_hidden, fc1_1_weight, expert_cumsum)
+    up = _grouped_linear_triton(sorted_hidden, fc1_2_weight, expert_cumsum)
+    intermediate = F.silu(gate) * up
+    intermediate.mul_(sorted_routing.unsqueeze(-1))
+    sorted_outputs = _grouped_linear_triton(intermediate.contiguous(), fc2_weight, expert_cumsum)
+    return _unroute_tokens(sorted_outputs, order, hidden_states.shape[0], selected_experts.shape[1])
+
+
+def _eager_moe_forward(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    sorted_hidden, sorted_routing, expert_cumsum, order = _route_tokens(
+        num_experts, routing_weights, selected_experts, hidden_states
+    )
+    expert_ends = expert_cumsum.to(device="cpu", dtype=torch.int64).tolist()
+    outputs: list[torch.Tensor] = []
+    start = 0
+    for expert, end in enumerate(expert_ends):
+        if end > start:
+            expert_inputs = sorted_hidden[start:end]
+            gate = F.linear(expert_inputs, fc1_1_weight[expert])
+            up = F.linear(expert_inputs, fc1_2_weight[expert])
+            intermediate = F.silu(gate) * up
+            intermediate.mul_(sorted_routing[start:end].unsqueeze(-1))
+            outputs.append(F.linear(intermediate, fc2_weight[expert]))
+        start = end
+
+    sorted_outputs = torch.cat(outputs, dim=0) if outputs else hidden_states.new_empty((0, hidden_states.shape[1]))
+    return _unroute_tokens(sorted_outputs, order, hidden_states.shape[0], selected_experts.shape[1])
+
+
+def fused_moe_forward(
+    module: torch.nn.Module,
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Run the VeOmni v0.1.0 split-weight MoE operation for inference.
+
+    ``module`` is retained for call-site compatibility. Like VeOmni's original
+    non-EP implementation, this function does not use it.
+
+    Set ``LLADA_MOE_BACKEND`` to ``auto`` (default), ``triton``, or ``eager``.
+    The ``triton`` setting fails loudly if its requirements are not met;
+    ``auto`` falls back to the PyTorch implementation.
+    """
+    del module
+    _validate_inputs(
+        num_experts,
+        routing_weights,
+        selected_experts,
+        hidden_states,
+        fc1_1_weight,
+        fc1_2_weight,
+        fc2_weight,
+    )
+
+    backend = os.getenv("LLADA_MOE_BACKEND", "auto").lower()
+    if backend not in {"auto", "triton", "eager"}:
+        raise ValueError(f"LLADA_MOE_BACKEND must be auto, triton, or eager; got {backend!r}")
+
+    compute_dtype = fc1_1_weight.dtype
+    if fc1_2_weight.dtype != compute_dtype or fc2_weight.dtype != compute_dtype:
+        raise TypeError("all expert weights must have the same dtype")
+    hidden_states = hidden_states.to(dtype=compute_dtype)
+    routing_weights = routing_weights.to(dtype=compute_dtype)
+
+    if hidden_states.shape[0] == 0:
+        return hidden_states
+
+    can_use_triton = (
+        triton is not None
+        and hidden_states.is_cuda
+        and compute_dtype in _SUPPORTED_TRITON_DTYPES
+        and fc1_1_weight.is_contiguous()
+        and fc1_2_weight.is_contiguous()
+        and fc2_weight.is_contiguous()
+    )
+    if backend == "triton" and not can_use_triton:
+        raise RuntimeError(
+            "The Triton backend requires Triton, CUDA tensors, contiguous expert weights, "
+            "and float16 or bfloat16 weights."
+        )
+    if backend != "eager" and can_use_triton:
+        return _triton_moe_forward(
+            num_experts,
+            routing_weights,
+            selected_experts,
+            hidden_states,
+            fc1_1_weight,
+            fc1_2_weight,
+            fc2_weight,
+        )
+    return _eager_moe_forward(
+        num_experts,
+        routing_weights,
+        selected_experts,
+        hidden_states,
+        fc1_1_weight,
+        fc1_2_weight,
+        fc2_weight,
+    )
+
+
+__all__ = ["fused_moe_forward"]

text_encoder/modeling_llada2uni_moe.py

@@ -41,7 +41,7 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_u
 from transformers.modeling_utils import PreTrainedModel
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import logging
-from veomni.ops import fused_moe_forward
+from .fused_moe_ops import fused_moe_forward
 
 from .configuration_llada2uni_moe import LLaDA2MoeConfig