fix(runtime): saturate FlushingPolicyState counters to prevent add-overflow panic (#1359)

[teddytennant]

What

FlushingPolicyState::register now accumulates its bytes_count and bytes_size counters with saturating_add instead of +=, and converts the allocation length with u32::try_from(bytes.len()).unwrap_or(u32::MAX) instead of as u32.

Why

PendingDropQueue::push() calls register() on every staged Bytes allocation, growing the two u32 counters. Flushing is advisory — the consumer decides when to call flush()/reset() based on should_flush() — so between flushes the counters can grow without bound. On a large workload (reported in #1359 while training a resnet model on a 24 GB GPU; more data-loader workers reproduce it faster) self.bytes_size += bytes.len() as u32 eventually exceeds u32::MAX:

• in debug (overflow-checks on) it panics with attempt to add with overflow at policy.rs, the site reported in the issue;
• in release it silently wraps, corrupting the flush decision.

saturating_add can never panic or wrap. Because should_flush() compares with >=, saturating at u32::MAX never changes the flush decision — once at/over threshold a flush+reset follows. The saturating usize -> u32 conversion also fixes a latent truncation bug: a single allocation >= 4 GiB previously truncated to a small wrong value and could suppress a needed flush; it now counts as u32::MAX, guaranteeing a flush.

No public API change (the u32 fields on FlushingPolicy are untouched), the default max_bytes_size of 64 MiB means normal configs are unaffected, and no-std/alloc compatibility is preserved.

Testing

Added register_saturates_instead_of_overflowing to the existing policy_tests module. It stages an 8-byte allocation on a state whose counters are already at/near u32::MAX. On the current code this panics with attempt to add with overflow (debug overflow-checks); with the fix the counters saturate and should_flush returns true.

cargo test -p cubecl-runtime -- memory_management::drop_queue::policy

All 6 policy tests pass. cargo fmt --check and cargo clippy -p cubecl-runtime --all-targets -- -D warnings are clean.

Fixes #1359

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.