TrNi/efficient-cubepart3d

CubePart INT4 (HQQ W4A16, group size 128)

INT4 weight-only quantization of the diffusion transformer in Roblox's CubePart (open-vocabulary, part-controllable 3D mesh generation, arXiv 2605.28763, SIGGRAPH 2026). Companion to TrNi/efficient-cube3d, which does the same for Cube3D.

Summary

• Checkpoint size: 1.479 GB, versus 8.582 GB for the BF16 diffusion transformer (82.8% smaller).
• Peak VRAM during generation: 27.09 GiB, versus 33.71 GiB for BF16 (19.6% lower).
• Quality: median chamfer distance between BF16 and INT4 part outputs is 0.0114 per part, which is near lossless on the typical part. Errors concentrate in a few small, thin parts.
• Latency: INT4 is slower than BF16 in this release (66.5 to 66.9 s versus 34.7 to 34.8 s per generation). torch.compile has to be disabled for CubePart, so INT4 dequantization runs in eager mode.

Efficiency

Measured on one NVIDIA L40S (44 GiB usable), PyTorch 2.8.0+cu128, torchao 0.11.0, eager mode, 50 diffusion steps, guidance scale 7.5, over 10 meshes, one from each of 10 categories. Two independent runs; ranges below are run 1 to run 2.

Metric	BF16	INT4
Peak VRAM during generation	33.71 GiB (both runs)	27.09 GiB (both runs)
Latency per generation, mean of 10 meshes	34.7 to 34.8 s	66.5 to 66.9 s
Setup time (pipeline load)	149.4 to 161.8 s	153.5 to 166.1 s
Quantization time	n/a	4.1 to 4.3 s
Diffusion transformer on disk	8.582 GB	1.479 GB

Notes:

• Latency covers input_to_part_shape (diffusion sampling and mesh extraction). Surface sampling and shape encoding are excluded.
• Peak VRAM is torch.cuda.max_memory_allocated() over the full generation call. It includes the Qwen3-VL text encoder, the shape VAE and the mesh extraction step.
• INT4 setup time is the BF16 pipeline load plus quantization, because the INT4 file is applied on top of an instantiated pipeline (see Usage). A direct INT4-only load path is not implemented.
• Not tested on GPUs smaller than the L40S. With a measured 27.09 GiB peak, a 24 GB card is not expected to fit this workload as configured.
• Units: file sizes are decimal GB, as the file page shows them (exact bytes: 1,478,846,193 for the INT4 checkpoint and 8,582,452,032 for the BF16 diffusion transformer). VRAM figures are GiB (2^30 bytes) from torch.cuda.max_memory_allocated().

Quantization details

• Method: HQQ through torchao.quantization.int4_weight_only(group_size=128, use_hqq=True), W4A16.
• Quantized: nn.Linear layers inside transformer_blocks and multi_transformer_blocks of diffusion_model (about 2134M of the 2145.6M Linear parameters in diffusion_model, 99.5%).
• Not quantized: img_in, txt_in, proj_out, pos_embed, time_text_embed, every layer whose name contains norm, the shape VAE (shape_model) and the Qwen3-VL text encoder (base_model.text_encoder).
• Why HQQ and not plain round to nearest: torchao's default tinygemm packing path raised a dtype mismatch ("Expected zeros dtype float32, got bfloat16") in the versions we tried (0.10 to 0.13). HQQ avoids that path and needs no calibration data.

Quality benchmark (chamfer distance)

250 prompts across 12 categories, 828 parts in total (2 to 4 parts per prompt, part schemas written per prompt for this benchmark). Prompts come from the 310 prompt master suite in TrNi/efficient-cube3d; the categories abstract_mathematical, geometric_primitive and symmetry_topology are excluded.

Metric: average of the forward and backward mean nearest neighbour distance between 5000 seeded surface samples per part, in the units of the input meshes. All sampling is seeded, so the distance between two identical configurations is exactly 0.

• CD(A,B): same BF16 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.
• CD(C,D): same INT4 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.

	min	median	mean	max
CD(A,B)	0.0011	0.0114	0.0193	0.8206
CD(C,D)	0.0014	0.0118	0.0239	1.6261

Per category (n is the number of parts):

Category	n	CD(A,B) median	CD(A,B) mean	CD(A,B) max	CD(C,D) median	CD(C,D) mean	CD(C,D) max
animal_domestic	77	0.0100	0.0307	0.8206	0.0093	0.0249	0.7246
animal_wild	78	0.0106	0.0158	0.1654	0.0106	0.0220	0.5996
architecture	61	0.0162	0.0173	0.0541	0.0193	0.0270	0.3630
electronics	60	0.0141	0.0165	0.0835	0.0135	0.0197	0.2594
fine_detail	58	0.0142	0.0161	0.0512	0.0161	0.0245	0.0999
furniture	62	0.0131	0.0356	0.4662	0.0125	0.0153	0.0617
musical_instrument	69	0.0096	0.0177	0.1520	0.0110	0.0697	1.6261
nature_plant	61	0.0138	0.0167	0.1059	0.0134	0.0143	0.0541
original_visuals	107	0.0118	0.0203	0.3210	0.0113	0.0226	0.6126
tool_hardware	56	0.0094	0.0168	0.2186	0.0086	0.0200	0.2537
vehicle_air_water	68	0.0093	0.0130	0.1371	0.0081	0.0124	0.1168
vehicle_land	71	0.0097	0.0132	0.1608	0.0104	0.0127	0.1108

