[5321981] fix: Fix the Llama-405B hanging issue.

Signed-off-by: Yukun He <[email protected]>

Author: hyukn

cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp

@@ -82,9 +82,6 @@ nvinfer1::DimsExprs FusedLayernormPlugin::getOutputDimensions(
         {
             ret.d[di] = inputs[0].d[di];
         }
-        // Div up by 16 as the storage type has 16 FP4 values per element.
-        ret.d[ret.nbDims - 1]
-            = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16));
         return ret;
     }
 
@@ -105,11 +102,8 @@ nvinfer1::DimsExprs FusedLayernormPlugin::getOutputDimensions(
             = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128));
         ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128));
         // Hidden size dimension.
-        // Div (rounding up) by 64.
-        // 16 elements share one SF, and we need number of SFs to be multiple of 4.
-        // The output data type is already int32_t (4 SFs), so just need to div (rounding up) by 64.
         ret.d[ret.nbDims - 1]
-            = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(64));
+            = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16));
         return ret;
     }
     catch (std::exception const& e)