Zhwang/llama

.vscode/settings.json

@@ -67,6 +67,33 @@
         "unordered_set": "cpp",
         "future": "cpp",
         "cfenv": "cpp",
-        "typeindex": "cpp"
+        "typeindex": "cpp",
+        "__bit_reference": "cpp",
+        "__bits": "cpp",
+        "__config": "cpp",
+        "__debug": "cpp",
+        "__errc": "cpp",
+        "__hash_table": "cpp",
+        "__locale": "cpp",
+        "__mutex_base": "cpp",
+        "__node_handle": "cpp",
+        "__split_buffer": "cpp",
+        "__threading_support": "cpp",
+        "__tree": "cpp",
+        "__tuple": "cpp",
+        "__verbose_abort": "cpp",
+        "bit": "cpp",
+        "ios": "cpp",
+        "locale": "cpp",
+        "queue": "cpp",
+        "stack": "cpp",
+        "variant": "cpp",
+        "__nullptr": "cpp",
+        "__string": "cpp",
+        "compare": "cpp",
+        "concepts": "cpp",
+        "filesystem": "cpp",
+        "__memory": "cpp",
+        "version": "cpp"
     }
-}
+}

CMakeLists.txt

@@ -16,6 +16,21 @@ project(FasterTransformer LANGUAGES CXX CUDA)
 
 find_package(CUDA 10.2 REQUIRED)
 
+include(FetchContent)
+
+FetchContent_Declare(
+  repo-cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG        cc85b64cf676c45f98a17e3a47c0aafcf817f088
+)
+
+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)
+
 if(${CUDA_VERSION_MAJOR} VERSION_GREATER_EQUAL "11")
   add_definitions("-DENABLE_BF16")
   message("CUDA_VERSION ${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR} is greater or equal than 11.0, enable -DENABLE_BF16 flag")
