[Triton] Shaoclee/085 triton fusion

vllm/envs.py

@@ -115,7 +115,8 @@
     VLLM_AITER_TRITON_FP8_BMM: bool = False
     VLLM_AITER_TRITON_FUSED_CONCAT_ZEROS: bool = False
     VLLM_AITER_TRITON_FUSED_ROPE_CACHE_CONCAT: bool = False
-    VLLM_MLA_FP8_PADDING: bool = False
+    VLLM_MLA_FP8_PADDING: bool = False,
+    VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT: bool = False
 
 
 def get_default_cache_root():
@@ -755,6 +756,9 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
     "VLLM_MLA_FP8_PADDING":
     lambda: bool(int(os.getenv("VLLM_MLA_FP8_PADDING", "0"))),
 
+    "VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT":
+    lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT", "0"))),
+
 }
 
 # end-env-vars-definition

vllm/model_executor/layers/linear.py

@@ -387,11 +387,17 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
         param.data.copy_(loaded_weight)
 
     def forward(
-        self, x: torch.Tensor
+        self, x: torch.Tensor, x_quant_scales: torch.Tensor = None,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
-        output = self.quant_method.apply(self, x, bias)
+        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
+        if isinstance(self.quant_method, CompressedTensorsLinearMethod):
+            output = self.quant_method.apply(self, x, bias, x_quant_scales=x_quant_scales)
+        else:
+            # assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            output = self.quant_method.apply(self, x, bias)
+
         output_bias = self.bias if self.skip_bias_add else None
         if not self.return_bias:
             return output
@@ -605,13 +611,18 @@ def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor):
         param.load_column_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
-        self, input_
+        self, input_, x_quant_scales: torch.Tensor = None,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         bias = self.bias if not self.skip_bias_add else None
 
         # Matrix multiply.
         assert self.quant_method is not None
-        output_parallel = self.quant_method.apply(self, input_, bias)
+        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
+        if isinstance(self.quant_method, CompressedTensorsLinearMethod):
+            output_parallel = self.quant_method.apply(self, input_, bias, x_quant_scales=x_quant_scales)
+        else:
+            # assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            output_parallel = self.quant_method.apply(self, input_, bias)
         if self.gather_output:
             # All-gather across the partitions.
             output = tensor_model_parallel_all_gather(output_parallel)
