vllm.models.dots3_note.nvidia.vision ¶

class MoESwiGLUFFN(nn.Module):
    """MoE FFN with per-expert SwiGLU experts, sigmoid/softmax gating, top-k routing."""

    def __init__(self, config: DotsMoEVitConfig, layer_number: int):
        super().__init__()
        self.config = config
        self.layer_number = layer_number
        self.hidden_size = config.embed_dim
        self.num_routed = config.pyramid_num_routed[layer_number]
        self.capacity_factor = config.capacity_factor
        self.router_scoring_func = config.router_scoring_func
        self.router_scale = config.router_scale

        self.register_buffer(
            "router_bias", torch.zeros(self.num_routed, dtype=torch.float32)
        )

        self.experts = nn.ModuleList(
            [
                DotsSwiGLUFFN(
                    self.hidden_size, config.moe_intermediate_size, bias=config.use_bias
                )
                for _ in range(self.num_routed)
            ]
        )

        self.gate_weight = nn.Parameter(
            torch.empty((self.num_routed, self.hidden_size), dtype=torch.float32)
        )
        nn.init.kaiming_uniform_(self.gate_weight, a=math.sqrt(5))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Mirror cybertron's ``AIMv2MoEGEMMSwiGLUFFN.forward``: keep ``gating_prob`` /
        # router-bias add in fp32 and run topk in fp32. The legacy ``.type_as(x_flat)``
        # path on bf16 made expert routing diverge whenever two routed-expert scores
        # tied within bf16 precision, which is the dominant source of numerical drift
        # observed against latest cybertron checkpoints.
        epsilon = 1e-9
        x_flat = x.contiguous()
        num_tokens = x_flat.shape[0]

        gate_logits = F.linear(x_flat.float(), self.gate_weight.float())

        if self.router_scoring_func == "sigmoid":
            gating_prob = torch.sigmoid(gate_logits)
        else:
            gating_prob = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)

        aggregated_output = torch.zeros_like(x_flat)
        aggregated_gate = torch.zeros(num_tokens, dtype=x.dtype, device=x.device)

        topk = min(int(self.capacity_factor), self.num_routed)

        gating_with_bias = gating_prob + self.router_bias.to(torch.float32).unsqueeze(0)
        _, topk_indices = torch.topk(gating_with_bias, k=topk, dim=-1, sorted=False)

        routed_weights = gating_prob.gather(1, topk_indices)
        if self.router_scoring_func == "sigmoid" and topk > 1:
            routed_weights = routed_weights / (
                routed_weights.sum(dim=-1, keepdim=True) + epsilon
            )
        routed_weights = (routed_weights * self.router_scale).to(x_flat.dtype)

        for expert_idx in range(self.num_routed):
            selected_mask = topk_indices == expert_idx
            if selected_mask.sum() == 0:
                continue
            n_idx, top = torch.where(selected_mask)
            # Fancy indexing can yield non-contiguous rows; cuBLAS bf16 GEMM may then fail
            # with ``CUBLAS_STATUS_INVALID_VALUE`` inside ``F.linear``.
            x_selected = x_flat[n_idx].contiguous()
            expert_output = self.experts[expert_idx](x_selected)
            contrib = expert_output * routed_weights[n_idx, top].unsqueeze(-1)
            aggregated_output[n_idx] = aggregated_output[n_idx] + contrib
            aggregated_gate[n_idx] = aggregated_gate[n_idx] + routed_weights[n_idx, top]

        aggregated_output = aggregated_output / (
            aggregated_gate.unsqueeze(-1) + epsilon
        )
        return aggregated_output

