diff --git a/src/core/common.h b/src/core/common.h
index b06989bc..9c6c8f6e 100644
--- a/src/core/common.h
+++ b/src/core/common.h
@@ -12,9 +12,6 @@
 #include <sstream>
 #include <vector>
 
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <cuda_runtime_api.h>
 #include <json.hpp>
 
 #include "config.h"	
@@ -82,6 +79,8 @@ typedef unsigned char uchar;
 # define KRR_DEVICE_LAMBDA(...) [ =, *this ] KRR_DEVICE(__VA_ARGS__) mutable 
 
 #if !defined(__CUDA_ARCH__)
+struct uint3;
+struct dim3;
 extern const uint3 threadIdx, blockIdx;
 extern const dim3 blockDim, gridDim;
 #endif	// eliminate intellisense warnings for these kernel built-in variables
diff --git a/src/core/device/buffer.h b/src/core/device/buffer.h
index 3b564590..f522988f 100644
--- a/src/core/device/buffer.h
+++ b/src/core/device/buffer.h
@@ -6,6 +6,7 @@
 #include <thrust/execution_policy.h>
 
 #include "common.h"
+#include "device/cuda.h"
 #include "util/check.h"
 
 KRR_NAMESPACE_BEGIN
diff --git a/src/core/device/cuda.h b/src/core/device/cuda.h
index e7e1ae60..9ba171e9 100644
--- a/src/core/device/cuda.h
+++ b/src/core/device/cuda.h
@@ -1,10 +1,13 @@
 #pragma once
 
-#include <cuda.h>
 #include <map>
 #include <typeinfo>
 #include <typeindex>
 
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
 #include "common.h"
 #include "render/color.h"
 #include "device/taggedptr.h"
diff --git a/src/core/device/gpustd.h b/src/core/device/gpustd.h
index e799d74a..76c71d67 100644
--- a/src/core/device/gpustd.h
+++ b/src/core/device/gpustd.h
@@ -17,6 +17,7 @@
 #include <utility>
 
 #include "util/check.h"
+#include "cuda.h"
 #include "logger.h"
 #include "common.h"
 
diff --git a/src/core/device/scene.cpp b/src/core/device/scene.cpp
index 189c4c1c..0ee7d058 100644
--- a/src/core/device/scene.cpp
+++ b/src/core/device/scene.cpp
@@ -154,27 +154,42 @@ void RTScene::uploadSceneMediumData() {
 			HomogeneousMedium gMedium(m->sigma_a, m->sigma_s, m->Le, m->g, KRR_DEFAULT_COLORSPACE);
 			mHomogeneousMedium.push_back(gMedium);
 		} else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium)) {
-			float maxDensity{0};
-			mNanoVDBMedium.emplace_back(m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s,
-										m->g, std::move(*m->densityGrid),
-										m->temperatureGrid ? std::move(*m->temperatureGrid)
-														   : NanoVDBGrid{},
-										m->LeScale, m->temperatureScale, m->temperatureOffset, 
-										KRR_DEFAULT_COLORSPACE);
+			if (std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid)) {
+				mNanoVDBMedium.emplace_back(
+					m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s, m->g,
+					std::move(*std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid)),
+					m->temperatureGrid ? std::move(*m->temperatureGrid) : NanoVDBGrid<float>{},
+					m->LeScale, m->temperatureScale, m->temperatureOffset, KRR_DEFAULT_COLORSPACE);
+			} else if (std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid)) {
+				mNanoVDBRGBMedium.emplace_back(
+					m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s, m->g,
+					std::move(*std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid)),
+					m->temperatureGrid ? std::move(*m->temperatureGrid) : NanoVDBGrid<float>{},
+					m->LeScale, m->temperatureScale, m->temperatureOffset, KRR_DEFAULT_COLORSPACE);
+			} else {
+				Log(Error, "Unsupported heterogeneous VDB medium data type");
+				continue;
+			}
 			mNanoVDBMedium.back().initializeFromHost();
-		}
+		} else
+			Log(Error, "Unknown medium type not uploaded to device memory.");
 	}
 	mHomogeneousMediumBuffer.alloc_and_copy_from_host(mHomogeneousMedium);
 	mNanoVDBMediumBuffer.alloc_and_copy_from_host(mNanoVDBMedium);
+	mNanoVDBRGBMediumBuffer.alloc_and_copy_from_host(mNanoVDBRGBMedium);
 
 	size_t homogeneousId = 0;
 	size_t nanoVDBId	 = 0;