@@ -1380,7 +1391,8 @@ def weight_loader_v2(self, param: BasevLLMParameter,
         param.load_row_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
-        self, input_
+        self, input_,
+        x_quant_scales = None
     ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
         if self.input_is_parallel:
             input_parallel = input_
@@ -1395,9 +1407,12 @@ def forward(
         # Only fuse bias add into GEMM for rank 0 (this ensures that
         # bias will not get added more than once in TP>1 case)
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
-        output_parallel = self.quant_method.apply(self,
-                                                  input_parallel,
-                                                  bias=bias_)
+        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
+        if isinstance(self.quant_method, CompressedTensorsLinearMethod):
+            output_parallel = self.quant_method.apply(self, input_parallel, bias_, x_quant_scales=x_quant_scales)
+        else:
+            # assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
+            output_parallel = self.quant_method.apply(self, input_parallel, bias_)
         if self.reduce_results and self.tp_size > 1:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -566,7 +566,8 @@ def create_weights(self, layer: torch.nn.Module,
     def apply(self,
               layer: torch.nn.Module,
               x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None):
+              bias: Optional[torch.Tensor] = None,
+              x_quant_scales: Optional[torch.Tensor] = None):
         """
         Use the output of create_weights and the CompressedTensorsScheme
         associated with the layer to apply the forward pass with the
@@ -577,7 +578,11 @@ def apply(self,
         scheme = layer.scheme
         if scheme is None:
             raise ValueError("A scheme must be defined for each layer")
-        return scheme.apply_weights(layer, x, bias=bias)
+        elif isinstance(scheme, CompressedTensorsW8A8Fp8):        
+            return scheme.apply_weights(layer, x, bias=bias, x_quant_scales=x_quant_scales)
+        else:
+            assert x_quant_scales is None, f"x_quant_scales is not supported for {scheme.__class__}"
+            return scheme.apply_weights(layer, x, bias=bias)
 
 
 class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -156,11 +156,13 @@ def create_weights(self, layer: torch.nn.Module,
     def apply_weights(self,
                       layer: torch.nn.Module,
                       x: torch.Tensor,
-                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+                      bias: Optional[torch.Tensor] = None,
+                      x_quant_scales: Optional[torch.Tensor] = None) -> torch.Tensor:
 
         return self.fp8_linear.apply(input=x,
                                      weight=layer.weight,
                                      weight_scale=layer.weight_scale,
                                      out_dtype=self.out_dtype,
                                      input_scale=layer.input_scale,
-                                     bias=bias)
+                                     bias=bias,
+                                     x_quant_scales=x_quant_scales)

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -599,4 +599,4 @@ def grid(META):
         **config,
     )
 
-    return C
+    return C

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

@@ -345,6 +345,7 @@ def apply(
         input_scale: Optional[torch.Tensor] = None,
         input_scale_ub: Optional[torch.Tensor] = None,
         bias: Optional[torch.Tensor] = None,
+        x_quant_scales: Optional[torch.Tensor] = None,
         # TODO(luka) remove this parameter in favor of __init__
         use_per_token_if_dynamic: Optional[bool] = None
     ) -> torch.Tensor:
@@ -404,6 +405,8 @@ def apply(
                         qinput, x_scale = (
                             torch.ops.vllm.rocm_aiter_per_tensor_quant_fp8(
                                 input_2d, scale=input_scale))
+            elif x_quant_scales is not None:
+                qinput, x_scale = input_2d, x_quant_scales
             else:
                 qinput, x_scale = input_2d, input_scale
 

vllm/model_executor/models/deepseek_v2.py

@@ -62,7 +62,14 @@
     is_rocm_aiter_fusion_shared_expert_enabled,
     is_rocm_aiter_fuse_routed_scaling_factor,
 )
+from vllm.platforms import current_platform
 
+import vllm.envs as envs
+
+if current_platform.is_rocm():
+    VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT=envs.VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT
+else:
+    VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT=False
 
 class DeepseekV2MLP(nn.Module):
 
@@ -93,7 +100,10 @@ def __init__(
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
-        gate_up, _ = self.gate_up_proj(x)
+        x_quant_scales = None
+        if isinstance(x, tuple):
+            x, x_quant_scales = x
+        gate_up, _ = self.gate_up_proj(x, x_quant_scales=x_quant_scales)
         x = self.act_fn(gate_up)
         x, _ = self.down_proj(x)
         return x
@@ -156,11 +166,16 @@ def __init__(
             )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if isinstance(hidden_states, tuple):
+            hidden_states_shared, hidden_states = hidden_states
+        else:
+            hidden_states_shared = hidden_states
+
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         shared_output = None
         if self.n_shared_experts is not None and not is_rocm_aiter_fusion_shared_expert_enabled():
-            shared_output = self.shared_experts(hidden_states)
+            shared_output = self.shared_experts(hidden_states_shared)
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
         if hidden_states.dtype != torch.float16:
@@ -484,21 +499,43 @@ def forward(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
+        hidden_states_quant = None
+        if isinstance(hidden_states, tuple):
+            hidden_states, hidden_states_quant = hidden_states
+
         if self.q_lora_rank is not None:
             # q_c = self.q_a_proj(hidden_states)[0]
             # kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
-            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
+            qkv_lora = self.fused_qkv_a_proj(hidden_states, x_quant_scales = hidden_states_quant)[0]
+            # qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
             q_c, kv_lora = qkv_lora.split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                 dim=-1,
             )
-            hidden_states_or_q_c = self.q_a_layernorm(q_c)
+            if VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
+                from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_token_quant
+                fp8_dtype = current_platform.fp8_dtype()
+                weight = self.q_a_layernorm.weight
+                eps = self.q_a_layernorm.variance_epsilon
+                weight2 = self.kv_a_layernorm.weight
+                eps2 = self.kv_a_layernorm.variance_epsilon
+                kv_c, k_pe = kv_lora.split(
+                    [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+                (_, _), hidden_states_or_q_c, kv_c_normed, _ = fused_rms_fp8_per_token_quant(q_c, weight, eps, 
+                                                        kv_c, weight2, eps2, 
+                                                        dtype_quant=fp8_dtype, 
+                                                        res1=None,
+                                                        output_unquantized_inp1=True,
+                                                        output_quantiezed_inp1=False)
+            else:
+                hidden_states_or_q_c = self.q_a_layernorm(q_c)
         else:
             hidden_states_or_q_c = hidden_states
             kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
-        kv_c, k_pe = kv_lora.split(
-            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
-        kv_c_normed = self.kv_a_layernorm(kv_c)
+        if not VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
+            kv_c, k_pe = kv_lora.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+            kv_c_normed = self.kv_a_layernorm(kv_c)
         return self.mla_attn(hidden_states_or_q_c,
                              kv_c_normed,
                              k_pe,
@@ -576,12 +613,30 @@ def forward(
         residual: Optional[torch.Tensor],
     ) -> torch.Tensor:
         # Self Attention
-        if residual is None:
-            residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+        if VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
+            weight = self.input_layernorm.weight
+            eps = self.input_layernorm.variance_epsilon
+            from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_token_quant
+            fp8_dtype = current_platform.fp8_dtype()
+            if residual is None:
+                residual = hidden_states
+                (hidden_states, hidden_states_quant), _, _, _ = fused_rms_fp8_per_token_quant(hidden_states, weight, eps, 
+                                                            None, None, eps, 
+                                                            dtype_quant=fp8_dtype, 
+                                                            res1=None)
+            else:
+                (hidden_states, hidden_states_quant), _, _, residual = fused_rms_fp8_per_token_quant(hidden_states, weight, eps, 
+                                                            None, None, eps, 
+                                                            dtype_quant=fp8_dtype, 
+                                                            res1=residual)
+            hidden_states = (hidden_states, hidden_states_quant)
         else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
+            if residual is None:
+                residual = hidden_states
+                hidden_states = self.input_layernorm(hidden_states)
+            else:
+                hidden_states, residual = self.input_layernorm(
+                    hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,