Fix: size the DSv4-Pro prefill rings that overflow the default heap

golden/runner.py

@@ -536,10 +536,29 @@ def _run_benchmark(
     stats = None
     with _Stage("benchmark"):
         try:
+            # Forward the caller's RunConfig when there is one, mirroring the
+            # L3 branch. The benchmark is a second, independent dispatch, so a
+            # program that sizes its rings for the correctness run must size
+            # them here too or it validates and then deadlocks. benchmark()
+            # takes config= or platform=/device_id=, never both, and a bare
+            # RunConfig still defaults to platform "a2a3sim" / device 0 — pin
+            # this run's real target onto a copy rather than trusting the
+            # caller to have restated it.
+            rc = runtime_cfg.get("config")
+            if rc is not None:
+                import dataclasses
+
+                bench_kwargs: dict[str, Any] = {
+                    "config": dataclasses.replace(
+                        rc, platform=platform, device_id=device_id
+                    )
+                }
+            else:
+                bench_kwargs = {"platform": platform, "device_id": device_id}
             stats = benchmark(
                 compiled, ordered,
                 rounds=rounds, warmup=warmup,
-                platform=platform, device_id=device_id,
+                **bench_kwargs,
             )
         except RuntimeError as e:
             if not _benchmark_unavailable(e):

models/deepseek_v4_pro/prefill_attention_csa.py

@@ -127,13 +127,13 @@
 
 # PRO's wider hidden/HC dims make one prefill attention layer's per-task args and
 # intermediates overflow the runtime's default ring-2 output heap, which surfaces as
-# `orch_error_code=2 HEAP_RING_DEADLOCK`. prefill_fwd.py fixes the same thing with
-# run()'s `ring_heap=` argument, which only the L3 dispatch path honours. A
-# single-chip leaf has no way to apply it, so the size below records what the
-# kernel needs and is NOT applied -- the run takes the 256 MiB compile default.
-# All four rings, not just ring 2: raising ring 2 alone (what prefill_fwd.py
-# needs) still deadlocks here at both 2 GiB and 4 GiB, measured on device.
-PREFILL_ATTN_RING_HEAP = (4 * 1024 * 1024 * 1024,) * 4
+# `orch_error_code=2 HEAP_RING_DEADLOCK`. The default is CHIP_HEAP_SIZE, 256 MiB per
+# ring; measured on a5, every ring at 512 MiB is already enough for all three prefill
+# attention variants, so 1 GiB is one doubling of headroom over the measured need.
+# All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not
+# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on
+# the ChipWorker route -- see golden/runner.py::_execute_via_runner.
+PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4
 
 # Cache and state pools are allocator-managed global tensors.  Their physical
 # capacities are runtime values; request-local tables carry the actual block
@@ -976,6 +976,8 @@ def _quant_w_per_output_channel_local(w):
     # elements), but keep the 0.5% fraction bar identical to full prefill.
     x_out_diff_thd, x_out_max_diff = (8e-3, 2) if args.start_pos else (5e-3, 1)
 
+    from pypto.runtime import RunConfig
+
     result = run_jit(
         fn=prefill_attention_csa_test,
         specs=build_tensor_specs(
@@ -989,6 +991,9 @@ def _quant_w_per_output_channel_local(w):
             device_id=args.device,
             enable_chip_swimlane=args.enable_chip_swimlane,
             enable_dep_gen=args.enable_dep_gen,
+            # Ring sizing lives on execute_compiled's `config`, not on its
+            # signature, so name it here rather than as a bare keyword.
+            config=RunConfig(ring_heap=PREFILL_ATTN_RING_HEAP),
         ),
         rtol=1e-2,
         atol=1e-2,

models/deepseek_v4_pro/prefill_attention_hca.py

@@ -110,13 +110,13 @@
 
 # PRO's wider hidden/HC dims make one prefill attention layer's per-task args and
 # intermediates overflow the runtime's default ring-2 output heap, which surfaces as
-# `orch_error_code=2 HEAP_RING_DEADLOCK`. prefill_fwd.py fixes the same thing with
-# run()'s `ring_heap=` argument, which only the L3 dispatch path honours. A
-# single-chip leaf has no way to apply it, so the size below records what the
-# kernel needs and is NOT applied -- the run takes the 256 MiB compile default.
-# All four rings, not just ring 2: raising ring 2 alone (what prefill_fwd.py
-# needs) still deadlocks here at both 2 GiB and 4 GiB, measured on device.
-PREFILL_ATTN_RING_HEAP = (4 * 1024 * 1024 * 1024,) * 4
+# `orch_error_code=2 HEAP_RING_DEADLOCK`. The default is CHIP_HEAP_SIZE, 256 MiB per
+# ring; measured on a5, every ring at 512 MiB is already enough for all three prefill
+# attention variants, so 1 GiB is one doubling of headroom over the measured need.
+# All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not
+# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on
+# the ChipWorker route -- see golden/runner.py::_execute_via_runner.
+PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4
 
 
 @pl.jit.inline
@@ -713,6 +713,8 @@ def init_wo_b():
     args = parser.parse_args()
     compare_tokens = args.num_tokens
 
+    from pypto.runtime import RunConfig
+
     result = run_jit(
         fn=prefill_attention_hca_test,
         specs=build_tensor_specs(
@@ -726,6 +728,9 @@ def init_wo_b():
             device_id=args.device,
             enable_chip_swimlane=args.enable_chip_swimlane,
             enable_dep_gen=args.enable_dep_gen,
+            # Ring sizing lives on execute_compiled's `config`, not on its
+            # signature, so name it here rather than as a bare keyword.
+            config=RunConfig(ring_heap=PREFILL_ATTN_RING_HEAP),
         ),
         rtol=1e-2,
         atol=1e-2,

models/deepseek_v4_pro/prefill_attention_swa.py

@@ -111,13 +111,13 @@
 
 # PRO's wider hidden/HC dims make one prefill attention layer's per-task args and
 # intermediates overflow the runtime's default ring-2 output heap, which surfaces as
-# `orch_error_code=2 HEAP_RING_DEADLOCK`. prefill_fwd.py fixes the same thing with
-# run()'s `ring_heap=` argument, which only the L3 dispatch path honours. A
-# single-chip leaf has no way to apply it, so the size below records what the
-# kernel needs and is NOT applied -- the run takes the 256 MiB compile default.
-# All four rings, not just ring 2: raising ring 2 alone (what prefill_fwd.py
-# needs) still deadlocks here at both 2 GiB and 4 GiB, measured on device.
-PREFILL_ATTN_RING_HEAP = (4 * 1024 * 1024 * 1024,) * 4
+# `orch_error_code=2 HEAP_RING_DEADLOCK`. The default is CHIP_HEAP_SIZE, 256 MiB per
+# ring; measured on a5, every ring at 512 MiB is already enough for all three prefill
+# attention variants, so 1 GiB is one doubling of headroom over the measured need.
+# All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not
+# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on
+# the ChipWorker route -- see golden/runner.py::_execute_via_runner.
+PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4
 
 
 @pl.jit.inline
@@ -690,6 +690,8 @@ def init_wo_b():
     args = parser.parse_args()
     compare_tokens = args.num_tokens
 
+    from pypto.runtime import RunConfig
+
     result = run_jit(
         fn=prefill_attention_swa_test,
         specs=build_tensor_specs(
@@ -703,6 +705,9 @@ def init_wo_b():
             device_id=args.device,
             enable_chip_swimlane=args.enable_chip_swimlane,
             enable_dep_gen=args.enable_dep_gen,
+            # Ring sizing lives on execute_compiled's `config`, not on its
+            # signature, so name it here rather than as a bare keyword.
+            config=RunConfig(ring_heap=PREFILL_ATTN_RING_HEAP),
         ),
         compile_only=args.compile_only,
         rtol=1e-2,

models/deepseek_v4_pro/prefill_layer.py

@@ -121,6 +121,11 @@
 # ``T`` is the fixed child-kernel token-tile capacity (Qwen's ``TOK_TILE``). It is
 # NOT the packed token total. The packed prefill only ever feeds the children a
 # fixed ``[T, ...]`` tile at a time.
+# One prefill layer overflows the runtime's default 256 MiB-per-ring output
+# heap and fails with `orch_error_code=2 HEAP_RING_DEADLOCK`. Size every ring:
+# ring 2 alone, which prefill_fwd.py sets, does not clear it.
+LAYER_RING_HEAP = 1024 * 1024 * 1024
+
 TOK_TILE = T
 PREFILL_CHUNK_TOKENS = T
 DEFAULT_CHUNK_LENS = (T, T + T // 2)
@@ -1724,6 +1729,7 @@ def compare(actual, expected, **kwargs):
         runtime_cfg=dict(
             platform=args.platform,
             enable_chip_swimlane=args.enable_chip_swimlane,
+            ring_heap=LAYER_RING_HEAP,
         ),
         rtol=1e-3,
         atol=1e-3,

models/deepseek_v4_pro/prefill_mtp.py

@@ -73,6 +73,12 @@
 MTP_LAYER_ID = M.num_hidden_layers
 MTP_MOE_EPOCH = 1
 
+# The MTP prefill layer overflows the runtime's default 256 MiB-per-ring output
+# heap and fails on device with `orch_error_code=2 HEAP_RING_DEADLOCK`. Unlike
+# prefill_fwd.py, sizing ring 2 alone does not clear it -- measured on a5, ring 2
+# at 2 GiB still deadlocks -- so size every ring.
+MTP_RING_HEAP = 1024 * 1024 * 1024
+
 
 @pl.jit
 def mtp_prefill_fwd(
@@ -593,6 +599,7 @@ def main():
             platform=args.platform,
             enable_chip_swimlane=args.enable_chip_swimlane,
             enable_scope_stats=args.enable_scope_stats,
+            ring_heap=MTP_RING_HEAP,
         ),
         rtol=1e-3,
         atol=1e-3,