Refactor process input for torch.compile

Author: jinzhen-lin

humming/csrc/launcher/launcher.cpp

@@ -348,14 +348,18 @@ COMMON_TORCH_LIBRARY(humming, m) {
   m.def("register_process_input_kernel(str cubin_path) -> (int, str)");
   m.def("get_kernel_smem_size(int kernel_id) -> int");
   m.def(
-      "launch_process_input(Tensor configs, Tensor inputs, Tensor? outputs, Tensor? group_scales, "
-      "Tensor? token_scales, Tensor? expert_layout, Tensor? indices, bool inplace) "
-      "-> (Tensor, Tensor?, Tensor?)");
+      "launch_process_input(Tensor configs, Tensor inputs, Tensor(a!) outputs, "
+      "Tensor(b!)? group_scales, Tensor(c!)? token_scales, Tensor? expert_layout, "
+      "Tensor? indices) -> ()");
+  m.def(
+      "launch_process_input.inplace(Tensor configs, Tensor(a!) inputs, "
+      "Tensor? expert_layout, Tensor? indices) -> ()");
 };
 
 COMMON_TORCH_LIBRARY_IMPL(humming, CUDA, m) {
   m.impl("launch_kernel", COMMON_TORCH_BOX(&launch_kernel));
   m.impl("launch_process_input", COMMON_TORCH_BOX(&launch_process_input));
+  m.impl("launch_process_input.inplace", COMMON_TORCH_BOX(&launch_process_input_inplace));
   m.impl("launch_kernel.out", COMMON_TORCH_BOX(&launch_kernel_out));
 };
 