class MoESwiGLUFFNFP8(MoESwiGLUFFN):
    """NOTE vision MoE using the checkpoint's local block-FP8 semantics."""

    @torch.no_grad()
    def process_weights_after_loading(self) -> None:
        if hasattr(self, "_fused_w13_fp8"):
            return

        w13_weights = []
        w13_scales = []
        w2_weights = []
        w2_scales = []
        for expert in self.experts:
            w1, s1 = _per_block_cast_to_fp8_padded(expert.fc1.weight)
            w3, s3 = _per_block_cast_to_fp8_padded(expert.fc3.weight)
            w2, s2 = _per_block_cast_to_fp8_padded(expert.fc2.weight)
            w13_weights.append(torch.cat((w1, w3), dim=0))
            w13_scales.append(torch.cat((s1, s3), dim=0))
            w2_weights.append(w2)
            w2_scales.append(s2)

        self.register_buffer(
            "_fused_w13_fp8",
            torch.stack(w13_weights).contiguous(),
            persistent=False,
        )
        self.register_buffer(
            "_fused_w13_scale",
            torch.stack(w13_scales).contiguous(),
            persistent=False,
        )
        self.register_buffer(
            "_fused_w2_fp8",
            torch.stack(w2_weights).contiguous(),
            persistent=False,
        )
        self.register_buffer(
            "_fused_w2_scale",
            torch.stack(w2_scales).contiguous(),
            persistent=False,
        )
        del self.experts

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not hasattr(self, "_fused_w13_fp8"):
            raise RuntimeError("NOTE vision FP8 weights were not initialized")

        gate_logits = F.linear(x.float(), self.gate_weight.float())
        if self.router_scoring_func == "sigmoid":
            scores = torch.sigmoid(gate_logits)
        else:
            scores = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)

        topk = min(int(self.capacity_factor), self.num_routed)
        biased_scores = scores + self.router_bias.float().unsqueeze(0)
        topk_ids = torch.topk(biased_scores, k=topk, dim=-1, sorted=False).indices
        topk_weights = scores.gather(1, topk_ids)
        if self.router_scoring_func == "sigmoid" and topk > 1:
            topk_weights = topk_weights / (
                topk_weights.sum(dim=-1, keepdim=True) + 1e-9
            )
        topk_weights = topk_weights * self.router_scale

        output = note_vision_fused_moe_fp8(
            x.contiguous(),
            self._fused_w13_fp8,
            self._fused_w2_fp8,
            topk_weights,
            topk_ids.to(torch.int32),
            self._fused_w13_scale,
            self._fused_w2_scale,
        )
        denominator = topk_weights.sum(dim=-1, keepdim=True).clamp_min(1e-9)
        return (output / denominator).type_as(x)

class PatchMergerAdapter(nn.Module):
    """Cybertron ``PatchMerger`` (``pool_kind='patch_merger', proj_kind='identity'``).

    Assumes the encoder output is already laid out in ``merge_size``x``merge_size`` groups
    (qwen ``pre_pixel_shuffle`` preprocessor + RoPE grouped accordingly), so merging is a
    simple ``view(-1, merge**2 * in_dim)`` of consecutive tokens. State-dict layout matches
    cybertron's ``PatchMerger`` (``ln_q`` over the per-token dim, ``mlp.0`` / ``mlp.2`` Linear).
    """

    def __init__(self, config: DotsMoEVitConfig):
        super().__init__()
        in_dim = config.adapter_in_dim
        out_dim = config.adapter_out_dim
        merge_size = config.adapter_merge_size
        merged_dim = in_dim * merge_size**2
        self.merge_size = merge_size
        self.merged_dim = merged_dim
        self.ln_q = LayerNorm(in_dim, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(merged_dim, merged_dim),
            nn.GELU(),
            nn.Linear(merged_dim, out_dim),
        )

    def forward(
        self,
        patch_embed: torch.Tensor,
        grid_thw: torch.Tensor,
    ) -> torch.Tensor:
        assert patch_embed.dim() == 2 and grid_thw is not None
        x = self.ln_q(patch_embed)
        x = x.reshape(-1, self.merged_dim)
        return self.mlp(x)

class PixelShuffleAdapter(nn.Module):
    """Legacy adapter: NHWC pixel-shuffle spatial merge + LayerNorm + 2-layer MLP.

    Mirrors ``cybertron`` ``FCAdapter(pool_kind='pixel_shuffle', proj_kind='mlp2x_ln_gelu')``.
    State-dict keys: ``proj.0`` (LayerNorm of in_dim*merge**2), ``proj.1`` / ``proj.3`` (Linear).
    """

    def __init__(self, config: DotsMoEVitConfig):
        super().__init__()
        in_dim = config.adapter_in_dim
        out_dim = config.adapter_out_dim
        merge_size = config.adapter_merge_size
        merged_dim = in_dim * merge_size**2
        self.proj = nn.Sequential(
            LayerNorm(merged_dim),
            nn.Linear(merged_dim, out_dim),
            nn.GELU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(
        self,
        patch_embed: torch.Tensor,
        grid_thw: torch.Tensor,
    ) -> torch.Tensor:
        assert patch_embed.dim() == 2 and grid_thw is not None
        image_features = []
        token_index = 0
        for i in range(grid_thw.shape[0]):
            grid_t, grid_h, grid_w = grid_thw[i]
            images_token_length = grid_t * grid_h * grid_w
            _pe = patch_embed[token_index : token_index + images_token_length]
            token_index += images_token_length
            if grid_t == 1:
                _pe = _pe.reshape(int(grid_h), int(grid_w), -1).unsqueeze(0)
            else:
                _pe = _pe.reshape(int(grid_t), int(grid_h), int(grid_w), -1)
            _pe = _pixel_shuffle(_pe, scale_factor=0.5)
            _pe = _pe.squeeze(0) if grid_t == 1 else _pe.reshape(-1, _pe.shape[-1])
            image_features.append(_pe.reshape(-1, _pe.shape[-1]))
        out = torch.cat(image_features, dim=0)
        out = self.proj(out)
        return out