Skip to content

vllm.models.dots3_note.nvidia.vision_moe

NOTE vision MoE execution matching the native encoder's FP8 semantics.

Functions:

note_vision_fused_moe_fp8(hidden_states, w13, w2, topk_weights, topk_ids, w13_scale, w2_scale)

Run top-k SwiGLU MoE with the encoder's block-FP8 formula.

Unlike vLLM's generic functional MoE path, the native NOTE encoder keeps dynamic activation scales as FP32 instead of rounding them to E8M0. Weight scales are block 128x128 and activations are quantized per token/group-128.

Source code in vllm/models/dots3_note/nvidia/vision_moe.py 
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def note_vision_fused_moe_fp8(
    hidden_states: torch.Tensor,
    w13: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
) -> torch.Tensor:
    """Run top-k SwiGLU MoE with the encoder's block-FP8 formula.

    Unlike vLLM's generic functional MoE path, the native NOTE encoder keeps
    dynamic activation scales as FP32 instead of rounding them to E8M0.  Weight
    scales are block 128x128 and activations are quantized per token/group-128.
    """
    if hidden_states.dtype != torch.bfloat16:
        raise TypeError(
            "NOTE vision FP8 MoE requires bfloat16 activations, got "
            f"{hidden_states.dtype}"
        )

    num_tokens = hidden_states.shape[0]
    num_experts, intermediate_size, _ = w13.shape
    output_size = w2.shape[1]
    topk = topk_ids.shape[1]
    config = try_get_optimal_moe_config(
        w13.shape,
        w2.shape,
        topk,
        "fp8_w8a8",
        num_tokens,
        block_shape=_BLOCK_SHAPE,
    )

    cache = torch.empty(
        num_tokens * topk * max(intermediate_size, output_size),
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )
    first_gemm = cache[: num_tokens * topk * intermediate_size].view(
        num_tokens, topk, intermediate_size
    )
    second_gemm = cache[: num_tokens * topk * output_size].view(
        num_tokens, topk, output_size
    )
    activated_size = intermediate_size // 2
    activated = torch.empty(
        (num_tokens * topk, activated_size),
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )

    quantized_input, input_scale = per_token_group_quant_fp8(
        hidden_states,
        _BLOCK_SHAPE[1],
        use_ue8m0=False,
    )
    sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
        topk_ids,
        config,
        num_tokens,
        topk,
        num_experts,
        None,
        block_shape=_BLOCK_SHAPE,
    )
    dispatch_fused_moe_kernel(
        quantized_input,
        w13,
        first_gemm,
        input_scale,
        w13_scale,
        None,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        False,
        topk,
        config,
        tl.bfloat16,
        True,
        False,
        False,
        False,
        False,
        _BLOCK_SHAPE,
    )

    apply_moe_activation(
        MoEActivation.SILU,
        activated,
        first_gemm.view(-1, intermediate_size),
    )
    quantized_activated, activated_scale = per_token_group_quant_fp8(
        activated,
        _BLOCK_SHAPE[1],
        use_ue8m0=False,
    )
    dispatch_fused_moe_kernel(
        quantized_activated,
        w2,
        second_gemm,
        activated_scale,
        w2_scale,
        None,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        True,
        1,
        config,
        tl.bfloat16,
        True,
        False,
        False,
        False,
        False,
        _BLOCK_SHAPE,
    )

    output = torch.empty_like(hidden_states)
    ops.moe_sum(second_gemm, output)
    return output