add FLAGS instead max_partition_size

PaddlePaddle · Dec 12, 2024 · ca4e474 · ca4e474
1 parent f445a7a
commit ca4e474
Show file tree

Hide file tree

Showing 24 changed files with 26 additions and 116 deletions.
diff --git a/csrc/gpu/append_attention.cu b/csrc/gpu/append_attention.cu
@@ -61,8 +61,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
     const float out_linear_in_scale,
     const int encoder_block_shape_q,
     const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
@@ -209,8 +207,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
             quant_max_bound,
             quant_min_bound,
             out_linear_in_scale,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             false,
@@ -248,8 +244,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
             quant_max_bound,
             quant_min_bound,
             out_linear_in_scale,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             false,
@@ -292,8 +286,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
           quant_max_bound,
           quant_min_bound,
           out_linear_in_scale,
-          max_partition_size,
-          encoder_max_partition_size,
           speculate_max_draft_token_num,
           causal,
           false,
@@ -440,8 +432,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
             quant_max_bound,
             quant_min_bound,
             out_linear_in_scale,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             !speculate_decoder,
@@ -479,8 +469,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
             quant_max_bound,
             quant_min_bound,
             out_linear_in_scale,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             !speculate_decoder,
@@ -524,8 +512,6 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
           quant_max_bound,
           quant_min_bound,
           out_linear_in_scale,
-          max_partition_size,
-          encoder_max_partition_size,
           speculate_max_draft_token_num,
           causal,
           !speculate_decoder,
@@ -585,8 +571,6 @@ std::vector<paddle::Tensor> AppendAttention(
     const float out_linear_in_scale,
     const int encoder_block_shape_q,
     const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
@@ -650,8 +634,6 @@ std::vector<paddle::Tensor> AppendAttention(
           out_linear_in_scale,
           encoder_block_shape_q,
           decoder_block_shape_q,
-          max_partition_size,
-          encoder_max_partition_size,
           speculate_max_draft_token_num,
           causal,
           speculate_decoder);
@@ -700,8 +682,6 @@ std::vector<paddle::Tensor> AppendAttention(
           out_linear_in_scale,
           encoder_block_shape_q,
           decoder_block_shape_q,
-          max_partition_size,
-          encoder_max_partition_size,
           speculate_max_draft_token_num,
           causal,
           speculate_decoder);
@@ -751,8 +731,6 @@ std::vector<paddle::Tensor> AppendAttention(
             out_linear_in_scale,
             encoder_block_shape_q,
             decoder_block_shape_q,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             speculate_decoder);
@@ -800,8 +778,6 @@ std::vector<paddle::Tensor> AppendAttention(
             out_linear_in_scale,
             encoder_block_shape_q,
             decoder_block_shape_q,
-            max_partition_size,
-            encoder_max_partition_size,
             speculate_max_draft_token_num,
             causal,
             speculate_decoder);
@@ -905,8 +881,6 @@ std::vector<paddle::DataType> AppendAttentionInferDtype(
     const float out_linear_in_scale,
     const int encoder_block_shape_q,
     const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
@@ -985,8 +959,6 @@ PD_BUILD_OP(append_attention)
             "out_linear_in_scale: float",
             "encoder_block_shape_q: int",
             "decoder_block_shape_q: int",
-            "max_partition_size: int",
-            "encoder_max_partition_size: int",
             "speculate_max_draft_token_num: int",
             "causal: bool",
             "speculate_decoder: bool"})

diff --git a/csrc/gpu/append_attn/append_attention_c16_impl.cuh b/csrc/gpu/append_attn/append_attention_c16_impl.cuh
@@ -786,8 +786,6 @@ void MultiQueryAppendAttention(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool is_decoder,
     cudaStream_t &stream,
@@ -839,9 +837,9 @@ void MultiQueryAppendAttention(
     int sm_count;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
 
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
     const int num_chunks = div_up(max_dec_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
@@ -1058,9 +1056,9 @@ void MultiQueryAppendAttention(
     int sm_count;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
 
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
     const int num_chunks = div_up(max_dec_len, chunk_size);
 
@@ -1301,8 +1299,6 @@ void CascadeAppendAttentionC16Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -1363,8 +1359,6 @@ void CascadeAppendAttentionC16Kernel(
                                 quant_max_bound,
                                 quant_min_bound,
                                 in_scale,
-                                max_partition_size,
-                                encoder_max_partition_size,
                                 speculate_max_draft_token_num,
                                 is_decoder,
                                 stream,

diff --git a/csrc/gpu/append_attn/append_attention_c4_impl.cuh b/csrc/gpu/append_attn/append_attention_c4_impl.cuh
@@ -973,8 +973,6 @@ void MultiQueryAppendC4Attention(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool is_decoder,
     cudaStream_t &stream,
@@ -1036,9 +1034,9 @@ void MultiQueryAppendC4Attention(
     const float ratio = static_cast<float>(num_blocks_need) /
                         static_cast<float>(num_blocks_per_wave);
 
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
     const int num_chunks = div_up(max_dec_len, chunk_size);
 
@@ -1282,9 +1280,9 @@ void MultiQueryAppendC4Attention(
                         static_cast<float>(num_blocks_per_wave);
 
 
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    static uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
     const int num_chunks = div_up(max_dec_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
@@ -1538,8 +1536,6 @@ void CascadeAppendAttentionC4Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -1604,8 +1600,6 @@ void CascadeAppendAttentionC4Kernel(
                                 quant_max_bound,
                                 quant_min_bound,
                                 in_scale,
-                                max_partition_size,
-                                encoder_max_partition_size,
                                 speculate_max_draft_token_num,
                                 is_decoder,
                                 stream,

diff --git a/csrc/gpu/append_attn/append_attention_c8_impl.cuh b/csrc/gpu/append_attn/append_attention_c8_impl.cuh
@@ -860,8 +860,6 @@ void MultiQueryAppendC8Attention(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool is_decoder,
     cudaStream_t &stream,
@@ -914,9 +912,9 @@ void MultiQueryAppendC8Attention(
     const int dev_id = 0;
     int sm_count;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
     const int num_chunks = div_up(max_dec_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
@@ -1136,9 +1134,9 @@ void MultiQueryAppendC8Attention(
     const int dev_id = 0;
     int sm_count;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
-    uint32_t chunk_size = static_cast<uint32_t>(max_partition_size);
+    uint32_t chunk_size = get_max_partition_size();
     if (!is_decoder) {
-      chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
+      chunk_size = get_encoder_max_partition_size();
     }
 
     const int num_chunks = div_up(max_dec_len, chunk_size);
@@ -1377,8 +1375,6 @@ void CascadeAppendAttentionC8Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -1441,8 +1437,6 @@ void CascadeAppendAttentionC8Kernel(
                                 quant_max_bound,
                                 quant_min_bound,
                                 in_scale,
-                                max_partition_size,
-                                encoder_max_partition_size,
                                 speculate_max_draft_token_num,
                                 is_decoder,
                                 stream,

diff --git a/csrc/gpu/append_attn/append_attention_kernel.h b/csrc/gpu/append_attn/append_attention_kernel.h
@@ -52,8 +52,6 @@ void CascadeAppendAttentionC16Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -97,8 +95,6 @@ void CascadeAppendAttentionC8Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -142,8 +138,6 @@ void CascadeAppendAttentionC4Kernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -188,8 +182,6 @@ void CascadeAppendAttentionKernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,
@@ -223,8 +215,6 @@ void CascadeAppendAttentionKernel(
                                              quant_max_bound,
                                              quant_min_bound,
                                              in_scale,
-                                             max_partition_size,
-                                             encoder_max_partition_size,
                                              speculate_max_draft_token_num,
                                              causal,
                                              is_decoder,
@@ -258,8 +248,6 @@ void CascadeAppendAttentionKernel(
                                             quant_max_bound,
                                             quant_min_bound,
                                             in_scale,
-                                            max_partition_size,
-                                            encoder_max_partition_size,
                                             speculate_max_draft_token_num,
                                             causal,
                                             is_decoder,
@@ -293,8 +281,6 @@ void CascadeAppendAttentionKernel(
                                             quant_max_bound,
                                             quant_min_bound,
                                             in_scale,
-                                            max_partition_size,
-                                            encoder_max_partition_size,
                                             speculate_max_draft_token_num,
                                             causal,
                                             is_decoder,
@@ -307,3 +293,17 @@ void CascadeAppendAttentionKernel(
         "cache_int4_zp]");
   }
 }
+
+inline uint32_t get_max_partition_size() {
+    static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
+    static const uint32_t max_partition_size =
+            max_partition_size_env == nullptr ? 128 : std::stoul(std::string(max_partition_size_env));
+    return max_partition_size;
+}
+
+inline uint32_t get_encoder_max_partition_size() {
+    static const char* encoder_max_partition_size_env = std::getenv("FLAGS_cascade_encoder_attention_max_partition_size");
+    static const uint32_t encoder_max_partition_size =
+            encoder_max_partition_size_env == nullptr ? 32768 : std::stoul(std::string(encoder_max_partition_size_env));
+    return encoder_max_partition_size;
+}
diff --git a/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_bfloat16_kernel.cu b/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_bfloat16_kernel.cu
@@ -49,8 +49,6 @@ template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, paddle::bfloat16
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,

diff --git a/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_fp8_kernel.cu b/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_fp8_kernel.cu
@@ -48,8 +48,6 @@ template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, paddle::float8_e
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,

diff --git a/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_int8_kernel.cu b/csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_int8_kernel.cu
@@ -48,8 +48,6 @@ template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, int8_t>(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,

diff --git a/csrc/gpu/append_attn/template_instantiation/append_attention_c16_float16_float16_kernel.cu b/csrc/gpu/append_attn/template_instantiation/append_attention_c16_float16_float16_kernel.cu
@@ -48,8 +48,6 @@ template void CascadeAppendAttentionC16Kernel<paddle::float16, paddle::float16>(
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
     const int speculate_max_draft_token_num,
     const bool causal,
     const bool is_decoder,