Repository navigation
[WS1][CUDA][MiniMax-H3] timestep_mlp_fp32: deterministic FP32 timestep MLP (RFC #420 stack 2/10) #484
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -505,6 +505,14 @@ at::Tensor prefix_shared_attention( | |||||||||||||||||
| #if defined(KERNEL_ALIGN_WITH_H3_KERNELS) | ||||||||||||||||||
| torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, | ||||||||||||||||||
| double max_period, bool check_range); | ||||||||||||||||||
| std::vector<torch::Tensor> h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, | ||||||||||||||||||
| c10::optional<torch::Tensor> bias, | ||||||||||||||||||
| int64_t activation, bool save_pre_activation); | ||||||||||||||||||
| torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, | ||||||||||||||||||
| c10::ScalarType out_dtype); | ||||||||||||||||||
| std::vector<torch::Tensor> h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, | ||||||||||||||||||
| c10::ScalarType w_dtype, | ||||||||||||||||||
| bool with_bias); | ||||||||||||||||||
| #endif | ||||||||||||||||||
|
|
||||||||||||||||||
| // PyBind11 Module Registration | ||||||||||||||||||
|
|
@@ -771,5 +779,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { | |||||||||||||||||
| "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", | ||||||||||||||||||
| py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, | ||||||||||||||||||
| py::arg("check_range") = true); | ||||||||||||||||||
| m.def("h3_det_linear_forward", &h3_det_linear_forward, | ||||||||||||||||||
| "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU", | ||||||||||||||||||
| py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), | ||||||||||||||||||
| py::arg("activation") = 0, py::arg("save_pre_activation") = false); | ||||||||||||||||||
| m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, | ||||||||||||||||||
| "Deterministic grad @ weight with fixed 64-row N chunks folded in order", | ||||||||||||||||||
| py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); | ||||||||||||||||||
| m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, | ||||||||||||||||||
|
Comment on lines
+786
to
+789
There was a problem hiding this comment.
Suggestion:
Suggested change
Contributor
Author
There was a problem hiding this comment. Not changed. |
||||||||||||||||||
| "Deterministic dW/dbias as ascending-row FP32 folds", | ||||||||||||||||||
| py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); | ||||||||||||||||||
| #endif | ||||||||||||||||||
| } | ||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
torch::wrap_pybind_functionis inconsistently applied: the threeh3_det_linear_*ops are registered with the raw&function pointer.torch::wrap_pybind_functionis what releases the GIL and re-acquires the current CUDA device around the call, so these ops run with the GIL held and (via pybind's implicittorch::Tensorcaster) do not do the device-guard bookkeeping of the wrapped path. Please wrap them too, matchingh3_timestep_sinusoid_forward.Suggestion:
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not changed.
torch::wrap_pybind_functiondoes not release the GIL and does not touch the CUDA device: in torch 2.13 (torch/csrc/Exceptions.h) it instantiateswrap_pybind_function_impl_withrelease_gil = false; the GIL-releasing variant iswrap_pybind_function_no_gil. What the wrapper adds is PyTorch's exception and warning translation. The device is handled inside the entry points:h3_det_linear_forward,h3_det_linear_backward_inputandh3_det_linear_backward_weighteach take ac10::cuda::CUDAGuardon their tensors' device before the first allocation or launch and read the stream after it. The sinusoid op keeps its wrapper for its own reason: it is the one H3 op that raises a typed error (TORCH_CHECK_VALUE, so the range check reaches Python asValueError); the det_linear checks are plainTORCH_CHECKs that reach Python asRuntimeErroreither way (the validation tests assert those messages), and every other native op inops.cppis registered with the raw pointer. Pinned in 2ee5abc:test_ops_run_on_the_device_of_their_tensors_not_the_current_onecalls forward and both backwards with the tensors on cuda:1 while cuda:0 is current, for FP32 and BF16, and compares the bytes with the same calls on cuda:0.