Make FP8 BMM output contiguous (pytorch#370)

jiawenliu64 · facebook-github-bot · commit 4e77e95a8e41 · 2024-10-23T09:39:50.000-07:00
Summary: X-link: pytorch#3270 Pull Request resolved: facebookresearch/FBGEMM#370 Make fp8 bmm output contiguous as [silu_mul](https://fburl.com/code/sa1faq0w) requests output tensor of fp8 bmm stride(-1) to be 1. This Diff fixes the issue Reviewed By: jspark1105 Differential Revision: D64811808 fbshipit-source-id: e0f213f24fbf8bf989576371af1e2ada4cafbfb1
diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched.cu b/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched.cu
@@ -448,7 +448,7 @@ at::Tensor handle_transposition(
           BIAS_DTYPE>(
           WQ.transpose(1, 2), XQ.transpose(1, 2), w_scale, x_scale, bias, out);
     }
-    return out_.transpose(1, 2);
+    return out_.transpose(1, 2).contiguous();
   }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -448,7 +448,7 @@ at::Tensor handle_transposition(`
`448`	`448`	`BIAS_DTYPE>(`
`449`	`449`	`WQ.transpose(1, 2), XQ.transpose(1, 2), w_scale, x_scale, bias, out);`
`450`	`450`	`}`
`451`		`- return out_.transpose(1, 2);`
	`451`	`+ return out_.transpose(1, 2).contiguous();`
`452`	`452`	`}`
`453`	`453`	`}`
`454`	`454`