[Feat][Kernel] Add TP Fused linear logp Triton - #208
Conversation
|
Warning Review limit reached
Next review available in: 12 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe Triton ChangesTriton tensor-parallel linear_logp
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Op as TritonLinearLogpOp.apply
participant Wrapper as _triton_tensor_parallel_linear_logp
participant Function as _TensorParallelTritonLinearLogpFunction
participant Kernel as _run_forward_kernel
participant TPGroup as Tensor-parallel group
Op->>Wrapper: dispatch tensor-parallel inputs
Wrapper->>Function: normalize metadata and invoke autograd path
Function->>Kernel: compute local logp and LSE
Function->>TPGroup: reduce target logits with SUM
Function->>TPGroup: reduce LSE with MAX then SUM
Function-->>Wrapper: return global target log-probabilities
Function->>Function: delegate backward computation
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.
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.
| local_vocab_size=weight.size(0), | ||
| global_vocab_size=global_vocab_size, | ||
| ) | ||
| _validate_global_targets(target_1d, global_vocab_size, tp_group) |
There was a problem hiding this comment.
Calling _validate_tp_vocab_partition and _validate_global_targets here introduces extra collectives (all_gather, all_reduce) and .item() host syncs on the hot path for every forward step.
Since the TP vocabulary partition is static after initialization, evaluating this on every pass will become a severe latency bottleneck in a distributed training loop. I recommend caching the validation state (e.g., executing it only once per tp_group) or gating these checks behind a validation/debug flag to keep the production forward pass as lightweight as possible.
|
please resolve the code conflicts first, then we will merge this PR. |
Resolve conflicts by keeping main's TP backward helpers (tensor_parallel_linear_logp_backward / _sm90_linear_logp_backward), which add needs_input_grad gating and avoid per-chunk host syncs, and dropping this branch's duplicate tensor_parallel_chunked_linear_logp_backward. The new Triton TP path now routes its backward through the shared helper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK
Use the memoized _validate_tp_vocab_partition_cached (the partition is static per tp_group after init) and gate _validate_global_targets behind the existing RL_KERNEL_LINEAR_LOGP_VALIDATE_TP_TARGETS debug flag, matching the native and SM90 TP paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rl_engine/kernels/ops/triton/loss/linear_logp.py (1)
204-217: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReject TP vocab shards with
local_vocab_size <= 0in_validate_tp_vocab_partition.
should_use_tensor_parallel_linear_logpalready rejects zero local vocab in non-TP or single-rank cases, but_validate_tp_vocab_partitionstill acceptslocal_vocab_size <= 0whenworld_size >= 2. A zero shard creates a [x, x) rank range; the partition loop then skips it, an out-of-range global target can be clamped into another rank’s local column, and the backward path can divide bylocal_vocab == 0.🤖 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 204 - 217, Update _validate_tp_vocab_partition to reject any local_vocab_size <= 0 for tensor-parallel groups, including world_size >= 2, before validating or iterating partition ranges. Preserve the existing validation behavior for positive shard sizes and ensure invalid zero-sized shards cannot reach the forward or backward paths.
🤖 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.
Outside diff comments:
In `@rl_engine/kernels/ops/triton/loss/linear_logp.py`:
- Around line 204-217: Update _validate_tp_vocab_partition to reject any
local_vocab_size <= 0 for tensor-parallel groups, including world_size >= 2,
before validating or iterating partition ranges. Preserve the existing
validation behavior for positive shard sizes and ensure invalid zero-sized
shards cannot reach the forward or backward paths.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ddd5b3f0-5d03-4f5f-b041-dccdd64c4301
📒 Files selected for processing (1)
rl_engine/kernels/ops/triton/loss/linear_logp.py
…_vocab_partition.
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM now, Thank you for update.
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