+	size_t nanoVDBRGBId	 = 0;
 	for (auto medium : mScene.lock()->getMedia()) {
 		if (auto m = std::dynamic_pointer_cast<HomogeneousVolume>(medium)) 
 			mMedium.push_back(Medium(&mHomogeneousMediumBuffer[homogeneousId++]));
-		else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium))
-			mMedium.push_back(Medium(&mNanoVDBMediumBuffer[nanoVDBId++]));
-		else Log(Error, "Unknown medium type not uploaded to device memory.");
+		else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium)) {
+			if (std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid))
+				mMedium.push_back(Medium(&mNanoVDBMediumBuffer[nanoVDBId++]));
+			else if (std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid))
+				mMedium.push_back(Medium(&mNanoVDBRGBMediumBuffer[nanoVDBRGBId++]));
+		}
 	}
 	mMediumBuffer.alloc_and_copy_from_host(mMedium);
 }
diff --git a/src/core/device/scene.h b/src/core/device/scene.h
index 7bcba790..e3bca0d3 100644
--- a/src/core/device/scene.h
+++ b/src/core/device/scene.h
@@ -1,11 +1,14 @@
 #pragma once
+#include <optix.h>
+#include <optix_stubs.h>
+
 #include "common.h"
 #include "mesh.h"
 #include "light.h"
 #include "camera.h"
 #include "texture.h"
-#include "device/gpustd.h"
 
+#include "device/gpustd.h"
 #include "device/buffer.h"
 #include "device/memory.h"
 #include "render/lightsampler.h"
