Add FP8 MoE for turbomind

CMakeLists.txt

@@ -38,21 +38,6 @@ option(BUILD_FAST_MATH "Build in fast math mode" ON)
 include(FetchContent)
 
 if (BUILD_TEST)
-  FetchContent_Declare(
-    repo-cutlass
-    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-    GIT_TAG        6f47420213f757831fae65c686aa471749fa8d60
-    GIT_SHALLOW ON
-  )
-
-  set(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-
-  FetchContent_MakeAvailable(repo-cutlass)
-
-  set(CUTLASS_HEADER_DIR ${PROJECT_SOURCE_DIR}/3rdparty/cutlass/include)
-  set(CUTLASS_EXTENSIONS_DIR ${PROJECT_SOURCE_DIR}/src/turbomind/cutlass_extensions/include)
-
-
   FetchContent_Declare(
     Catch2
     GIT_REPOSITORY https://github.com/catchorg/Catch2.git
@@ -61,6 +46,19 @@ if (BUILD_TEST)
   FetchContent_MakeAvailable(Catch2)
 endif()
 
+
+FetchContent_Declare(
+  repo-cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG        v3.9.2
+  GIT_SHALLOW ON
+)
+
+set(CUTLASS_ENABLE_SM90_EXTENDED_MMA_SHAPES ON CACHE BOOL "Enable extended GMMA shapes")
+set(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
+
+FetchContent_MakeAvailable(repo-cutlass)
+
 FetchContent_Declare(
   yaml-cpp
   GIT_REPOSITORY https://github.com/jbeder/yaml-cpp.git
@@ -129,10 +127,13 @@ if (NOT CMAKE_CUDA_ARCHITECTURES)
     list(APPEND CMAKE_CUDA_ARCHITECTURES 86-real)
   endif ()
   if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.8")
-    list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real 90-real)
+    list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real)
+  endif ()
+  if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.0")
+    list(APPEND CMAKE_CUDA_ARCHITECTURES 90a-real)
   endif ()
   if (MSVC)
-    list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES 80-real 90-real)
+    list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES 80-real 90a-real)
   endif ()
 endif ()
 

lmdeploy/cli/utils.py

@@ -124,7 +124,7 @@ def model_format(parser, default: str = None):
         return parser.add_argument('--model-format',
                                    type=str,
                                    default=default,
-                                   choices=['hf', 'awq', 'gptq'],
+                                   choices=['hf', 'awq', 'gptq', 'fp8'],
                                    help='The format of input model. `hf` means `hf_llama`, '
                                    '`awq` represents the quantized model by AWQ,'
                                    ' and `gptq` refers to the quantized model by GPTQ')

lmdeploy/turbomind/deploy/converter.py

@@ -20,7 +20,7 @@
 from .source_model.base import INPUT_MODELS
 from .target_model.base import OUTPUT_MODELS, BaseOutputModel
 
-SUPPORTED_FORMATS = ['hf', 'awq', 'gptq', None]
+SUPPORTED_FORMATS = ['hf', 'awq', 'gptq', 'fp8', None]
 logger = get_logger('lmdeploy')
 
 
