FP4 grouped API for torch (#4958)

Summary:
Pull Request resolved: #4958

X-link: https://github.com/facebookresearch/FBGEMM/pull/1979

We upgrade the FP4 grouped kernel with new API `f4f4bf16_grouped_mm`  that could be used in torch and vLLM. The kernel will support both MX and NV FP4, determined based on the scale dtypes (E4M3 vs E8M0). The API largely matches existing one we added for MXFP8. We also add unit tests for these new APIs.

Next steps:
- Full re-tune of the kernel
- Add other layouts to better support backwards

Reviewed By: q10

Differential Revision: D83171662

fbshipit-source-id: ba0abe5e1adf151e1b98b19b0abb14c9325d7966

Author: cthi

fbgemm_gpu/experimental/gemm/triton_gemm/fp4_quantize.py

@@ -5540,3 +5540,53 @@ def calculate_group_max(
         USE_INT64=use_int64,
     )
     return out, tensor_idx
+
+
+def get_nvfp4_global_scales_naive(
+    xs: list[torch.Tensor], ws: list[torch.Tensor]
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    """
+    Get global scales for each tensor in xs and ws.
+    This is done "naively" (not efficiently with a kernel). This function is used in unit tests or debugging.
+    """
+    global_scales = []
+    x_global_scales = []
+    w_global_scales = []
+
+    for x, w in zip(xs, ws):
+        # pyre-ignore
+        x_global_scale: torch.Tensor = (448.0 * 6.0) / torch.amax(
+            torch.abs(x.flatten()), dim=-1
+        ).to(torch.float32)
+        # pyre-ignore
+        w_global_scale: torch.Tensor = (448.0 * 6.0) / torch.amax(
+            torch.abs(w.flatten()), dim=-1
+        ).to(torch.float32)
+        # pyre-ignore
+        global_scale: torch.Tensor = 1 / (x_global_scale * w_global_scale)
+
+        global_scales.append(global_scale)
+        x_global_scales.append(x_global_scale)
+        w_global_scales.append(w_global_scale)
+
+    return global_scales, x_global_scales, w_global_scales
+
+
+def quantize_nvfp4_naive(
+    xs: list[torch.Tensor], global_scales: list[torch.Tensor]
+) -> tuple[
+    list[torch.Tensor],
+    list[torch.Tensor],
+]:
+    """
+    Quantize A to NVFP4 format.
+    This is done "naively" using a kernel for each group. This function is largely used in unit tests or debugging.
+    """
+    xqs, x_scales = zip(
+        *(
+            triton_scale_nvfp4_quant(x, global_scale)
+            for x, global_scale in zip(xs, global_scales)
+        )
+    )
+
+    return xqs, x_scales