[WS1][kernels] Batch-invariant logprob (CUDA)#204
Open
KJLdefeated wants to merge 6 commits into
Open
Conversation
Implements batch_invariant_logp for selected-token log probabilities from materialized logits with row-local, batch-invariant semantics. - PyTorch NativeBatchInvariantLogpOp: FP32 row-wise reference with ignore_index handling and target validation. - Triton TritonBatchInvariantLogpOp: online-softmax forward with fixed vocab tiling and tile-wise backward using saved per-row lse. - Registry dispatch, PyTorch/Triton tests, and operator docs.
KJLdefeated
requested review from
EthanZero2Hero,
Flink-ddd,
bitborne and
inaniloquentee
as code owners
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a new ChangesBatch-invariant logp operator
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant KernelRegistry
participant BatchInvariantLogpSM90Op
participant _BatchInvariantLogpSM90Function
participant TritonBatchInvariantLogpOp
participant NativeBatchInvariantLogpOp
Caller->>KernelRegistry: get_op("batch_invariant_logp")
KernelRegistry-->>Caller: selected backend op
Caller->>BatchInvariantLogpSM90Op: __call__(logits, target_ids, ignore_index)
BatchInvariantLogpSM90Op->>BatchInvariantLogpSM90Op: check shape, dtype, alignment
alt SM90 supported
BatchInvariantLogpSM90Op->>_BatchInvariantLogpSM90Function: apply()
_BatchInvariantLogpSM90Function->>_BatchInvariantLogpSM90Function: call compiled SM90 kernel forward
_BatchInvariantLogpSM90Function-->>BatchInvariantLogpSM90Op: logp, lse
else unsupported input
BatchInvariantLogpSM90Op->>TritonBatchInvariantLogpOp: apply() (fallback)
alt Triton unavailable
TritonBatchInvariantLogpOp->>NativeBatchInvariantLogpOp: apply() (fallback)
NativeBatchInvariantLogpOp-->>BatchInvariantLogpSM90Op: logp
else Triton available
TritonBatchInvariantLogpOp-->>BatchInvariantLogpSM90Op: logp
end
end
BatchInvariantLogpSM90Op-->>Caller: selected log-probability
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
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: 9
🧹 Nitpick comments (3)
tests/test_batch_invariant_logp.py (2)
45-56: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNarrow the broad exception catch.
Catching bare
Exceptionto detect kernel availability masks unrelated failures (e.g., a broken_Cmodule import due to an ABI mismatch) as "SM90 not available," silently skipping tests instead of surfacing the real problem.🔧 Suggested narrowing
try: from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - except Exception: + except ImportError: return False🤖 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 `@tests/test_batch_invariant_logp.py` around lines 45 - 56, The _sm90_kernel_available helper is swallowing unrelated import failures by catching a broad Exception when importing rl_engine.kernels.ops.base. Narrow this to the specific import error(s) expected for missing kernel bindings, and keep the availability check in _sm90_kernel_available focused on genuine “not installed/not built” cases so ABI or module breakages still surface during tests.Source: Linters/SAST tools
831-1020: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing validation-error test coverage for the SM90 backend.
Every other backend has a dedicated validation test class (
TestValidationfor native at lines 294-324,TestTritonValidation/TestTritonCPUValidationat lines 788-824), but there is no equivalent forBatchInvariantLogpSM90Op(e.g., rejecting 1D logits, shape mismatches, or out-of-range targets). Given the PR objective explicitly calls out "expanded tests covering correctness, backward, ignore-index, alignment, and fallback behavior," validation parity across backends would close a gap.🤖 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 `@tests/test_batch_invariant_logp.py` around lines 831 - 1020, Add a SM90 validation test class for BatchInvariantLogpSM90Op to match the existing backend coverage in TestValidation and TestTritonValidation/TestTritonCPUValidation. In tests/test_batch_invariant_logp.py, create cases that assert invalid inputs are rejected for the BatchInvariantLogpSM90Op path, such as 1D logits, mismatched logits/target shapes, and targets outside the vocabulary range. Keep the tests grouped near the other TestCudaSM90* classes and use _get_op() so the new coverage clearly targets the SM90 backend.rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py (1)
22-35: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the fallback exception handling.
Only the import path should fall back for missing Triton. Catching every
Exceptionaround the backward kernel can mask real kernel failures and silently take the full[N, V]PyTorch fallback.Suggested shape
- try: + try: import triton from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( _BLOCK_V, _batch_invariant_logp_bwd_kernel, ) + except (ImportError, ModuleNotFoundError): # pragma: no cover - Triton missing + ... + else: grid = (num_tokens, triton.cdiv(vocab_size, _BLOCK_V)) _batch_invariant_logp_bwd_kernel[grid]( ... ) - except Exception: # pragma: no cover - Triton missing + # PyTorch fallback for the missing-Triton branch only.Also applies to: 67-89
🤖 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/cuda/loss/batch_invariant_logp.py` around lines 22 - 35, The _fallback_op helper currently catches every Exception, which can hide real Triton/kernel failures and incorrectly route to the PyTorch fallback. Narrow the exception handling in batch_invariant_logp so only Triton import/module-missing cases fall back to NativeBatchInvariantLogpOp, and let actual kernel errors propagate; update the same logic wherever the fallback is duplicated in the backward path referenced by TritonBatchInvariantLogpOp and NativeBatchInvariantLogpOp.Source: Linters/SAST tools
🤖 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_batch_invariant_logp.py`:
- Around line 77-88: The VRAM helper currently hardcodes its warmup and
iteration counts, so CLI overrides are ignored for memory measurement. Update
`_peak_vram_gb` to accept the same `warmup` and `iters` inputs used by
`_time_ms`, and change the call sites in `benchmark_batch_invariant_logp` to
pass `args.warmup` and `args.iters` so timing and VRAM runs stay consistent.
- Around line 63-92: The timing and memory helpers are hardcoded to torch.cuda
even though run_benchmark supports xpu, so the benchmark will break on Intel
GPUs. Update _time_ms and _peak_vram_gb to dispatch through the active backend
in device_ctx.device_type, using torch.xpu for xpu and torch.cuda for cuda/hip,
including synchronize, Event creation, empty_cache, memory_allocated,
reset_peak_memory_stats, and max_memory_allocated calls.
In `@csrc/cuda/batch_invariant_logp_kernel_sm90.cu`:
- Around line 179-200: The SM90 entrypoint in batch_invariant_logp_sm90_forward
currently assumes the current CUDA context matches logits, and it also only
validates target shape, not device. Add a CUDA device guard based on logits
before any allocation or launch, and add a TORCH_CHECK that target is on the
same CUDA device as logits (and CUDA at all) alongside the existing logits
checks. Keep the validation and setup near batch_invariant_logp_sm90_forward so
direct calls through the exported binding are safe even when the Python wrapper
is bypassed.
In `@docs/operators/batch-invariant-logp.md`:
- Around line 85-100: The “Forward + backward” measurements in the
batch-invariant logp docs are not reproducible from the current benchmark path.
Update the benchmark source used by run_benchmark so it actually exercises and
times a backward pass (including VRAM tracking), or else remove/qualify this
table as coming from a different benchmark variant. Use the existing benchmark
identifiers in benchmark_batch_invariant_logp.py and its run_benchmark flow to
keep the documentation aligned with the script.
In `@rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py`:
- Around line 114-130: The SM90 wrapper currently disables target validation by
default in BatchInvariantLogp.__call__ and apply, which can allow unsafe CUDA
gathers for bad non-ignored targets. Update the default in these methods so
validate matches the native operator’s behavior (true by default), and keep an
explicit opt-out path for callers that already validated targets. Preserve the
existing ignore_index handling and ensure the wrapper forwards the validate flag
consistently through BatchInvariantLogp.
- Around line 12-19: The `_sm90_supported` guard only checks tensor device type
and alignment, so tensors from non-Hopper CUDA devices can still route into the
SM90 TMA path. Update `_sm90_supported` to also verify the input logits’ CUDA
capability (cc_major == 9) on the tensor’s current device, alongside the
existing dtype and stride-multiple-of-16 checks, so cached registry entries
correctly fall back for non-SM90 GPUs.
In `@rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py`:
- Around line 115-124: The shape validator in _validate_shapes currently allows
non-integral target_ids and empty vocabularies to slip through. Update
_validate_shapes in batch_invariant_logp.py to first reject target_ids unless
they are an integer/long type before any cast happens in apply(), and add an
explicit check that logits.size(-1) is greater than zero before proceeding. Keep
the existing logits/target_ids shape checks in place after these new
validations.
In `@rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py`:
- Around line 202-212: The Triton batch-invariant logp path defaults validation
off, which makes it inconsistent with the native backend and can hide invalid
target IDs. Update the default in the BatchInvariantLogp call/apply flow so
validation is enabled by default, matching the native implementation; adjust the
default value on the validate parameter in the relevant method(s) and ensure
BatchInvariantLogp.apply preserves that default behavior.
- Around line 219-245: `_BatchInvariantLogpFunction.apply` currently validates
shapes and index ranges but not the dtype of `target_ids`, so non-integer inputs
can be silently converted later. Add an explicit integer-type check in the
validation path of the batch invariant logp function before
`target_ids.reshape(-1)` is used, and raise a clear `ValueError` for float,
bool, complex, or other non-integer dtypes. Keep the existing shape and range
checks in place and ensure the new check is near the `validate` block in
`batch_invariant_logp.py`.
---
Nitpick comments:
In `@rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py`:
- Around line 22-35: The _fallback_op helper currently catches every Exception,
which can hide real Triton/kernel failures and incorrectly route to the PyTorch
fallback. Narrow the exception handling in batch_invariant_logp so only Triton
import/module-missing cases fall back to NativeBatchInvariantLogpOp, and let
actual kernel errors propagate; update the same logic wherever the fallback is
duplicated in the backward path referenced by TritonBatchInvariantLogpOp and
NativeBatchInvariantLogpOp.
In `@tests/test_batch_invariant_logp.py`:
- Around line 45-56: The _sm90_kernel_available helper is swallowing unrelated
import failures by catching a broad Exception when importing
rl_engine.kernels.ops.base. Narrow this to the specific import error(s) expected
for missing kernel bindings, and keep the availability check in
_sm90_kernel_available focused on genuine “not installed/not built” cases so ABI
or module breakages still surface during tests.
- Around line 831-1020: Add a SM90 validation test class for
BatchInvariantLogpSM90Op to match the existing backend coverage in
TestValidation and TestTritonValidation/TestTritonCPUValidation. In
tests/test_batch_invariant_logp.py, create cases that assert invalid inputs are
rejected for the BatchInvariantLogpSM90Op path, such as 1D logits, mismatched
logits/target shapes, and targets outside the vocabulary range. Keep the tests
grouped near the other TestCudaSM90* classes and use _get_op() so the new
coverage clearly targets the SM90 backend.
🪄 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: b1bd5092-d130-4cca-b342-e7dd0294cf08
📒 Files selected for processing (14)
.gitignorebenchmarks/benchmark_batch_invariant_logp.pycsrc/cuda/batch_invariant_logp_kernel_sm90.cucsrc/ops.cppdocs/.nav.ymldocs/operators/README.mddocs/operators/batch-invariant-logp.mdrl_engine/_C.pyirl_engine/kernels/ops/cuda/loss/batch_invariant_logp.pyrl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.pyrl_engine/kernels/ops/triton/loss/batch_invariant_logp.pyrl_engine/kernels/registry.pysetup.pytests/test_batch_invariant_logp.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[WS1][kernels] SM90 Batch-Invariant logprob CUDA Kernel
Summary
Adds a Hopper (SM90) TMA CUDA kernel for the batch-invariant selected-token log-probability operator introduced in #199 plus a Native/Triton/CUDA benchmark. The operator computes, from already materialized logits:
The kernel streams the vocab through shared memory with TMA bulk-tensor copies and folds it into a running online log-sum-exp, so it never materializes the
[N, V]softmax and its extra forward memory is ~0. It is registered as thetop-priority CUDA backend on Hopper and falls back cleanly everywhere else.
csrc/cuda/batch_invariant_logp_kernel_sm90.cu)BatchInvariantLogpSM90Op)lse(no[N, V]materialization).batch_invariant_logplist only when the SM90 symbol is compiled oncc_major == 9.benchmarks/benchmark_batch_invariant_logp.pycompares Native/Triton/CUDA. New.This builds on the open branch for #199 (Native + Triton). Reviewed after #199 merges.
Implementation
csrc/cuda/batch_invariant_logp_kernel_sm90.cu.csrc/ops.cppand build source insetup.py'ssm90_srcs; stub inrl_engine/_C.pyi.rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py:_BatchInvariantLogpSM90Functionforward calls the compiled kernel;_batch_invariant_logp_bwd_kernel.rl_engine/kernels/registry.py.docs/operators/batch-invariant-logp.md; added the benchmark script.Validation environment
nvccnvidiachannel)KERNEL_ALIGN_FORCE_SM90=1,_C.batch_invariant_logp_sm90availableTORCH_CUDA_ARCH_LIST=9.0a,KERNEL_ALIGN_DEV_RPATH=1Correctness / Tests
Unit coverage
Result: 67 passed (14 new SM90-kernel tests: correctness fp32/bf16, large and unaligned vocab, single-token, 3-D, bitwise batch-invariance, backward, ignore-index, fp16 fallback; plus the registry-dispatch test now accepting the CUDA op).
Result: 14 passed. Registry selects
BatchInvariantLogpSM90Opon the H200:Vocab sweep (bf16, vs
log_softmax + gather)Correct on both TMA-aligned and unaligned (tail-path) vocab sizes:
Benchmarks
python benchmarks/benchmark_batch_invariant_logp.py— bf16, H200, 20 iters + 5 warmup. Memory is peak extra above baseline (MB).Forward
Forward + backward
Files
csrc/cuda/batch_invariant_logp_kernel_sm90.cu(new)rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py(new)benchmarks/benchmark_batch_invariant_logp.py(new)csrc/ops.cpp,setup.py,rl_engine/_C.pyi,rl_engine/kernels/registry.pytests/test_batch_invariant_logp.py,docs/operators/batch-invariant-logp.mdSummary by CodeRabbit
ignore_indexhandling, validation, and graceful fallback when an optimized backend isn’t available.