[Draft] Add PyTorch engine support for Hy3 BF16 and static FP8 (NO MTP)

lmdeploy/pytorch/backends/base.py

@@ -28,7 +28,9 @@ class OpType(Enum):
     SoftmaxTopK = auto()
     FusedMoE = auto()
     FusedMoEW8A8 = auto()
+    FusedMoEStaticF8 = auto()
     LinearBlockedF8 = auto()
+    LinearStaticF8 = auto()
     FusedMoEBlockedF8 = auto()
     NSAIndexFP8 = auto()
     Embedding = auto()

lmdeploy/pytorch/backends/cuda/moe/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .blocked_fp8 import TritonFusedMoEBlockedF8Builder  # noqa: F401
 from .default import TritonFusedMoEBuilder  # noqa: F401
+from .static_fp8 import TritonFusedMoEStaticF8Builder  # noqa: F401
 from .w8a8 import TritonFusedMoEW8A8Builder  # noqa: F401

lmdeploy/pytorch/backends/cuda/moe/static_fp8.py

@@ -0,0 +1,96 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+import torch
+
+from lmdeploy.pytorch.backends.moe import (
+    FusedMoEStaticF8Builder,
+    FusedMoEStaticF8Impl,
+)
+from lmdeploy.pytorch.kernels.cuda.w8a8_fused_moe import (
+    fused_moe_static_fp8,
+)
+
+
+class TritonFusedMoEStaticF8Impl(FusedMoEStaticF8Impl):
+    """Triton static FP8 fused MoE implementation."""
+
+    def __init__(
+        self,
+        top_k: int,
+        num_experts: int,
+        renormalize: bool = False,
+        out_dtype: torch.dtype = torch.float16,
+        quant_dtype: torch.dtype = torch.float8_e4m3fn,
+    ):
+        self.top_k = top_k
+        self.num_experts = num_experts
+        self.renormalize = renormalize
+        self.out_dtype = out_dtype
+        self.quant_dtype = quant_dtype
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.LongTensor,
+        gate_up_weights: torch.Tensor,
+        gate_up_weight_scale: torch.Tensor,
+        gate_up_input_scale: torch.Tensor,
+        down_weights: torch.Tensor,
+        down_weight_scale: torch.Tensor,
+        down_input_scale: torch.Tensor,
+        expert_list: list[int] = None,
+    ):
+        """Forward."""
+        hidden_states = hidden_states.contiguous()
+
+        expert_offset = 0
+        num_experts = None
+
+        if (
+            expert_list is not None
+            and len(expert_list) != self.num_experts
+        ):
+            expert_offset = expert_list[0]
+            num_experts = self.num_experts
+
+        return fused_moe_static_fp8(
+            hidden_states,
+            gate_up_input_scale,
+            gate_up_weights,
+            gate_up_weight_scale,
+            down_input_scale,
+            down_weights,
+            down_weight_scale,
+            topk_weights,
+            topk_ids,
+            topk=self.top_k,
+            out_dtype=self.out_dtype,
+            quant_dtype=self.quant_dtype,
+            expert_offset=expert_offset,
+            num_experts=num_experts,
+            renormalize=self.renormalize,
+        )
+
+
+class TritonFusedMoEStaticF8Builder(
+    FusedMoEStaticF8Builder,
+):
+    """Triton static FP8 fused MoE builder."""
+
+    @staticmethod
+    def build(
+        top_k: int,
+        num_experts: int,
+        renormalize: bool = False,
+        out_dtype: torch.dtype = torch.float16,
+        quant_dtype: torch.dtype = torch.float8_e4m3fn,
+    ):
+        """Build static FP8 fused MoE."""
+        return TritonFusedMoEStaticF8Impl(
+            top_k=top_k,
+            num_experts=num_experts,
+            renormalize=renormalize,
+            out_dtype=out_dtype,
+            quant_dtype=quant_dtype,
+        )

lmdeploy/pytorch/backends/cuda/op_backend.py

@@ -58,9 +58,17 @@ def get_layer_impl_builder(cls, layer_type: OpType):
         elif layer_type == OpType.FusedMoEW8A8:
             from .moe import TritonFusedMoEW8A8Builder
             return TritonFusedMoEW8A8Builder
+        elif layer_type == OpType.FusedMoEStaticF8:
+            from .moe import TritonFusedMoEStaticF8Builder
+            return TritonFusedMoEStaticF8Builder
         elif layer_type == OpType.FusedMoEBlockedF8:
             from .moe import TritonFusedMoEBlockedF8Builder
             return TritonFusedMoEBlockedF8Builder
+        elif layer_type == OpType.LinearStaticF8:
+            from .static_fp8_modules import (
+                TritonLinearStaticF8Builder,
+            )
+            return TritonLinearStaticF8Builder
         elif layer_type == OpType.LinearBlockedF8:
             from .blockedf8_modules import TritonLinearBlockedF8Builder
             return TritonLinearBlockedF8Builder

lmdeploy/pytorch/backends/cuda/static_fp8_modules.py

