fix zero work cublas gemm

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

Author: vthumbe1503

tests/pytorch/test_numerics.py

@@ -3036,7 +3036,7 @@ def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -
 
     def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a):
         """Create a GroupedTensor with non-zero logical_shape but zero first_dims."""
-        buf = torch.randn(k, logical_last_dim, dtype=dtype, device=device)
+        buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device)
         if use_mxfp8:
             if is_a:
                 rowwise, columnwise = transa, not transa

transformer_engine/pytorch/csrc/type_converters.cpp

@@ -221,6 +221,8 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) {
     DType data_dtype =
         quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype;
     ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data));
+  } else if (quantizer_dtype != DType::kNumTypes) {
+    ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector<size_t>{0});
   }
 
   // Columnwise data
@@ -229,6 +231,8 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) {
     DType data_dtype =
         quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype;
     ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data));
+  } else if (quantizer_dtype != DType::kNumTypes) {
+    ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector<size_t>{0});
   }
 
   // Scale

transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py

@@ -543,7 +543,7 @@ def make_grouped_tensor(
         all_same_last = last_dims is None
 
         assert all_same_last, "Last dim must be uniform for GroupedTensor"
-        assert logical_first_dim > 0, "Logical first dim must be positive for GroupedTensor"
+        assert logical_first_dim >=0 , "Logical first dim must be non-negative for GroupedTensor"
         assert logical_last_dim > 0, "Logical last dim must be positive for GroupedTensor"
 
         # assert (