@@ -102,6 +102,9 @@ def get_output_model_registered_name_and_config(model_path: str, model_format: s
     if model_format in ['awq', 'gptq']:
         weight_type = 'int4'
         group_size = 128 if group_size == 0 else group_size
+    elif model_format == 'fp8':
+        weight_type = 'fp8'
+        group_size = 128
     else:
         torch_dtype = getattr(model_config, 'torch_dtype', 'float16')
         TORCH_DTYPE_MAP = {torch.bfloat16: 'bfloat16', torch.float16: 'float16'}
@@ -112,7 +115,7 @@ def get_output_model_registered_name_and_config(model_path: str, model_format: s
             weight_type = 'bfloat16'
 
     if dtype == 'auto':
-        weight_type = weight_type if weight_type in ['float16', 'bfloat16', 'int4'] else 'float16'
+        weight_type = weight_type if weight_type in ['float16', 'bfloat16', 'int4', 'fp8'] else 'float16'
     elif dtype in ['float16', 'bfloat16']:
         if weight_type == 'int4':
             logger.warning(f'The model {model_path} is a quantized model, so the '
@@ -197,6 +200,8 @@ def get_tm_model(model_path,
             assert not quant_config.get('desc_act', False) and \
                 quant_config.get('sym', True), \
                 f'unsupported quant config: {quant_config}'
+        elif quant_method == 'fp8':
+            pass
         else:
             assert 0, f'unsupported quant_config: {quant_config}'
 

lmdeploy/turbomind/deploy/module.py

@@ -100,14 +100,15 @@ def __init__(self, model: BaseOutputModel):
         self.inter_size = model.model_config.inter_size
         self.group_size = max(1, model.model_config.group_size)
 
-    def _export(self, inter_size: int, fmt: str, idx: int, w123, kind: str, pack_fn, apply_gs=False):
+    def _export(self, inter_size: int, fmt: str, idx: int, w123, kind: str, pack_fn, apply_gs=False, block_size=1, **kwargs):
         is_lora_a, is_lora_b = get_lora_flags(kind)
         w1, w2, w3 = map(transpose, w123)
 
-        if not is_lora_a:
+        # TODO: handle padding for block_size != 1
+        if not is_lora_a and block_size == 1:
             w1 = pad_out_dims(w1, inter_size)
             w3 = pad_out_dims(w3, inter_size)
-        if not is_lora_b:
+        if not is_lora_b and block_size == 1:
             group_size = self.group_size if apply_gs else 1
             w2 = pad_in_dims(w2, inter_size // group_size)
 
@@ -171,12 +172,15 @@ def __init__(self, model: BaseOutputModel):
         self.attn_bias = model.model_config.attn_bias
         self.qk_norm = model.model_config.qk_norm
 
-    def _reorder_and_merge(self, qkvo):
+    def _reorder_and_merge(self, qkvo, block_size):
         q, k, v, o = qkvo
         # reorder output dim for tm's rotary embedding layout
         if self.model.permute_qk:
-            q = permute_v2(q, self.head_dim)
-            k = permute_v2(k, self.head_dim)
+            if block_size == 1:
+                q = permute_v2(q, self.head_dim)
+                k = permute_v2(k, self.head_dim)
+            else:
+                assert block_size % self.head_dim == 0
         qkv = merge_qkv_v2(q, k, v, self.tp)
         # zero bias for `wo` when `w_qkv` has bias but `wo` doesn't
         if o is None and q.dim() == 1:
@@ -204,7 +208,7 @@ def _repeat(x):
 
         return (q, k, v, o)
 
-    def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs):
+    def _export(self, idx: int, qkvo, kind: str, pack_fn, block_size=1, **kwargs):
         if all(x is None for x in qkvo):
             return
         is_lora_a, is_lora_b = get_lora_flags(kind)
@@ -214,7 +218,7 @@ def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs):
             qkvo = tuple(map(transpose, qkvo))
             if self.model.repeat_kv:
                 qkvo = self._repeat_kv(qkvo, kind)
-            qkv, o = self._reorder_and_merge(qkvo)
+            qkv, o = self._reorder_and_merge(qkvo, block_size)
         self.model.save_split(pack_fn(qkv),
                               self._attn.format(idx, 'w_qkv', kind),
                               split_dim=-1,

lmdeploy/turbomind/deploy/parameter.py

@@ -13,6 +13,15 @@ def to_half(x: torch.Tensor):
     return x.to(torch.half)
 
 
+def to_float(x: torch.Tensor):
+    return x.to(torch.float)
+
+
+def to_fp8(x: torch.Tensor):
+    assert x.dtype == torch.uint8
+    return x.view(dtype=torch.float8_e4m3fn)
+
+
 def pack_u4_row(x: torch.Tensor) -> torch.Tensor:
     assert x.dtype == torch.uint8
     xs = x.view(*x.shape[:-1], -1, 8).split(1, dim=-1)
@@ -51,6 +60,15 @@ def __call__(self, f, g, i):
         f(i, g('qzeros'), 'zeros', to_half, apply_gs=True)
 
 
+class WeightScaleInv(Parameter):
+    KEYS =  '.weight_scale_inv', '.weight'
+
+    # TODO: flag any operations crossing the quant blocks as illegal
+    def __call__(self, f, g, i):
+        f(i, g('weight_scale_inv'), 'scales', to_float, block_size=128)
+        f(i, g('weight'), 'weight', identity)
+
+
 class Weight(Parameter):
     KEYS = '.weight',
 
@@ -79,6 +97,8 @@ def get_params(keys: List[str], bias=0):
         ps.append(PLora())
     if QuantWeightOnly.take(keys):
         ps.append(QuantWeightOnly())
+    if WeightScaleInv.take(keys):
+        ps.append(WeightScaleInv())
     if Weight.take(keys):
         ps.append(Weight())
     if bias and Bias.take(keys):

lmdeploy/turbomind/deploy/policy.py

@@ -46,10 +46,23 @@ def process_gptq(x: torch.Tensor, kind: str):
     return x
 
 
+def process_fp8(x: torch.Tensor, kind: str):
+    x = x.cuda()
+    if x.dtype == torch.float8_e4m3fn:
+        # some ops (e.g. torch.cat) for fp8 is not implemented in pytorch
+        return x.view(dtype=torch.uint8)
+    elif kind != 'weight_scale_inv' and x.dtype == torch.float:
+        return x.to(dtype=torch.bfloat16)
+    else:
+        return x
+
+
 def get_input_policy(model_format):
     if model_format == 'awq':
         return process_awq_gemm
     elif model_format == 'gptq':
         return process_gptq
+    elif model_format == 'fp8':
+        return process_fp8
     else:
         return to_cuda

lmdeploy/turbomind/deploy/source_model/llama.py

@@ -45,15 +45,15 @@ def filter(self, pattern: str):
 
     def tok_embeddings(self):
         """Get embeddings."""
-        return self.params.get(self.tok_embeddings_key, None)
+        return self.transform(self.params.get(self.tok_embeddings_key, None), 'weight')
 
     def norm_weight(self):
         """Get norm."""
-        return self.params.get(self.norm_weight_key, None)
+        return self.transform(self.params.get(self.norm_weight_key, None), 'weight')
 
     def output_weight(self):
         """Get output."""
-        return self.params.get(self.output_weight_key, None)
+        return self.transform(self.params.get(self.output_weight_key, None), 'weight')
 
     def _transform(self, x: torch.Tensor, kind: str):
         return self.processor(x, kind)
@@ -74,7 +74,7 @@ def attn(self, i: int, kind: str):
 
     def attn_norm(self, i: int):
         """Get attn norm for layer i."""
-        return self.params[f'{self.attn_layer_prefix}.{i}.input_layernorm.weight']
+        return self.transform(self.params[f'{self.attn_layer_prefix}.{i}.input_layernorm.weight'], 'weight')
 
     def _ffn(self, i: int, kind: str):
         """Get ffn kind for layer i."""
@@ -94,7 +94,7 @@ def ffn(self, i: int, kind: str):
 
     def ffn_norm(self, i: int):
         """Get ffn norm for layer i."""
-        return self.params[f'{self.attn_layer_prefix}.{i}.post_attention_layernorm.weight']
+        return self.transform(self.params[f'{self.attn_layer_prefix}.{i}.post_attention_layernorm.weight'], 'weight')
 
 
 @INPUT_MODELS.register_module(name='llama')

lmdeploy/turbomind/deploy/source_model/qwen.py

@@ -131,7 +131,7 @@ def moe_ffn_expert(self, e=None, i=None, kind=None):
         return (*result, )
 
     def moe_ffn_gate(self, i):
-        return self.params.get(f'model.layers.{i}.mlp.gate.weight')
+        return self.transform(self.params.get(f'model.layers.{i}.mlp.gate.weight'), 'weight')
 
     def _ffn(self, i: int, kind: str):
         """Get ffn kind for layer i."""
@@ -172,7 +172,7 @@ def qk_norm(self, i: int):
         result = []
         for x in ['q', 'k']:
             name = f'{self.attn_layer_prefix}.{i}.self_attn.{x}_norm.weight'
-            result.append(self.params.get(name))
+            result.append(self.transform(self.params.get(name), 'weight'))
         return (*result, )
 
 
@@ -193,7 +193,7 @@ def qk_norm(self, i: int):
         result = []
         for x in ['q', 'k']:
             name = f'{self.attn_layer_prefix}.{i}.self_attn.{x}_norm.weight'
-            result.append(self.params.get(name))
+            result.append(self.transform(self.params.get(name), 'weight'))
         return (*result, )
 
 

lmdeploy/turbomind/deploy/target_model/base.py

@@ -142,17 +142,19 @@ def _tofile(tensor, path):
         elif len(self.tm_params) > 0:
             tm_params = self.tm_params
             weight_type = self.model_config.weight_type
-            assert weight_type in ['float16', 'bfloat16', 'int4']
+            assert weight_type in ['float16', 'bfloat16', 'int4', 'fp8']
 
             # currently, the tensor type should in
             # [torch.float, torch.half, torch.bfloat16, torch.int32]
             torch_tensor = param.cuda().contiguous()
-            assert torch_tensor.dtype in [torch.int32, torch.float, torch.half, torch.bfloat16]
+            assert torch_tensor.dtype in [torch.int32, torch.float, torch.half, torch.bfloat16, torch.uint8]
             if torch_tensor.dtype != torch.int32:
                 if weight_type in ['float16', 'int4']:
                     torch_tensor = torch_tensor.half()
                 elif weight_type == 'bfloat16':
                     torch_tensor = torch_tensor.bfloat16()
+                elif weight_type == 'fp8':
+                    pass
                 else:
                     torch_tensor = torch_tensor.half()
             for tm_tensor in tm_params[name]:

src/turbomind/core/CMakeLists.txt

@@ -13,7 +13,7 @@ add_library(core STATIC
         tensor.cu
         module.cc)
 
-target_link_libraries(core PUBLIC cuda_utils CUDA::cudart CUDA::cuda_driver)
+target_link_libraries(core PUBLIC cuda_utils logger CUDA::cudart CUDA::cuda_driver)
 
 set_property(TARGET core PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET core PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)

src/turbomind/core/cuda_data_type.h

@@ -1,6 +1,10 @@
-#include <cublas_v2.h>
+
+
+#include <cuda.h>
 #include <cuda_runtime.h>
 
+#include <cublas_v2.h>
+
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_fp8.h>
@@ -54,6 +58,22 @@ constexpr DataType from_cuda_dtype(cudaDataType type) {
     }
 }
 
+constexpr CUtensorMapDataType to_CUtensorMap_dtype(DataType type) {
+    switch (type) {
+        case kFloat32: 
+            return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
+        case kFloat16: 
+            return CU_TENSOR_MAP_DATA_TYPE_FLOAT16;
+        case kBfloat16: 
+            return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
+        case kFloat8_e4m3:
+        case kFloat8_e5m2:
+            return CU_TENSOR_MAP_DATA_TYPE_UINT8;
+        default:
+            throw std::runtime_error("Not supported " + std::string{to_string(type)});
+    }
+}
+
 // clang-format on
 
 }  // namespace turbomind

src/turbomind/core/data_type.h

@@ -218,8 +218,8 @@ constexpr const char* to_string(DataType type) {
         case kFloat32: return "f32";
         case kFloat64: return "f64";
         case kBfloat16: return "bf16";
-        case kFloat8_e4m3: return "f8_e4m3";
-        case kFloat8_e5m2: return "f8_e5m2";
+        case kFloat8_e4m3: return "e4m3";
+        case kFloat8_e5m2: return "e5m2";
         case kUint2: return "u2";
         case kUint4: return "u4";
         case kUint6: return "u8";