Merge branch 'develop_v2.5' of http://10.16.6.30/dcutoolkit/deeplearing/TransformerEngine

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -267,12 +267,12 @@ def general_gemm(
             )[0]
             if out_dtype is torch.bfloat16:
                 if accumulate:
-                    out = channelwise_dequantize_transA_add(dy_scales, x_scales, dw_int32, out)
+                    channelwise_dequantize_transA_add(dy_scales, x_scales, dw_int32, out)
                 else:
                     out = channelwise_dequantize_transA(dy_scales, x_scales, dw_int32)
             else:
                 if accumulate:
-                    out = channelwise_dequantize_transA_float_add(dy_scales, x_scales, dw_int32, out)     
+                    channelwise_dequantize_transA_float_add(dy_scales, x_scales, dw_int32, out)     
                 else:
                     out = channelwise_dequantize_transA_float(dy_scales, x_scales, dw_int32)     
             return out, None, None, None
@@ -572,14 +572,14 @@ def general_grouped_gemm(
             if out_dtype is torch.bfloat16:
                 if accumulate:
                     for i in num_gemms:
-                        out[i] = channelwise_dequantize_transA_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])
+                        channelwise_dequantize_transA_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])
                 else:
                     for i in num_gemms:
                         out[i] = channelwise_dequantize_transA(scales_dout_list[i], scales_x_list[i], dw_int32[i])
             else:
                 if accumulate:
                     for i in num_gemms:
-                        out[i] = channelwise_dequantize_transA_float_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])     
+                        channelwise_dequantize_transA_float_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])     
                 else:
                     for i in num_gemms:
                         out[i] = channelwise_dequantize_transA_float(scales_dout_list[i], scales_x_list[i], dw_int32[i])    

transformer_engine/pytorch/triton/per_token_group_quant.py

@@ -331,12 +331,12 @@ def channelwise_dequantize_transA_float(A, B, C):
 @torch.compile(mode="max-autotune-no-cudagraphs")
 def channelwise_dequantize_transA_add(A, B, C, D):
     out_scales = A.T * B
-    return (out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16) + D
+    D.add_((out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16))
 
 @torch.compile(mode="max-autotune-no-cudagraphs")
 def channelwise_dequantize_transA_float_add(A, B, C, D):
     out_scales = A.T * B
-    return out_scales * C.to(dtype=torch.float32) + D
+    D.add_(out_scales * C.to(dtype=torch.float32))
 
 @torch.compile(mode="max-autotune-no-cudagraphs")
 def channelwise_dequantize_transB(A, B, C):