[Feat][Kernel] Add TP Fused linear logp Triton#208
Conversation
KJLdefeated
requested review from
EthanZero2Hero,
Flink-ddd and
inaniloquentee
as code owners
📝 WalkthroughWalkthroughThe tensor-parallel backward computation for linear_logp was consolidated into a new shared helper function, ChangesTensor-parallel linear_logp backward consolidation
Estimated code review effort: 4 (Complex) | ~50 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
rl_engine/kernels/ops/triton/loss/linear_logp.py (1)
201-208: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftPer-forward TP validation adds collectives + host syncs on the hot path.
_validate_tp_vocab_partition(2×all_gather+.item()syncs) and_validate_global_targets(all_reduceMIN/MAX/MAX +.item()syncs) run on every forward call. As flagged in the PR description, these extra collectives and host syncs can become a latency bottleneck at higher TP degree or higher inter-rank latency. Consider gating them behind a one-time/debug validation (e.g. validate once per (tp_group, shard-layout) and cache the result, or behind an env/debug flag) since the vocab partition is static across steps.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/kernels/ops/triton/loss/linear_logp.py` around lines 201 - 208, The per-forward validation in linear_logp is adding expensive collectives and host syncs on the hot path. Update the validation flow around _validate_tp_vocab_partition and _validate_global_targets to avoid running on every forward call, e.g. by guarding them with a debug/env flag or caching a one-time result per tp_group and shard layout. Keep the existing checks in place for validation mode, but make the default forward path skip the repeated all_gather/all_reduce and .item() syncs.rl_engine/kernels/ops/pytorch/loss/linear_logp.py (1)
458-474: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDrop the per-chunk host sync in the backward loop.
bool(owns_target.any().item())forces a device→host synchronization on every chunk iteration, in a hot backward path. The guard is only an optimization: when no rows own a target,rowsis empty and the scatter-add is a safe no-op — which is exactly how the non-TPchunked_linear_logp_backwardhandles it (it scatters unconditionally). Removing the sync keeps correctness and improves throughput on higher-rank/higher-latency setups.♻️ Proposed change to remove the per-chunk sync
dz = -torch.exp(logits.float() - lse[i0:i1].unsqueeze(1)) local_idx = target_1d[i0:i1] - vocab_start_index owns_target = (local_idx >= 0) & (local_idx < local_vocab) - if bool(owns_target.any().item()): - rows = torch.arange(i1 - i0, device=dz.device)[owns_target] - dz[rows, local_idx[owns_target].long()] += 1.0 + rows = torch.arange(i1 - i0, device=dz.device)[owns_target] + dz[rows, local_idx[owns_target].long()] += 1.0 dz *= g[i0:i1].unsqueeze(1)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/kernels/ops/pytorch/loss/linear_logp.py` around lines 458 - 474, The backward loop in linear_logp’s chunked path is doing an unnecessary device-to-host sync via bool(owns_target.any().item()) on every iteration. Update the chunk handling in the linear_logp backward logic to avoid the host-side guard and perform the target scatter-add unconditionally, matching chunked_linear_logp_backward behavior. Keep the existing owns_target/local_idx logic in place, but remove the .item()-based branching so the hot path stays fully on device.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@rl_engine/kernels/ops/pytorch/loss/linear_logp.py`:
- Around line 458-474: The backward loop in linear_logp’s chunked path is doing
an unnecessary device-to-host sync via bool(owns_target.any().item()) on every
iteration. Update the chunk handling in the linear_logp backward logic to avoid
the host-side guard and perform the target scatter-add unconditionally, matching
chunked_linear_logp_backward behavior. Keep the existing owns_target/local_idx
logic in place, but remove the .item()-based branching so the hot path stays
fully on device.
In `@rl_engine/kernels/ops/triton/loss/linear_logp.py`:
- Around line 201-208: The per-forward validation in linear_logp is adding
expensive collectives and host syncs on the hot path. Update the validation flow
around _validate_tp_vocab_partition and _validate_global_targets to avoid
running on every forward call, e.g. by guarding them with a debug/env flag or
caching a one-time result per tp_group and shard layout. Keep the existing
checks in place for validation mode, but make the default forward path skip the
repeated all_gather/all_reduce and .item() syncs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 19b2aa5f-d05f-4d12-9075-ab9d4fea8163
📒 Files selected for processing (3)
rl_engine/kernels/ops/cuda/loss/linear_logp.pyrl_engine/kernels/ops/pytorch/loss/linear_logp.pyrl_engine/kernels/ops/triton/loss/linear_logp.py
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
Great！Here's a small suggestion for your consideration.
2 participants
[FEAT][kernels] Add tensor-parallel path to the Triton linear_logp kernel
Summary
Gives the Triton
linear_logpop a real vocab-parallel (tensor-parallel)path, so a vocab-sharded LM head runs the fused Triton forward on each rank instead of silently falling back to the pure-PyTorch TP implementation._TensorParallelTritonLinearLogpFunction)TritonLinearLogpOpnow routes TP calls.tensor_parallel_chunked_linear_logp_backward)tests/linear_logp_tp.py --op-source tritonBuilds on the merged #122 (fused linear_logp) and #189 (native + SM90 TP).
Implementation
Triton TP forward —
rl_engine/kernels/ops/triton/loss/linear_logp.py_run_forward_kernel(...) -> (logp, lse),reused by both the plain and TP forwards (no kernel change).
_TensorParallelTritonLinearLogpFunctionmirrors the SM90 TP path(
_FusedTensorParallelLinearLogpSM90Function):fused forward over its shard, giving the local
lseand (for the ownedtoken) the target logit, recovered as
target_logit = logp + lse;all_reduce(MAX)the running max, rescale +all_reduce(SUM)the runningsum,
all_reduce(SUM)the target logit;logp = target_logit - global_lse;_validate_tp_vocab_partition+_validate_global_targets(from [FEAT][kernels] Add tensor-parallel linear_logp path #189)guarantee a contiguous
[0, V)partition and in-range targets, so eachtarget has exactly one owner — matching the SM90 path (no extra owner-count
all-reduce).
TritonLinearLogpOp.applynow dispatches TP calls to_triton_tensor_parallel_linear_logp.Shared TP backward —
rl_engine/kernels/ops/pytorch/loss/linear_logp.pytensor_parallel_chunked_linear_logp_backward.Validation
Validated base on
torchrun --standalone --nproc_per_node=4 tests/linear_logp_tp.py --op-source tritonreference from docs/operators/linear-logp-tp-test.mdResults (backend=nccl, world_size=4)
--atol 1e-3 --rtol 1e-3)Backend equivalence (bf16, same harness/config, 4 NCCL ranks):
Unit tests —
tests/test_linear_logp.py, 27 passed (incl. the SM90 TP tests that exercise the shared backward).Files
rl_engine/kernels/ops/triton/loss/linear_logp.py— Triton TP forward + dispatch.rl_engine/kernels/ops/pytorch/loss/linear_logp.py— shared TP backward helper.rl_engine/kernels/ops/cuda/loss/linear_logp.py— SM90 TP backward now uses the shared helper.Potential Issues: per-forward validation collectives
The exactly-one-owner guarantee comes from
_validate_tp_vocab_partitionand_validate_global_targets, which run on every forward and issue collectives on top of the 3 all-reduces the TP reduction itself needs (target logit, runningmax, running sum):
_validate_tp_vocab_partition→ 2all_gathers (shard ranges + declared global sizes). The vocab partition is fixed at model construction, so this re-proves a static invariant every step._validate_global_targets→ up to 3all_reduces (invalid-flag MAX, target MIN, target MAX) plus.item()host-device syncs. Data-dependent, but the syncs serialize the stream.So validation roughly doubles the collective count per forward. The messages are tiny (scalar / 2-element), so this is latency-bound, not bandwidth-bound — but on high-rank or high-latency interconnects the extra round-trips (and the host
syncs) can become a connection/throughput bottleneck in a tight training loop. This behavior is inherited unchanged from #189 (native + SM90 run the same checks); the Triton path mirrors it for consistency rather than diverging.
Summary by CodeRabbit
New Features
Bug Fixes