@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 149 lines 3.81 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """NOTE vision MoE execution matching the native encoder's FP8 semantics."""
4
5 import torch
6
7 from vllm import _custom_ops as ops
8 from vllm.model_executor.layers.fused_moe.activation import (
9 MoEActivation,
10 apply_moe_activation,
11 )
12 from vllm.model_executor.layers.fused_moe.fused_moe import (
13 _prepare_expert_assignment,
14 dispatch_fused_moe_kernel,
15 try_get_optimal_moe_config,
16 )
17 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 per_token_group_quant_fp8,
19 )
20 from vllm.triton_utils import tl
21
22 _BLOCK_SHAPE = [128, 128]
23
24
25 def note_vision_fused_moe_fp8(
26 hidden_states: torch.Tensor,
27 w13: torch.Tensor,
28 w2: torch.Tensor,
29 topk_weights: torch.Tensor,
30 topk_ids: torch.Tensor,
31 w13_scale: torch.Tensor,
32 w2_scale: torch.Tensor,
33 ) -> torch.Tensor:
34 """Run top-k SwiGLU MoE with the encoder's block-FP8 formula.
35
36 Unlike vLLM's generic functional MoE path, the native NOTE encoder keeps
37 dynamic activation scales as FP32 instead of rounding them to E8M0. Weight
38 scales are block 128x128 and activations are quantized per token/group-128.
39 """
40 if hidden_states.dtype != torch.bfloat16:
41 raise TypeError(
42 "NOTE vision FP8 MoE requires bfloat16 activations, got "
43 f"{hidden_states.dtype}"
44 )
45
46 num_tokens = hidden_states.shape[0]
47 num_experts, intermediate_size, _ = w13.shape
48 output_size = w2.shape[1]
49 topk = topk_ids.shape[1]
50 config = try_get_optimal_moe_config(
51 w13.shape,
52 w2.shape,
53 topk,
54 "fp8_w8a8",
55 num_tokens,
56 block_shape=_BLOCK_SHAPE,
57 )
58
59 cache = torch.empty(
60 num_tokens * topk * max(intermediate_size, output_size),
61 dtype=hidden_states.dtype,
62 device=hidden_states.device,
63 )
64 first_gemm = cache[: num_tokens * topk * intermediate_size].view(
65 num_tokens, topk, intermediate_size
66 )
67 second_gemm = cache[: num_tokens * topk * output_size].view(
68 num_tokens, topk, output_size
69 )
70 activated_size = intermediate_size // 2
71 activated = torch.empty(
72 (num_tokens * topk, activated_size),
73 dtype=hidden_states.dtype,
74 device=hidden_states.device,
75 )
76
77 quantized_input, input_scale = per_token_group_quant_fp8(
78 hidden_states,
79 _BLOCK_SHAPE[1],
80 use_ue8m0=False,
81 )
82 sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
83 topk_ids,
84 config,
85 num_tokens,
86 topk,
87 num_experts,
88 None,
89 block_shape=_BLOCK_SHAPE,
90 )
91 dispatch_fused_moe_kernel(
92 quantized_input,
93 w13,
94 first_gemm,
95 input_scale,
96 w13_scale,
97 None,
98 topk_weights,
99 sorted_token_ids,
100 expert_ids,
101 num_tokens_post_padded,
102 False,
103 topk,
104 config,
105 tl.bfloat16,
106 True,
107 False,
108 False,
109 False,
110 False,
111 _BLOCK_SHAPE,
112 )
113
114 apply_moe_activation(
115 MoEActivation.SILU,
116 activated,
117 first_gemm.view(-1, intermediate_size),
118 )
119 quantized_activated, activated_scale = per_token_group_quant_fp8(
120 activated,
121 _BLOCK_SHAPE[1],
122 use_ue8m0=False,
123 )
124 dispatch_fused_moe_kernel(
125 quantized_activated,
126 w2,
127 second_gemm,
128 activated_scale,
129 w2_scale,
130 None,
131 topk_weights,
132 sorted_token_ids,
133 expert_ids,
134 num_tokens_post_padded,
135 True,
136 1,
137 config,
138 tl.bfloat16,
139 True,
140 False,
141 False,
142 False,
143 False,
144 _BLOCK_SHAPE,
145 )
146
147 output = torch.empty_like(hidden_states)
148 ops.moe_sum(second_gemm, output)
149 return output