Commit a7f4e460 authored by Vijay Korthikanti's avatar Vijay Korthikanti
Browse files

build_directory argument

parent 37d77ca0

megatron/fused_kernels/__init__.py

+4 −2
Original line number Diff line number Diff line
@@ -51,6 +51,7 @@ def load_scaled_upper_triang_masked_softmax_fusion_kernel(): 
        name='scaled_upper_triang_masked_softmax_cuda', 

        sources=[srcpath / 'scaled_upper_triang_masked_softmax.cpp', 

                 srcpath / 'scaled_upper_triang_masked_softmax_cuda.cu'],

        build_directory=srcpath,

        extra_cflags=['-O3',],

        extra_cuda_cflags=['-O3',

                           '-gencode', 'arch=compute_70,code=sm_70',
@@ -75,6 +76,7 @@ def load_scaled_masked_softmax_fusion_kernel(): 
        name='scaled_masked_softmax_cuda', 

        sources=[srcpath / 'scaled_masked_softmax.cpp', 

                 srcpath / 'scaled_masked_softmax_cuda.cu'],

        build_directory=srcpath,

        extra_cflags=['-O3',],

        extra_cuda_cflags=['-O3',

                           '-gencode', 'arch=compute_70,code=sm_70',