Merge branch 'v0.15.1-fusion' into 'v0.15.1-dev'

feat: Support rms+quant fusion in minimax_m2 series model.

See merge request dcutoolkit/deeplearing/vllm!538

vllm/model_executor/models/minimax_m2.py

@@ -39,7 +39,7 @@ from vllm.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from vllm.model_executor.layers.fused_moe import FusedMoE
-from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.layernorm import RMSNorm, FusedRMSNormQuant
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
     ReplicatedLinear,
@@ -58,6 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import (
     maybe_remap_kv_scale_name,
 )
 from vllm.sequence import IntermediateTensors
+from vllm import envs
 
 from .interfaces import SupportsLoRA, SupportsPP
 from .utils import (
@@ -229,8 +230,10 @@ class MiniMaxM2Attention(nn.Module):
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
+        iqis: tuple[torch.Tensor, torch.Tensor] | None = None
     ) -> torch.Tensor:
-        qkv, _ = self.qkv_proj(hidden_states)
+        # iqis: (input_quant, input_scale)
+        qkv, _ = self.qkv_proj(hidden_states, iqis = iqis)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = MiniMaxText01RMSNormTP.forward_qk(
             self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
@@ -282,17 +285,54 @@ class MiniMaxM2DecoderLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.mlp",
         )
+        if envs.USE_FUSED_RMS_QUANT:
+            self.input_layernorm = FusedRMSNormQuant(config.hidden_size, eps=config.rms_norm_eps)
+        else:
             self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
         self.post_attention_layernorm = RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
 
+    def rms_quant_fusion_forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states.clone()
+            input_quant, input_scale, _ = self.input_layernorm(
+                x = hidden_states, 
+                residual = None, 
+                quant_dtype = torch.int8,
+                update_input = False
+            )
+        else:
+            input_quant, input_scale, residual = self.input_layernorm(
+                x = hidden_states, 
+                residual = residual, 
+                quant_dtype = torch.int8,
+                update_input = False
+            )
+        hidden_states = self.self_attn(
+            positions = positions,
+            hidden_states = hidden_states,
+            iqis = (input_quant, input_scale)
+        )
+        return hidden_states, residual
+
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None,
     ) -> torch.Tensor:
+        if envs.USE_FUSED_RMS_QUANT:
+            hidden_states, residual = self.rms_quant_fusion_forward(
+                positions, hidden_states, residual
+            )
+        else:
             # Self Attention
             if residual is None:
                 residual = hidden_states