Skip to content

fix(rocm): avoid indirect host pointers during HIP graph capture - #316

Open
ErikBPF wants to merge 1 commit into
FlashML-org:mainfrom
ErikBPF:fix/rocm-hip-graph-copy
Open

fix(rocm): avoid indirect host pointers during HIP graph capture#316
ErikBPF wants to merge 1 commit into
FlashML-org:mainfrom
ErikBPF:fix/rocm-hip-graph-copy

Conversation

@ErikBPF

@ErikBPF ErikBPF commented Sep 1, 2026

Copy link
Copy Markdown

Summary

ROCm graph capture can fault when the fused multi-bank expert copy reads pinned-host source addresses indirectly through a device pointer table. During HIP capture only, use the existing per-bank copy path whose source is a direct tensor argument.

  • eager ROCm keeps the fused path
  • CUDA keeps the fused path
  • only HIP graph capture falls back

Reproducer and evidence

Validated on Orion ROCm gfx1201 with realistic Qwen expert-row geometry:

  • 40 layers, 256 experts, 2,117 slots, top-k 8
  • 1,179,648-byte Q4_K gate/up rows and 720,896-byte Q5_K down rows
  • 512 graph replays completed with exact fingerprints in 17.96 seconds using direct per-bank arguments
  • the fused indirect source-pointer path reproduced the GPU page fault
  • full endpoint completed 3 × 1,024-token generations with graph-off canary/output hashes and an empty kernel log
  • later 10-sample soak remained exact and fault-free

The historical endpoint throughput varied between runs, so this PR claims graph correctness, not a speedup.

Tests

pytest -q tests/moe/test_fused_copy.py::test_rocm_graph_capture_uses_per_bank_copy
pytest -q tests/moe/test_fused_copy.py::test_qwen36_sized_pinned_rows_survive_rocm_graph_replay

Packaging validation: Python compile, Ruff on the changed test, and git diff --check pass.

Scope and overlap

Two files only. This complements #260, whose MVP intentionally disables HIP graph capture, and does not duplicate its ROCm bring-up changes. Qwen model support, FTW changes, tuning, and local research docs are excluded.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant