[cherry-pick releasev0.18.0post1] Codex/release v0.18.0.post1 cherry-pick #2735/#2736/#2555 (#2813)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Samit <285365963@qq.com>

tests/diffusion/test_diffusion_step_pipeline.py

@@ -443,14 +443,42 @@ class TestWorker:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        log_probs = torch.ones(300_000, dtype=torch.float32)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor, trajectory_log_probs=log_probs),
+        )
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)
         assert packed.result.output["__tensor_shm__"] is True
+        assert isinstance(packed.result.trajectory_log_probs, dict)
+        assert packed.result.trajectory_log_probs["__tensor_shm__"] is True
 
         unpacked = unpack_diffusion_output_shm(packed)
         assert isinstance(unpacked.result.output, torch.Tensor)
+        assert isinstance(unpacked.result.trajectory_log_probs, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+        torch.testing.assert_close(unpacked.result.trajectory_log_probs, log_probs)
+
+    def test_pack_unpack_runner_output_shm_bfloat16(self):
+        tensor = torch.arange(600_000, dtype=torch.bfloat16)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor),
+        )
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+        assert packed.result.output["torch_dtype"] == str(torch.bfloat16)
+        assert packed.result.output["numpy_dtype"] == "uint16"
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        assert unpacked.result.output.dtype == torch.bfloat16
         torch.testing.assert_close(unpacked.result.output, tensor)