Fix FP8 offload alignment for cuBLAS 12.9/13.0

cosmos_framework/inference/offloading_storage.py

@@ -90,6 +90,9 @@ def __init__(self, tensors: list[torch.Tensor]):
         for tensor in physical:
             entries = self.views.setdefault(tensor.dtype, [])
             offset = entries[-1].offset + entries[-1].span if entries else 0
+            # Work around cuBLAS 12.9/13.0 requiring 16-byte-aligned FP8 scalar scales.
+            alignment = max(1, 16 // tensor.element_size())
+            offset = (offset + alignment - 1) // alignment * alignment
             span = 1 + sum((size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride(), strict=True))
             entries.append(_WeightView(tensor, offset, span if tensor.numel() else 0, tensor.shape, tensor.stride()))
         # Complete the copies before changing any module storage.

cosmos_framework/inference/offloading_test.py

@@ -174,6 +174,23 @@ def init_weights(self, buffer_device):
         assert getattr(weight, part).device.type == "cpu"
 
 
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float8_e4m3fn])
+def test_weight_binding_preserves_alignment_and_values(dtype):
+    # Scalar scales and odd-sized payloads must stay aligned after packing.
+    tensors = [torch.arange(size, dtype=torch.float32).to(dtype) for size in (1, 1, 3, 5)]
+    tensors.append(torch.arange(6, dtype=torch.float32).reshape(2, 3).t().to(dtype))
+    expected = [tensor.clone() for tensor in tensors]
+    strides = [tensor.stride() for tensor in tensors]
+    weights = _LayerWeights(tensors)
+    for device in ("cpu", "cuda", "cpu"):
+        backing = {dtype: host.to(device) for dtype, host in weights.host.items()}
+        weights.bind(backing)
+        for tensor, reference, stride in zip(tensors, expected, strides, strict=True):
+            assert tensor.data_ptr() % 16 == 0
+            assert tensor.stride() == stride
+            torch.testing.assert_close(tensor.cpu().float(), reference.float(), rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("inference_parameter", [False, True])
 def test_weight_binding_preserves_inference_status(inference_parameter):
     with torch.inference_mode(inference_parameter):