feat: add gradient buffer tracking for entry RMSNorm layers in path-c fusion runtime

Author: apstenku123

.services.pid

@@ -1,2 +1,2 @@
-backend=30112
-frontend=30113
+backend=41057
+frontend=41058

backend.log

@@ -451,9 +451,9 @@ def _gemm_softmax_factory(*, M: int, N: int, K: int,
                 f"{exc.__class__.__name__}: {exc}") from exc
 
         @T.prim_func
-        def gemm_softmax(Q: T.Buffer((M, K), in_dtype),
-                          K_buf: T.Buffer((N, K), in_dtype),
-                          Out: T.Buffer((M, N), out_dtype)):
+        def gemm_softmax(Q: T.Tensor((M, K), in_dtype),
+                          K_buf: T.Tensor((N, K), in_dtype),
+                          Out: T.Tensor((M, N), out_dtype)):
             # Fused: Out = softmax(Q @ K^T) along axis=-1.
             with T.Kernel(M) as i:
                 row = T.alloc_fragment((N,), "float32")

cppmega_mlx/nn/_tilelang/_engine_dispatch.py

@@ -263,6 +263,10 @@ def __init__(self, compiled: Any) -> None:
         self._kernel = compiled.compiled.artifact
         self._prim_func = self._kernel.prim_func
         self.manifest = load_path_c_abi_manifest(self._prim_func)
+        self._result_param_indices = tuple(
+            int(idx)
+            for idx in self._prim_func.attrs.get("tilelang_out_idx", ())
+        )
 
     # --- inspection ------------------------------------------------------
 
@@ -420,8 +424,42 @@ def __call__(
                     f"not present in banks, dtype scratch, or top-level scratch"
                 )
 
-        # Launch.
-        self._kernel(*positional)
+        # Launch. Descriptor schedules mark every ABI parameter as a native
+        # owner-output/inout buffer; keep the returned aliases in the manifest
+        # maps so extraction reads the MLX graph value produced by TileLang.
+        returned = self._kernel(*positional)
+        if returned is not None:
+            returned_outputs = (
+                [returned]
+                if isinstance(returned, mx.array)
+                else list(returned)
+            )
+            if returned_outputs:
+                result_indices = self._result_param_indices
+                if len(returned_outputs) != len(result_indices):
+                    raise ValueError(
+                        "Path C TileLang kernel returned "
+                        f"{len(returned_outputs)} owner outputs for "
+                        f"{len(result_indices)} result ABI parameters"
+                    )
+                for param_index, value in zip(
+                    result_indices,
+                    returned_outputs,
+                    strict=True,
+                ):
+                    param_name = manifest.param_order[param_index]
+                    logical_name = (
+                        param_name[: -len("_handle")]
+                        if param_name.endswith("_handle")
+                        else param_name
+                    )
+                    if logical_name in banks:
+                        banks[logical_name] = value
+                    elif logical_name in scratch_buffers:
+                        scratch_buffers[logical_name] = value
+                    elif logical_name in top_level_buffers:
+                        top_level_buffers[logical_name] = value
+                    positional[param_index] = value
         # Force the lazy MLX graph to materialise so subsequent reads see the writes.
         mx.eval(*positional)
 

cppmega_mlx/runtime/path_c_fusion_launcher.py

@@ -1275,6 +1275,15 @@ def build_path_c_descriptor_prim_func(
             "tilelang_pass_configs",
             compile_pass_configs,
         )
+    owner_output_param_indices = tuple(range(len(physical_abi_plan.param_lines)))
+    if owner_output_param_indices:
+        prim_func = prim_func.with_attr(
+            "tilelang_out_idx",
+            list(owner_output_param_indices),
+        ).with_attr(
+            "tilelang_metal_zero_init_output_positions",
+            [],
+        )
     prim_func._cppmega_path_c_schedule_generator = PATH_C_DESCRIPTOR_SCHEDULE_GENERATOR
     prim_func._cppmega_path_c_brick_ops = tuple(
         descriptor.op_name for descriptor in descriptors
@@ -1546,7 +1555,7 @@ def _descriptor_prim_func_source(
     param_lines = list(physical_abi_plan.param_lines)
     if not param_lines:
         param_lines = [
-            f"{indent}_dummy: T.Buffer(({buffer_extent},), \"float32\"),"
+            f"{indent}_dummy: T.Tensor(({buffer_extent},), \"float32\"),"
         ]
     access_by_buffer = {
         buffer_name: _internal_buffer_ref(
@@ -2745,7 +2754,7 @@ def _spill_large_shared_scratch_to_abi(
         raise ValueError("descriptor source did not contain a function signature close")
     param_indent = " " * 4
     spill_param_lines = [
-        f'{param_indent}{info["param_name"]}: T.Buffer({_shape_literal(info["shape"])}, '
+        f'{param_indent}{info["param_name"]}: T.Tensor({_shape_literal(info["shape"])}, '
         f'"{info["dtype"]}"),'
         for name, info in spilled.items()
     ]
@@ -2866,7 +2875,7 @@ def _physical_abi_plan(
     }
     param_lines = tuple(
         f"{indent}{_physical_abi_bank_name(dtype)}: "
-        f"T.Buffer({_shape_literal(physical_shapes[_physical_abi_bank_name(dtype)])}, "
+        f"T.Tensor({_shape_literal(physical_shapes[_physical_abi_bank_name(dtype)])}, "
         f"\"{dtype}\"),"
         for dtype in bank_order
     )
@@ -2920,7 +2929,7 @@ def _direct_physical_abi_plan(
         for buffer_name in external_buffers
     }
     param_lines = tuple(
-        f"{indent}{name}: T.Buffer({_shape_literal(direct_shape_by_buffer[name])}, "
+        f"{indent}{name}: T.Tensor({_shape_literal(direct_shape_by_buffer[name])}, "
         f"\"{dtype_by_buffer[name]}\"),"
         for name in external_buffers
     )