diff --git a/cpp/kernels/contextAttentionKernels/cuteDslFMHARunner.cpp b/cpp/kernels/contextAttentionKernels/cuteDslFMHARunner.cpp index f849741c..83bfee41 100644 --- a/cpp/kernels/contextAttentionKernels/cuteDslFMHARunner.cpp +++ b/cpp/kernels/contextAttentionKernels/cuteDslFMHARunner.cpp @@ -38,7 +38,7 @@ namespace bool isSupportedBlackwellFmha(int32_t smVersion) { - return smVersion == 100 || smVersion == 101 || smVersion == 110; + return smVersion == 100 || smVersion == 101 || smVersion == 103 || smVersion == 110; } template diff --git a/cpp/kernels/contextAttentionKernels/cuteDslFMHAV2Runner.cpp b/cpp/kernels/contextAttentionKernels/cuteDslFMHAV2Runner.cpp index 40cb3700..39a79ffd 100644 --- a/cpp/kernels/contextAttentionKernels/cuteDslFMHAV2Runner.cpp +++ b/cpp/kernels/contextAttentionKernels/cuteDslFMHAV2Runner.cpp @@ -64,7 +64,8 @@ namespace bool isFMHAV2SM(int32_t smVersion) { return smVersion == 80 || smVersion == 86 || smVersion == 87 || smVersion == 89 || smVersion == 90 - || smVersion == 100 || smVersion == 101 || smVersion == 110 || smVersion == 120 || smVersion == 121; + || smVersion == 100 || smVersion == 101 || smVersion == 103 || smVersion == 110 || smVersion == 120 + || smVersion == 121; } bool isFMHAV2PagedLlmHeadSize(int32_t headSize) diff --git a/cpp/kernels/decodeAttentionKernels/decoderXQAJitCompiler.cpp b/cpp/kernels/decodeAttentionKernels/decoderXQAJitCompiler.cpp index 9922199c..c33a8013 100644 --- a/cpp/kernels/decodeAttentionKernels/decoderXQAJitCompiler.cpp +++ b/cpp/kernels/decodeAttentionKernels/decoderXQAJitCompiler.cpp @@ -251,7 +251,7 @@ bool canCompileXQAKernel(int32_t numQHeads, int32_t numKVHeads, int32_t headSize nvinfer1::DataType dataType, nvinfer1::DataType kvDataType) noexcept { // SM versions with an XQA decode kernel at all. - constexpr std::array kALLOWED_SM_VERSIONS{80, 86, 87, 89, 90, 100, 101, 120, 121}; + constexpr std::array kALLOWED_SM_VERSIONS{80, 86, 87, 89, 90, 100, 101, 103, 120, 121}; // Of those, the ones with native FP8 converts. Listed explicitly rather // than tested as `smVersion >= 89`: a numeric threshold silently admits // every future SM added to kALLOWED_SM_VERSIONS, and whether that SM has diff --git a/kernelSrcs/build_cutedsl.py b/kernelSrcs/build_cutedsl.py index b8fd3535..953bc89d 100644 --- a/kernelSrcs/build_cutedsl.py +++ b/kernelSrcs/build_cutedsl.py @@ -264,14 +264,14 @@ class KernelVariant: KernelVariant( name="fmha_d64", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM, ), KernelVariant( name="fmha_d128", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM, ), @@ -282,7 +282,7 @@ class KernelVariant: KernelVariant( name="fmha_d64_skipsoftmax", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM + ["--skip_softmax_threshold", "1e-6"], @@ -290,7 +290,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_skipsoftmax", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM + ["--skip_softmax_threshold", "1e-6"], @@ -298,7 +298,7 @@ class KernelVariant: KernelVariant( name="fmha_d256", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM, ), @@ -307,14 +307,14 @@ class KernelVariant: KernelVariant( name="fmha_d512_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_PAGED, ), KernelVariant( name="fmha_d512_paged_bidirectional", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_PAGED + ["--window_size", "4096,-1", "--bidirectional"], @@ -322,7 +322,7 @@ class KernelVariant: KernelVariant( name="fmha_d512_sw_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_PAGED + ["--window_size", "4096,-1"], @@ -330,28 +330,28 @@ class KernelVariant: KernelVariant( name="fmha_d512_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_FP8_PAGED, ), KernelVariant( name="fmha_d512_dense_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_DENSE_PAGED, ), KernelVariant( name="fmha_d512_dense_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_DENSE_FP8_PAGED, ), KernelVariant( name="fmha_d512_sw_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,8,512", "--k_shape", "1,1024,1,512"] + _LLM_FP8_PAGED + ["--window_size", "4096,-1"], @@ -359,7 +359,7 @@ class KernelVariant: KernelVariant( name="fmha_d64_sw", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM + ["--window_size", "4096,-1"], @@ -367,7 +367,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_sw", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM + ["--window_size", "4096,-1"], @@ -375,7 +375,7 @@ class KernelVariant: KernelVariant( name="fmha_d256_sw", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM + ["--window_size", "4096,-1"], @@ -384,28 +384,28 @@ class KernelVariant: KernelVariant( name="fmha_d64_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_FP8, ), KernelVariant( name="fmha_d128_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_FP8, ), KernelVariant( name="fmha_d256_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_FP8, ), KernelVariant( name="fmha_d64_sw_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_FP8 + ["--window_size", "4096,-1"], @@ -413,7 +413,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_sw_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_FP8 + ["--window_size", "4096,-1"], @@ -421,7 +421,7 @@ class KernelVariant: KernelVariant( name="fmha_d256_sw_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_FP8 + ["--window_size", "4096,-1"], @@ -430,14 +430,14 @@ class KernelVariant: KernelVariant( name="fmha_d64_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_PAGED, ), KernelVariant( name="fmha_d128_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_PAGED, ), @@ -447,7 +447,7 @@ class KernelVariant: KernelVariant( name="fmha_d64_skipsoftmax_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_PAGED + ["--skip_softmax_threshold", "1e-6"], @@ -455,7 +455,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_skipsoftmax_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_PAGED + ["--skip_softmax_threshold", "1e-6"], @@ -463,21 +463,21 @@ class KernelVariant: KernelVariant( name="fmha_d256_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_PAGED, ), KernelVariant( name="fmha_d256_dense_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_DENSE_PAGED, ), KernelVariant( name="fmha_d64_sw_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_PAGED + ["--window_size", "4096,-1"], @@ -485,7 +485,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_sw_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_PAGED + ["--window_size", "4096,-1"], @@ -494,7 +494,7 @@ class KernelVariant: KernelVariant( name="fmha_d256_sw_paged", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_PAGED + ["--window_size", "4096,-1"], @@ -502,35 +502,35 @@ class KernelVariant: KernelVariant( name="fmha_d64_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_FP8_PAGED, ), KernelVariant( name="fmha_d128_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_FP8_PAGED, ), KernelVariant( name="fmha_d256_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_FP8_PAGED, ), KernelVariant( name="fmha_d256_dense_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_DENSE_FP8_PAGED, ), KernelVariant( name="fmha_d64_sw_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,1,64"] + _LLM_FP8_PAGED + ["--window_size", "4096,-1"], @@ -538,7 +538,7 @@ class KernelVariant: KernelVariant( name="fmha_d128_sw_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,1,128"] + _LLM_FP8_PAGED + ["--window_size", "4096,-1"], @@ -546,7 +546,7 @@ class KernelVariant: KernelVariant( name="fmha_d256_sw_paged_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,16,256", "--k_shape", "1,1024,2,256"] + _LLM_FP8_PAGED + ["--window_size", "4096,-1"], @@ -554,35 +554,35 @@ class KernelVariant: KernelVariant( name="vit_fmha_d64", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,14,64"] + _VIT, ), KernelVariant( name="vit_fmha_d72", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,72", "--k_shape", "1,1024,14,72"] + _VIT, ), KernelVariant( name="vit_fmha_d80", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,80", "--k_shape", "1,1024,14,80"] + _VIT, ), KernelVariant( name="vit_fmha_d96", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,96", "--k_shape", "1,1024,14,96"] + _VIT, ), KernelVariant( name="vit_fmha_d128", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,14,128"] + _VIT, ), @@ -594,14 +594,14 @@ class KernelVariant: KernelVariant( name="vit_fmha_d64_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,64", "--k_shape", "1,1024,14,64"] + _VIT_FP8 + ["--kv_stage", "2"], ), KernelVariant( name="vit_fmha_d80_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", # Pinned to the measured optimum: the padded K=96 smem tiles need a # deeper KV pipeline than the narrower dims (shallow staging costs @@ -611,14 +611,14 @@ class KernelVariant: KernelVariant( name="vit_fmha_d96_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,96", "--k_shape", "1,1024,14,96"] + _VIT_FP8, ), KernelVariant( name="vit_fmha_d128_fp8", group="fmha", - supported_sms=[100, 101, 110], + supported_sms=[100, 101, 103, 110], script="fmha_cutedsl_blackwell/fmha.py", script_args=["--q_shape", "1,1024,14,128", "--k_shape", "1,1024,14,128"] + _VIT_FP8 + ["--kv_stage", "3"], ), @@ -631,7 +631,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -644,7 +644,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64_small", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -655,7 +655,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d128", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "128", @@ -666,7 +666,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -677,7 +677,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64_sw", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -689,7 +689,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d128_sw", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "128", @@ -701,7 +701,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256_sw", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -713,7 +713,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -724,7 +724,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64_small_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -735,7 +735,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d128_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "128", @@ -746,7 +746,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -757,7 +757,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d512_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "512", @@ -768,7 +768,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d64_sw_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -780,7 +780,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d128_sw_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "128", @@ -792,7 +792,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256_sw_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -804,7 +804,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d512_sw_paged", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "512", @@ -816,7 +816,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256_padding", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -827,7 +827,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_vit_d64", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "64", @@ -838,7 +838,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_vit_d72", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "72", @@ -849,7 +849,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_vit_d80", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "80", @@ -860,7 +860,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_vit_d128", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "128", @@ -871,7 +871,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d256_bidirectional", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "256", @@ -889,7 +889,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d512", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "512", @@ -901,7 +901,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d512_sw", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "512", @@ -919,7 +919,7 @@ class KernelVariant: KernelVariant( name="fmha_v2_d512_bidirectional", group="fmha", - supported_sms=[80, 86, 87, 89, 90, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 90, 100, 101, 103, 110, 120, 121], script="fmha_v2_cutedsl/fmha.py", script_args=[ "--head_dim", "512", @@ -1755,7 +1755,7 @@ class KernelVariant: KernelVariant( name=f"int4_fp16_gemm_{_tile_tag}_s{_stages}_sk{_sk}", group="int4_fp16_gemm", - supported_sms=[80, 86, 87, 89, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 100, 101, 103, 110, 120, 121], script="int4_fp16_gemm_cutedsl/int4_fp16_gemm_ampere.py", script_args=[ "--mnk", "256,512,1024", @@ -1788,7 +1788,7 @@ class KernelVariant: KernelVariant( name=f"int4_fp16_gemv_m{_m}", group="int4_fp16_gemm", - supported_sms=[80, 86, 87, 89, 100, 101, 110, 120, 121], + supported_sms=[80, 86, 87, 89, 100, 101, 103, 110, 120, 121], script="int4_fp16_gemm_cutedsl/int4_fp16_gemv_ampere.py", script_args=[ "--mnk", f"{_m},512,1024", diff --git a/kernelSrcs/xqa/mha.cu b/kernelSrcs/xqa/mha.cu index 159f2fc3..7248547a 100644 --- a/kernelSrcs/xqa/mha.cu +++ b/kernelSrcs/xqa/mha.cu @@ -102,7 +102,7 @@ constexpr uint32_t preferedKHeadPartBytes = tiledQKVStagingHeadDim512 ? 32 : 64; __constant__ constexpr uint32_t cacheVTileSeqLen = tiledQKVStagingHeadDim512 ? tiledQKVStagingHeadDim512VTileSeqLen : 32; #elif __CUDA_ARCH__ == 800 || __CUDA_ARCH__ == 870 || __CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000 \ - || __CUDA_ARCH__ == 1010 || __CUDA_ARCH__ == 1100 + || __CUDA_ARCH__ == 1010 || __CUDA_ARCH__ == 1030 || __CUDA_ARCH__ == 1100 // This path tiles Q/K/V shared-memory staging to keep shared memory within the per-CTA budget. constexpr uint32_t preferedKHeadPartBytes = tiledQKVStagingHeadDim512 ? 32 : (useSmallHeadDim512Tiles ? 64 : 128); constexpr bool useWide2CtaHeadDim512VTile = twoCtaHeadDim512 && cacheElemSize < inputElemSize; diff --git a/kernelSrcs/xqa/utils.cuh b/kernelSrcs/xqa/utils.cuh index 40669db1..aa2e7513 100644 --- a/kernelSrcs/xqa/utils.cuh +++ b/kernelSrcs/xqa/utils.cuh @@ -51,7 +51,8 @@ inline constexpr bool isLowPrecCacheElem = mha::is_same_v || isFp8E4M constexpr uint32_t kMAX_SMEM_SIZE = (99u << 10); #elif __CUDA_ARCH__ == 800 || __CUDA_ARCH__ == 870 constexpr uint32_t kMAX_SMEM_SIZE = (163u << 10); -#elif __CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1010 || __CUDA_ARCH__ == 1100 +#elif __CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1010 || __CUDA_ARCH__ == 1030 \ + || __CUDA_ARCH__ == 1100 constexpr uint32_t kMAX_SMEM_SIZE = (227u << 10); #endif #endif