PR Category

Execute Infrastructure

PR Types

Bug fixes

Description

修复paddle.incubate.softmax_mask_fuse使其能支持big tensor。
对softmax_mask_fuse的前向和反向的kernel都进行了修改；

• 反向kernel只存在访存越界的问题，修复后对整体性能影响较小。
• 前向kernel存在访存越界和launch kernel时的非法配置问题，原有的实现为了兼顾更好的性能对block的分配比较固定，不能应对规模较大的情况。修复后的kernel更加通用，但是性能比不上原有的kernel。
• 为了兼顾性能和能够应对比较大规模的Tensor，当前采用两种cuda kernel并存的方案。当Tensor数据量比较小时仍采用原来的实现（命名为SoftmaxMaskFuseV1GPUKernel）以保证较好的性能。数据量较大时采用更通用的SoftmaxMaskFuseV2GPUKernel

存在问题1: 访存越界

成因:
前向和反向kernel中访问全局显存时使用的是int类型的offset，导致访问超过2^31以上的offset时出现溢出，导致出现非法地址。
修复方案：
将关键的索引部分替换为int64类型

存在问题2:launch kernel时的非法配置

成因：
launch kernel时 gird的各个维度的配置为dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);
由于gridDim.y/z的最大值为65535，因此当attn_heads或者batches超过此值时会报cuda error 9
修复方案：
计算总共需要的block数，然后只使用gridDim中的x维度（最大为2^32 - 1），在cuda kernel中重新计算mask Tensor对应的内存索引。
Grid中block的配置修复前后的对比

//修复前launch kernel时Grid中block的配置
//attn_head和batches都不能超过65536
 dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);
//修复后launch kernel时Grid中block的配置
//total_blocks = (query_seq_len / batches_per_block)*attn_heads*batches
int64_t total_blocks = batch_count / batches_per_block;
dim3 blocks(total_blocks);

CUDA Kernel中重新计算每个block的任务分配

//修复前的一些关键代码
  int data_first_idx =
      (blockDim.y *
           (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) +
       threadIdx.y) *
      kLocalBatchSize;
//(blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z)) 是在计算当前block 是所有Grid中所有block中的第几个block
  int mask_fist_idx =
      (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) *
      kLocalBatchSize;
//修复后对上述代码的一些等价转换
  u_int32_t blockInGrid = blockIdx.x;
  u_int32_t indexInMaskDim0 = blockInGrid / (attn_heads * query_seqs);
  u_int32_t indexInMaskDim2 = blockInGrid % (query_seqs);
  int64_t data_first_idx =
      (blockDim.y * static_cast<int64_t>(blockInGrid) + threadIdx.y) *
      kLocalBatchSize;
  int64_t mask_fist_idx =
      (blockDim.y * (indexInMaskDim2 +
                     static_cast<int64_t>(query_seqs) * indexInMaskDim0) +
       threadIdx.y) *
      kLocalBatchSize;
  //修复前后的一些对应关系
  // query_seqs <-> gridDim.x
  // attn_heads <-> gridDim.y
  // indexInMaskDim0 <-> blockIdx.z
  // indexInMaskDim2 <-> blockIdx.x

当前支持的Tensor的最大numel为 batches * attn_heads * query_seq_len * key_seq_len = （2^32 - 1 ）*8192 ，batches * attn_heads * query_seq_len最大值为（2^32 - 1 ），以float16 为例，此时需要的显存约为65536GB，可以应对绝大多数场景。
性能测试

通过nsight system和nsight compute进行深入的profiling发现。修改后的cuda kernel 每个thread所需的register由30涨到了37，进而导致warp的occupancy 下降，从而影响了性能。但是修复后的cuda kernel更加通用可以应对更大规模的Tensor。为了保证性能和支持Big Tensor，对Tensor中数据规模较小时，仍采用原有的cuda kernel（SoftmaxMaskFuseV1GPUKernel）以保证性能。当Tensor中数据规模较大时采用修复后的SoftmaxMaskFuseV2GPUKernel。

正确性测试

• PaddleAPITest回测
  
• 添加单测case用以验证mask和x 在shape不同时的计算结果的正确性

pcard-67164