Skip to content

tpmoe_sm90: double-buffer the LL all-reduce peer buffers by epoch parity - #1

Merged
AichenF merged 1 commit into
AichenF:sm90/tpmoefrom
liz-badada:sm90/tpmoe
Oct 19, 2026
Merged

AichenF merged 1 commit into
AichenF:sm90/tpmoefrom
liz-badada:sm90/tpmoe

Conversation

@liz-badada

Copy link
Copy Markdown

rs[]/ag[] were single buffers tagged by epoch. A rank preempted between pushing its T1 messages and reading its peers' does not hold the peers up: they run the next call, overwrite its pending messages, and it spins forever on the old tag.

Split both buffers into two epoch-parity halves (RS_PAR/AG_PAR in the kernel, RS_BYTES/AG_BYTES x2 on the host, +11 MiB per rank). A rank is never two calls ahead, so two halves suffice. Numerics unchanged; sources.json re-pinned.

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit (or used your preferred method).
  • I have installed the hooks with pre-commit install.
  • I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

🔬 Experimental Track

  • This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
    • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
    • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
    • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
    • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
    • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.
# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

rs[]/ag[] were single buffers tagged by epoch. A rank preempted between
pushing its T1 messages and reading its peers' does not hold the peers
up: they run the next call, overwrite its pending messages, and it spins
forever on the old tag.

Split both buffers into two epoch-parity halves (RS_PAR/AG_PAR in the
kernel, RS_BYTES/AG_BYTES x2 on the host, +11 MiB per rank). A rank is
never two calls ahead, so two halves suffice. Numerics unchanged;
sources.json re-pinned.
@AichenF
AichenF merged commit cd11a11 into AichenF:sm90/tpmoe Oct 19, 2026
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.

2 participants