@@ -367,4 +371,6 @@ COMMON_TORCH_LIBRARY_IMPL(humming, Undefined, m) {
 
 COMMON_TORCH_LIBRARY_IMPL(humming, Meta, m) {
   m.impl("launch_kernel.out", COMMON_TORCH_BOX(&launch_kernel_out));
+  m.impl("launch_process_input", COMMON_TORCH_BOX(&launch_process_input));
+  m.impl("launch_process_input.inplace", COMMON_TORCH_BOX(&launch_process_input_inplace));
 };

humming/csrc/launcher/process_input.h

@@ -251,19 +251,6 @@ inline std::vector<int64_t> process_input_output_shape(
   return {shape.num_output_rows, columns};
 }
 
-inline std::vector<int64_t> process_input_leading_shape(
-    const ProcessInputKernelData &data, const Tensor &inputs, const ProcessInputShape &shape) {
-  if (data.layout == 0) {
-    std::vector<int64_t> leading_shape;
-    leading_shape.reserve(inputs.dim() - 1);
-    for (int64_t dimension = 0; dimension + 1 < inputs.dim(); dimension++)
-      leading_shape.push_back(inputs.size(dimension));
-    return leading_shape;
-  }
-  if (data.layout == 3) return {shape.num_experts, shape.max_tokens_per_expert};
-  return {shape.num_output_rows};
-}
-
 inline Tensor prepare_process_input_output(
     const ProcessInputKernelData &data,
     const Tensor &inputs,
@@ -278,7 +265,7 @@ inline Tensor prepare_process_input_output(
     ASSERT_CHECK(!outputs.has_value() || outputs->data_ptr() == inputs.data_ptr(), "inplace output must alias inputs");
     outputs = inputs;
   }
-  if (!outputs.has_value()) return torch_empty(expected_shape, dtype, inputs.device());
+  ASSERT_CHECK(outputs.has_value(), "outputs must be allocated by prepare_process_input");
   check_process_input_tensor(*outputs, "outputs", inputs.get_device(), dtype);
   ASSERT_CHECK(outputs->dim() == static_cast<int64_t>(expected_shape.size()), "invalid output rank");
   for (size_t dimension = 0; dimension < expected_shape.size(); dimension++)
@@ -302,18 +289,7 @@ inline std::optional<Tensor> prepare_process_input_group_scales(
   if (data.scale_layout == 2) elements = shape.group_scale_stride * CEIL_DIV(groups, 4) * 4;
   ScalarType dtype = dtype_id_to_tensor_dtype(data.group_scale_dtype_id);
   bool allow_byte = data.group_scale_dtype_id == 20080800;
-  if (!scales.has_value()) {
-    std::vector<int64_t> scale_shape;
-    if (data.scale_layout == 0) {
-      scale_shape = process_input_leading_shape(data, inputs, shape);
-      scale_shape.push_back(groups);
-    } else if (data.scale_layout == 1) {
-      scale_shape = {groups, shape.group_scale_stride};
-    } else {
-      scale_shape = {CEIL_DIV(groups, 4), shape.group_scale_stride, 4};
-    }
-    return torch_empty(scale_shape, dtype, inputs.device());
-  }
+  ASSERT_CHECK(scales.has_value(), "group_scales must be allocated by prepare_process_input");
   check_process_input_tensor(*scales, "group_scales", inputs.get_device(), dtype, allow_byte);
   ASSERT_CHECK(scales->numel() == elements, "invalid group_scales size");
   return scales;
@@ -331,11 +307,7 @@ inline std::optional<Tensor> prepare_process_input_token_scales(
     return std::nullopt;
   }
   int64_t elements = static_scale ? shape.num_experts : shape.num_output_rows;
-  if (!scales.has_value()) {
-    ASSERT_CHECK(dynamic_scale, "static token_scales is required");
-    auto scale_shape = process_input_leading_shape(data, inputs, shape);
-    return torch_empty(scale_shape, ScalarType::Float, inputs.device());
-  }
+  ASSERT_CHECK(scales.has_value(), "token_scales must be allocated by prepare_process_input");
   check_process_input_tensor(*scales, "token_scales", inputs.get_device(), ScalarType::Float);
   ASSERT_CHECK(scales->numel() == elements, "invalid token_scales size");
   return scales;
@@ -457,7 +429,7 @@ inline void launch_process_input_finalizer(
   check_curesult(cuLaunchKernelEx(&config, data.func, kernel_args, nullptr), "cuLaunchKernelEx");
 }
 
-inline std::tuple<Tensor, std::optional<Tensor>, std::optional<Tensor>> launch_process_input(
+inline void launch_process_input_impl(
     Tensor configs_tensor,
     Tensor inputs,
     std::optional<Tensor> outputs,
@@ -524,5 +496,29 @@ inline std::tuple<Tensor, std::optional<Tensor>, std::optional<Tensor>> launch_p
           secondary, intermediate, *group_scales, *token_scales, expert_layout, indices, shape);
     }
   }
-  return std::make_tuple(prepared_output, group_scales, token_scales);
+}
+
+inline void launch_process_input(
+    Tensor configs_tensor,
+    Tensor inputs,
+    Tensor outputs,
+    std::optional<Tensor> group_scales,
+    std::optional<Tensor> token_scales,
+    std::optional<Tensor> expert_layout,
+    std::optional<Tensor> indices) {
+  if (!inputs.is_cuda()) return;
+  launch_process_input_impl(
+      configs_tensor, inputs, outputs, group_scales, token_scales,
+      expert_layout, indices, false);
+}
+
+inline void launch_process_input_inplace(
+    Tensor configs_tensor,
+    Tensor inputs,
+    std::optional<Tensor> expert_layout,
+    std::optional<Tensor> indices) {
+  if (!inputs.is_cuda()) return;
+  launch_process_input_impl(
+      configs_tensor, inputs, inputs, std::nullopt, std::nullopt,
+      expert_layout, indices, true);
 }

humming/forward.py

@@ -25,7 +25,9 @@ def _resolve_use_pdl(
 
 
 def _prepare_input_scale(config: LayerConfig, input_scale: torch.Tensor) -> torch.Tensor:
-    if str(config.as_dtype) == "float8e8m0" and input_scale.dtype != torch.int32:
+    mx_scale_dtype = str(config.as_dtype) in ("float8e4m3", "float8e8m0")
+    grouped_mxmma = config.mma_type == MmaType.MXMMA and config.input_scale_group_size > 0
+    if mx_scale_dtype and grouped_mxmma and input_scale.dtype != torch.int32:
         packed_scale = input_scale.view(torch.int32)
         if input_scale.ndim == 3:
             packed_scale = packed_scale.reshape(input_scale.size(0), input_scale.size(1))

humming/kernel/process_input.py

@@ -227,16 +227,9 @@ def spec_key(spec):
                 secondary_id = kernels[spec_key(secondary)].kernel_id
             plan_kernel_ids[plan] = primary_id, secondary_id
 
-        compiled_intervals = []
-        for first, last, plan in intervals:
-            kernel_ids = plan_kernel_ids[plan]
-            if compiled_intervals and compiled_intervals[-1][2] == kernel_ids:
-                compiled_intervals[-1] = compiled_intervals[-1][0], last, kernel_ids
-            else:
-                compiled_intervals.append((first, last, kernel_ids))
-
         launch_configs = []
-        for first, last, (primary_id, secondary_id) in compiled_intervals:
+        for first, last, plan in intervals:
+            primary_id, secondary_id = plan_kernel_ids[plan]
             launch_configs.extend((first - 1, last, primary_id, secondary_id))
         result = torch.tensor(launch_configs, dtype=torch.int64, device="cpu")
         if cache_key is not None:

humming/ops/input/hadamard.py

@@ -44,13 +44,26 @@ def hadamard_quant_input(
     global_scale: torch.Tensor | None = None,
     use_pdl: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    if group_size is None or group_size == 0:
+    channelwise = group_size is None or group_size == 0
+    if channelwise and scales is None and scale_dtype == "float32" and global_scale is None:
+        quantized, _, result_token_scales = process_input(
+            inputs,
+            outputs=outputs,
+            quant_mode=QuantizationMode.DynamicToken,
+            quant_dtype=quant_dtype,
+            hadamard_block_size=block_size,
+            use_pdl=use_pdl,
+        )
+        assert result_token_scales is not None
+        return quantized, result_token_scales.unsqueeze(-1)
+
+    if channelwise:
         group_size = inputs.size(-1)
 
     group_scale_layout = GroupScaleLayout.RowMajor
     if m_major_scale:
         group_scale_layout = GroupScaleLayout.MMajor
-        if scale_dtype == "float8e8m0":
+        if scale_dtype in ("float8e4m3", "float8e8m0"):
             group_scale_layout = GroupScaleLayout.MxPacked
 
     quant_mode = QuantizationMode.DynamicGroup