fix(rng): make default_generators iterable instead of looping forever

[lvyufeng]

What

The three default_generators shims — torch_fl/accelerator/metax/_metax_compat.py, torch_fl/accelerator/cuda/_cuda_compat.py, and torch_fl/flagos/__init__.py — stand in for a value upstream types as a tuple (torch/cuda/__init__.py:158), so callers are entitled to iterate it, slice it, or wrap it in list(). All three were list-like proxies defining only __getitem__/__len__, and __getitem__ was unbounded.

With no __iter__, Python falls back to the legacy iteration protocol: call __getitem__(0, 1, 2, ...) until IndexError. An unbounded __getitem__ never raises, so:

for g in torch.cuda.default_generators:   # never terminates

became an infinite loop that allocated a fresh CUDA generator on every step. It presented as a hang rather than an error, which is why nothing surfaced it — found while walking the generator list on an 8-card C550, where the process sat at 200% CPU indefinitely (34 min before I killed it). The flagos variant did not hang but was equally wrong: out-of-range indices surface as a RuntimeError from C++, which aborts iteration instead of ending it.

All three now define __iter__ and bounds-check __getitem__, raising IndexError past the device count and wrapping negative indices the way the tuple they replace does. Slices are supported for the same reason.

Testing

New test_default_generators_iterable is parametrized over both the torch.cuda and flagos shims and asserts all three properties (iteration terminates at device_count, slicing works, out-of-range raises IndexError). Verified to hang on the unfixed tree (stashed the fix, pytest timed out) and pass on the fixed one — not merely green-on-green.

Verified on 8xC550 (MetaX boxing, torch 2.10.0+metax3.8.1.0, triton-metax 3.6.0):

suite	vendor (FLAGOS_METAX_BOXING=1)	FlagGems (+FLAGOS_USE_FLAGGEMS=1)
test_rng_dispatch.py	104 passed / 2 skipped / 1 xfailed	106 passed / 1 xfailed
-m main_ops	104 passed / 28 skipped / 1 xfailed	—

This branch also confirms the unified-RNG work from #39/#49 is complete on MetaX — the boxing build compiles 103 GetFlagosDefaultCudaGenerator injection sites, and python_op_caller.cc rides in via FLAGGEMS_PYTHON=ON. Beyond the repo suite I ran an independent probe of 91 RNG ops covering every overload family (.out / .names_out / .generator / *_like), fp16/bf16/fp64/int32, and the nn.init.* entry points, checking run + same-seed-same-draw + different-seed-differs for each: 89/91 vendor, 91/91 FlagGems. The 2 vendor misses are exactly the documented native_dropout schema gap (already an xfail).

Distribution moments were checked against theory rather than just reproducibility, since a mis-bound overload can be perfectly reproducible and still wrong: randn skew +0.004 / excess kurtosis −0.0005, poisson(4) mean 4.0001 var 3.9925, randint chi² 6.99 (df=9), gamma(3) mean 2.997 var 2.992, dirichlet rows sum to 1, cauchy median +0.001, consecutive draws uncorrelated (r = −0.009), get_state/set_state round-trip exact. All 8 cards pass independently, and seeding card 0 does not perturb card 1.

ruff==0.15.12: ruff check . and ruff format --check . both pass.

Not caused by this change

Under -m "flaggems and main_ops" on MetaX, test_mm_dispatch / test_mean_dispatch test_dispatch_log_flaggems_runtime fail: they assert -> flagos_python, but backends_metax_flaggems.conf deliberately routes mm/bmm/mean.dim to cuda. Confirmed reproducing on the unmodified tree.

🤖 Generated with Claude Code