feat(kernels): add fused masking + variable-length pack-and-pad op #182
feat(kernels): add fused masking + variable-length pack-and-pad op #182Chen-BUPT wants to merge 9 commits into
Conversation
…AM) (RL-Align#42) Measure TritonPackOp vs a PyTorch boolean-index baseline for pack latency, and the end-to-end peak VRAM of dense logp vs pack->logp to quantify the memory saving on sparse masks (the motivation behind RL-Align#42). Follows the existing benchmark_ratio_kl.py conventions (CUDA-event median timing, max_memory_allocated, CSV output, --smoke).
Pack hidden states before the vocab projection so the dense [B,S,V] logits are never materialized for masked-out tokens, which is the actual RL-Align#42 saving. Drop the fp32 upcast in selected-logp (use logits - logsumexp) so the dense path's peak memory reflects the logits, not an fp32 copy. Add --hidden-dim.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a PyTorch pack-and-pad operator with autograd support, registry dispatch across platforms, correctness tests, a configurable latency and memory benchmark, and documentation. ChangesPack-and-pad operator
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant KernelRegistry
participant NativePackOp
participant _PackFunction
Caller->>KernelRegistry: get_op("pack")
KernelRegistry-->>Caller: NativePackOp
Caller->>NativePackOp: __call__(x, mask)
NativePackOp->>_PackFunction: apply(x, mask)
_PackFunction-->>NativePackOp: packed, cu_seqlens
NativePackOp-->>Caller: packed, cu_seqlens
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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.
Actionable comments posted: 2
🧹 Nitpick comments (2)
rl_engine/kernels/ops/pytorch/packing/pack.py (1)
27-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicate dimension check.
The
mask.dim() < 1guard at Lines 35-36 is dead code — it repeats the identical check at Lines 28-29, and the shape comparison in between never altersmask.dim().♻️ Proposed cleanup
def _validate(x: torch.Tensor, mask: torch.Tensor) -> None: if mask.dim() < 1: raise ValueError("mask must have at least one dimension.") if mask.shape != x.shape[: mask.dim()]: raise ValueError( f"mask shape {tuple(mask.shape)} must match the leading dims of " f"x.shape {tuple(x.shape)} (expected {tuple(x.shape[: mask.dim()])})." ) - if mask.dim() < 1: - raise ValueError("mask must have at least one dimension.")🤖 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/packing/pack.py` around lines 27 - 37, The _validate function contains a duplicate dimension check for mask.dim() < 1 that appears after the shape validation logic. Remove the second occurrence of the identical check (the one appearing after the shape comparison) since the intermediate validation logic does not modify mask.dim() and therefore makes the repeated check dead code. Keep only the initial dimension check and remove the redundant check that follows the shape validation.rl_engine/kernels/ops/triton/packing/pack.py (1)
99-102: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffGather grid launches a program for every source row, including inactive ones.
The grid is sized over
n_rows = B*S, and each program loadsdestand early-exits when the row is inactive (dest < 0). For the low-density packing this op targets (e.g. 0.05), the large majority of launched programs do no work. Launching over then_activerows via a packed→source inverse index would avoid the wasted launches on the hot path. This is the design tradeoff already acknowledged in the PR, so feel free to defer.🤖 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/packing/pack.py` around lines 99 - 102, The grid in the `_pack_gather_kernel` launch is currently sized over `n_rows` which includes all source rows, causing programs to be launched for inactive rows that do no work (they just early-exit when dest < 0). To optimize this for sparse packing scenarios, resize the grid to be sized over `n_active` instead of `n_rows`, and introduce a packed-to-source inverse index mapping that the kernel can use to look up which source rows are actually active. This eliminates wasted kernel launches on inactive rows while keeping the same kernel logic.
🤖 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.
Inline comments:
In `@benchmarks/benchmark_pack.py`:
- Around line 214-252: The broad except Exception block that catches all
exceptions and sets status to "blocked" masks real execution errors
(kernel/runtime/math failures) as if the candidate backend is simply
unavailable. Replace the generic except Exception handler with specific
exception handling that only catches exceptions that genuinely indicate
unavailability (such as ImportError for missing kernel_registry or RuntimeError
for unavailable backends), and allow other exceptions like runtime/math/kernel
errors to propagate or be handled separately so they surface as actual failures
rather than being silently masked as blocked candidates.
In `@tests/test_pack.py`:
- Around line 18-31: The TritonPackOp import on line 18 happens unconditionally
before the Triton availability check, causing test collection to fail in
environments without Triton. Move the TritonPackOp import into the try block
alongside the triton import, and add a fallback assignment in the except block
(such as TritonPackOp = None) so the name can be safely referenced elsewhere in
the code. This ensures tests are properly skipped by the requires_triton_cuda
marker instead of failing during collection.
---
Nitpick comments:
In `@rl_engine/kernels/ops/pytorch/packing/pack.py`:
- Around line 27-37: The _validate function contains a duplicate dimension check
for mask.dim() < 1 that appears after the shape validation logic. Remove the
second occurrence of the identical check (the one appearing after the shape
comparison) since the intermediate validation logic does not modify mask.dim()
and therefore makes the repeated check dead code. Keep only the initial
dimension check and remove the redundant check that follows the shape
validation.
In `@rl_engine/kernels/ops/triton/packing/pack.py`:
- Around line 99-102: The grid in the `_pack_gather_kernel` launch is currently
sized over `n_rows` which includes all source rows, causing programs to be
launched for inactive rows that do no work (they just early-exit when dest < 0).
To optimize this for sparse packing scenarios, resize the grid to be sized over
`n_active` instead of `n_rows`, and introduce a packed-to-source inverse index
mapping that the kernel can use to look up which source rows are actually
active. This eliminates wasted kernel launches on inactive rows while keeping
the same kernel logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b134abab-ef3b-4adc-8ca3-c50dd21ef7ae
📒 Files selected for processing (7)
benchmarks/benchmark_pack.pyrl_engine/kernels/ops/pytorch/packing/__init__.pyrl_engine/kernels/ops/pytorch/packing/pack.pyrl_engine/kernels/ops/triton/packing/__init__.pyrl_engine/kernels/ops/triton/packing/pack.pyrl_engine/kernels/registry.pytests/test_pack.py
The second mask.dim() < 1 guard was dead code: the intervening shape check does not alter mask.dim(). Addresses CodeRabbit review on RL-Align#182.
| packed = flat_x.index_select(0, index) | ||
|
|
||
| # cu_seqlens: prefix-sum of per-row active counts, for varlen consumers. | ||
| per_row_active = mask.reshape(mask.shape[0], -1).to(torch.int64).sum(dim=1) |
There was a problem hiding this comment.
packed uses mask.to(bool), but cu_seqlens sums the raw mask. Non-bool masks can produce wrong prefix sums. Please either require bool masks or compute counts from the bool mask.
Packing selects rows via mask.to(bool) (nonzero == active), but cu_seqlens
summed the raw mask, so a non-bool mask (e.g. values in {0, 2}) inflated the
prefix sum beyond the number of rows actually packed. Count from the same
bool mask so cu_seqlens always matches the packed row count. Adds a
regression test. Addresses review feedback on RL-Align#182.
Co-authored-by: Cursor <cursoragent@cursor.com>
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
Excellent work! but need add related doc file and Please revise the following review comments. Everything else is very well done. Once these are done, I will approve it. Thanks.
| counts = active.to(torch.int64) | ||
| # Exclusive prefix sum: position of each active row in the packed buffer. | ||
| excl = torch.cumsum(counts, dim=0) - counts | ||
| n_active = int(counts.sum().item()) |
There was a problem hiding this comment.
(_dest_index, n_active = int(counts.sum().item())) .item() forces a device-to-host sync that blocks the stream, since torch.empty(n_active, ...) depends on this Python int. It's hard to avoid given the output shape depends on n_active, but please add a comment marking this as an intentional sync point so it isn't duplicated on the hot path later.
| grid = (n_rows, triton.cdiv(T, _BLOCK_T)) | ||
| _pack_gather_kernel[grid](src, packed, dest, T, BLOCK_T=_BLOCK_T) | ||
|
|
||
| per_row_active = flat_mask.to(torch.bool).reshape(mask.shape[0], -1).to(torch.int64).sum(1) |
There was a problem hiding this comment.
(per_row_active = ...reshape(mask.shape[0], -1)...) cu_seqlens treats mask.shape[0] as the batch dim, but packing supports an arbitrary-dim leading mask. For mask.dim() > 2 the per-row grouping no longer matches the row-major pack order. Please document that cu_seqlens is only well-defined for a 2D [B, S] mask.
| """forward: gather active rows; backward: scatter grads to active rows.""" | ||
|
|
||
| @staticmethod | ||
| def forward(ctx, x: torch.Tensor, mask: torch.Tensor): |
There was a problem hiding this comment.
(_PackFunction.forward / _validate) If x and mask are on different devices, index_select throws a low-level device-mismatch error. Please add a device check in _validate for a clear message, consistent with the shape check.
| return _PackFunction.apply(x, mask) | ||
|
|
||
| @staticmethod | ||
| def unpack( |
There was a problem hiding this comment.
(NativePackOp.unpack) When tail_shape is passed explicitly and the element count happens to match but dims differ, reshape silently produces a wrong shape. Consider asserting tail_shape against packed.shape[1:].
| return parser | ||
|
|
||
|
|
||
| def main() -> None: |
There was a problem hiding this comment.
(main(), OOM handler) The OOM-row dict duplicates the structure of _pack_row's return value, so it's easy to miss when CSV_COLUMNS changes. Consider a _blank_row(config, status, notes) helper.
| assert torch.equal(packed_t, packed_n) | ||
| assert cu_t.tolist() == [0, 2, 4] | ||
|
|
||
|
|
There was a problem hiding this comment.
Add a negative test where x and mask are on different devices, asserting a clear error is raised (matches the native/triton device-check suggestion above).
| packed_n, cu_n = NativePackOp()(x, batch.completion_mask) | ||
| assert torch.equal(packed_t, packed_n) | ||
| assert torch.equal(cu_t, cu_n) | ||
|
|
There was a problem hiding this comment.
(Triton tests) Add a dedicated n_active == 0 test asserting the Triton path doesn't launch the kernel (the if n_active > 0 guard) and returns the correct [0, ...] shape.
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
Correction: my previous review was approved by mis-click. The doc and the review comments above still need to be addressed before this can be merged re-submitting as Request changes. Will approve once they're done. Thanks!
KJLdefeated
left a comment
KJLdefeated
left a comment
There was a problem hiding this comment.
The native op is clean and correct gather/scatter, race-free (unique dest indices, no atomics), drift-free by construction. Aligning with Flink's existing requested-changes (device check, unpack tail_shape assert, cu_seqlens-is-2D-only doc, _blank_row helper, the negative-device test — all fair, should land).
| class TritonPackOp: | ||
| """Triton fused masking + variable-length packing (pack-and-pad). | ||
|
|
||
| Forward packs the active rows of ``x`` (selected by ``mask``) into a | ||
| contiguous ``[Total_Active, *tail]`` tensor and returns the per-row | ||
| ``cu_seqlens`` prefix-sum. Backward scatters the upstream gradient back to the | ||
| original ``[*mask.shape, *tail]`` layout, leaving zeros at inactive positions. | ||
| Numerically identical to ``NativePackOp``; CUDA & ROCm via Triton. | ||
| """ |
There was a problem hiding this comment.
What does the Triton path buy over index_select today? Reading the benchmark honestly: TritonPackOp runs 1.3–4× slower than the PyTorch boolean index baseline (speedup 0.23–0.78), and its output and peak memory are identical to NativePackOp. It's the same copy. So as it stands the Triton kernel is strictly dominated by x[mask] / the native op on every axis.
The only justification I can see is future fusion, folding this gather into the downstream lm_head/logp kernel so the packed hidden is never materialized as a standalone tensor. But that fused consumer isn't in this PR, so right now we'd be committing to maintain a slower-than-native CUDA/ROCm kernel with no caller that needs it.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
rl_engine/kernels/ops/pytorch/packing/pack.py (2)
94-96: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify redundant conditional expression for
tail_shape.In Python, unpacking an empty tuple inside a function argument list naturally expands to nothing. The explicit
if tail_shape:conditional evaluates to the exact same result whethertail_shapeis empty or not, and can be safely removed for cleaner code. (This simplification also applies to thex.reshapeconditional in the forward pass).♻️ Proposed refactoring
- grad_flat = grad_packed.new_zeros( - (ctx.flat_rows, *tail_shape) if tail_shape else (ctx.flat_rows,) - ) + grad_flat = grad_packed.new_zeros((ctx.flat_rows, *tail_shape))🤖 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/packing/pack.py` around lines 94 - 96, Remove the redundant tail_shape conditional in the grad_flat allocation so the tuple uses direct unpacking. Apply the same simplification to the x.reshape conditional in the forward pass, preserving identical behavior for both empty and non-empty tail_shape values.
74-76: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider cloning the empty view to strictly preserve the non-aliasing contract.
Although an empty tensor naturally prevents data corruption upon mutation, returning a direct view of
flat_xskips the strict non-aliasing intent established in theall_activefast path. Appending.clone()keeps the returned tensor fully detached fromx's storage without incurring any memory cost.♻️ Proposed refactoring
elif ctx.none_active: - packed = flat_x[:0] + packed = flat_x[:0].clone() ctx.save_for_backward()🤖 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/packing/pack.py` around lines 74 - 76, Update the none_active branch of the packing operation to clone the empty tensor produced from flat_x before assigning it to packed. Preserve the existing empty shape and ctx.save_for_backward() behavior while ensuring the result does not alias flat_x storage.
🤖 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/packing/pack.py`:
- Around line 94-96: Remove the redundant tail_shape conditional in the
grad_flat allocation so the tuple uses direct unpacking. Apply the same
simplification to the x.reshape conditional in the forward pass, preserving
identical behavior for both empty and non-empty tail_shape values.
- Around line 74-76: Update the none_active branch of the packing operation to
clone the empty tensor produced from flat_x before assigning it to packed.
Preserve the existing empty shape and ctx.save_for_backward() behavior while
ensuring the result does not alias flat_x storage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f246fdae-274a-4e14-b50c-dfdc1849064e
📒 Files selected for processing (8)
benchmarks/benchmark_pack.pydocs/.nav.ymldocs/benchmarking/README.mddocs/operators/README.mddocs/operators/pack-and-pad.mdrl_engine/kernels/ops/pytorch/packing/pack.pyrl_engine/kernels/registry.pytests/test_pack.py
🚧 Files skipped from review as they are similar to previous changes (1)
- benchmarks/benchmark_pack.py
4 participants
Fused masking + variable-length pack-and-pad op (#42)
What this adds
A
packoperator that compacts the active rows of a dense[B, S, *tail]tensor (selected by a
[B, S]mask) into a contiguous[Total_Active, *tail]tensor, returns per-row
cu_seqlens, and scatters gradients back to the denselayout on the backward pass.
NativePackOpSyntheticRLKernelBatch.compact_completion_values.TritonPackOpRegistry dispatch for
"pack": Triton on GPU, PyTorch fallback on CPU.Correctness
tests/test_pack.pypass on an NVIDIA H20 (SM90, CUDA 13.0).gradcheck.drift =
0.000e+00across all benchmarked shapes.Why it matters: end-to-end VRAM
In RL training only the response / non-padding tokens contribute to the loss.
Packing the hidden states before the vocab projection means the full
[B, S, V]logits are never materialized for masked-out tokens — exactly thesaving #42 targets ("saved memory can be used for larger batches or longer
CoT").
Benchmark:
hidden=4096 -> lm_head -> logits -> selected logp, bf16,B=32, S=1024, comparing dense (full logits) vs pack-then-project.vocab = 131072
vocab = 32768
(Full data:
pack_h20.csv.)How to read the numbers
response — the common RL case), packing before the projection cuts peak logp
memory by up to ~88 % (density 0.05). The
17 GB -> 2 GBdrop lets thesame GPU fit a much larger batch or longer CoT.
pack overhead), as expected — there is nothing to compact, so nothing is
saved. This confirms the measurement is honest.
speedup< 1) is reported as-is. The pack op itself ismemory-bound and a touch slower than PyTorch's boolean indexing
(
index_selectis already highly tuned). Its absolute cost is 0.06–0.35 ms,negligible against the multi-GB memory it saves. This PR targets the VRAM
win ([FEAT][kernels]: implement Fused Masking and Variable-Length Sequence Packing (Pack-and-Pad) #42's stated motivation);
Reproduce
PYTHONPATH=. python -m pytest tests/test_pack.py -v # correctness (needs CUDA for Triton cases) PYTHONPATH=. python benchmarks/benchmark_pack.py \ --num-prompts 4 --g-sizes 8 --hidden-dim 4096 \ --mask-densities 0.05,0.1,0.3,0.5,1.0 \ --completion-lens 1024 --vocab-sizes 32768,131072 \ --output benchmarks/results/pack_h20.csvNotes for reviewers
NativePackOp(CPU path); Triton tests skip cleanly.they run on ROCm via Triton without CUDA-specific intrinsics.
needs-gpu-cilabel so the GPU CI exercises the Tritonpath.
Summary by CodeRabbit
New Features
Documentation
Tests