Support zero tensor in FP4 GEMM (#4460)

Summary:
Pull Request resolved: #4460

X-link: facebookresearch/FBGEMM#1520

As title. Handle corner cases when M, N, or K = 0 in E2E

Reviewed By: jianyuh

Differential Revision: D77986266

fbshipit-source-id: 517ba131ba783320162d7955cebc0f1e4b1fc2be

Author: jiawenliu64

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16.cu

@@ -34,6 +34,13 @@ at::Tensor dispatch_f4f4bf16_kernel(
       N % BLOCK_SIZE == 0 && K % BLOCK_SIZE == 0,
       "Weight dimensions N and K must be multiples of block size 16");
 
+  auto out_sizes = XQ.sizes().vec();
+  out_sizes.back() = N;
+  if (M == 0 || N == 0 || K == 0) {
+    // Use zeros instead of empty for special case where K=0.
+    return at::zeros(out_sizes, XQ.options().dtype(at::kBFloat16));
+  }
+
   // MXFP4
   if (use_mx) {
     if (M <= 128) {