Skip to content

GitLab

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

build_directory argument

parent 37d77ca0
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
megatron/fused_kernels/__init__.py
View file @ a7f4e460
......@@ -50,7 +50,8 @@ def load_scaled_upper_triang_masked_softmax_fusion_kernel():
scaled_upper_triang_masked_softmax_cuda = cpp_extension.load(
name='scaled_upper_triang_masked_softmax_cuda',
sources=[srcpath / 'scaled_upper_triang_masked_softmax.cpp',
srcpath / 'scaled_upper_triang_masked_softmax_cuda.cu'],
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',
......@@ -74,7 +75,8 @@ def load_scaled_masked_softmax_fusion_kernel():
scaled_upper_triang_masked_softmax_cuda = cpp_extension.load(
name='scaled_masked_softmax_cuda',
sources=[srcpath / 'scaled_masked_softmax.cpp',
srcpath / 'scaled_masked_softmax_cuda.cu'],
srcpath / 'scaled_masked_softmax_cuda.cu'],
build_directory=srcpath,
extra_cflags=['-O3',],
extra_cuda_cflags=['-O3',
'-gencode', 'arch=compute_70,code=sm_70',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment