refactor: set path_b as the default backward route for Mamba3 in Path C training and add support for explicit full-Path-C backward overrides.

apstenku123 · apstenku123 · commit b269b61ad0ae · 2026-05-18T15:14:32.000+02:00
diff --git a/reports/v4/v4_parquet_smoke_h1024_L12_1779106370.json b/reports/v4/v4_parquet_smoke_h1024_L12_1779106370.json
@@ -0,0 +1,36 @@
+{
+  "template_name": "v4_smoke_h1024_L12",
+  "hidden_size": 1024,
+  "layers": 12,
+  "seq_len": 256,
+  "batch_size": 1,
+  "steps": 3,
+  "parquet_path": "/Volumes/external/sources/cppmega.mlx/data/parquet_samples/gb10/clang_semantic_4k_v10/val_00000.parquet",
+  "fwd_median_ms": 59.37258299672976,
+  "bwd_median_ms": 27.799417002825066,
+  "total_median_ms": 90.63329198397696,
+  "tokens_per_sec": 2824.569144473514,
+  "per_step": [
+    {
+      "step": 0,
+      "fwd_ms": 96.77291702246293,
+      "bwd_ms": 7.3932919767685235,
+      "total_ms": 104.16620899923146,
+      "loss": 11.183396339416504
+    },
+    {
+      "step": 1,
+      "fwd_ms": 56.313292006962,
+      "bwd_ms": 34.31999997701496,
+      "total_ms": 90.63329198397696,
+      "loss": 11.018841743469238
+    },
+    {
+      "step": 2,
+      "fwd_ms": 59.37258299672976,
+      "bwd_ms": 27.799417002825066,
+      "total_ms": 87.17199999955483,
+      "loss": 10.956659317016602
+    }
+  ]
+}
diff --git a/reports/v4/v4_parquet_smoke_h1024_L8_1779106283.json b/reports/v4/v4_parquet_smoke_h1024_L8_1779106283.json
@@ -0,0 +1,36 @@
+{
+  "template_name": "v4_smoke_h1024_L8",
+  "hidden_size": 1024,
+  "layers": 8,
+  "seq_len": 256,
+  "batch_size": 1,
+  "steps": 3,
+  "parquet_path": "/Volumes/external/sources/cppmega.mlx/data/parquet_samples/gb10/clang_semantic_4k_v10/val_00000.parquet",
+  "fwd_median_ms": 6.824791984399781,
+  "bwd_median_ms": 0.4123759863432497,
+  "total_median_ms": 7.097541994880885,
+  "tokens_per_sec": 36068.82497978039,
+  "per_step": [
+    {
+      "step": 0,
+      "fwd_ms": 34.21054198406637,
+      "bwd_ms": 31.339375040261075,
+      "total_ms": 65.54991702432744,
+      "loss": 11.423054695129395
+    },
+    {
+      "step": 1,
+      "fwd_ms": 6.824791984399781,
+      "bwd_ms": 0.1975830236915499,
+      "total_ms": 7.0223750080913305,
+      "loss": 11.195006370544434
+    },
+    {
+      "step": 2,
+      "fwd_ms": 6.685166008537635,
+      "bwd_ms": 0.4123759863432497,
+      "total_ms": 7.097541994880885,
+      "loss": 11.354619026184082
+    }
+  ]
+}
diff --git a/reports/v4/v4_parquet_smoke_h2048_L12_1779108259.json b/reports/v4/v4_parquet_smoke_h2048_L12_1779108259.json
@@ -0,0 +1,29 @@
+{
+  "template_name": "v4_smoke_h2048_L12",
+  "hidden_size": 2048,
+  "layers": 12,
+  "seq_len": 256,
+  "batch_size": 1,
+  "steps": 2,
+  "parquet_path": "/Volumes/external/sources/cppmega.mlx/data/parquet_samples/gb10/clang_semantic_4k_v10/val_00000.parquet",
+  "fwd_median_ms": 222.840249989531,
+  "bwd_median_ms": 28.563021012814716,
+  "total_median_ms": 154.23662500688806,
+  "tokens_per_sec": 1659.7873558797548,
+  "per_step": [
+    {
+      "step": 0,
+      "fwd_ms": 352.65866699046455,
+      "bwd_ms": 0.0,
+      "total_ms": 158.32537499954924,
+      "loss": 11.305686950683594
+    },
+    {
+      "step": 1,
+      "fwd_ms": 93.02183298859745,
+      "bwd_ms": 57.12604202562943,
+      "total_ms": 150.14787501422688,
+      "loss": 11.046894073486328
+    }
+  ]
+}
diff --git a/scripts/bench_1b_training_matrix.py b/scripts/bench_1b_training_matrix.py
@@ -178,6 +178,16 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--dtypes", default=",".join(DTYPE_CHOICES))
     parser.add_argument("--optimizers", default=",".join(OPTIMIZER_CHOICES))
     parser.add_argument("--paths", default=",".join(PATH_CHOICES))
