feat: switch Sparse-MLA FP8 backward default to native Path C using graph outputs and remove manual allocation helpers

Author: apstenku123

cppmega_mlx/nn/_tilelang/sparse_mla_fp8_path_c.py

@@ -2546,58 +2546,6 @@ def _clear_fp8_bwd_dkv_buffer(dkv_buffer: mx.array) -> mx.array:
     return cleared_flat.reshape(shape)
 
 
-def _zero_fp8_bwd_output(shape: tuple[int, int, int, int]) -> mx.array:
-    debug_alloc = os.environ.get("CPPMEGA_DEBUG_SPARSE_MLA_BWD_ALLOC") == "1"
-    if debug_alloc:
-        print(
-            "DEBUG_BWD_ZERO_BEGIN",
-            shape,
-            int(mx.get_active_memory()) if hasattr(mx, "get_active_memory") else None,
-        )
-    total = 1
-    for dim in shape:
-        total *= int(dim)
-    if total <= 0:
-        return mx.zeros(shape, dtype=mx.float32)
-    kernel = _fp8_bwd_clear_dkv_tvm_ffi_kernel_for(
-        shape[0],
-        shape[1],
-        shape[2],
-        shape[3],
-        min(_SMFP8_BWD_CLEAR_THREADS, max(1, total)),
-    )
-    returned = kernel()
-    if isinstance(returned, mx.array):
-        zero_flat = returned
-    elif isinstance(returned, (list, tuple)) and len(returned) == 1:
-        zero_flat = cast(mx.array, returned[0])
-    else:
-        raise SparseMLAFp8PathCDirectError(
-            "direct tvm-ffi FP8 Sparse-MLA backward lazy zero did not return "
-            "one graph output"
-        )
-    if tuple(zero_flat.shape) != (total,) or zero_flat.dtype != mx.float32:
-        raise SparseMLAFp8PathCDirectError(
-            "direct tvm-ffi FP8 Sparse-MLA backward lazy zero returned "
-            f"shape/dtype {tuple(zero_flat.shape)}/{zero_flat.dtype}, expected "
-            f"{(total,)}/float32"
-        )
-    if debug_alloc:
-        print(
-            "DEBUG_BWD_ZERO_END",
-            shape,
-            int(mx.get_active_memory()) if hasattr(mx, "get_active_memory") else None,
-        )
-    return zero_flat.reshape(shape)
-
-
-def _empty_fp8_bwd_output(shape: tuple[int, ...]) -> mx.array:
-    empty = getattr(mx, "empty", None)
-    if empty is None:
-        return mx.zeros(shape, dtype=mx.float32)
-    return cast(mx.array, empty(shape, dtype=mx.float32))
-
-
 def _dispatch_fp8_bwd_owner_output_path_c(
     *,
     q_fp8: mx.array,
@@ -2628,9 +2576,11 @@ def _dispatch_fp8_bwd_owner_output_path_c(
     dkv_shape = (batch, seq_len_kv, kv_group, K)
     index_dtype = _index_dtype_name(indices, op_name="FP8 Sparse-MLA Path C backward")
     dkv_needs_clear = True
+    graph_output_route = False
     if dq_buffer is None and dkv_buffer is None:
-        dq_owner = _zero_fp8_bwd_output(dq_shape)
-        dkv_owner = _zero_fp8_bwd_output(dkv_shape)
+        graph_output_route = True
+        dq_owner = None
+        dkv_owner = None
         dkv_needs_clear = False
     elif (dq_buffer is None) != (dkv_buffer is None):
         raise SparseMLAFp8PathCDirectError(
@@ -2664,32 +2614,54 @@ def _dispatch_fp8_bwd_owner_output_path_c(
             d_out_dtype,
             index_dtype,
         )
-        if dkv_needs_clear:
-            dkv_owner = _clear_fp8_bwd_dkv_buffer(dkv_owner)
-        dq_flat = _flat_1d_view(dq_owner)
-        dkv_flat = _flat_1d_view(dkv_owner)
-        returned = kernel(
-            _flat_1d_view(q_fp8),
-            _flat_1d_view(q_scale),
-            _flat_1d_view(kv_fp8),
-            _flat_1d_view(kv_scale),
-            _flat_1d_view(d_out),
-            _flat_1d_view(indices),
-            sm_scale_buf,
-            out=(dq_flat, dkv_flat),
-        )
+        if graph_output_route:
+            dq_flat = None
+            dkv_flat = None
+            returned = kernel(
+                _flat_1d_view(q_fp8),
+                _flat_1d_view(q_scale),
+                _flat_1d_view(kv_fp8),
+                _flat_1d_view(kv_scale),
+                _flat_1d_view(d_out),
+                _flat_1d_view(indices),
+                sm_scale_buf,
+            )
+        else:
+            if dkv_needs_clear:
+                dkv_owner = _clear_fp8_bwd_dkv_buffer(cast(mx.array, dkv_owner))
+            dq_flat = _flat_1d_view(cast(mx.array, dq_owner))
+            dkv_flat = _flat_1d_view(cast(mx.array, dkv_owner))
+            returned = kernel(
+                _flat_1d_view(q_fp8),
+                _flat_1d_view(q_scale),
+                _flat_1d_view(kv_fp8),
+                _flat_1d_view(kv_scale),
+                _flat_1d_view(d_out),
+                _flat_1d_view(indices),
+                sm_scale_buf,
+                out=(dq_flat, dkv_flat),
+            )
     except Exception as exc:
         if force_path_c:
             raise RuntimeError(
                 "sparse_mla_fp8_bwd_path_c: direct tvm-ffi owner-output "
                 f"dispatch failed: {type(exc).__name__}: {exc}"
             ) from exc
         return None
-    dq_flat, dkv_flat = _owner_output_graph_tuple(
-        returned,
-        expected=(dq_flat, dkv_flat),
-        op_name="direct tvm-ffi FP8 Sparse-MLA backward",
-    )
+    if graph_output_route:
+        if not isinstance(returned, (list, tuple)) or len(returned) != 2:
+            raise SparseMLAFp8PathCDirectError(
+                "direct tvm-ffi FP8 Sparse-MLA backward graph-output route "
+                "did not return dq/dkv"
+            )
+        dq_flat = cast(mx.array, returned[0])
+        dkv_flat = cast(mx.array, returned[1])
+    else:
+        dq_flat, dkv_flat = _owner_output_graph_tuple(
+            returned,
+            expected=(cast(mx.array, dq_flat), cast(mx.array, dkv_flat)),
+            op_name="direct tvm-ffi FP8 Sparse-MLA backward",
+        )
     return dq_flat.reshape(dq_shape), dkv_flat.reshape(dkv_shape)
 
 
@@ -3331,8 +3303,6 @@ def _prepared_fp8_bwd_ste(
         dq = dq.astype(q_dtype)
     if dkv.dtype != kv_dtype:
         dkv = dkv.astype(kv_dtype)
-    if os.environ.get("CPPMEGA_SPARSE_MLA_FP8_BWD_EVAL") == "1":
-        mx.eval(dq, dkv)
     return dq, dkv
 
 

scripts/bench_1b_training_matrix.py

@@ -378,7 +378,7 @@ def path_env_and_support(
         }
         if dtype == "fp8":
             env[SPARSE_MLA_FP8_ROUTE_ENV] = "path_c"
-            env[SPARSE_MLA_FP8_BWD_ENV] = "path_b"
+            env[SPARSE_MLA_FP8_BWD_ENV] = "path_c"
         return (env, True, None, "cold", cache_dir)
     if path == "path_c_warm":
         env = {
@@ -390,7 +390,7 @@ def path_env_and_support(
         }
         if dtype == "fp8":
             env[SPARSE_MLA_FP8_ROUTE_ENV] = "path_c"
-            env[SPARSE_MLA_FP8_BWD_ENV] = "path_b"
+            env[SPARSE_MLA_FP8_BWD_ENV] = "path_c"
         return (env, True, None, "warm", cache_dir)
     return {}, False, f"unknown path {path!r}", "not_applicable", None
 

scripts/m04_train_step.py

@@ -271,7 +271,7 @@
 MAMBA3_PATH_C_BWD_ENV = "CPPMEGA_MAMBA3_PATH_C_BWD"
 FP8_PATH_C_RUNTIME_ENV: dict[str, str] = {
     SPARSE_MLA_FP8_ROUTE_ENV: "path_c",
-    SPARSE_MLA_FP8_BWD_ENV: "path_b",
+    SPARSE_MLA_FP8_BWD_ENV: "path_c",
     MAMBA3_PATH_C_BWD_ENV: "path_b",
 }
 FP8_PATH_B_RUNTIME_ENV: dict[str, str] = {SPARSE_MLA_FP8_ROUTE_ENV: "path_b"}
@@ -2780,16 +2780,17 @@ def fp8_path_c_training_route_payload(
                 "training_surface": fp8_producer_configured,
                 "producer_required": True,
                 "producer_status": producer["status"],
-                "backward_surface": "prepared_fp8_path_b_reference_vjp",
-                "path_c_backward_surface": "native_tvm_ffi_graph_output_scatter",
-                "default_backward_route": "path_b",
+                "backward_surface": "native_tvm_ffi_graph_output_scatter",
+                "fallback_backward_surface": "prepared_fp8_path_b_reference_vjp",
+                "default_backward_route": "path_c",
                 "default_backward_reason": (
-                    "native Sparse-MLA Path C backward is kept opt-in until "
-                    "peak-memory non-regression beats Path B"
+                    "native Sparse-MLA Path C backward uses graph outputs for "
+                    "the no-owner VJP path; caller-owned buffers remain explicit "
+                    "for fused runtimes"
                 ),
                 "kernel_policy_env": {
                     SPARSE_MLA_FP8_ROUTE_ENV: "path_c",
-                    SPARSE_MLA_FP8_BWD_ENV: "path_b",
+                    SPARSE_MLA_FP8_BWD_ENV: "path_c",
                 },
             },
             {

tests/test_bench_1b_training_matrix.py

@@ -65,7 +65,7 @@ def test_bench_1b_matrix_plan_covers_dtype_optimizer_path_cells(
     assert by_case["fp8_adamw_path_b"].env["CPPMEGA_KERNEL_PATH"] == "path_b"
     assert by_case["fp8_adamw_path_b"].env["CPPMEGA_KERNEL_PATH__SPARSE_MLA"] == "path_b"
     assert by_case["fp8_adamw_path_c_warm"].env["CPPMEGA_SPARSE_MLA_FP8_ROUTE"] == "path_c"
-    assert by_case["fp8_adamw_path_c_warm"].env["CPPMEGA_SPARSE_MLA_FP8_BWD"] == "path_b"
+    assert by_case["fp8_adamw_path_c_warm"].env["CPPMEGA_SPARSE_MLA_FP8_BWD"] == "path_c"
     assert "--use-path-c-direct-chain-runtime" not in by_case[
         "bf16_adamw_path_c_warm"
     ].command

tests/test_m04_train_step.py

@@ -450,7 +450,7 @@ def test_fp8_path_policies_set_explicit_runtime_routes(tmp_path: Path) -> None:
         ):
             assert os.environ["CPPMEGA_KERNEL_PATH__SPARSE_MLA"] == "path_c"
             assert os.environ["CPPMEGA_SPARSE_MLA_FP8_ROUTE"] == "path_c"
-            assert os.environ["CPPMEGA_SPARSE_MLA_FP8_BWD"] == "path_b"
+            assert os.environ["CPPMEGA_SPARSE_MLA_FP8_BWD"] == "path_c"
             assert os.environ["CPPMEGA_MAMBA3_PATH_C_BWD"] == "path_b"
         assert "CPPMEGA_SPARSE_MLA_FP8_ROUTE" not in os.environ
         assert "CPPMEGA_SPARSE_MLA_FP8_BWD" not in os.environ
@@ -1591,13 +1591,13 @@ def test_fp8_path_c_training_dtype_route_blocks_missing_sparse_mla_producer(
     )
     assert (
         surfaces["sparse_mla_fp8_path_c_apply"]["backward_surface"]
-        == "prepared_fp8_path_b_reference_vjp"
+        == "native_tvm_ffi_graph_output_scatter"
     )
-    assert surfaces["sparse_mla_fp8_path_c_apply"]["path_c_backward_surface"] == (
-        "native_tvm_ffi_graph_output_scatter"
+    assert surfaces["sparse_mla_fp8_path_c_apply"]["fallback_backward_surface"] == (
+        "prepared_fp8_path_b_reference_vjp"
     )
     assert surfaces["sparse_mla_fp8_path_c_apply"]["default_backward_route"] == (
-        "path_b"
+        "path_c"
     )
     assert (
         "FP8 parameter/weight producers that create the required dtype/layout "

tests/test_tilelang_sparse_mla_fp8.py

@@ -1018,6 +1018,21 @@ def test_fp8_path_c_backward_vjp_does_not_force_internal_eval() -> None:
         assert "mx.synchronize" not in source, name
 
 
+def test_fp8_path_c_backward_no_owner_route_uses_graph_outputs() -> None:
+    import cppmega_mlx.nn._tilelang.sparse_mla_fp8_path_c as path_c_module
+
+    source = inspect.getsource(path_c_module._dispatch_fp8_bwd_owner_output_path_c)
+    no_owner_block = source.split(
+        "if dq_buffer is None and dkv_buffer is None:", maxsplit=1
+    )[1].split("elif (dq_buffer is None) != (dkv_buffer is None):", maxsplit=1)[0]
+
+    assert "_empty_fp8_bwd_output" not in source
+    assert "_zero_fp8_bwd_output" not in source
+    assert "graph_output_route = True" in no_owner_block
+    assert "returned = kernel(" in source
+    assert "out=(dq_flat, dkv_flat)" in source
+
+
 def test_fp8_path_c_prepared_float_can_force_memory_safe_reference_backward() -> None:
     import cppmega_mlx.nn._tilelang.sparse_mla_fp8_path_c as path_c_module