Allow MHA plugin to run on SM_86 as well

rajeevsrao · rajeevsrao · commit 1fe5532c0c78 · 2021-04-12T14:56:06.000-07:00
Signed-off-by: Rajeev Rao &lt;rajeevrao@nvidia.com&gt;
diff --git a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp
@@ -62,7 +62,7 @@ QKVToContextInterleavedPlugin::QKVToContextInterleavedPlugin(
     mSM = getSMVersion();
     // variable sequence length is only supported with the fused MHA kernels
     // we should not override mS!
-    assert((mSM == kSM_AMPERE || mSM == kSM_TURING || mSM == kSM_XAVIER)
+    assert((mSM == kSM_AMPERE_100 || mSM == kSM_AMPERE_10X || mSM == kSM_TURING || mSM == kSM_XAVIER)
         && "requesting maxSeqlen not compatible with GPU arch");
     // the layout changes: SxB will be a combined \sum_i s_i and hdim will be the 2nd dimension instead of the third
     mXmmaKernel = getXMMAKernelsV2(DATA_TYPE_INT8, mSM);
diff --git a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.h b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.h
@@ -29,7 +29,8 @@ namespace bert
 {
 static constexpr int32_t kSM_XAVIER = 72;
 static constexpr int32_t kSM_TURING = 75;
-static constexpr int32_t kSM_AMPERE = 80;
+static constexpr int32_t kSM_AMPERE_100 = 80;
+static constexpr int32_t kSM_AMPERE_10X = 86;
 
 class QKVToContextInterleavedPlugin : public nvinfer1::IPluginV2DynamicExt
 {

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,8 @@ namespace bert`
`29`	`29`	`{`
`30`	`30`	`static constexpr int32_t kSM_XAVIER = 72;`
`31`	`31`	`static constexpr int32_t kSM_TURING = 75;`
`32`		`-static constexpr int32_t kSM_AMPERE = 80;`
	`32`	`+static constexpr int32_t kSM_AMPERE_100 = 80;`
	`33`	`+static constexpr int32_t kSM_AMPERE_10X = 86;`
`33`	`34`
`34`	`35`	`class QKVToContextInterleavedPlugin : public nvinfer1::IPluginV2DynamicExt`
`35`	`36`	`{`