perf(comm): add opt-in Triton AR+RMSNorm backend - #968
Conversation
|
Figures to support above claims in PR summary. Complete data available in full companion branch. GLM-5.2 on MI355, fused triton-shmem vs. unfused baseline and default fused Iris |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bb2a93b328
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if ts_state is not None: | ||
| _ts_mod.TRITON_SHMEM_AR_RMSNORM_STATES[ts_key] = ts_state |
There was a problem hiding this comment.
When state creation returns None—for example on a non-gfx950 AMD GPU or after allocation/setup failure—the key remains absent, so every subsequent fused norm retries collective validation and allocation. In normal graph-enabled serving, a retry can therefore execute all_gather_object and memory setup during CUDA graph capture instead of taking a stable unfused path, potentially breaking capture; eager serving also incurs this work at every layer. Cache a permanent decline sentinel or disable this fusion after the first failed initialization.
Useful? React with 👍 / 👎.
| import torch | ||
| import torch.distributed as dist | ||
|
|
||
| _HIP = ctypes.CDLL("libamdhip64.so") |
There was a problem hiding this comment.
This new module loads and binds the vendor HIP runtime directly from ops/communication, bypassing the repository's required third-party boundary and registration path. Move the ABI wrapper under tokenspeed_kernel/thirdparty and expose it to ops through the established kernel registration/import mechanism so vendor-specific integration remains isolated.
AGENTS.md reference: AGENTS.md:L36-L37
Useful? React with 👍 / 👎.
| return eligible | ||
|
|
||
|
|
||
| def triton_shmem_allreduce_residual_rmsnorm( |
There was a problem hiding this comment.
This function is exported through __all__ as a new public API but has no docstring describing its state and tensor arguments, optional output-buffer ownership, or returned tensors. Add an Args/Returns docstring so callers can safely use the output and lifetime contract.
AGENTS.md reference: AGENTS.md:L44-L44
Useful? React with 👍 / 👎.
14f6fdc to
aba2ba6
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aba2ba623e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Slot 0 is scheduler-owned and can contain live request state. Every | ||
| # padded graph row must use the reserved sink row, which remains neutral | ||
| # across capture and replay for attention and sampling consumers. | ||
| sentinel = int(self.config.max_req_pool_size) |
There was a problem hiding this comment.
When a padded CUDA graph is replayed on a single-table attention backend, this sentinel is later passed to gather_page_table_with_padding, whose kernel uses each req_pool_indices value as a row index into the supplied page_table. That table is the batch-ordered DraftPageStaging.table, allocated with only max_bs rows, while config.max_req_pool_size is the reserved sink after the scheduler-owned range (max_bs + 1), so any replay with padded_bs > bs can read past the page-table buffer instead of using the dummy row. Keep the metadata-facing padding index in range (or pad the page table to include the sink row) while still using the sink row only for pool-indexed runtime state.
Useful? React with 👍 / 👎.
| except Exception as exc: # noqa: BLE001 - decline rather than crash forward | ||
| logger.warning("triton_shmem AR+RMSNorm state creation failed: %s", exc) | ||
| return None |
There was a problem hiding this comment.
If the constructor fails on only one rank after the eligibility all_gather_object—for example from a rank-local CUDA allocation/OOM or HIP runtime error—this catch returns None only there, while peers can cache a live triton_shmem state. The next fused norm then splits the group between the NCCL all_reduce fallback on the failed rank and Triton shmem barriers on the others, which can hang the process group; either let the exception abort all ranks or exchange a success flag so every rank takes the same fallback path.
Useful? React with 👍 / 👎.
Add a graph-safe MI350X backend with coarse HIP-IPC buffers while preserving Iris-first defaults and complete unfused fallback. Harden capture and topology guards, and cover the supported multi-GPU paths with focused correctness tests and a reproducible benchmark. Signed-off-by: Jeremy Wang <w.jeremy220@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Jeremy Wang <w.jeremy220@gmail.com>
Signed-off-by: Jeremy Wang <w.jeremy220@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Jeremy Wang <w.jeremy220@gmail.com>
Signed-off-by: Jeremy Wang <w.jeremy220@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
aba2ba6 to
c204cbc
Compare
|
@Jeremy-wj can you give a bit more detail behind choosing PyTorch symmetric memory over Iris? Mostly for my own edification since behind the scenes Iris does a lot of the IPC setup for us. |
|
@qedawkins For IPC setup the new backend uses the existing pytorch symmetric memory API to set up symmetric memory, which already had functions set up in If you were referring to motivation, the general idea is that pytorch symmetric memory is directly integrated as part of the torch+rocm stack and doesn't introduce external dependencies. |
qedawkins
left a comment
qedawkins
left a comment
There was a problem hiding this comment.
Sorry for the late reply, a few codex driven comments below. There was some offline discussion about motivation behind a new comms backend so I'll defer to those folks if they want to add details here
| """ | ||
| world_size = group.size() | ||
| rank = dist.get_rank(group) | ||
| tensor = torch.empty(shape, dtype=dtype, device=device) |
There was a problem hiding this comment.
The initial torch.empty allocation happens before the first cross-rank error exchange. If it fails on only one rank, that rank enters the fine-grained fallback while the other ranks block in all_gather_object, leaving the ranks in different collectives and deadlocking. Could we include allocation failure in the collective status protocol so every rank decides on the fallback together? A test injecting a rank-local allocation failure would also cover this path.
| help="Enable allreduce fusion for improved decode performance. Auto-enabled on supported single-node TP configurations.", | ||
| ) | ||
| parser.add_argument( | ||
| "--disable-allreduce-fusion", |
There was a problem hiding this comment.
Could this be represented as a single tri-state option, such as --allreduce-fusion={auto,on,off}, instead of two independent booleans? The current representation has three valid states plus the invalid case where both flags are set. A tri-state would model the intended behavior directly and match the existing auto-mode configuration pattern in this file.
| "are mutually exclusive" | ||
| ) | ||
|
|
||
| arnorm_backend = os.environ.get("TS_ARNORM_BACKEND", "auto").strip().lower() |
There was a problem hiding this comment.
Could we avoid duplicating ARNorm backend selection and triton_shmem policy in ServerArgs? tokenspeed-kernel already owns _arnorm_backend, backend dispatch, state eligibility, and per-call capability checks, while ordinary communication uses the existing AutoBackend/can_run pattern for this responsibility. The missing information appears to be runtime scheduling context such as DP or speculative execution guarantees. Could that context be passed to an ARNorm backend preflight API instead of parsing TS_ARNORM_BACKEND and hard-coding one backend's policy here? This would keep backend names and capabilities behind the backend boundary and avoid changing ServerArgs for every implementation.
2 participants
Summary (Human)
Add an alternative triton-native fused backend for AR+RMSNorm op using PyTorch symmetric memory. Results show improved performance across a majority of tested configs compared to Iris, the current default fused backend. Initial end-to-end results compared to unfused are mixed but show a promising improvement for the WS4 case. Benchmarking performed on GLM-5.2 and GPT-OSS-120B. Project effort has been paused, so defaults remain unchanged. Changes also made to unfused fallback behavior, AR+RMSNorm backend selection, and graph padding behavior.
Full project documentation, implementation history, and benchmarking results are available on companion branch
jeremwan/triton-shmem-experimentsin the source fork.Summary (AI)
TS_ARNORM_BACKEND=triton_shmemimplementation of fusedall-reduce + residual + RMSNorm for single-node MI350X tensor parallelism.
collectively fall back to symmetric-memory buffers if IPC setup fails.
autobehavior, add an explicit fusion-disable control,and restore the complete unfused operation whenever the backend declines.
captured-output configurations.
ordering.
Validated scope: contiguous BF16
(tokens, hidden)tensors and BF16 weights onsingle-node MI350X (
gfx950), TP 2/4/8. This is explicit opt-in only; no modelprofile or default dispatch policy is added.
The full investigation, profiles, and campaign artifacts are in the
finalized research snapshot.
Correctness note
Universal graph-padding sink routing previously landed in #531 and was narrowed
to DFLASH in temporary performance rollback #594. Ordinary non-speculative
GPT-OSS serving later reproduced slot-zero metadata corruption and passed after
the universal sink was restored. Current upstream is safer (reserved row 0,
batch-ordered page tables), so this PR treats terminal-sink routing as invariant
hardening rather than claiming the historical KV-underflow path is unchanged.
Performance evidence
The finalized research campaigns support an experimental backend, not default
promotion:
five-pair screen; WS2 was inconclusive and WS8 lost.
captured-operator matrix on MI355X; full-model serving was not qualified.
These results were collected on the research snapshot, not rerun after the port
to upstream
c0f41e0. Model-specific profiles remain reference material.Test Plan
From the repository root.
CPU/runtime integration:
MI350X communication and fallback:
Two-rank benchmark smoke:
Both paths completed; unfused remained faster for this smoke row, so it is not
used as performance evidence.