test: unify cross-backend RNG integration coverage

[lvyufeng]

AI Agent Information

• Agent/Tool: Claude Code CLI
• Model: Claude Opus 5
• Human Reviewer: @lvyufeng
• Session Summary: Consolidated duplicated MUSA RNG integration coverage into the shared cross-backend dispatch suite. Unified the public RNG test contract around torch.flagos, synchronized CUDA and MetaX compatibility state, and updated CI and documentation.

Summary

This PR makes tests/integration/ops/test_rng_dispatch.py the single cross-backend RNG integration suite. Shared behavior is exercised through the torch.flagos public API, while genuine backend differences use existing pytest platform marks; the redundant MUSA-only integration file is removed. CUDA and MetaX compatibility layers now synchronize their internal CUDA-shaped Philox generators with the public torch.flagos seed and state APIs, and CUDA, DCU, MetaX, and documentation references are updated accordingly.

Change Type

• Bug Fix
• New Feature
• Performance Optimization
• Refactoring
• Documentation
• Testing
• CI/Infrastructure
• Breaking Change

Platforms Affected

• CUDA
• MetaX
• Ascend
• PPU
• Platform-agnostic (all platforms)

Problem Analysis

What was broken/missing?

RNG integration coverage was split between the shared dispatch test and a duplicate MUSA-only test file. The split allowed public RNG behavior to drift between hardware-specific suites and caused the shared tests to encode CUDA generator implementation details even though the user-facing backend is torch.flagos.

Why did it happen?

CUDA-compatible FlagGems and boxing paths use CUDA-shaped Philox generators internally, while native PrivateUse1 paths expose torch.flagos generators. The test consolidation initially treated those internal representations as the public contract. In addition, CUDA and MetaX compatibility layers seeded their internal Philox generators without synchronizing the torch.flagos state API.

Investigation process:

1. Read test_rng_dispatch.py, the deleted MUSA RNG suite, platform marker setup, all affected CI manifests, and MUSA/operator-support documentation.
2. Ran the unified suite and reproduced state round-trip and multi-device failures when torch.flagos state was separate from CUDA-shaped generators.
3. Traced generator setup through torch_fl/flagos/random.py, csrc/runtime/generator.cc, the CUDA/MetaX compatibility shims, and FlagGems integration paths.
4. Verified the corrected public API behavior with the unified suite, unit tests, integration tests, collection checks, and repository lint checks.

Solution Design

Implementation approach:

• Move reusable MUSA assertions into test_rng_dispatch.py and select MUSA-only reservation behavior with @pytest.mark.musa.
• Express shared RNG state, seeding, factory, out, like, dropout, integer-range, explicit-generator, and multi-device contracts through torch.flagos.
• Keep CUDA-shaped generators as internal compatibility details, but synchronize them from torch.flagos.manual_seed, manual_seed_all, get_rng_state, and set_rng_state on CUDA and MetaX.
• Add explicit unified RNG steps to CUDA, DCU, and MetaX manifests; Ascend already had the dedicated command.
• Guard the MUSA bridge unit case when the MUSA-only C++ reservation symbol is not present in another build.

Key design decisions:

The public cross-backend contract uses torch.flagos, not torch.cuda. CUDA and MetaX still require Philox-shaped generator state internally for FlagGems, so those representations remain behind compatibility shims rather than being exposed as test prerequisites. The duplicate MUSA integration file is removed instead of maintaining another hardware-specific source of truth.

Code changes by file:

• .github/configs/cuda.yml: run the unified RNG suite explicitly.
• .github/configs/dcu.yml: run the unified RNG suite explicitly.
• .github/configs/metax.yml: run the unified RNG suite explicitly.
• docs/reference/operator-support.md: replace the standalone MUSA RNG test reference and retain the upstream MUPTI history entry.
• docs/vendors/musa/installation.md: document the unified RNG test path and MUSA mark selection.
• tests/integration/ops/test_musa_rng.py: remove duplicate MUSA integration coverage.
• tests/integration/ops/test_rng_dispatch.py: consolidate shared RNG contracts and isolate backend-specific behavior with marks and capability checks.
• tests/unit/test_musa_rng_bridge.py: skip the MUSA-symbol-dependent case on builds without the bridge symbol.
• torch_fl/accelerator/cuda/_cuda_compat.py: synchronize public torch.flagos seed/state APIs with internal CUDA Philox generators.
• torch_fl/accelerator/metax/_metax_compat.py: synchronize public torch.flagos seed/state APIs with internal MetaX Philox generators.