@@ -77,8 +80,10 @@ class RTScene {
 
 	std::vector<HomogeneousMedium> mHomogeneousMedium;
 	TypedBuffer<HomogeneousMedium> mHomogeneousMediumBuffer;
-	std::vector<NanoVDBMedium> mNanoVDBMedium;
-	TypedBuffer<NanoVDBMedium> mNanoVDBMediumBuffer;
+	std::vector<NanoVDBMedium<float>> mNanoVDBMedium;
+	TypedBuffer<NanoVDBMedium<float>> mNanoVDBMediumBuffer;
+	std::vector<NanoVDBMedium<Array3f>> mNanoVDBRGBMedium;
+	TypedBuffer<NanoVDBMedium<Array3f>> mNanoVDBRGBMediumBuffer;
 
 	std::vector<Medium> mMedium;
 	TypedBuffer<Medium> mMediumBuffer;
diff --git a/src/core/medium.h b/src/core/medium.h
index c6882f5b..cdb8e86c 100644
--- a/src/core/medium.h
+++ b/src/core/medium.h
@@ -11,9 +11,9 @@ class Ray;
 class HGPhaseFunction;
 class PhaseFunctionSample;
 class HomogeneousMedium;
-class NanoVDBMedium;
 class MediumProperties;
 class MajorantIterator;
+template <typename DataType> class NanoVDBMedium;
 
 class PhaseFunction : public TaggedPointer<HGPhaseFunction> {
 public:
@@ -32,7 +32,7 @@ struct MediumProperties {
 	Spectrum Le;
 };
 
-class Medium : public TaggedPointer<HomogeneousMedium, NanoVDBMedium> {
+class Medium : public TaggedPointer<HomogeneousMedium, NanoVDBMedium<float>, NanoVDBMedium<Array3f>> {
 public:
 	using TaggedPointer::TaggedPointer;
 
diff --git a/src/core/scenegraph.cpp b/src/core/scenegraph.cpp
index ccfa23a0..8414a7f7 100644
--- a/src/core/scenegraph.cpp
+++ b/src/core/scenegraph.cpp
@@ -611,12 +611,11 @@ void HomogeneousVolume::renderUI() {
 	if (isEmissive()) ui::Text("Le: %s", Le.string().c_str());
 }
 
-void VDBVolume::renderUI() {
-	ui::Text("VDB volume"); 
+void Volume::renderUI() {
+	ui::Text("Volume Data");
 	ui::Text(("Sigma_a: " + sigma_a.string()).c_str());
 	ui::Text(("Sigma_s: " + sigma_s.string()).c_str());
 	ui::Text("g: %f", g);
-	ui::Text("Max density: %f", densityGrid->getMaxDensity());
 }
 
 void SceneGraphNode::renderUI() { 
diff --git a/src/core/scenegraph.h b/src/core/scenegraph.h
index dde71d41..3c32f5ea 100644
--- a/src/core/scenegraph.h
+++ b/src/core/scenegraph.h
@@ -131,10 +131,15 @@ class Volume : public SceneGraphLeaf {
 public:
 	using SharedPtr = std::shared_ptr<Volume>;
 	Volume()		= default;
+	Volume(RGB sigma_a, RGB sigma_s, float g) : sigma_a(sigma_a), sigma_s(sigma_s), g(g) {}
 
 	int getMediumId() const { return mediumId; }
 	void setMediumId(int id) { mediumId = id; }
+	virtual void renderUI() override;
 
+	RGB sigma_a;
+	RGB sigma_s;
+	float g;
 protected:
 	friend class SceneGraph;
 	int mediumId{-1};
@@ -145,18 +150,14 @@ class HomogeneousVolume : public Volume {
 	using SharedPtr = std::shared_ptr<HomogeneousVolume>;
 
 	HomogeneousVolume(RGB sigma_a, RGB sigma_s, float g, RGB Le = RGB::Zero()) :
-		sigma_a(sigma_a), sigma_s(sigma_s), g(g), Le(Le) {}
+		Volume(sigma_a, sigma_s, g), Le(Le) {}
 
 	virtual std::shared_ptr<SceneGraphLeaf> clone() override;
 	virtual void renderUI() override;
 
 	bool isEmissive() const { return !Le.isZero(); }
 
-	RGB sigma_a;
-	RGB sigma_s;
 	RGB Le;
-	float g;
-
 private:
 	friend class SceneGraph;
 };
@@ -165,24 +166,20 @@ class VDBVolume : public Volume {
 public:
 	using SharedPtr = std::shared_ptr<VDBVolume>;
 
-	VDBVolume(RGB sigma_a, RGB sigma_s, float g, 
-		NanoVDBGrid::SharedPtr density, NanoVDBGrid::SharedPtr temperature = nullptr, 
-		float LeScale = 1, float temperatureScale = 1, float temperatureOffset = 0) :
-		sigma_a(sigma_a), sigma_s(sigma_s), g(g), densityGrid(density), temperatureGrid(temperature), 
+	VDBVolume(RGB sigma_a, RGB sigma_s, float g, NanoVDBGridBase::SharedPtr density,
+			  NanoVDBGrid<float>::SharedPtr temperature = nullptr, float LeScale = 1,
+			  float temperatureScale = 1, float temperatureOffset = 0) :
+		Volume(sigma_a, sigma_s, g), densityGrid(density), temperatureGrid(temperature),
 		LeScale(LeScale), temperatureScale(temperatureScale), temperatureOffset(temperatureOffset) {}
 
-	virtual std::shared_ptr<SceneGraphLeaf> clone() override;
+	virtual SceneGraphLeaf::SharedPtr clone() override;
 	virtual AABB getLocalBoundingBox() const override { return densityGrid->getBounds(); }
-	virtual void renderUI() override;
 
-	RGB sigma_a;
-	RGB sigma_s;
-	float g;
 	float LeScale;
 	float temperatureScale;
 	float temperatureOffset;
-	NanoVDBGrid::SharedPtr densityGrid;
-	NanoVDBGrid::SharedPtr temperatureGrid;
+	NanoVDBGridBase::SharedPtr densityGrid;
+	NanoVDBGrid<float>::SharedPtr temperatureGrid;
 
 protected:
 	friend class SceneGraph;
diff --git a/src/core/vulkan/cuvk.h b/src/core/vulkan/cuvk.h
index e0b2c0ca..4902841a 100644
--- a/src/core/vulkan/cuvk.h
+++ b/src/core/vulkan/cuvk.h
@@ -16,6 +16,7 @@
 
 #include <common.h>
 #include <logger.h>
+#include <device/cuda.h>
 #include <util/check.h>
 
 KRR_NAMESPACE_BEGIN
diff --git a/src/render/color.h b/src/render/color.h
index 7e685156..46cf9fcb 100644
--- a/src/render/color.h
+++ b/src/render/color.h
@@ -1,6 +1,5 @@
 #pragma once
 
-#include "common.h"
 #include "krrmath/math.h"
 #include "util/math_utils.h"
 
diff --git a/src/render/media.cpp b/src/render/media.cpp
index e4663880..bd8257a9 100644
--- a/src/render/media.cpp
+++ b/src/render/media.cpp
@@ -8,8 +8,8 @@
 KRR_NAMESPACE_BEGIN
 /* Note that the function qualifier (e.g. inline) should be consistent between declaration and definition. */
 
-void initializeMajorantGrid(MajorantGrid& majorantGrid,
-							nanovdb::FloatGrid *floatGrid) {
+template <typename GridType>
+void initializeMajorantGrid(MajorantGrid &majorantGrid, GridType *densityGrid) {
 	auto res = majorantGrid.res;
 	cudaDeviceSynchronize();
 	// [TODO] This device memory is not properly freed
@@ -30,14 +30,14 @@ void initializeMajorantGrid(MajorantGrid& majorantGrid,
 														float(z + 1) / majorantGrid.res.z())));
 
 			// Compute corresponding NanoVDB index-space bounds in floating-point.
-			nanovdb::Vec3R i0 = floatGrid->worldToIndexF(
+			nanovdb::Vec3R i0 = densityGrid->worldToIndexF(
 				nanovdb::Vec3R(wb.min().x(), wb.min().y(), wb.min().z()));
-			nanovdb::Vec3R i1 = floatGrid->worldToIndexF(
+			nanovdb::Vec3R i1 = densityGrid->worldToIndexF(
 				nanovdb::Vec3R(wb.max().x(), wb.max().y(), wb.max().z()));
 
 			// Now find integer index-space bounds, accounting for both
 			// filtering and the overall index bounding box.
-			auto bbox	= floatGrid->indexBBox();
+			auto bbox	= densityGrid->indexBBox();
 			float delta = 1.f; // Filter slop
 			int nx0		= max(int(i0[0] - delta), bbox.min()[0]);
 			int nx1		= min(int(i1[0] + delta), bbox.max()[0]);
@@ -47,33 +47,30 @@ void initializeMajorantGrid(MajorantGrid& majorantGrid,
 			int nz1		= min(int(i1[2] + delta), bbox.max()[2]);
 
 			float maxValue = 0;
-			auto accessor  = floatGrid->getAccessor();
+			auto accessor  = densityGrid->getAccessor();
 
 			for (int nz = nz0; nz <= nz1; ++nz)
 				for (int ny = ny0; ny <= ny1; ++ny)
 					for (int nx = nx0; nx <= nx1; ++nx)
-						maxValue = max(maxValue, accessor.getValue({nx, ny, nz}));
-
+						if constexpr (std::is_same_v<GridType, nanovdb::FloatGrid>) {
+							maxValue = max(maxValue, accessor.getValue({nx, ny, nz}));
+						} else if constexpr (std::is_same_v<GridType, nanovdb::Vec3fGrid>) {
+							auto value = accessor.getValue({nx, ny, nz});
+							RGB color  = {value[0], value[1], value[2]};
+							RGBUnboundedSpectrum spectrum(color, *KRR_DEFAULT_COLORSPACE_GPU);
+							maxValue = max(maxValue, spectrum.maxValue());
+						} else {
+							static_assert(false, "Unsupported grid type!");
+						}
 			majorantGrid.set(x, y, z, maxValue);
 		}, 0);
 }
 
-NanoVDBMedium::NanoVDBMedium(const Affine3f &transform, RGB sigma_a, RGB sigma_s, float g,
-							 NanoVDBGrid density, NanoVDBGrid temperature, float LeScale,
-							 float temperatureScale, float temperatureOffset, 
-							 const RGBColorSpace *colorSpace) :
-	transform(transform), phase(g), sigma_a(sigma_a), sigma_s(sigma_s), densityGrid(std::move(density)), 
-	temperatureGrid(std::move(temperature)), LeScale(LeScale), temperatureScale(temperatureScale), 
-	temperatureOffset(temperatureOffset), colorSpace(colorSpace) {
-	inverseTransform = transform.inverse();
-	const Vector3f majorantGridRes{64, 64, 64};
-	majorantGrid	 = MajorantGrid(densityGrid.getBounds(), majorantGridRes);
-}
-
-void NanoVDBMedium::initializeFromHost() {
-	densityGrid.toDevice();
-	initializeMajorantGrid(majorantGrid, densityGrid.getFloatGrid());
-}
+// explicit instantiation of enable data types of vdb grids
+template void initializeMajorantGrid<nanovdb::FloatGrid>(MajorantGrid &majorantGrid,
+														 nanovdb::FloatGrid *densityGrid);
+template void initializeMajorantGrid<nanovdb::Vec3fGrid>(MajorantGrid &majorantGrid,
+														 nanovdb::Vec3fGrid *densityGrid);
 
 KRR_HOST_DEVICE PhaseFunctionSample HGPhaseFunction::sample(const Vector3f &wo,
 														 const Vector2f &u) const {
diff --git a/src/render/media.h b/src/render/media.h
index 8e330ab2..6528b467 100644
--- a/src/render/media.h
+++ b/src/render/media.h
@@ -33,6 +33,9 @@ class MajorantGrid : public Grid<float, 3> {
 	AABB3f bounds;
 };
 
+template <typename GridType>
+void initializeMajorantGrid(MajorantGrid &majorantGrid, GridType *floatGrid);
+
 class MajorantIterator {
 public:
 	MajorantIterator() = default;
@@ -135,10 +138,11 @@ class HomogeneousMedium {
 	}
 };
 
+template <typename DataType>
 class NanoVDBMedium {
 public:
-	NanoVDBMedium(const Affine3f &transform, RGB sigma_a, RGB sigma_s, float g, NanoVDBGrid density,
-				  NanoVDBGrid temperature, float LeScale, float temperatureScale, float temperatureOffset,
+	NanoVDBMedium(const Affine3f &transform, RGB sigma_a, RGB sigma_s, float g, NanoVDBGrid<DataType> density,
+				  NanoVDBGrid<float> temperature, float LeScale, float temperatureScale, float temperatureOffset,
 				  const RGBColorSpace *colorSpace = KRR_DEFAULT_COLORSPACE);
 
 	KRR_HOST void initializeFromHost();
@@ -147,9 +151,12 @@ class NanoVDBMedium {
 	
 	KRR_CALLABLE MediumProperties samplePoint(Vector3f p, const SampledWavelengths &lambda) const { 
 		p = inverseTransform * p;
-		float d = densityGrid.getDensity(p);
-		return {Spectrum::fromRGB(sigma_a, SpectrumType::RGBUnbounded, lambda, *colorSpace) * d,
-				Spectrum::fromRGB(sigma_s, SpectrumType::RGBUnbounded, lambda, *colorSpace) * d,
+		DataType d = densityGrid.getDensity(p);
+		Spectrum color;
+		if constexpr (std::is_same_v<DataType, float>) color = Spectrum::Constant(d);
+		else color = Spectrum::fromRGB(d, SpectrumType::RGBUnbounded, lambda, *colorSpace);
+		return {Spectrum::fromRGB(sigma_a, SpectrumType::RGBUnbounded, lambda, *colorSpace) * color,
+				Spectrum::fromRGB(sigma_s, SpectrumType::RGBUnbounded, lambda, *colorSpace) * color,
 				&phase, Le(p, lambda)};
 	}
 
@@ -166,8 +173,8 @@ class NanoVDBMedium {
 			&majorantGrid};
 	}
 
-	NanoVDBGrid densityGrid;
-	NanoVDBGrid temperatureGrid;
+	NanoVDBGrid<DataType> densityGrid;
+	NanoVDBGrid<float> temperatureGrid;
 	MajorantGrid majorantGrid;
 	Affine3f transform, inverseTransform;
 	HGPhaseFunction phase;
@@ -189,6 +196,25 @@ class NanoVDBMedium {
 	}
 };
 
+template <typename DataType>
+NanoVDBMedium<DataType>::NanoVDBMedium(const Affine3f &transform, RGB sigma_a, RGB sigma_s, float g,
+							 NanoVDBGrid<DataType> density, NanoVDBGrid<float> temperature, float LeScale,
+							 float temperatureScale, float temperatureOffset, 
+							 const RGBColorSpace *colorSpace) :
+	transform(transform), phase(g), sigma_a(sigma_a), sigma_s(sigma_s), densityGrid(std::move(density)), 
+	temperatureGrid(std::move(temperature)), LeScale(LeScale), temperatureScale(temperatureScale), 
+	temperatureOffset(temperatureOffset), colorSpace(colorSpace) {
+	inverseTransform = transform.inverse();
+	const Vector3f majorantGridRes{64, 64, 64};
+	majorantGrid	 = MajorantGrid(densityGrid.getBounds(), majorantGridRes);
+}
+
+template <typename DataType> 
+void NanoVDBMedium<DataType>::initializeFromHost() {
+	densityGrid.toDevice();
+	initializeMajorantGrid(majorantGrid, densityGrid.getNativeGrid());
+}
+
 /* Put these definitions here since the optix kernel will need them... */
 /* Definitions of inline functions should be put into header files. */
 
diff --git a/src/scene/openvdb.cpp b/src/scene/openvdb.cpp
index 99be5c92..8d9e9d35 100644
--- a/src/scene/openvdb.cpp
+++ b/src/scene/openvdb.cpp
@@ -25,8 +25,8 @@ bool OpenVDBImporter::import(const fs::path filepath, Scene::SharedPtr scene,
 	float temperatureOffset = params.value<float>("offset_temperature", 0);
 	float LeScale			= params.value<float>("scale_le", 1);
 
-	NanoVDBGrid::SharedPtr densityGrid	   = loadNanoVDB(filepath, key_density);
-	NanoVDBGrid::SharedPtr temperatureGrid = loadNanoVDB(filepath, key_temperature);
+	auto densityGrid	 = loadNanoVDB(filepath, key_density);
+	auto temperatureGrid = std::dynamic_pointer_cast<NanoVDBGrid<float>>(loadNanoVDB(filepath, key_temperature));
 	auto volume = std::make_shared<VDBVolume>(sigma_a, sigma_s, g, densityGrid, temperatureGrid,
 											  LeScale, temperaetureScale, temperatureOffset);
 
diff --git a/src/util/check.h b/src/util/check.h
index ededf717..01afd77f 100644
--- a/src/util/check.h
+++ b/src/util/check.h
@@ -1,31 +1,10 @@
 #pragma once
 
 #include "common.h"
-#include <cuda.h>
-#include <optix.h>
-#include <optix_stubs.h>
 #include "logger.h"
 
 KRR_NAMESPACE_BEGIN
 
-#define OPTIX_CHECK(call)                                                                          \
-	do {                                                                                           \
-		OptixResult res = call;                                                                    \
-		if (res != OPTIX_SUCCESS) {                                                                \
-			Log(Fatal, "Optix call (%s) failed with code %d (line %d)\n", #call, res,              \
-					__LINE__);                                                                     \
-		}                                                                                          \
-	} while (false)
-
-#define OPTIX_CHECK_WITH_LOG(EXPR, LOG)                                                            \
-	do {                                                                                           \
-		OptixResult res = EXPR;                                                                    \
-		if (res != OPTIX_SUCCESS) {                                                                \
-			Log(Fatal, "Optix call " #EXPR " failed with code %d: \"%s\"\nLogs: %s",               \
-					int(res), optixGetErrorString(res), LOG);                                      \
-		}                                                                                          \
-	} while (false) /* eat semicolon */
-
 #define CUDA_CHECK(call)                                                                           \
 	do {                                                                                           \
 		cudaError_t rc = call;                                                                     \
@@ -53,6 +32,25 @@ KRR_NAMESPACE_BEGIN
 		}                                                                                          \
 	} while (0)
 
+#define OPTIX_CHECK(call)                                                                          \
+	do {                                                                                           \
+		OptixResult res = call;                                                                    \
+		if (res != OPTIX_SUCCESS) {                                                                \
+			Log(Fatal, "Optix call (%s) failed with code %d (line %d)\n", #call, res,              \
+					__LINE__);                                                                     \
+		}                                                                                          \
+	} while (false)
+
+#define OPTIX_CHECK_WITH_LOG(EXPR, LOG)                                                            \
+	do {                                                                                           \
+		OptixResult res = EXPR;                                                                    \
+		if (res != OPTIX_SUCCESS) {                                                                \
+			Log(Fatal, "Optix call " #EXPR " failed with code %d: \"%s\"\nLogs: %s",               \
+					int(res), optixGetErrorString(res), LOG);                                      \
+		}                                                                                          \
+	} while (false) /* eat semicolon */
+
+
 #define CHECK_LOG(EXPR, LOG, ...)                                                                  \
 	do {                                                                                           \
 		if (!(EXPR)) {                                                                             \
diff --git a/src/util/volume.cpp b/src/util/volume.cpp
index 06dd4029..716cf40f 100644
--- a/src/util/volume.cpp
+++ b/src/util/volume.cpp
@@ -17,7 +17,7 @@
 
 KRR_NAMESPACE_BEGIN
 
-NanoVDBGrid::SharedPtr loadNanoVDB(std::filesystem::path path, std::string key) {
+NanoVDBGridBase::SharedPtr loadNanoVDB(std::filesystem::path path, string key) {
 	openvdb::initialize();
 	Log(Info, "Loading openvdb file from %s", path.string().c_str());
 	openvdb::io::File file(path.generic_string());
@@ -32,14 +32,24 @@ NanoVDBGrid::SharedPtr loadNanoVDB(std::filesystem::path path, std::string key)
 		return nullptr;
 	}
 
-	auto grid							  = openvdb::gridPtrCast<openvdb::FloatGrid>(baseGrid);
-	auto transform						  = grid->transform();
-	auto handle							  = nanovdb::openToNanoVDB<nanovdb::CudaDeviceBuffer>(grid);
+	auto transform						  = baseGrid->transform();
+	auto handle = nanovdb::openToNanoVDB<nanovdb::CudaDeviceBuffer>(baseGrid);
 	const nanovdb::GridMetaData *metadata = handle.gridMetaData();
-	if (metadata->gridType() != nanovdb::GridType::Float) Log(Fatal, "only support float grid!");
-	float minValue, maxValue;
-	grid->evalMinMax(minValue, maxValue);
-	return std::make_shared<NanoVDBGrid>(std::move(handle), maxValue);
+	if (metadata->gridType() == nanovdb::GridType::Float) {
+		float minValue, maxValue;
+		auto grid = openvdb::gridPtrCast<openvdb::FloatGrid>(baseGrid);
+		grid->evalMinMax(minValue, maxValue);
+		return std::make_shared<NanoVDBGrid<float>>(std::move(handle), maxValue);
+	} else if (metadata->gridType() == nanovdb::GridType::Vec3f) {
+		auto grid = openvdb::gridPtrCast<openvdb::Vec3fGrid>(baseGrid);
+		openvdb::Vec3f minValue, maxValue;
+		grid->evalMinMax(minValue, maxValue);
+		return std::make_shared<NanoVDBGrid<Array3f>>(std::move(handle), 
+			Array3f{maxValue[0], maxValue[1], maxValue[2]});
+	} else {
+		Log(Fatal, "Unsupported data type for openvdb grid!");
+		return nullptr;
+	}
 }
 
 KRR_NAMESPACE_END
\ No newline at end of file
diff --git a/src/util/volume.h b/src/util/volume.h
index bd022ff0..81a9227e 100644
--- a/src/util/volume.h
+++ b/src/util/volume.h
@@ -9,19 +9,61 @@
 #include <nanovdb/util/SampleFromVoxels.h>
 #include <nanovdb/util/CudaDeviceBuffer.h>
 
+#include "common.h"
 #include "krrmath/math.h"
 
 KRR_NAMESPACE_BEGIN
 
-class NanoVDBGrid {
+/* We did not use virtual device functions for VDB grids on purpose.
+	For cuda device routines to access virtual functions, the class (i.e. virtual table)
+	must also be created on device. Otherwise, each call to a virtual function will 
+	attempt to access host memory (which causes the illegal memory access error). */
+
+class NanoVDBGridBase {
+public:
+	using SharedPtr = std::shared_ptr<NanoVDBGridBase>;
+	NanoVDBGridBase()  = default;
+	~NanoVDBGridBase() = default;
+
+	NanoVDBGridBase(nanovdb::GridHandle<nanovdb::CudaDeviceBuffer> &&density) :
+		densityHandle(std::move(density)) {}
+
+	KRR_HOST virtual void toDevice() { densityHandle.deviceUpload(); }
+
+	KRR_CALLABLE AABB3f getBounds() const { return bounds; }
+
+protected:
+	AABB3f bounds;
+	nanovdb::GridHandle<nanovdb::CudaDeviceBuffer> densityHandle{};
+};
+
+template <typename DataType>
+class NanoVDBGridBaseImpl: public NanoVDBGridBase {
+public:
+	NanoVDBGridBaseImpl() = default;
+	~NanoVDBGridBaseImpl() = default;
+
+	KRR_HOST NanoVDBGridBaseImpl(nanovdb::GridHandle<nanovdb::CudaDeviceBuffer> &&density,
+		DataType maxDensity) : NanoVDBGridBase(std::move(density)), maxDensity(maxDensity) {}
+
+	KRR_CALLABLE DataType getMaxDensity() const { return maxDensity; }
+
+protected:
+	DataType maxDensity{0};
+};
+
+template <typename DataType> 
+class NanoVDBGrid : public NanoVDBGridBaseImpl<DataType> {};
+
+template<> class NanoVDBGrid<float> : public NanoVDBGridBaseImpl<float> {
 public:
-	using SharedPtr  = std::shared_ptr<NanoVDBGrid>;
+	using SharedPtr  = std::shared_ptr<NanoVDBGrid<float>>;
 	using VDBSampler = nanovdb::SampleFromVoxels<nanovdb::FloatGrid::TreeType, 1, false>;
 
 	NanoVDBGrid() = default;
 
 	NanoVDBGrid(nanovdb::GridHandle<nanovdb::CudaDeviceBuffer>&& density, float maxDensity) :
-		densityHandle(std::move(density)), maxDensity(maxDensity) {
+		NanoVDBGridBaseImpl(std::move(density), maxDensity) {
 		densityGrid						   = densityHandle.grid<float>();
 		nanovdb::BBox<nanovdb::Vec3R> bbox = densityGrid->worldBBox();
 		bounds = AABB3f{Vector3f{bbox.min()[0], bbox.min()[1], bbox.min()[2]},
@@ -35,29 +77,58 @@ class NanoVDBGrid {
 		maxDensity	  = other.maxDensity;
 	}
 
-	void toDevice() { densityHandle.deviceUpload(); }
-
 	KRR_CALLABLE operator bool() const { return densityGrid != nullptr; }
 
-	KRR_CALLABLE AABB3f getBounds() const { return bounds; }
-
 	KRR_CALLABLE float getDensity(const Vector3f &p) const {
 		nanovdb::Vec3<float> pIndex =
 			densityGrid->worldToIndexF(nanovdb::Vec3<float>(p.x(), p.y(), p.z()));
 		return VDBSampler(densityGrid->tree())(pIndex);
 	}
 
-	KRR_CALLABLE float getMaxDensity() const { return maxDensity; }
+	KRR_CALLABLE nanovdb::FloatGrid *getNativeGrid() const { return densityGrid; }
 
-	KRR_CALLABLE nanovdb::FloatGrid *getFloatGrid() const { return densityGrid; }
-
-private:
-	AABB3f bounds;
-	nanovdb::GridHandle<nanovdb::CudaDeviceBuffer> densityHandle{};
+protected:
 	nanovdb::FloatGrid *densityGrid{nullptr};
-	float maxDensity{0};
 };
 
-NanoVDBGrid::SharedPtr loadNanoVDB(std::filesystem::path path, std::string key);
+template<> class NanoVDBGrid<Array3f> : public NanoVDBGridBaseImpl<Array3f> {
+public:
+	using NativeType = nanovdb::Vec3f;
+	using SharedPtr  = std::shared_ptr<NanoVDBGrid<Array3f>>;
+	using VDBSampler = nanovdb::SampleFromVoxels<nanovdb::Vec3fGrid::TreeType, 1, false>;
+
+	NanoVDBGrid() = default;
+
+	NanoVDBGrid(nanovdb::GridHandle<nanovdb::CudaDeviceBuffer> &&density, Array3f maxDensity) :
+		NanoVDBGridBaseImpl(std::move(density), maxDensity) {
+		densityGrid						   = densityHandle.grid<NativeType>();
+		nanovdb::BBox<nanovdb::Vec3R> bbox = densityGrid->worldBBox();
+		bounds = AABB3f{Vector3f{bbox.min()[0], bbox.min()[1], bbox.min()[2]},
+						Vector3f{bbox.max()[0], bbox.max()[1], bbox.max()[2]}};
+	}
+
+	NanoVDBGrid(NanoVDBGrid&& other) noexcept {
+		std::swap(densityHandle, other.densityHandle);
+		densityGrid	  = densityHandle.grid<NativeType>();
+		bounds		  = other.bounds;
+		maxDensity	  = other.maxDensity;
+	}
+
+	KRR_CALLABLE operator bool() const { return densityGrid != nullptr; }
+
+	KRR_CALLABLE Array3f getDensity(const Vector3f &p) const {
+		nanovdb::Vec3<float> pIndex =
+			densityGrid->worldToIndexF(nanovdb::Vec3<float>(p.x(), p.y(), p.z()));
+		auto color = VDBSampler(densityGrid->tree())(pIndex);
+		return {color[0], color[1], color[2]};
+	}
+
+	KRR_CALLABLE nanovdb::Vec3fGrid *getNativeGrid() const { return densityGrid; }
+
+protected:
+	nanovdb::Vec3fGrid *densityGrid{nullptr};
+};
+
+NanoVDBGridBase::SharedPtr loadNanoVDB(std::filesystem::path path, std::string key);
 
 KRR_NAMESPACE_END
\ No newline at end of file