Skip to content

Commit 77c154c

Browse files
committed
fix(path-c): reflect CUDA bridge buffer writes back to owner (fused forward was silently dropped)
run_path_c_direct_fusion_chain_route runs segment artifacts against a route-local buffers dict. On Metal the kernel mutates caller-owned arrays in-place (zero-copy); on CUDA the MLX<->torch bridge writes results as NEW mlx arrays into the route-local dict only, and the route never reflected them back to logical_owner.buffers. So value_and_grad re-read the owner (still seeded zeros) after the forward route -> fused forward outputs silently dropped -> degenerate loss -> near-zero cotangents -> near-zero fused grads. This made fused-CUDA training silently produce GARBAGE gradients. Fix: after the segment loop, reflect the route-local buffers back into logical_owner.buffers / logical_buffers. CUDA-gated (_path_c_default_target()=="cuda") so Metal is byte-for-byte unchanged; env CPPMEGA_PATH_C_DISABLE_CUDA_WRITEBACK=1 disables it for verification calibration. Found + verified by a standalone b0 ABI verifier (drives the real fused mamba3 backward via the production banked-ABI runtime at seq=8): with the fix the fused forward delta matches eager to ~2% and loss to 0.002 (was: total changed 0/71, output exactly 0.0). Metal direct-chain route tests 7 passed.
1 parent 8798113 commit 77c154c

1 file changed

Lines changed: 24 additions & 0 deletions

File tree

scripts/m04_train_step.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6419,6 +6419,30 @@ def run_path_c_direct_fusion_chain_route(
64196419
"result_type": type(result).__name__,
64206420
}
64216421
)
6422+
# CUDA-only write-back fix: the MLX<->torch bridge replaces caller-owned bank
6423+
# arrays with NEW mlx arrays holding the kernel results, living only in the
6424+
# route-local ``buffers`` dict. Metal mutates the owner's arrays in-place
6425+
# (zero-copy) and needs none of this. The caller
6426+
# (PathCDirectFusionChainTrainingRuntime.value_and_grad) re-reads
6427+
# ``logical_owner.buffers`` after the forward route to seed the loss bridge
6428+
# and backward, so without reflecting the writes back the fused forward
6429+
# outputs are silently dropped -> degenerate loss -> near-zero fused grads.
6430+
# CUDA-gated so Metal is byte-for-byte unchanged.
6431+
# (env CPPMEGA_PATH_C_DISABLE_CUDA_WRITEBACK=1 disables it to reproduce the
6432+
# pre-fix broken state for verification calibration.)
6433+
if (
6434+
_path_c_default_target() == "cuda"
6435+
and os.environ.get("CPPMEGA_PATH_C_DISABLE_CUDA_WRITEBACK", "") != "1"
6436+
):
6437+
for _target_map in (getattr(logical_owner, "buffers", None), logical_buffers):
6438+
if _target_map is None:
6439+
continue
6440+
try:
6441+
for _name, _value in buffers.items():
6442+
_target_map[_name] = _value
6443+
except (TypeError, AttributeError):
6444+
# immutable mapping (e.g. MappingProxyType): cannot write back.
6445+
pass
64226446
return {
64236447
"status": "ok",
64246448
"runtime_uses_direct_fusion_chain": True,

0 commit comments

Comments
 (0)