+    parser.add_argument(
+        "--mamba3-bwd",
+        choices=("path_b", "path_c"),
+        default="path_b",
+        help=(
+            "Backward route for Mamba3 inside Path C cells. The default keeps "
+            "the performance-safe Path C forward + Path B backward route; use "
+            "--mamba3-bwd path_c only for an explicit full-Path-C experiment."
+        ),
+    )
     parser.add_argument("--fresh-process", action="store_true", default=False)
     parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
     parser.add_argument("--csv", type=Path, default=DEFAULT_CSV)
@@ -231,6 +241,7 @@ def path_env_and_support(
     dtype: str,
     path: str,
     cache_dir: Path,
+    mamba3_bwd: str,
 ) -> tuple[dict[str, str], bool, str | None, str, Path | None]:
     if dtype == "fp8" and path == "path_b":
         return (
@@ -265,7 +276,7 @@ def path_env_and_support(
             "CPPMEGA_KERNEL_PATH__MAMBA3_MIMO": "path_c",
             "CPPMEGA_KERNEL_PATH__M2RNN": "path_c",
             "CPPMEGA_KERNEL_PATH__SPARSE_MLA": "path_c",
-            MAMBA3_PATH_C_BWD_ENV: "path_c",
+            MAMBA3_PATH_C_BWD_ENV: mamba3_bwd,
         }
         if dtype == "fp8":
             env[SPARSE_MLA_FP8_ROUTE_ENV] = "path_c"
@@ -276,7 +287,7 @@ def path_env_and_support(
             "CPPMEGA_KERNEL_PATH__MAMBA3_MIMO": "path_c",
             "CPPMEGA_KERNEL_PATH__M2RNN": "path_c",
             "CPPMEGA_KERNEL_PATH__SPARSE_MLA": "path_c",
-            MAMBA3_PATH_C_BWD_ENV: "path_c",
+            MAMBA3_PATH_C_BWD_ENV: mamba3_bwd,
         }
         if dtype == "fp8":
             env[SPARSE_MLA_FP8_ROUTE_ENV] = "path_c"
@@ -301,6 +312,7 @@ def build_cell(
         dtype=dtype,
         path=path,
         cache_dir=cache_dir,
+        mamba3_bwd=args.mamba3_bwd,
     )
     if not path_supported:
         supported = False
@@ -858,6 +870,7 @@ def main(argv: list[str] | None = None) -> int:
                 "dtypes": list(parse_csv_list(args.dtypes, DTYPE_CHOICES)),
                 "optimizers": list(parse_csv_list(args.optimizers, OPTIMIZER_CHOICES)),
                 "paths": list(parse_csv_list(args.paths, PATH_CHOICES)),
+                "mamba3_bwd": args.mamba3_bwd,
                 "fresh_process": bool(args.fresh_process),
                 "dry_run": bool(args.dry_run),
                 "reuse_existing_ok": bool(args.reuse_existing_ok),
diff --git a/scripts/m04_train_step.py b/scripts/m04_train_step.py
@@ -161,7 +161,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",
-    MAMBA3_PATH_C_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"}
 FP8_PATH_C_SPLIT_GRAD_UPDATE_EVAL_REASON = (
@@ -1061,10 +1061,13 @@ def fp8_path_c_training_route_payload(
                 "fp8_route_auto_selected": True,
                 "fp8_route_reason": (
                     "Path C now uses direct tvm-ffi owner-output for contiguous "
-                    "selective-scan buffers, including Path C backward, and "
-                    "matches Path B dispatch geometry on that explicit boundary"
+                    "selective-scan forward buffers and keeps the performance-safe "
+                    "Path B backward route until full Path C backward wins the "
+                    "1B training matrix"
                 ),
-                "training_surface": True,
+                "training_surface": False,
+                "full_path_c_backward_available": True,
+                "default_backward_route": "path_b",
             },
             {
                 "name": "m2rnn_path_c",
diff --git a/scripts/render_1b_training_matrix_html.py b/scripts/render_1b_training_matrix_html.py
@@ -550,15 +550,15 @@ def render_route_legend() -> str:
           </div>
           <div class="route-card">
             <h3>BF16 Path C</h3>
-            <p>Runs <code>--dtype bfloat16</code> with <code>CPPMEGA_KERNEL_PATH=path_c</code> and <code>CPPMEGA_MAMBA3_PATH_C_BWD=path_c</code>. Mamba3, M2RNN, and Sparse-MLA are all requested through Path C; no Path B backward fallback is allowed in these cells.</p>
+            <p>Runs <code>--dtype bfloat16</code> with <code>CPPMEGA_KERNEL_PATH=path_c</code>. The default matrix keeps Mamba3 on Path C forward plus Path B backward via <code>--mamba3-bwd path_b</code>; pass <code>--mamba3-bwd path_c</code> only for an explicit full-Path-C experiment.</p>
           </div>
           <div class="route-card">
             <h3>FP8 Path B</h3>
             <p>Runs <code>--dtype fp8_path_b</code> and forces <code>CPPMEGA_KERNEL_PATH__SPARSE_MLA=path_b</code> plus <code>CPPMEGA_SPARSE_MLA_FP8_ROUTE=path_b</code>. The DSA Sparse-MLA baseline dispatch is recorded as <code>sparse_mla_fp8_reference_path_b</code>; it is an honest non-Path-C FP8 training baseline, not a Path C fallback.</p>
           </div>
           <div class="route-card">
             <h3>FP8 Path C</h3>
-            <p>Runs <code>--dtype fp8_path_c</code>, <code>CPPMEGA_SPARSE_MLA_FP8_ROUTE=path_c</code>, and the same Mamba3 Path C forward + Path B backward training policy. Sparse-MLA consumes prepared <code>q_fp8/q_scale/kv_fp8/kv_scale</code> buffers through the TileLang/tvm-ffi route.</p>
+            <p>Runs <code>--dtype fp8_path_c</code>, <code>CPPMEGA_SPARSE_MLA_FP8_ROUTE=path_c</code>, and the selected Mamba3 backward policy from the matrix command. Sparse-MLA consumes prepared <code>q_fp8/q_scale/kv_fp8/kv_scale</code> buffers through the TileLang/tvm-ffi route.</p>
           </div>
         </div>
       </section>
diff --git a/tests/test_bench_1b_training_matrix.py b/tests/test_bench_1b_training_matrix.py
@@ -49,7 +49,7 @@ def test_bench_1b_matrix_plan_covers_dtype_optimizer_path_cells(
     assert by_case["bf16_adamw_path_c_cold"].env["CPPMEGA_KERNEL_PATH"] == "path_c"
     assert by_case["bf16_adamw_path_c_cold"].env["CPPMEGA_KERNEL_PATH__MAMBA3_MIMO"] == "path_c"
     assert by_case["bf16_adamw_path_c_cold"].env["CPPMEGA_KERNEL_PATH__M2RNN"] == "path_c"
-    assert by_case["bf16_adamw_path_c_cold"].env["CPPMEGA_MAMBA3_PATH_C_BWD"] == "path_c"
+    assert by_case["bf16_adamw_path_c_cold"].env["CPPMEGA_MAMBA3_PATH_C_BWD"] == "path_b"
     assert by_case["bf16_adamw_path_c_cold"].cache_mode == "cold"
     assert by_case["bf16_adamw_path_c_warm"].cache_mode == "warm"
     assert "--seq-len" in by_case["bf16_adamw_path_b"].command
@@ -67,6 +67,26 @@ def test_bench_1b_matrix_plan_covers_dtype_optimizer_path_cells(
     assert by_case["fp8_adamw_path_c_warm"].env["CPPMEGA_SPARSE_MLA_FP8_ROUTE"] == "path_c"
 
 
+def test_bench_1b_matrix_can_explicitly_force_full_mamba3_path_c_bwd(
+    tmp_path: Path,
+) -> None:
+    args = _args(
+        tmp_path,
+        "--dtypes",
+        "bf16",
+        "--optimizers",
+        "adamw",
+        "--paths",
+        "path_c_warm",
+        "--mamba3-bwd",
+        "path_c",
+    )
+    cell = matrix.build_cells(args)[0]
+
+    assert cell.env["CPPMEGA_KERNEL_PATH__MAMBA3_MIMO"] == "path_c"
+    assert cell.env["CPPMEGA_MAMBA3_PATH_C_BWD"] == "path_c"
+
+
 def test_bench_1b_matrix_dry_run_writes_markdown_csv_and_json(
     tmp_path: Path,
 ) -> None:
@@ -106,6 +126,7 @@ def test_bench_1b_matrix_dry_run_writes_markdown_csv_and_json(
     payload = json.loads((tmp_path / "matrix.json").read_text(encoding="utf-8"))
     assert payload["scope"] == "cppmega_1b_path_matrix"
     assert payload["config"]["block_size"] == 2048
+    assert payload["config"]["mamba3_bwd"] == "path_b"
 
 
 def test_bench_1b_matrix_extracts_m04_receipt_metrics(
diff --git a/tests/test_m04_train_step.py b/tests/test_m04_train_step.py
@@ -114,7 +114,7 @@ def test_fp8_path_policies_set_explicit_runtime_routes(
     with m04_train_step.fp8_path_c_kernel_policy(path_c_args):
         assert os.environ["CPPMEGA_KERNEL_PATH__SPARSE_MLA"] == "path_c"
         assert os.environ["CPPMEGA_SPARSE_MLA_FP8_ROUTE"] == "path_c"
-        assert os.environ["CPPMEGA_MAMBA3_PATH_C_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_MAMBA3_PATH_C_BWD" not in os.environ
 
@@ -1191,7 +1191,9 @@ def test_fp8_path_c_training_dtype_route_blocks_missing_sparse_mla_producer(
         surface["name"]: surface for surface in route["available_path_c_surfaces"]
     }
     assert surfaces["matmul_tl_fp8_scaled_matmul"]["kernel_surface_available"] is True
-    assert surfaces["mamba3_mimo_path_c"]["training_surface"] is True
+    assert surfaces["mamba3_mimo_path_c"]["training_surface"] is False
+    assert surfaces["mamba3_mimo_path_c"]["full_path_c_backward_available"] is True
+    assert surfaces["mamba3_mimo_path_c"]["default_backward_route"] == "path_b"
     assert surfaces["mamba3_mimo_path_c"]["fp8_route_auto_selected"] is True
     assert surfaces["m2rnn_path_c"]["training_surface"] is True
     assert surfaces["m2rnn_path_c"]["fallback_to_path_b_allowed"] is False
