fix(neuron): Wan and HunyuanVideo 1.0 DiTs get the scheduler timestep in fp32

difflet/backends/neuron/hunyuan_video/transformer.py

@@ -189,7 +189,9 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]:
         pooled_projections, guidance)`` at the static shape: the latents ``[B, C, F', H/8,
         W/8]``, the per-sample timestep, the right-padded Llama embeddings ``[B, text_seq_len,
         text_embed_dim]`` with an int64 mask, the CLIP pooled projection and the guidance
-        (scale x 1000). Exactly the per-step shapes the orchestrator sends, so the warm-up
+        (scale x 1000); the timestep in fp32, as diffusers passes it, the guidance in ``dtype``
+        (diffusers builds it in the transformer's dtype). Exactly the per-step shapes the
+        orchestrator sends, so the warm-up
         compiles the only block graphs the run needs. Under CP or SP every rank receives the
         whole latent and the module scatters it.
         """
@@ -200,7 +202,7 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]:
         return (
             torch.zeros(batch, int(cfg.in_channels), cfg.latent_frames, cfg.latent_height,
                         cfg.latent_width, dtype=dt),
-            torch.ones(batch, dtype=dt),
+            torch.ones(batch, dtype=torch.float32),
             torch.zeros(batch, text, int(cfg.text_embed_dim), dtype=dt),
             mask,
             torch.zeros(batch, int(cfg.pooled_projection_dim), dtype=dt),

difflet/backends/neuron/wan/transformer.py

@@ -118,7 +118,8 @@ def teacache_mod_input(self, hidden_states: torch.Tensor, timestep: torch.Tensor
                                         timestep.to(self.device), self._teacache_text)
 
     def get_example_inputs(self) -> tuple[torch.Tensor, ...]:
-        """``(latents, timestep, text embeddings)`` at the static shape, in ``dtype``.
+        """``(latents, timestep, text embeddings)`` at the static shape, in ``dtype`` except the
+        timestep, which the DiT takes in fp32 as diffusers passes it.
 
         These are exactly the per-step shapes the orchestrator sends (two batch-1 forwards per
         CFG step), so the warm-up compiles the only block graph the run needs. With CFG
@@ -136,7 +137,7 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]:
                 cfg.latent_width,
                 dtype=self.dtype,
             ),
-            torch.zeros(batch, dtype=self.dtype),
+            torch.zeros(batch, dtype=torch.float32),
             torch.zeros(batch, int(cfg.text_seq_len), int(cfg.text_dim), dtype=self.dtype),
         )
 

difflet/models/hunyuan_video/neuron_application.py

@@ -210,6 +210,10 @@ class _OnDeviceHunyuanVideoTransformer(OnDeviceDiT):
     probe's per-step signal: the orchestrator's fused branch.
     """
 
+    # As diffusers passes it: the scheduler's timestep unrounded. bf16 holds only multiples of 4
+    # between 512 and 1024, which shifts the sinusoidal timestep embedding at every step.
+    timestep_dtype = torch.float32
+
     def __init__(self, module, config: NeuronHunyuanVideoConfig, dtype, device,
                  probe: _HunyuanVideoTeaCacheProbe | None = None) -> None:
         super().__init__(module, config, dtype, device)
@@ -223,7 +227,12 @@ def __call__(self, bundle: HunyuanVideoDiTInputBundle) -> torch.Tensor:
         if not isinstance(bundle, HunyuanVideoDiTInputBundle):
             raise TypeError(f"expected a HunyuanVideoDiTInputBundle, got {type(bundle).__name__}")
         _check_dit_shapes(self.config, bundle)
-        return self._forward(bundle.as_model_inputs())
+        # The timestep goes to the device in fp32, not through _forward's cast to the model
+        # dtype; the other inputs follow it already on the device, keeping the module's order.
+        hidden_states, timestep, *rest = bundle.as_model_inputs()
+        timestep = timestep.to(self.timestep_dtype).contiguous().to(self.device)
+        return self._forward((hidden_states,),
+                             extra=(timestep, *(self._device_input(t) for t in rest)))
 
     def _device_input(self, tensor: torch.Tensor) -> torch.Tensor:
         if not tensor.is_floating_point():
@@ -233,8 +242,8 @@ def _device_input(self, tensor: torch.Tensor) -> torch.Tensor:
     def teacache_delta(self, bundle: HunyuanVideoDiTInputBundle) -> torch.Tensor:
         if self.probe is None:
             raise RuntimeError("no TeaCache probe: adaptive TeaCache (teacache_speedup) is off")
-        return self.probe.teacache_delta(
-            *(t.contiguous() for t in (bundle.hidden_states, bundle.timestep,
+        return self.probe.teacache_delta(  # the probe keeps its model-dtype timestep
+            *(t.contiguous() for t in (bundle.hidden_states, bundle.timestep.to(self.dtype),
                                        bundle.pooled_projections, bundle.guidance)))
 
 
@@ -399,7 +408,7 @@ def phase_seconds(self) -> dict[str, float]:
 
     def dit_input_contract(self) -> dict[str, dict[str, Any]]:
         """The Trainium contract's six tensors at the static shape (``dit_input_shapes``)."""
-        dtypes = {"encoder_attention_mask": torch.int64}
+        dtypes = {"encoder_attention_mask": torch.int64, "timestep": torch.float32}
         return {name: {"shape": shape, "dtype": dtypes.get(name, self.dtype)}
                 for name, shape in dit_input_shapes(self.config).items()}
 

difflet/models/hunyuan_video/pipeline.py

@@ -233,11 +233,11 @@ def _denoise(
         prev_mod_handle: torch.Tensor | None = None
         for step_index, timestep in enumerate(timesteps):
             model_dtype = _component_dtype(self.transformer, self.dtype)
-            timestep_batch = _batch_timestep(
+            timestep_batch = _batch_timestep(  # the DiT's own timestep dtype, if it declares one
                 timestep,
                 batch_size=latents.shape[0],
                 device=latents.device,
-                dtype=model_dtype,
+                dtype=getattr(self.transformer, "timestep_dtype", model_dtype),
             )
             model_bundle = HunyuanVideoDiTInputBundle(
                 hidden_states=latents.to(dtype=model_dtype),

difflet/models/wan/neuron_application.py

@@ -224,13 +224,21 @@ class _OnDeviceTransformer(OnDeviceDiT):
     expert's adapter reports to the high-noise one's, so one hook times every DiT forward).
     """
 
+    # As diffusers passes it: the scheduler's timestep unrounded. bf16 holds only multiples of 4
+    # between 512 and 1024, which shifts the sinusoidal timestep embedding at every step.
+    timestep_dtype = torch.float32
+
     def __init__(self, module, config, dtype, device, hook_from=None):
         super().__init__(module, config, dtype, device)
         self.hook_from = hook_from
 
     def __call__(self, hidden_states, timestep, encoder_hidden_states) -> torch.Tensor:
         _check_dit_shapes(self.config, hidden_states, timestep, encoder_hidden_states)
-        return self._forward((hidden_states, timestep, encoder_hidden_states))
+        # The timestep goes to the device in fp32, not through _forward's cast to the model
+        # dtype; the text states follow it already on the device, keeping the module's order.
+        timestep = timestep.to(self.timestep_dtype).contiguous().to(self.device)
+        return self._forward((hidden_states,),
+                             extra=(timestep, self._device_input(encoder_hidden_states)))
 
     def _forward_hook(self):
         return self.forward_hook if self.hook_from is None else self.hook_from.forward_hook
@@ -529,7 +537,7 @@ def dit_input_contract(self) -> dict[str, dict[str, Any]]:
                 ),
                 "dtype": self.dtype,
             },
-            "timestep": {"shape": (batch,), "dtype": self.dtype},
+            "timestep": {"shape": (batch,), "dtype": torch.float32},
             "encoder_hidden_states": {
                 "shape": (batch, int(cfg.text_seq_len), int(cfg.text_dim)),
                 "dtype": self.dtype,

tests/unit/backends/test_neuron_hv1_app.py

@@ -190,6 +190,7 @@ def test_application_accepts_the_trainium_cli_kwargs(tiny_pipeline):
     assert app.shape == SHAPE
     contract = app.dit_input_contract()
     assert contract["encoder_attention_mask"] == {"shape": (1, TEXT), "dtype": torch.int64}
+    assert contract["timestep"]["dtype"] == torch.float32
     assert contract["hidden_states"]["shape"] == (1, 4, 5, 8, 8)
 
 

tests/unit/backends/test_neuron_hv1_dit.py

@@ -99,6 +99,7 @@ def test_example_inputs_are_the_static_shape():
     assert latents.shape == (1, 16, 16, 40, 64) and latents.dtype == torch.bfloat16
     assert text.shape == (1, 256, 4096) and pooled.shape == (1, 768)
     assert timestep.shape == guidance.shape == (1,)
+    assert timestep.dtype == torch.float32  # unrounded, as diffusers passes it
     assert torch.equal(guidance, torch.tensor([6000.0], dtype=torch.bfloat16))  # 6016 in bf16
     assert mask.dtype == torch.int64 and mask.shape == (1, 256) and int(mask.sum()) == 64
 

tests/unit/backends/test_neuron_wan_w1.py

@@ -286,7 +286,7 @@ def test_example_inputs_match_the_contract(tiny):
     )
     assert _shapes_and_dtypes(real.get_example_inputs()) == [
         ((1, 16, 3, 60, 104), torch.bfloat16),
-        ((1,), torch.bfloat16),
+        ((1,), torch.float32),  # the timestep unrounded, as diffusers passes it
         ((1, 512, 4096), torch.bfloat16),
     ]
 
@@ -423,9 +423,11 @@ def __call__(self, *args):
         sent.append((latents, timestep, text))
     assert len(received) == 2
     for args, originals in zip(received, sent, strict=True):
-        for got, original in zip(args, originals, strict=True):
-            assert got.is_contiguous() and got.dtype is torch.bfloat16
-            assert torch.equal(got, original.to(torch.bfloat16))
+        # The timestep (index 1) arrives in fp32, as diffusers passes it; the rest in bf16.
+        for i, (got, original) in enumerate(zip(args, originals, strict=True)):
+            want = torch.float32 if i == 1 else torch.bfloat16
+            assert got.is_contiguous() and got.dtype is want
+            assert torch.equal(got, original.to(want))
     assert received[1][1].stride() == (1,)  # the batch-2 timestep was materialised
     assert len(seconds) == 2 and all(s >= 0.0 for s in seconds)