Reading the numbers: the median is about 0.011 for both pairs, and the mean sits above it because of a thin tail. The largest errors are small or thin parts next to a dominant part (for example a peacock head at 0.8206 in CD(A,B), a double bass endpin at 1.6261 and sitar frets at 1.2668 in CD(C,D)). CD(C,D) is higher than CD(A,B) on the overall statistics but lower in some categories, so it is a mild aggregate trend and not a per mesh rule.

Comparison to Cube3D INT4

The efficient-cube3d card reports a median chamfer distance of 67.7e-3 for Cube3D INT4 over 15 categories and 310 prompts. Recomputing Cube3D's own BF16 versus INT4 error with the seeded protocol above, on the 250 meshes of the 12 shared categories, gives a median of 45.5e-3. CubePart INT4 under the same protocol has a per mesh median of 13.0e-3 (median of each mesh's mean over its parts), about 3.5 times lower, and lower in every one of the 12 categories (ratio 2.2 to 5.3).

These measure different things: whole mesh drift for Cube3D and per part decomposition drift on a fixed input mesh for CubePart. Treat the ratio as indicative, not exact. The two models also fail on different meshes (per mesh rank correlation 0.27, and 2 of the 15 worst meshes in common).

Usage

You need the base weights and code from Roblox/cubepart and the cube_part package from Roblox/cube. The INT4 file replaces the diffusion transformer weights on an instantiated pipeline. The order matters: quantize first, then load the state dict (no assign=True).

import torch
from torchao.quantization import quantize_, int4_weight_only
from cube_part.pipelines import PartShapeDenoiserPipeline

pipe = PartShapeDenoiserPipeline(
    config_path="configs/shape_denoiser_multimesh.yaml",
    checkpoint_path="weights/multi_part_dit.safetensors",   # base BF16 weights
    vae_checkpoint_path="weights/vae.safetensors",
    extract_geometry_fn_name="extract_geometry_coarse_to_fine",
)

SKIP = ("img_in", "txt_in", "proj_out", "pos_embed", "time_text_embed", "norm")
def only_blocks(module, fqn):
    return isinstance(module, torch.nn.Linear) and not any(s in fqn for s in SKIP)

quantize_(pipe.system.diffusion_model,
          int4_weight_only(group_size=128, use_hqq=True),
          filter_fn=only_blocks)

sd = torch.load("multi_part_dit_int4_hqq_g128.pt", map_location="cuda", weights_only=False)
pipe.system.diffusion_model.load_state_dict(sd)

Generating parts from a mesh (this is the call used for every number above):

import numpy as np, torch, trimesh
from cube_part.pipelines import ShapeInput
from cube_part.utils.mesh import sample_surface

mesh = trimesh.load("input.glb", force="mesh")
parts = ["circular seat", "vertical legs", "footrest ring"]

np.random.seed(0)
surface = sample_surface(mesh, num_samples=128_000)
surface = torch.from_numpy(surface).to(pipe.device).unsqueeze(0).float()
latents, _ = pipe.encode_shape(surface)

part_meshes = pipe.input_to_part_shape(
    ShapeInput(prompt=[parts], latents=latents),
    guidance_scale=7.5, num_inference_steps=50, seed=0,
)
# part_meshes[i] is a (vertices, faces) pair for parts[i]

Environment notes: set TORCHDYNAMO_DISABLE=1 (CubePart's compiled path raises a Dynamo error in hijack_qwenimage_multi.py). On our setup PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True also avoided a Warp marching cubes out of memory error.

The checkpoint is a pickle and needs weights_only=False because torchao's INT4 tensors are tensor subclasses. Only load it from a source you trust.

Reload check: a fresh process that quantizes and then loads this file reproduces a live quantized model's output exactly (vertices and faces bit for bit on all 3 parts of a bar stool prompt).

Files

File	Size	Description
multi_part_dit_int4_hqq_g128.pt	1.479 GB	INT4 diffusion_model state dict, saved after in place quantization
quant_config.json	small	Quantization settings and load recipe

The shape VAE and the text encoder are unchanged, so they are not re-hosted. Use the base repository for them.

Limitations

• INT4 is slower than BF16 here (eager mode). We expect torch.compile accounts for much of the gap but have not tested compiled execution.
• Peak VRAM is 27.09 GiB; smaller GPUs were not tested.
• The base BF16 weights are needed to instantiate the pipeline. There is no INT4 only loading path yet.
• One quantization setting only (HQQ W4A16, group size 128). No other bit widths or methods were compared.
• The error of a fully INT4 pipeline (INT4 Cube3D input and INT4 CubePart against a BF16 reference) was not measured. CD(C,D) only measures what CubePart's quantization adds on an INT4 input.
• Chamfer distance is a point cloud metric. Watertightness and part reassembly were not checked.

License

Derived from Roblox/cubepart, which is released under the OpenRAIL license. This derivative carries the same license.