Apply suggestions from code review

Co-authored-by: Masaki <mkozuki@nvidia.com>

Author: 2 people

torchao/prototype/mx_formats/kernels.py

@@ -711,7 +711,7 @@ def triton_quantize_nvfp4(
         return scales, xq.view(torch.uint8)
 
     @triton_quantize_nvfp4.register_fake
-    def _(x, per_tensor_scale=None, rounding_mode=0, seed=None):
+    def _(x, per_tensor_scale=None, rounding_mode=RoundingMode.RN, seed=None):
         M, N = x.shape
         num_scales = N // 16
         n_row_blocks = triton.cdiv(M, 128)
@@ -749,7 +749,7 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
     def triton_quantize_nvfp4(
         x: torch.Tensor,
         tensor_scale: Optional[torch.Tensor] = None,
-        rounding_mode: int = 0,
+        rounding_mode: RoundingMode = RoundingMode.RN,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         raise AssertionError("needs torch version 2.8+ and triton")