@@ -0,0 +1,82 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+import torch
+
+import lmdeploy.pytorch.distributed as dist
+from lmdeploy.pytorch.kernels.cuda.w8a8_triton_kernels import (
+    matmul_kernel_static_quant,
+)
+
+from ..static_fp8_modules import (
+    LinearStaticF8Builder,
+    LinearStaticF8Impl,
+)
+
+
+class TritonLinearStaticF8Impl(LinearStaticF8Impl):
+    """Triton static per-tensor FP8 linear implementation."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        out_dtype: torch.dtype = torch.float16,
+    ):
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        input_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        all_reduce: bool = False,
+        group: dist.ProcessGroup | None = None,
+        rank: int = 0,
+        scatter_size: list[int] | None = None,
+    ):
+        """Run static FP8 linear."""
+        output_dtype = self.out_dtype or x.dtype
+
+        output = matmul_kernel_static_quant(
+            x,
+            weight,
+            input_scale,
+            weight_scale,
+            bias=bias,
+            output_dtype=output_dtype,
+        )
+
+        if all_reduce:
+            if scatter_size is not None:
+                output = dist.reduce_scatter_by_tp_sizes(
+                    output,
+                    rank,
+                    scatter_size,
+                    group=group,
+                )
+            else:
+                dist.all_reduce(output, group=group)
+
+        return output
+
+
+class TritonLinearStaticF8Builder(LinearStaticF8Builder):
+    """Triton static per-tensor FP8 linear builder."""
+
+    @staticmethod
+    def build(
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        dtype: torch.dtype | None = None,
+    ):
+        """Build static FP8 linear implementation."""
+        return TritonLinearStaticF8Impl(
+            in_features,
+            out_features,
+            out_dtype=dtype,
+        )

lmdeploy/pytorch/backends/moe.py

@@ -116,6 +116,60 @@ def build(top_k: int,
         """Build from mlp."""
         raise NotImplementedError
 
+class FusedMoEStaticF8Impl(ABC):
+    """Fused MoE static FP8 implementation."""
+
+    def update_weights(
+        self,
+        gate_up_weights: torch.Tensor,
+        gate_up_weight_scale: torch.Tensor,
+        gate_up_input_scale: torch.Tensor,
+        down_weights: torch.Tensor,
+        down_weight_scale: torch.Tensor,
+        down_input_scale: torch.Tensor,
+    ):
+        """Update weights and scales."""
+        return (
+            gate_up_weights,
+            gate_up_weight_scale,
+            gate_up_input_scale,
+            down_weights,
+            down_weight_scale,
+            down_input_scale,
+        )
+
+    @abstractmethod
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.LongTensor,
+        gate_up_weights: torch.Tensor,
+        gate_up_weight_scale: torch.Tensor,
+        gate_up_input_scale: torch.Tensor,
+        down_weights: torch.Tensor,
+        down_weight_scale: torch.Tensor,
+        down_input_scale: torch.Tensor,
+        expert_list: list[int] = None,
+    ):
+        """Forward."""
+        raise NotImplementedError
+
+
+class FusedMoEStaticF8Builder(ABC):
+    """Fused MoE static FP8 builder."""
+
+    @staticmethod
+    @abstractmethod
+    def build(
+        top_k: int,
+        num_experts: int,
+        renormalize: bool = False,
+        out_dtype: torch.dtype = torch.float16,
+        quant_dtype: torch.dtype = torch.float8_e4m3fn,
+    ):
+        """Build static FP8 MoE implementation."""
+        raise NotImplementedError
 
 class FusedMoEBlockedF8Impl(ABC):
     """Fused moe blocked f8 implementation."""

lmdeploy/pytorch/backends/static_fp8_modules.py

@@ -0,0 +1,50 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from abc import ABC, abstractmethod
+
+import torch
+import torch.distributed as dist
+
+
+class LinearStaticF8Impl(ABC):
+    """Static per-tensor FP8 linear implementation API."""
+
+    def update_weights(
+        self,
+        weight: torch.Tensor,
+        input_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ):
+        """Update weights."""
+        return weight, input_scale, weight_scale, bias
+
+    @abstractmethod
+    def forward(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        input_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        all_reduce: bool = False,
+        group: dist.ProcessGroup | None = None,
+        rank: int = 0,
+        scatter_size: list[int] | None = None,
+    ):
+        """Run static FP8 linear."""
+        raise NotImplementedError
+
+
+class LinearStaticF8Builder(ABC):
+    """Static per-tensor FP8 linear builder."""
+
+    @staticmethod
+    @abstractmethod
+    def build(
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        dtype: torch.dtype | None = None,
+    ):
+        """Build static FP8 linear implementation."""
+        raise NotImplementedError

lmdeploy/pytorch/configurations/hy3.py

@@ -0,0 +1,12 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+from .default import DefaultModelConfigBuilder
+
+
+class Hy3ModelConfigBuilder(DefaultModelConfigBuilder):
+    """Model config builder for Hy3."""
+
+    @classmethod
+    def condition(cls, hf_config):
+        """Match Hy3 configurations."""
+        return hf_config.model_type == 'hy_v3'