@@ -346,6 +361,9 @@ add_library(transformer-shared SHARED
   $<TARGET_OBJECTS:GptNeoXTritonBackend>
   $<TARGET_OBJECTS:GptNeoXWeight>
   $<TARGET_OBJECTS:LinearAdapterLayer>
+  $<TARGET_OBJECTS:llama_fmha>
+  $<TARGET_OBJECTS:Llama>
+  $<TARGET_OBJECTS:LlamaTritonBackend>
   $<TARGET_OBJECTS:OnlineBeamSearchLayer>
   $<TARGET_OBJECTS:ParallelGpt>
   $<TARGET_OBJECTS:ParallelGptContextDecoder>
@@ -466,6 +484,7 @@ set_target_properties(transformer-shared PROPERTIES POSITION_INDEPENDENT_CODE ON
 set_target_properties(transformer-shared PROPERTIES CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 set_target_properties(transformer-shared PROPERTIES LINKER_LANGUAGE CXX)
 target_link_libraries(transformer-shared PUBLIC -lcudart -lcublas -lcublasLt -lcurand)
+target_link_libraries(transformer-shared PUBLIC stdc++fs)
 
 include(GNUInstallDirs)
 set(INSTALL_CONFIGDIR ${CMAKE_INSTALL_LIBDIR}/cmake/FasterTransformer)

src/fastertransformer/kernels/decoder_masked_multihead_attention.h

@@ -117,6 +117,9 @@ struct Multihead_attention_params_base {
     const float* qkv_scale_out       = nullptr;
     const float* attention_out_scale = nullptr;
     int          int8_mode           = 0;
+
+    float        attention_k_scale   = 0.f;
+    float        attention_v_scale   = 0.f;
 };
 
 template<typename T, bool CROSS_ATTENTION>
@@ -135,6 +138,12 @@ struct Multihead_attention_params: public Multihead_attention_params_base<T> {
 
     // required in case of masked attention with different length
     const int* length_per_sample = nullptr;
+
+    T**    k_cache_per_sample         = nullptr;
+    T**    v_cache_per_sample         = nullptr;
+    size_t kv_cache_per_sample_offset = 0;
+    bool   k_cache_interleaved        = true;
+    int    num_kv_heads               = 0;
 };
 
 template<typename T>
@@ -152,6 +161,12 @@ struct Multihead_attention_params<T, true>: public Multihead_attention_params_ba
 
     // required in case of masked attention with different length
     const int* length_per_sample = nullptr;
+
+    T**    k_cache_per_sample         = nullptr;
+    T**    v_cache_per_sample         = nullptr;
+    size_t kv_cache_per_sample_offset = 0;
+    bool   k_cache_interleaved        = true;
+    int    num_kv_heads               = 0;
 };
 
 template<class T>

src/fastertransformer/kernels/unfused_attention_kernels.cu

@@ -1556,6 +1556,42 @@ void invokeAddFusedQKVBiasTranspose(T*                               q_buf,
     }
 }
 
+template<typename T>
+void invokeAddFusedQKVBiasTranspose(T*                               q_buf,
+                                    T*                               k_buf,
+                                    T*                               v_buf,
+                                    PrefixPromptBatchWeightsParam<T> param,
+                                    T*                               QKV,
+                                    const T*                         qkv_bias,
+                                    const int*                       padding_offset,
+                                    const int*                       history_length,
+                                    const int                        batch_size,
+                                    const int                        seq_len,
+                                    const int                        token_num,
+                                    const int                        head_num,
+                                    const int                        kv_head_num,
+                                    const int                        size_per_head,
+                                    const int                        rotary_embedding_dim,
+                                    const int                        neox_rotary_style,
+                                    const float*                     scale,
+                                    const int                        int8_mode,
+                                    cudaStream_t                     stream)
+{
+    FT_CHECK(rotary_embedding_dim);
+    FT_CHECK_WITH_INFO(int8_mode != 2, "w8a8 not yet implemented with prefix prompt");  // TODO(mseznec)
+    // To implement rotary embeddings, each thread processes two QKV elems:
+    dim3   block((size_per_head / Vec_t<T>::size + 31) / 32 * 32);
+    dim3   grid(token_num + batch_size * param.max_prefix_prompt_length, head_num);
+    size_t smem_size = neox_rotary_style ? 2 * rotary_embedding_dim * sizeof(T) : 0;
+    // NOTE: add offset for rotary embedding
+    if (param.max_prefix_prompt_length == 0) {
+        FUSED_QKV_BIAS_TRANSPOSE_LAUNCH(T, false);
+    }
+    else {
+        FUSED_QKV_BIAS_TRANSPOSE_LAUNCH(T, true);
+    }
+}
+
 #define INSTANTIATEADDFUSEDQKVBIASTRANSPOSE(T)                                                                         \
     template void invokeAddFusedQKVBiasTranspose(T*                               q_buf,                               \
                                                  T*                               k_buf,                               \
@@ -1573,6 +1609,25 @@ void invokeAddFusedQKVBiasTranspose(T*                               q_buf,
                                                  const int                        neox_rotary_style,                   \
                                                  const float*                     scale,                               \
                                                  const int                        int8_mode,                           \
+                                                 cudaStream_t                     stream);                             \
+    template void invokeAddFusedQKVBiasTranspose(T*                               q_buf,                               \
+                                                 T*                               k_buf,                               \
+                                                 T*                               v_buf,                               \
+                                                 PrefixPromptBatchWeightsParam<T> param,                               \
+                                                 T*                               QKV,                                 \
+                                                 const T*                         qkv_bias,                            \
+                                                 const int*                       padding_offset,                      \
+                                                 const int*                       history_length,                      \
+                                                 const int                        batch_size,                          \
+                                                 const int                        seq_len,                             \
+                                                 const int                        token_num,                           \
+                                                 const int                        head_num,                            \
+                                                 const int                        kv_head_num,                         \
+                                                 const int                        size_per_head,                       \
+                                                 const int                        rotary_embedding_dim,                \
+                                                 const int                        neox_rotary_style,                   \
+                                                 const float*                     scale,                               \
+                                                 const int                        int8_mode,                           \
                                                  cudaStream_t                     stream)
 INSTANTIATEADDFUSEDQKVBIASTRANSPOSE(float);
 INSTANTIATEADDFUSEDQKVBIASTRANSPOSE(half);

src/fastertransformer/kernels/unfused_attention_kernels.h

@@ -113,6 +113,27 @@ struct PrefixPromptBatchWeightsParam {
     const size_t prefix_prompt_layer_offset_per_seq = 0;
 };
 
+template<typename T>
+void invokeAddFusedQKVBiasTranspose(T*                               q_buf,
+                                    T*                               k_buf,
+                                    T*                               v_buf,
+                                    PrefixPromptBatchWeightsParam<T> param,
+                                    T*                               QKV,
+                                    const T*                         qkv_bias,
+                                    const int*                       padding_offset,
+                                    const int*                       history_length,
+                                    const int                        batch_size,
+                                    const int                        seq_len,
+                                    const int                        token_num,
+                                    const int                        head_num,
+                                    const int                        kv_head_num,
+                                    const int                        size_per_head,
+                                    const int                        rotary_embedding_dim,
+                                    const int                        neox_rotary_style,
+                                    const float*                     scale,
+                                    const int                        int8_mode,
+                                    cudaStream_t                     stream);
+
 template<typename T>
 void invokeAddFusedQKVBiasTranspose(T*           q_buf,
                                     T*           k_buf,

src/fastertransformer/layers/DynamicDecodeLayer.h

@@ -56,6 +56,14 @@ class DynamicDecodeLayer: public BaseLayer {
     int* h_pinned_finished_sum_ = nullptr;
 
 public:
+    curandState_t* topk_curandstate_buf()
+    {
+        return static_cast<BaseSamplingLayer<T>*>(topk_decode_)->curandstate_buf();
+    }
+    curandState_t* topp_curandstate_buf()
+    {
+        return static_cast<BaseSamplingLayer<T>*>(topp_decode_)->curandstate_buf();
+    }
     DynamicDecodeLayer(size_t           vocab_size,
                        size_t           vocab_size_padded,
                        int              end_id,

src/fastertransformer/layers/sampling_layers/BaseSamplingLayer.h

@@ -59,6 +59,11 @@ class BaseSamplingLayer: public DynamicDecodeBaseLayer {
     virtual void allocateBuffer(size_t batch_size, Tensor top_k, Tensor top_p);
 
 public:
+    curandState_t* curandstate_buf()
+    {
+        return curandstate_buf_;
+    }
+
     BaseSamplingLayer(size_t             max_batch_size,
                       size_t             vocab_size,
                       size_t             vocab_size_padded,

src/fastertransformer/models/CMakeLists.txt

@@ -19,6 +19,7 @@ add_subdirectory(bert_fp8)
 endif()
 add_subdirectory(deberta)
 add_subdirectory(decoder)
+add_subdirectory(llama)
 add_subdirectory(longformer)
 add_subdirectory(decoding)
 add_subdirectory(xlnet)

src/fastertransformer/models/llama/Barrier.h

@@ -0,0 +1,37 @@
+// Copyright (c) OpenMMLab. All rights reserved.
+
+#pragma once
+
+#include "src/fastertransformer/utils/logger.h"
+#include <pthread.h>
+
+namespace fastertransformer {
+
+class Barrier {
+public:
+    Barrier(unsigned count)
+    {
+        FT_LOG_INFO("Barrier(%d)", (int)count);
+        pthread_barrier_init(&barrier_, nullptr, count);
+    }
+
+    Barrier(const Barrier&) = delete;
+    Barrier& operator=(const Barrier&) = delete;
+    Barrier(Barrier&&) noexcept        = delete;
+    Barrier& operator=(Barrier&&) noexcept = delete;
+
+    void wait()
+    {
+        pthread_barrier_wait(&barrier_);
+    }
+
+    ~Barrier()
+    {
+        pthread_barrier_destroy(&barrier_);
+    }
+
+private:
+    pthread_barrier_t barrier_{};
+};
+
+}  // namespace fastertransformer

src/fastertransformer/models/llama/CMakeLists.txt

@@ -0,0 +1,43 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+cmake_minimum_required(VERSION 3.8)
+
+add_subdirectory(fused_multi_head_attention)
+
+add_library(Llama STATIC
+        LlamaV2.cc
+        LlamaBatch.cc
+        LlamaCacheManager.cc
+        LlamaContextDecoder.cc
+        LlamaContextAttentionLayer.cc
+        LlamaDecoderSelfAttentionLayer.cc
+        LlamaDecoder.cc
+        LlamaWeight.cc
+        LlamaDecoderLayerWeight.cc
+        LlamaFfnLayer.cc
+        llama_kernels.cu
+        llama_decoder_kernels.cu
+        llama_utils.cu)
+set_property(TARGET Llama PROPERTY POSITION_INDEPENDENT_CODE  ON)
+set_property(TARGET Llama PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS  ON)
+target_link_libraries(Llama PUBLIC -lcudart
+        cublasMMWrapper
+        DynamicDecodeLayer
+        activation_kernels
+        decoder_masked_multihead_attention
+        bert_preprocess_kernels
+        decoding_kernels
+        unfused_attention_kernels
+        custom_ar_kernels
+        custom_ar_comm
+        gpt_kernels
+        tensor
+        memory_utils
+        nccl_utils
+        cuda_utils
+        logger
+        stdc++fs
+        llama_fmha)
+
+add_executable(llama_gemm llama_gemm.cc)
+target_link_libraries(llama_gemm PUBLIC -lcudart gpt_gemm_func memory_utils cuda_utils logger)