From 6afc6126b2c8a12daffccd20ed2149aec93f1421 Mon Sep 17 00:00:00 2001 From: "Y.Hisaki" Date: Tue, 14 Jul 2026 17:26:25 +0900 Subject: [PATCH 1/3] Revert "Revert "fix: update DiTBlock forward method to include cross attention mask and modify Encoder to mask outputs based on encoding mask (#156)"" This reverts commit df1543a6d721a105db0be45837f7061023e2fe37. --- .../diffusion_planner/model/module/dit.py | 16 +++++++++++++--- .../diffusion_planner/model/module/encoder.py | 1 + 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/diffusion_planner/diffusion_planner/model/module/dit.py b/diffusion_planner/diffusion_planner/model/module/dit.py index eb5b2f61f..65ccf956e 100644 --- a/diffusion_planner/diffusion_planner/model/module/dit.py +++ b/diffusion_planner/diffusion_planner/model/module/dit.py @@ -34,7 +34,7 @@ def __init__(self, dim=192, heads=6, dropout=0.1, mlp_ratio=4.0): in_features=dim, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0 ) - def forward(self, x, cross_c, y, attn_mask): + def forward(self, x, cross_c, y, attn_mask, cross_attn_mask): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation( y ).chunk(6, dim=2) @@ -55,7 +55,16 @@ def forward(self, x, cross_c, y, attn_mask): modulated_x = modulate(self.norm2(x), shift_mlp, scale_mlp) x = x + gate_mlp * self.mlp1(modulated_x) - x = x + self.cross_attn(self.norm3(x), cross_c, cross_c, need_weights=False)[0] + x = ( + x + + self.cross_attn( + self.norm3(x), + cross_c, + cross_c, + key_padding_mask=cross_attn_mask, + need_weights=False, + )[0] + ) x = x + self.mlp2(self.norm4(x)) return x @@ -154,9 +163,10 @@ def forward(self, x, t, cross_c, neighbor_current_mask): ego_mask = torch.zeros((B, 1), dtype=torch.bool, device=x.device) attn_mask = torch.cat([ego_mask, neighbor_current_mask], dim=1) + cross_attn_mask = torch.all(cross_c == 0, dim=-1) for block in self.blocks: - x = block(x, cross_c, t, attn_mask) + x = block(x, cross_c, t, attn_mask, cross_attn_mask) x = self.final_layer(x, t) # (B, P, output_dim) x = x.reshape(B, P, T, D) diff --git a/diffusion_planner/diffusion_planner/model/module/encoder.py b/diffusion_planner/diffusion_planner/model/module/encoder.py index e0e541e79..bf631d869 100644 --- a/diffusion_planner/diffusion_planner/model/module/encoder.py +++ b/diffusion_planner/diffusion_planner/model/module/encoder.py @@ -305,6 +305,7 @@ def forward(self, inputs): encoding_input = encoding_input + encoding_pos_result.view(B, self.token_num, -1) encoder_outputs = self.fusion(encoding_input, encoding_mask.view(B, self.token_num)) + encoder_outputs = encoder_outputs.masked_fill(encoding_mask.view(B, self.token_num, 1), 0.0) return encoder_outputs From d500b75bcd67d8835521a6e80c14f20ba323dcb3 Mon Sep 17 00:00:00 2001 From: Yukihiro Saito Date: Fri, 24 Jul 2026 19:16:03 +0900 Subject: [PATCH 2/3] feat: force strict FP32 in ONNX export validation torch2onnx.py compares a CPU (strict FP32) torch reference against an ORT session whose CUDA EP defaults to TF32 matmuls on Ampere+, so the validation diff was dominated by TF32's 10-bit mantissa (~1e-3), not by export correctness. Disable TF32 on both sides for validation only: - pass use_tf32=0 to CUDAExecutionProvider and set NVIDIA_TF32_OVERRIDE=0 in the ORT subprocess - explicitly disable torch TF32 backends flags Deployed sessions keep the (faster) TF32 default. --- ros_scripts/torch2onnx.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/ros_scripts/torch2onnx.py b/ros_scripts/torch2onnx.py index 50b42c0fb..493e7c7bb 100644 --- a/ros_scripts/torch2onnx.py +++ b/ros_scripts/torch2onnx.py @@ -1,4 +1,5 @@ import argparse +import os import random import subprocess import sys @@ -27,6 +28,11 @@ torch.backends.cuda.enable_math_sdp(True) torch.backends.mha.set_fastpath_enabled(False) +# Validation must compare strict-FP32 numbers: TF32 matmuls (10-bit mantissa) introduce ~1e-3 +# per-layer error that makes torch-vs-ORT comparison meaningless on Ampere+ GPUs. +torch.backends.cuda.matmul.allow_tf32 = False +torch.backends.cudnn.allow_tf32 = False + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() @@ -123,6 +129,9 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda output_path = f"{tmpdir}/outputs.npz" np.savez(input_path, **np_inputs) + # use_tf32=0: ORT's CUDA EP defaults to TF32 matmuls on Ampere+, which would put ~1e-3 + # error between these outputs and the strict-FP32 torch reference. Validation-only — + # deployed sessions keep the (faster) TF32 default. script = f""" import numpy as np import onnxruntime as ort @@ -132,7 +141,7 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda sess_options.log_severity_level = 3 sess = ort.InferenceSession( "{model_path}", sess_options, - providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + providers=[("CUDAExecutionProvider", {{"use_tf32": 0}}), "CPUExecutionProvider"]) print("ORT providers:", sess.get_providers()) outputs = sess.run(None, inputs) np.savez("{output_path}", **{{f"out_{{i}}": o for i, o in enumerate(outputs)}}) @@ -142,6 +151,7 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda capture_output=True, text=True, timeout=300, + env={**os.environ, "NVIDIA_TF32_OVERRIDE": "0"}, ) if result.returncode != 0: raise RuntimeError(f"ORT subprocess failed:\n{result.stderr[-1000:]}") From f5118a0587133aec372c1be5ed46485ee3de9aba Mon Sep 17 00:00:00 2001 From: Shintaro Sakoda Date: Tue, 28 Jul 2026 20:49:31 +0900 Subject: [PATCH 3/3] Revert "feat: force strict FP32 in ONNX export validation" This reverts commit d500b75bcd67d8835521a6e80c14f20ba323dcb3. --- ros_scripts/torch2onnx.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/ros_scripts/torch2onnx.py b/ros_scripts/torch2onnx.py index 493e7c7bb..50b42c0fb 100644 --- a/ros_scripts/torch2onnx.py +++ b/ros_scripts/torch2onnx.py @@ -1,5 +1,4 @@ import argparse -import os import random import subprocess import sys @@ -28,11 +27,6 @@ torch.backends.cuda.enable_math_sdp(True) torch.backends.mha.set_fastpath_enabled(False) -# Validation must compare strict-FP32 numbers: TF32 matmuls (10-bit mantissa) introduce ~1e-3 -# per-layer error that makes torch-vs-ORT comparison meaningless on Ampere+ GPUs. -torch.backends.cuda.matmul.allow_tf32 = False -torch.backends.cudnn.allow_tf32 = False - def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() @@ -129,9 +123,6 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda output_path = f"{tmpdir}/outputs.npz" np.savez(input_path, **np_inputs) - # use_tf32=0: ORT's CUDA EP defaults to TF32 matmuls on Ampere+, which would put ~1e-3 - # error between these outputs and the strict-FP32 torch reference. Validation-only — - # deployed sessions keep the (faster) TF32 default. script = f""" import numpy as np import onnxruntime as ort @@ -141,7 +132,7 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda sess_options.log_severity_level = 3 sess = ort.InferenceSession( "{model_path}", sess_options, - providers=[("CUDAExecutionProvider", {{"use_tf32": 0}}), "CPUExecutionProvider"]) + providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) print("ORT providers:", sess.get_providers()) outputs = sess.run(None, inputs) np.savez("{output_path}", **{{f"out_{{i}}": o for i, o in enumerate(outputs)}}) @@ -151,7 +142,6 @@ def run_ort_in_subprocess(model_path: Path, np_inputs: NumpyDict) -> list[np.nda capture_output=True, text=True, timeout=300, - env={**os.environ, "NVIDIA_TF32_OVERRIDE": "0"}, ) if result.returncode != 0: raise RuntimeError(f"ORT subprocess failed:\n{result.stderr[-1000:]}")