Changes by commit:

1. 3f32be5 - test: unify cross-backend RNG integration coverage: consolidate tests, synchronize public and internal RNG state, and update CI/docs.

Verification

Pre-submission Checklist

• Linting passed (ruff check, ruff format --check)
• Type checking passed (not applicable; no type-checking workflow for these changes)
• All tests pass (unit + integration; vendor FlagGems-only environment has an existing backend-registration gap documented below)
• Manual testing completed (reproduced and verified state round-trip and multi-device behavior)
• No debug/temporary code (no print statements, commented code, TODOs)
• Documentation updated (operator-support and MUSA installation docs)
• Commit messages follow conventions (type: description format)
• All text in English (required per CLAUDE.md)

Linting Results

$ ruff check .
All checks passed!

$ ruff format --check .
176 files already formatted

Test Results

$ python -m pytest tests/unit/ -q --tb=short
96 passed, 29 skipped, 1 warning in 9.49s

$ python -m pytest tests/integration/ops/ -m "anyplatform or cuda" -q --tb=short
667 passed, 1 skipped, 200 deselected, 1 xfailed, 3 xpassed in 142.03s

$ python -m pytest tests/integration/ops/test_rng_dispatch.py -m main_ops -q --tb=short
111 passed, 3 skipped, 1 xfailed in 0.85s

$ python -m pytest tests/unit/test_musa_rng_bridge.py tests/unit/test_vendor_routing.py -q --tb=short
25 passed, 1 skipped in 4.56s

The FlagGems-only command was also attempted:

$ FLAGOS_USE_FLAGGEMS=1 python -m pytest tests/integration/ops/test_rng_dispatch.py -m "flaggems and main_ops" -q --tb=short
2 failed, 113 deselected in 0.15s

The two failures occur before RNG assertions because ones is not registered in the current local backends_flaggems.conf environment. Hardware CI with the vendor FlagGems runtime is required for that path.

Manual Verification

# Reproduction before the fix:
python -m pytest tests/integration/ops/test_rng_dispatch.py -m main_ops -q --tb=short

# Output before the fix:
2 failed, 109 passed, 3 skipped, 1 xfailed
# Failures: torch.flagos state round-trip and manual_seed_all replay.

# Output after the fix:
111 passed, 3 skipped, 1 xfailed in 0.85s

Collection verification:

$ python -m pytest tests/integration/ops/test_rng_dispatch.py --collect-only -q
115 tests collected in 0.11s

Code Quality Verification

Style Consistency

• Matched existing code style in modified files
• Followed naming conventions (checked similar code)
• Comment density matches surrounding code
• Used project's existing utilities/helpers (no reinventing)

Edge Cases Considered

1. RNG state round trips through torch.flagos.get_rng_state and set_rng_state.
2. manual_seed_all replay and independent per-device generators.
3. Explicit generator isolation, integer/out variants, like variants, dropout, full-width int64 ranges, and MUSA shared reservations.

Potential Risks

1. CUDA and MetaX compatibility shims now replace the public torch.flagos seed/state callables; regressions would affect FlagGems reproducibility on those platforms.
2. The FlagGems-only path still depends on vendor runtime registrations not available in this local environment and must be validated by hardware CI.

Rollback Plan

Revert commit 3f32be5. This restores the previous MUSA-specific integration file and compatibility behavior without changing unrelated commits.

Related Work

• Related to the existing native MUSA RNG and FlagGems hybrid support.
• Related to the existing GCU and CUDA generator compatibility work.

Explicitly Not Included

• No changes to vendor RNG kernels or generated operator routes.
• No claim of local execution for unavailable Ascend, MetaX, DCU, or MUSA hardware.
• No expansion of the FlagGems operator route set.

Human Review Notes

Areas needing special attention:

1. Confirm that CUDA and MetaX torch.flagos state synchronization matches their internal Philox consumption model.
2. Review whether explicit flagos generator acceptance remains correctly capability-conditioned for native backends.
3. Verify the unified RNG suite on vendor hardware, especially the FlagGems runtime path.

Questions for reviewer:

1. Should any additional backend-specific generator capability be represented by a pytest mark rather than the current build capability helper?
2. Should the existing operator-support evidence be refreshed after the next physical-hardware run?

Additional Context

The branch is based on the latest flagos/main and contains only the RNG unification commit.

---

🤖 Generated with Claude Code