[draft][WS1][kernels] Batch-invariant logprob Ascend C Kernel - #297
Conversation
Signed-off-by: chenyang <2082464740@qq.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR adds an Ascend CANN implementation for batch-invariant selected-token log probabilities. It integrates NPU detection, custom extension compilation, runtime dispatch, fallback behavior, tests, documentation, and accelerator-aware benchmarking. ChangesAscend batch-invariant log-probability
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: 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: 9
🧹 Nitpick comments (6)
rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py (1)
30-37: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueCache the fallback operator.
_fallback_op()performs a module import and constructs a newNativeBatchInvariantLogpOpon every call. For fp16 inputs this runs on each forward. Build it once.♻️ Proposed fix
+_FALLBACK_OP = None + + def _fallback_op(): """Portable op for inputs the Ascend forward cannot take. Triton rejects non-CUDA devices, so on NPU the only fallback is native. """ - from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp - - return NativeBatchInvariantLogpOp() + global _FALLBACK_OP + if _FALLBACK_OP is None: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, + ) + + _FALLBACK_OP = NativeBatchInvariantLogpOp() + return _FALLBACK_OP🤖 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/ascend/loss/batch_invariant_logp.py` around lines 30 - 37, Cache the operator returned by _fallback_op so the NativeBatchInvariantLogpOp import and construction occur only once, while preserving the existing fallback behavior for fp16 inputs.setup.py (2)
218-221: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnchor the source glob to the setup.py directory.
Path("csrc/ascend")resolves against the current working directory. If a build front-end invokes setup.py from another directory, the glob returns no files and the build raises the "no .asc sources" error. Use a path derived from__file__.♻️ Proposed fix
- asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + root = Path(__file__).parent.resolve() + asc_srcs = sorted(str(p) for p in (root / "csrc" / "ascend").glob("*.asc"))🤖 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 `@setup.py` around lines 218 - 221, Update the source discovery in the Ascend extension setup flow around the asc_srcs expression to resolve csrc/ascend relative to setup.py via __file__, rather than the current working directory. Preserve the existing sorting, missing-source RuntimeError, and Extension construction behavior.
6-8: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the
distutilsimports.The project supports Python 3.10 and later. Python 3.12 removed
distutils, and directsetup.pyexecution imports it beforesetuptoolscan install its compatibility shim. Usesetuptools.errors.CompileErrorandshutil.whichinstead. Thesetuptools>=64build requirement providesCompileError.🤖 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 `@setup.py` around lines 6 - 8, Replace the distutils imports in setup.py: import CompileError from setuptools.errors and use shutil.which instead of distutils.spawn.find_executable, while retaining sysconfig and updating any corresponding call sites to the new executable lookup symbol.tests/test_batch_invariant_logp.py (1)
1195-1213: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the empty batch and for
validate=True.Two reachable paths have no test:
- An empty batch (
logitsshaped[0, V]). The kernel computesblockNum = 0for this input. See the related comment oncsrc/ascend/batch_invariant_logp_ascend.asc.validate=Truewith an out-of-range target. The Ascend wrapper defaults tovalidate=False, and the kernel returns a silent-lsefor out-of-range targets rather than raising.Do you want me to generate these test cases?
🤖 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 1195 - 1213, Extend TestAscendIgnoreIndex with coverage for an empty logits batch shaped [0, _VC] and matching empty targets, asserting the operation returns an empty output. Add a validate=True case using an out-of-range target and assert the Ascend operation raises the expected validation error, while preserving the existing ignore-index and default validate=False coverage.csrc/ascend/batch_invariant_logp_ascend.asc (1)
164-172: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe 1-element
Logreads seven uninitialized floats.
scalar.SetValue(0, sumExp)initializes only element 0, andAscendC::Log(scalar, scalar, 8)processes elements 0 through 7. Elements 1 through 7 hold whatever remained in UB. Only element 0 is read afterwards, so the result is correct, but the operation computeslogon undefined data on every row. Initialize the padding once inInitto keep the behavior defined.♻️ Proposed fix
AscendC::LocalTensor<float> scalar = scalarBuf_.Get<float>(); scalar.SetValue(0, sumExp); + for (uint32_t i = 1; i < 8; ++i) { + scalar.SetValue(i, 1.0f); + } AscendC::SetFlag<AscendC::HardEvent::S_V>(eventSV_); // scalar write -> vector op🤖 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 `@csrc/ascend/batch_invariant_logp_ascend.asc` around lines 164 - 172, Initialize all eight elements of scalarBuf_ during Init before any Log operation, ensuring the padding elements have defined values. Update the scalar setup around scalarBuf_, scalar, and AscendC::Log so per-row processing only overwrites element 0 while preserving defined padding for the length-8 vector operation.docs/operators/batch-invariant-logp.md (1)
59-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the Ascend
validatedefault and the out-of-range behavior.Lines 154-157 state that the PyTorch backend validates target ranges by default and that Triton defaults to
validate=False.BatchInvariantLogpAscendOp.applyalso defaults tovalidate=False, and the Ascend kernel returns-lsefor an out-of-range target instead of raising. Add that to this section so callers know the reference operator and the Ascend operator disagree on invalid input.🤖 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 `@docs/operators/batch-invariant-logp.md` around lines 59 - 66, Update the Ascend backend documentation to state that BatchInvariantLogpAscendOp.apply defaults to validate=False and that the Ascend kernel returns -lse for out-of-range targets rather than raising. Explicitly contrast this with the PyTorch reference backend’s default target-range validation so callers understand the invalid-input behavior difference.
🤖 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`:
- Line 11: Update the description of NativeBatchInvariantLogpOp to state that it
materializes FP32 intermediate tensors for the explicit row-wise logsumexp
calculation, removing the claim that it materializes a full [N, V] log_softmax
tensor.
- Around line 88-95: Update _maybe_accelerated_op to accept or access device_ctx
and only call/select _maybe_sm90_op when device_ctx.device_type is "cuda", and
only call/select _maybe_ascend_op when it is "npu". Preserve the existing
(label, operator) return values and the (None, None) fallback for unsupported or
unavailable devices.
In `@csrc/ascend/batch_invariant_logp_ascend.asc`:
- Around line 141-143: Update the kernel header comment in
csrc/ascend/batch_invariant_logp_ascend.asc around the target selection in
BatchInvariantLogpAscendOp to state that invalid non-ignore targets remain
unvalidated and produce -lse rather than an error; no kernel logic change is
requested. Update docs/operators/batch-invariant-logp.md around lines 59-66 to
document the Ascend backend’s validate=False default and its -lse behavior for
out-of-range targets.
- Around line 270-285: In the batch setup before stream acquisition and kernel
launch, handle numRows == 0 by returning the already allocated empty logp and
lse tensors. Keep the existing validation and non-empty blockNum/kernel-launch
path unchanged.
In `@rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py`:
- Around line 62-80: Reduce the backward memory usage in the static method
backward by reusing a single [N, V] buffer: transform the logits buffer in place
into probabilities, apply the one-hot target and gradient scaling without
allocating separate probs, onehot, or grad tensors, and zero ignored rows in the
same buffer. Preserve the existing dtype conversion and reshape the final
gradient using tuple concatenation (ctx.lead_shape + (vocab_size,)) to avoid the
RUF005 warning.
In `@rl_engine/kernels/registry.py`:
- Around line 293-300: Update the npu platform map in the registry to include
every operator defined in the cpu map, preserving each CPU backend mapping by
default and overriding only batch_invariant_logp with its Ascend and PyTorch
backends. Ensure get_op resolves rms_norm, attention, grpo_loss, rope, matmul,
silu, swiglu, pack, lm_head, and embedding to their corresponding native
implementations on NPU hosts.
In `@rl_engine/platforms/device.py`:
- Around line 10-16: Make the shared _npu_available helper resilient by catching
Exception around both torch_npu import and torch.npu.is_available(), returning
False on any failure; update rl_engine/platforms/device.py:10-16 accordingly. In
tests/test_batch_invariant_logp.py:64-82, remove the duplicated local helper and
import _npu_available from rl_engine.platforms.device, while widening
_ascend_kernel_available’s handler to except Exception.
In `@setup.py`:
- Around line 256-279: Update the command construction in the extension build
function so include paths and library search paths precede ext.sources, followed
by the required -lascendcl, -ltorch_npu, -ltorch, -ltorch_cpu, -ltorch_python,
and -lc10 options for left-to-right linker resolution. Add runtime rpath entries
for the Ascend and PyTorch library directories when needed to allow _C_npu.so to
import without sourced CANN or torch_npu environments.
In `@tests/test_batch_invariant_logp.py`:
- Around line 64-82: Remove the local _npu_available implementation and import
the shared _npu_available helper from rl_engine.platforms.device. Update
_ascend_kernel_available to use that imported helper, preserving its existing
kernel-extension checks and avoiding duplicate import/error-handling logic.
---
Nitpick comments:
In `@csrc/ascend/batch_invariant_logp_ascend.asc`:
- Around line 164-172: Initialize all eight elements of scalarBuf_ during Init
before any Log operation, ensuring the padding elements have defined values.
Update the scalar setup around scalarBuf_, scalar, and AscendC::Log so per-row
processing only overwrites element 0 while preserving defined padding for the
length-8 vector operation.
In `@docs/operators/batch-invariant-logp.md`:
- Around line 59-66: Update the Ascend backend documentation to state that
BatchInvariantLogpAscendOp.apply defaults to validate=False and that the Ascend
kernel returns -lse for out-of-range targets rather than raising. Explicitly
contrast this with the PyTorch reference backend’s default target-range
validation so callers understand the invalid-input behavior difference.
In `@rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py`:
- Around line 30-37: Cache the operator returned by _fallback_op so the
NativeBatchInvariantLogpOp import and construction occur only once, while
preserving the existing fallback behavior for fp16 inputs.
In `@setup.py`:
- Around line 218-221: Update the source discovery in the Ascend extension setup
flow around the asc_srcs expression to resolve csrc/ascend relative to setup.py
via __file__, rather than the current working directory. Preserve the existing
sorting, missing-source RuntimeError, and Extension construction behavior.
- Around line 6-8: Replace the distutils imports in setup.py: import
CompileError from setuptools.errors and use shutil.which instead of
distutils.spawn.find_executable, while retaining sysconfig and updating any
corresponding call sites to the new executable lookup symbol.
In `@tests/test_batch_invariant_logp.py`:
- Around line 1195-1213: Extend TestAscendIgnoreIndex with coverage for an empty
logits batch shaped [0, _VC] and matching empty targets, asserting the operation
returns an empty output. Add a validate=True case using an out-of-range target
and assert the Ascend operation raises the expected validation error, while
preserving the existing ignore-index and default validate=False coverage.
🪄 Autofix
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 Plus
Run ID: b87b57d4-6126-4a63-9086-6fff9b9aed60
📒 Files selected for processing (13)
benchmarks/benchmark_batch_invariant_logp.pycsrc/ascend/batch_invariant_logp_ascend.ascdocs/operators/batch-invariant-logp.mdenvs.pyrl_engine/_C_npu.pyirl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/ascend/__init__.pyrl_engine/kernels/ops/ascend/loss/__init__.pyrl_engine/kernels/ops/ascend/loss/batch_invariant_logp.pyrl_engine/kernels/registry.pyrl_engine/platforms/device.pysetup.pytests/test_batch_invariant_logp.py
| (batch-invariant). The comparison here is latency and peak device memory across | ||
| a vocab sweep: | ||
|
|
||
| - Native materializes ``log_softmax`` over the full ``[N, V]`` tensor. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the native-backend description.
NativeBatchInvariantLogpOp does not materialize log_softmax. It materializes FP32 intermediate tensors for the explicit row-wise logsumexp calculation. Update this line to describe that implementation.
🤖 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 `@benchmarks/benchmark_batch_invariant_logp.py` at line 11, Update the
description of NativeBatchInvariantLogpOp to state that it materializes FP32
intermediate tensors for the explicit row-wise logsumexp calculation, removing
the claim that it materializes a full [N, V] log_softmax tensor.
| def _maybe_accelerated_op(): | ||
| """(label, op) for the fastest hardware kernel on this host, else (None, None).""" | ||
| sm90_op = _maybe_sm90_op() | ||
| if sm90_op is not None: | ||
| return "cuda", sm90_op | ||
| ascend_op = _maybe_ascend_op() | ||
| if ascend_op is not None: | ||
| return "ascend", ascend_op |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Select the accelerated operator for the active device.
Lines 90-92 can select BatchInvariantLogpSM90Op on a host that has SM90 CUDA and an active NPU device. run_benchmark then creates NPU tensors and passes them to the CUDA operator. Restrict SM90 selection to device_ctx.device_type == "cuda" and Ascend selection to device_ctx.device_type == "npu".
Proposed fix
def _maybe_accelerated_op():
- sm90_op = _maybe_sm90_op()
- if sm90_op is not None:
- return "cuda", sm90_op
- ascend_op = _maybe_ascend_op()
- if ascend_op is not None:
- return "ascend", ascend_op
+ if device_ctx.device_type == "cuda":
+ sm90_op = _maybe_sm90_op()
+ if sm90_op is not None:
+ return "cuda", sm90_op
+ if device_ctx.device_type == "npu":
+ ascend_op = _maybe_ascend_op()
+ if ascend_op is not None:
+ return "ascend", ascend_op
return None, None🤖 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 `@benchmarks/benchmark_batch_invariant_logp.py` around lines 88 - 95, Update
_maybe_accelerated_op to accept or access device_ctx and only call/select
_maybe_sm90_op when device_ctx.device_type is "cuda", and only call/select
_maybe_ascend_op when it is "npu". Preserve the existing (label, operator)
return values and the (None, None) fallback for unsupported or unavailable
devices.
| if (valid && target >= start && target < start + count) { | ||
| selected = fLocal.GetValue(static_cast<uint32_t>(target - start)); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The Ascend path skips target-range validation by default and does not report invalid targets. BatchInvariantLogpAscendOp.apply defaults to validate=False, and the kernel has no device-side range check. For a target outside [0, vocabSize) that is not ignoreIndex, no tile matches the condition on line 141, selected keeps 0.0f, and the row returns -lse. The native reference raises ValueError for the same input.
csrc/ascend/batch_invariant_logp_ascend.asc#L141-L143: state the out-of-range contract in the kernel header comment, or clamp the result to match the reference.docs/operators/batch-invariant-logp.md#L59-L66: document that the Ascend backend defaults tovalidate=Falseand that an out-of-range target yields-lseinstead of an error.
📍 Affects 2 files
csrc/ascend/batch_invariant_logp_ascend.asc#L141-L143(this comment)docs/operators/batch-invariant-logp.md#L59-L66
🤖 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 `@csrc/ascend/batch_invariant_logp_ascend.asc` around lines 141 - 143, Update
the kernel header comment in csrc/ascend/batch_invariant_logp_ascend.asc around
the target selection in BatchInvariantLogpAscendOp to state that invalid
non-ignore targets remain unvalidated and produce -lse rather than an error; no
kernel logic change is requested. Update docs/operators/batch-invariant-logp.md
around lines 59-66 to document the Ascend backend’s validate=False default and
its -lse behavior for out-of-range targets.
| TORCH_CHECK(logits.size(-1) > 0, "vocab size must be positive"); | ||
| TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as logits"); | ||
| TORCH_CHECK(target.numel() == logits.size(0), "target must have one entry per row"); | ||
|
|
||
| const int64_t numRows = logits.size(0); | ||
| const int64_t vocabSize = logits.size(1); | ||
|
|
||
| torch::Tensor targetI32 = target.to(torch::kInt32).contiguous(); | ||
| torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat)); | ||
| torch::Tensor lse = at::empty({numRows}, logits.options().dtype(at::kFloat)); | ||
|
|
||
| // stream(true): flush the task queue before launch so the kernel cannot | ||
| // overtake earlier NPU work; outputs were allocated with at::empty (no | ||
| // queued initializer) for the same reason. | ||
| auto aclStream = c10_npu::getCurrentNPUStream().stream(true); | ||
| const uint32_t blockNum = static_cast<uint32_t>(std::min(numRows, MAX_BLOCKS)); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard the empty-batch case before the launch.
TORCH_CHECK validates the vocab dimension but not the row dimension. If logits.size(0) == 0, then blockNum becomes 0 and the kernel launches with a zero block count. Return the empty outputs before the launch.
🛡️ Proposed fix
torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat));
torch::Tensor lse = at::empty({numRows}, logits.options().dtype(at::kFloat));
+
+ if (numRows == 0) {
+ return {logp, lse};
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| TORCH_CHECK(logits.size(-1) > 0, "vocab size must be positive"); | |
| TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as logits"); | |
| TORCH_CHECK(target.numel() == logits.size(0), "target must have one entry per row"); | |
| const int64_t numRows = logits.size(0); | |
| const int64_t vocabSize = logits.size(1); | |
| torch::Tensor targetI32 = target.to(torch::kInt32).contiguous(); | |
| torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat)); | |
| torch::Tensor lse = at::empty({numRows}, logits.options().dtype(at::kFloat)); | |
| // stream(true): flush the task queue before launch so the kernel cannot | |
| // overtake earlier NPU work; outputs were allocated with at::empty (no | |
| // queued initializer) for the same reason. | |
| auto aclStream = c10_npu::getCurrentNPUStream().stream(true); | |
| const uint32_t blockNum = static_cast<uint32_t>(std::min(numRows, MAX_BLOCKS)); | |
| TORCH_CHECK(logits.size(-1) > 0, "vocab size must be positive"); | |
| TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as logits"); | |
| TORCH_CHECK(target.numel() == logits.size(0), "target must have one entry per row"); | |
| const int64_t numRows = logits.size(0); | |
| const int64_t vocabSize = logits.size(1); | |
| torch::Tensor targetI32 = target.to(torch::kInt32).contiguous(); | |
| torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat)); | |
| torch::Tensor lse = at::empty({numRows}, logits.options().dtype(at::kFloat)); | |
| if (numRows == 0) { | |
| return {logp, lse}; | |
| } | |
| // stream(true): flush the task queue before launch so the kernel cannot | |
| // overtake earlier NPU work; outputs were allocated with at::empty (no | |
| // queued initializer) for the same reason. | |
| auto aclStream = c10_npu::getCurrentNPUStream().stream(true); | |
| const uint32_t blockNum = static_cast<uint32_t>(std::min(numRows, MAX_BLOCKS)); |
🤖 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 `@csrc/ascend/batch_invariant_logp_ascend.asc` around lines 270 - 285, In the
batch setup before stream acquisition and kernel launch, handle numRows == 0 by
returning the already allocated empty logp and lse tensors. Keep the existing
validation and non-empty blockNum/kernel-launch path unchanged.
| @staticmethod | ||
| def backward(ctx, grad_output): | ||
| logits_2d, target_1d, lse = ctx.saved_tensors | ||
| ignore_index = ctx.ignore_index | ||
| vocab_size = ctx.vocab_size | ||
|
|
||
| grad_flat = grad_output.reshape(-1).contiguous().to(torch.float32) | ||
|
|
||
| valid = target_1d != ignore_index | ||
| safe_target = torch.where(valid, target_1d, torch.zeros_like(target_1d)) | ||
| probs = torch.exp(logits_2d.float() - lse.unsqueeze(1)) | ||
| onehot = torch.zeros_like(probs) | ||
| onehot.scatter_(1, safe_target.unsqueeze(1), 1.0) | ||
| grad = grad_flat.unsqueeze(1) * (onehot - probs) | ||
| grad = torch.where(valid.unsqueeze(1), grad, torch.zeros_like(grad)) | ||
| grad_logits = grad.to(logits_2d.dtype) | ||
|
|
||
| grad_logits = grad_logits.reshape(ctx.lead_shape + (vocab_size,)) | ||
| return grad_logits, None, None |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
The backward pass allocates three full [N, V] fp32 tensors.
probs, onehot, and grad are each [N, V] in fp32. For a 8192-token batch with a 128k vocabulary that is roughly 12 GB of peak device memory for a single backward call, on top of the saved logits_2d. The stated goal is to avoid materializing extra [N, V] buffers. You can produce the same result with one buffer.
♻️ Proposed lower-memory backward
`@staticmethod`
def backward(ctx, grad_output):
logits_2d, target_1d, lse = ctx.saved_tensors
ignore_index = ctx.ignore_index
vocab_size = ctx.vocab_size
grad_flat = grad_output.reshape(-1).contiguous().to(torch.float32)
valid = target_1d != ignore_index
safe_target = torch.where(valid, target_1d, torch.zeros_like(target_1d))
- probs = torch.exp(logits_2d.float() - lse.unsqueeze(1))
- onehot = torch.zeros_like(probs)
- onehot.scatter_(1, safe_target.unsqueeze(1), 1.0)
- grad = grad_flat.unsqueeze(1) * (onehot - probs)
- grad = torch.where(valid.unsqueeze(1), grad, torch.zeros_like(grad))
- grad_logits = grad.to(logits_2d.dtype)
+ # grad = g * (onehot - softmax); built in a single [N, V] buffer.
+ grad = torch.exp(logits_2d.float() - lse.unsqueeze(1))
+ grad.neg_()
+ grad.scatter_add_(1, safe_target.unsqueeze(1), torch.ones_like(lse).unsqueeze(1))
+ grad.mul_(torch.where(valid, grad_flat, torch.zeros_like(grad_flat)).unsqueeze(1))
+ grad_logits = grad.to(logits_2d.dtype)
- grad_logits = grad_logits.reshape(ctx.lead_shape + (vocab_size,))
+ grad_logits = grad_logits.reshape((*ctx.lead_shape, vocab_size))
return grad_logits, None, NoneThe reshape change also resolves the Ruff RUF005 hint on line 79.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @staticmethod | |
| def backward(ctx, grad_output): | |
| logits_2d, target_1d, lse = ctx.saved_tensors | |
| ignore_index = ctx.ignore_index | |
| vocab_size = ctx.vocab_size | |
| grad_flat = grad_output.reshape(-1).contiguous().to(torch.float32) | |
| valid = target_1d != ignore_index | |
| safe_target = torch.where(valid, target_1d, torch.zeros_like(target_1d)) | |
| probs = torch.exp(logits_2d.float() - lse.unsqueeze(1)) | |
| onehot = torch.zeros_like(probs) | |
| onehot.scatter_(1, safe_target.unsqueeze(1), 1.0) | |
| grad = grad_flat.unsqueeze(1) * (onehot - probs) | |
| grad = torch.where(valid.unsqueeze(1), grad, torch.zeros_like(grad)) | |
| grad_logits = grad.to(logits_2d.dtype) | |
| grad_logits = grad_logits.reshape(ctx.lead_shape + (vocab_size,)) | |
| return grad_logits, None, None | |
| @staticmethod | |
| def backward(ctx, grad_output): | |
| logits_2d, target_1d, lse = ctx.saved_tensors | |
| ignore_index = ctx.ignore_index | |
| vocab_size = ctx.vocab_size | |
| grad_flat = grad_output.reshape(-1).contiguous().to(torch.float32) | |
| valid = target_1d != ignore_index | |
| safe_target = torch.where(valid, target_1d, torch.zeros_like(target_1d)) | |
| # grad = g * (onehot - softmax); built in a single [N, V] buffer. | |
| grad = torch.exp(logits_2d.float() - lse.unsqueeze(1)) | |
| grad.neg_() | |
| grad.scatter_add_(1, safe_target.unsqueeze(1), torch.ones_like(lse).unsqueeze(1)) | |
| grad.mul_(torch.where(valid, grad_flat, torch.zeros_like(grad_flat)).unsqueeze(1)) | |
| grad_logits = grad.to(logits_2d.dtype) | |
| grad_logits = grad_logits.reshape((*ctx.lead_shape, vocab_size)) | |
| return grad_logits, None, None |
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 79-79: Consider (*ctx.lead_shape, vocab_size) instead of concatenation
Replace with (*ctx.lead_shape, vocab_size)
(RUF005)
🤖 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/ascend/loss/batch_invariant_logp.py` around lines 62 -
80, Reduce the backward memory usage in the static method backward by reusing a
single [N, V] buffer: transform the logits buffer in place into probabilities,
apply the one-hot target and gradient scaling without allocating separate probs,
onehot, or grad tensors, and zero ignored rows in the same buffer. Preserve the
existing dtype conversion and reshape the final gradient using tuple
concatenation (ctx.lead_shape + (vocab_size,)) to avoid the RUF005 warning.
Source: Linters/SAST tools
| # Ascend NPU: op types without an entry fall back to PYTORCH_NATIVE | ||
| # (see get_op's default), so only Ascend-accelerated ops are listed. | ||
| "npu": { | ||
| "batch_invariant_logp": [ | ||
| OpBackend.ASCEND_BATCH_INVARIANT_LOGP, | ||
| OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, | ||
| ], | ||
| }, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
The npu platform map breaks every operator except batch_invariant_logp.
get_op on line 392 uses [OpBackend.PYTORCH_NATIVE] as the default, and PYTORCH_NATIVE maps to rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp. It is not a generic fallback. The cpu map lists every operator explicitly for that reason.
Before this change, an Ascend host reported device_ctx.device_type == "cpu" and resolved through the cpu map. After this change, line 421 routes it to npu, and kernel_registry.get_op("rms_norm") now returns NativeLogpOp instead of NativeRMSNormOp. The same applies to attention, grpo_loss, rope, matmul, silu, swiglu, pack, lm_head, and embedding. Each of these fails at call time with a signature or shape error.
Derive the npu map from the cpu map so only batch_invariant_logp differs.
🐛 Proposed fix
- # Ascend NPU: op types without an entry fall back to PYTORCH_NATIVE
- # (see get_op's default), so only Ascend-accelerated ops are listed.
- "npu": {
- "batch_invariant_logp": [
- OpBackend.ASCEND_BATCH_INVARIANT_LOGP,
- OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
- ],
- },
}
+ # Ascend NPU: start from the portable PyTorch map, then override the
+ # ops that have an Ascend kernel. get_op's default is NativeLogpOp,
+ # not a generic fallback, so every op needs an explicit entry.
+ self._priority_map["npu"] = {
+ **self._priority_map["cpu"],
+ "batch_invariant_logp": [
+ OpBackend.ASCEND_BATCH_INVARIANT_LOGP,
+ OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
+ ],
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Ascend NPU: op types without an entry fall back to PYTORCH_NATIVE | |
| # (see get_op's default), so only Ascend-accelerated ops are listed. | |
| "npu": { | |
| "batch_invariant_logp": [ | |
| OpBackend.ASCEND_BATCH_INVARIANT_LOGP, | |
| OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, | |
| ], | |
| }, | |
| }, | |
| } | |
| # Ascend NPU: start from the portable PyTorch map, then override the | |
| # ops that have an Ascend kernel. get_op's default is NativeLogpOp, | |
| # not a generic fallback, so every op needs an explicit entry. | |
| self._priority_map["npu"] = { | |
| **self._priority_map["cpu"], | |
| "batch_invariant_logp": [ | |
| OpBackend.ASCEND_BATCH_INVARIANT_LOGP, | |
| OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, | |
| ], | |
| } |
🤖 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/registry.py` around lines 293 - 300, Update the npu
platform map in the registry to include every operator defined in the cpu map,
preserving each CPU backend mapping by default and overriding only
batch_invariant_logp with its Ascend and PyTorch backends. Ensure get_op
resolves rms_norm, attention, grpo_loss, rope, matmul, silu, swiglu, pack,
lm_head, and embedding to their corresponding native implementations on NPU
hosts.
| def _npu_available() -> bool: | ||
| """torch.npu only exists after torch_npu is imported; probe defensively.""" | ||
| try: | ||
| import torch_npu # noqa: F401 | ||
| except ImportError: | ||
| return False | ||
| return hasattr(torch, "npu") and torch.npu.is_available() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
One _npu_available helper, duplicated, with an over-narrow except ImportError. Both copies assume import torch_npu fails only with ImportError. A partially installed CANN stack raises OSError for a missing libascendcl.so and RuntimeError for a missing driver, and torch.npu.is_available() can raise as well. The rl_engine/platforms/device.py copy runs at package import time, so the failure aborts import rl_engine; the test copy aborts pytest collection.
rl_engine/platforms/device.py#L10-L16: widen the handler toexcept Exceptionand move thetorch.npu.is_available()call inside thetryblock.tests/test_batch_invariant_logp.py#L64-L82: delete the local copy and import_npu_availablefromrl_engine.platforms.device; widen the_ascend_kernel_availablehandler toexcept Exception.
📍 Affects 2 files
rl_engine/platforms/device.py#L10-L16(this comment)tests/test_batch_invariant_logp.py#L64-L82
🤖 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/platforms/device.py` around lines 10 - 16, Make the shared
_npu_available helper resilient by catching Exception around both torch_npu
import and torch.npu.is_available(), returning False on any failure; update
rl_engine/platforms/device.py:10-16 accordingly. In
tests/test_batch_invariant_logp.py:64-82, remove the duplicated local helper and
import _npu_available from rl_engine.platforms.device, while widening
_ascend_kernel_available’s handler to except Exception.
| cmd = [ | ||
| "bisheng", | ||
| "-x", | ||
| "asc", | ||
| f"--npu-arch={soc}", | ||
| "-shared", | ||
| "-fPIC", | ||
| "-std=c++17", | ||
| "-O2", | ||
| f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", | ||
| f"-DTORCH_EXTENSION_NAME={module_name}", | ||
| "-lascendcl", | ||
| "-ltorch_npu", | ||
| "-ltorch", | ||
| "-ltorch_cpu", | ||
| "-ltorch_python", | ||
| "-lc10", | ||
| *ext.sources, | ||
| "-o", | ||
| ext_fullpath, | ||
| ] | ||
| cmd += [f"-I{d}" for d in include_dirs if d] | ||
| cmd += [f"-L{d}" for d in lib_dirs if d] | ||
| return cmd |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check how the existing CUDA extension handles library dirs and rpath, for consistency.
rg -n 'rpath|library_dirs|extra_link_args|torch_lib_dir' setup.pyRepository: RL-Align/RL-Kernel
Length of output: 659
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- setup.py outline ---'
ast-grep outline setup.py --lang python || true
printf '%s\n' '--- relevant setup.py sections ---'
sed -n '35,90p' setup.py
sed -n '130,205p' setup.py
sed -n '230,290p' setup.py
printf '%s\n' '--- all relevant symbols and link flags ---'
rg -n -C 3 'def .*build|build_ext|include_dirs|library_dirs|extra_link_args|bisheng|asc|lib_dirs|ext\.sources' setup.pyRepository: RL-Align/RL-Kernel
Length of output: 13362
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- Ascend build references ---'
rg -n -i -C 3 'bisheng|KERNEL_ALIGN_FORCE_ASCEND|ASCEND_HOME_PATH|torch_npu|rpath|LD_LIBRARY_PATH|libascendcl|libtorch_npu' --glob '!setup.py' .
printf '%s\n' '--- source and extension declarations ---'
rg -n -C 2 'language="asc"|\.asc|Extension\(|_bisheng_compile_cmd|build_extension' setup.py csrc README.md docs 2>/dev/null || true
printf '%s\n' '--- argument-order verifier ---'
python3 - <<'PY'
from pathlib import Path
import ast
tree = ast.parse(Path("setup.py").read_text())
fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_bisheng_compile_cmd")
cmd_assign = next(n for n in ast.walk(fn) if isinstance(n, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "cmd" for t in n.targets))
items = cmd_assign.value.elts
for i, item in enumerate(items):
if isinstance(item, ast.Starred):
label = "STARRED " + ast.unparse(item.value)
elif isinstance(item, ast.Constant):
label = repr(item.value)
else:
label = ast.unparse(item)
print(f"{i:02d}: {label}")
PYRepository: RL-Align/RL-Kernel
Length of output: 12074
🌐 Web query:
Bisheng compiler command line link order -l source files Ascend C
💡 Result:
When using the BiSheng compiler for Ascend C development, the command-line compilation and linking order generally follows conventions similar to standard compilers like g++ [1][2]. 1. Compilation of Individual Units: You typically compile source files (.cpp,.asc) into object files (.o) before linking them. If you are using the separate compilation mode (e.g., when making cross-file device code calls), you must use the -dc option for all.asc files [3]. Example: bisheng -dc add_compute.asc -o add_compute.o --npu-arch=dav-2201 bisheng -dc add_kernel.asc -o add_kernel.o --npu-arch=dav-2201 bisheng -c main.cpp -o main.o -I${ASCEND_HOME_PATH}/include 2. Linking Order: After generating the object files, you link them together to produce the final executable. The order in which you pass the object files and library flags matters; object files containing device code or dependencies should generally be placed before the host-side code or where they are referenced [4][3]. Example: bisheng add_compute.o add_kernel.o main.o -o demo 3. Unified Compilation: You can also compile and link in a single command, where the order of source files should reflect the dependencies (e.g., placing kernel source files before the main host code) [4]. Example: bisheng main.cpp add_kernel.asc -o demo --npu-arch=dav-2201 -I${ASCEND_HOME_PATH}/include Key considerations for Ascend C compilation: - --npu-arch: This flag is required to specify the AI processor architecture (e.g., dav-2201) [4][1][2]. - -x asc: Use this flag when compiling.cpp or.c files that contain Ascend C code [1][2]. - Default Linking: The compiler handles default libraries, but if you are linking Ascend C-generated static libraries using a different compiler (like g++), you may need to manually link the default libraries [1][2].
Citations:
- 1: https://www.hiascend.com/document/detail/zh/canncommercial/850/opdevg/Ascendcopdevg/atlas_ascendc_10_00037.html
- 2: https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1/opdevg/Ascendcopdevg/atlas_ascendc_10_00037.html
- 3: https://gitcode.com/cann/asc-devkit/blob/master/examples/01_simd_cpp_api/02_features/04_compile/01_separate_compile/README_en.md
- 4: https://gitcode.com/cann/asc-devkit/blob/master/examples/01_simd_cpp_api/02_features/04_compile/00_basic_compile/README_en.md
🌐 Web query:
torch_npu cpp_extension_asc bisheng rpath library_dirs link flags
💡 Result:
In the context of developing C++ extensions for torch_npu and utilizing the BiSheng compiler for Ascend-based projects, configuration is typically handled through setuptools (setup.py) or manual compiler commands [1][2][3]. Library Paths and Linking When building extensions, you must ensure that the compiler can locate the necessary torch_npu, PyTorch, and CANN/ACL libraries [1][4][2]. 1. Library Paths (library_dirs): You typically need to provide paths to the Python library directory, the PyTorch lib directory, and the torch_npu lib directory [1][2]. These are often retrieved dynamically using sysconfig and the installation paths of the installed packages [1][2]. 2. Link Flags (extra_link_args): Linking requires specifying the libraries to link against, such as -ltorch_npu, -ltorch, and -lc10 [1][2][3]. When using shared libraries, adding RPATH settings is common to ensure the runtime can find these dependencies without explicit environment variable configuration [5]. A standard practice is to use -Wl,-rpath,$ORIGIN/lib or relative paths to the library directory [5]. BiSheng Compiler Integration When integrating the BiSheng compiler (often invoked as the 'bisheng' command), projects frequently implement a custom build_ext class to replace or extend the default setuptools behavior [2][3]. - Compiler Invocation: The build process checks for the existence of the bisheng executable [2][3]. - Compilation Commands: Instead of standard setuptools flags, you may manually construct a command list for bisheng, including arguments such as -shared, -fPIC, -std=c++17, -I for includes, and -L for library paths [2][3]. - NPU-Specific Flags: When compiling, you may need architecture-specific flags like --npu-arch=dav-2201 [2][3]. Reference Implementation Pattern A common pattern seen in Ascend-optimized projects is: - Discovering dependencies: Use a helper function (often named get_dependency_paths) to locate torch, torch_npu, and Python headers/libs [1][2]. - Constructing the command: Iterate through collected include and library directories to append -I and -L flags [2][3]. - Building: Pass the constructed command to the subprocess or self.spawn method within a custom build_ext class [2][3]. For official guidance, refer to the Ascend Extension for PyTorch documentation, which provides details on environment setup and compiler optimizations [6][7]. [1][4][8][2][5][3][7]
Citations:
- 1: https://gitcode.com/Ascend/op-plugin/blob/26.1.0/examples/cpp_extension_structured/setup.py
- 2: https://gitcode.com/Ascend/torchair/blob/skill_install/examples/_kernel_extension_aclgraph/pybind/setup.py
- 3: https://gitcode.com/Ascend/op-plugin/blob/a988d82e5ca295f9f68d864f75e802a311b6744e/examples/cpp_extension/setup.py
- 4: https://gitcode.com/Ascend/pytorch/blob/v2.12.0-26.1.0/torch_npu/utils/cpp_extension.py
- 5: https://gitcode.com/Ascend/pytorch/blob/07a3ec02ec359b432cfd931a8e616e25b9c959dd/setup.py
- 6: https://www.hiascend.com/document/detail/en/Pytorch/2600/index/index.html
- 7: https://www.huaweicloud.com/guide/productsdesc-bms_f4fc4f651968cf7def95d259f2b40dedsupport7
- 8: https://github.com/Ascend/pytorch/blob/15c68ef6/torch_npu/_inductor/cpp_builder.py
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- runtime library-path setup ---'
sed -n '300,365p' rl_engine/executors/deepspeed_trainer.py
printf '%s\n' '--- all environment path mutations ---'
rg -n -C 4 '_prepend_env_path|ASCEND|LD_LIBRARY_PATH|PATH' rl_engine setup.py
printf '%s\n' '--- repository tests for the Bisheng command ---'
rg -n -i -C 3 'bisheng_compile|AscendBuildExtension|_bisheng|lib_dirs|rpath|link' tests setup.py
printf '%s\n' '--- cited reference snippets, if network access is available ---'
for url in \
'https://raw.githubusercontent.com/Ascend/pytorch/07a3ec02ec359b432cfd931a8e616e25b9c959dd/setup.py' \
'https://raw.githubusercontent.com/Ascend/pytorch/v2.12.0-26.1.0/torch_npu/utils/cpp_extension.py'
do
echo "URL: $url"
curl -L --max-time 10 -s "$url" | rg -n -C 3 'rpath|library_dirs|extra_link_args|torch_npu|ascendcl' | head -80 || true
doneRepository: RL-Align/RL-Kernel
Length of output: 28686
Place the source files and library search paths before the libraries.
BiSheng uses left-to-right linker resolution. With --as-needed, the current order can discard the -l options before ext.sources introduce references, which can leave _C_npu.so with unresolved symbols.
Add rpath entries for the required Ascend and PyTorch library directories if the extension must import without a sourced CANN or torch_npu runtime environment.
🤖 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 `@setup.py` around lines 256 - 279, Update the command construction in the
extension build function so include paths and library search paths precede
ext.sources, followed by the required -lascendcl, -ltorch_npu, -ltorch,
-ltorch_cpu, -ltorch_python, and -lc10 options for left-to-right linker
resolution. Add runtime rpath entries for the Ascend and PyTorch library
directories when needed to allow _C_npu.so to import without sourced CANN or
torch_npu environments.
| def _npu_available() -> bool: | ||
| try: | ||
| import torch_npu # noqa: F401 | ||
| except ImportError: | ||
| return False | ||
| return hasattr(torch, "npu") and torch.npu.is_available() | ||
|
|
||
|
|
||
| def _ascend_kernel_available() -> bool: | ||
| if not _npu_available(): | ||
| return False | ||
| try: | ||
| from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( | ||
| _NPU_EXT_AVAILABLE, | ||
| _C_npu, | ||
| ) | ||
| except ImportError: | ||
| return False | ||
| return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "batch_invariant_logp_ascend") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Reuse _npu_available from rl_engine.platforms.device.
This function is a byte-for-byte copy of _npu_available in rl_engine/platforms/device.py lines 10-16, including the narrow except ImportError. A broken CANN install raises OSError or RuntimeError from import torch_npu, which aborts pytest collection for the whole file instead of skipping the Ascend tests. Import the shared helper so one fix covers both call sites. See the related comment on rl_engine/platforms/device.py.
♻️ Proposed fix
-def _npu_available() -> bool:
- try:
- import torch_npu # noqa: F401
- except ImportError:
- return False
- return hasattr(torch, "npu") and torch.npu.is_available()
+from rl_engine.platforms.device import _npu_available
def _ascend_kernel_available() -> bool:
if not _npu_available():
return False
try:
from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import (
_NPU_EXT_AVAILABLE,
_C_npu,
)
- except ImportError:
+ except Exception:
return False
return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "batch_invariant_logp_ascend")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _npu_available() -> bool: | |
| try: | |
| import torch_npu # noqa: F401 | |
| except ImportError: | |
| return False | |
| return hasattr(torch, "npu") and torch.npu.is_available() | |
| def _ascend_kernel_available() -> bool: | |
| if not _npu_available(): | |
| return False | |
| try: | |
| from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( | |
| _NPU_EXT_AVAILABLE, | |
| _C_npu, | |
| ) | |
| except ImportError: | |
| return False | |
| return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "batch_invariant_logp_ascend") | |
| from rl_engine.platforms.device import _npu_available | |
| def _ascend_kernel_available() -> bool: | |
| if not _npu_available(): | |
| return False | |
| try: | |
| from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( | |
| _NPU_EXT_AVAILABLE, | |
| _C_npu, | |
| ) | |
| except Exception: | |
| return False | |
| return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "batch_invariant_logp_ascend") |
🤖 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 64 - 82, Remove the local
_npu_available implementation and import the shared _npu_available helper from
rl_engine.platforms.device. Update _ascend_kernel_available to use that imported
helper, preserving its existing kernel-extension checks and avoiding duplicate
import/error-handling logic.
There was a problem hiding this comment.
LGTM. I have tested this new kernel on NPU and it works. Please fix CI conflicts, thanks for your contribution!
You can use this command,
python scripts/check_operator.py --op batch_invariant_logp --candidate ascend \
--device npu --dtype fp32 --batch 4 --seq 8 --vocab 2048 --check-grad
And then the ouptut for fp32 is
INFO 08-13 21:34:55 [RL-Kernel]: Successfully linked to precompiled _C_npu.batch_invariant_logp_ascend kernel.
/home/ma-user/work/z84450661/rl-kernel/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py:105: UserWarning: Cannot create tensor with interal format while
allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp))
suite=batch_invariant_logp passed=True pass_rate=1.0000
candidate=ascend-batch_invariant_logp backend=ascend passed=True pass_rate=1.0000
case=batch_invariant_logp-torch.float32-4x8x2048 output=0 shape=(4, 8) dtype=torch.float32 max_abs=9.53674316e-07 mean_abs=1.78813934e-07 max_rel=1.45964989e-07
tol=(atol=1.000e-05, rtol=0.000e+00) passed=True
case=batch_invariant_logp-torch.float32-4x8x2048 output=1 gradient:logits shape=(4, 8, 2048) dtype=torch.float32 max_abs=1.32247806e-07 mean_abs=1.55981117e-10
max_rel=1.23174477e-05 tol=(atol=1.000e-05, rtol=0.000e+00) passed=True
The output for bf16 is
INFO 08-13 21:35:49 [RL-Kernel]: Successfully linked to precompiled _C_npu.batch_invariant_logp_ascend kernel.
...（同样的 UserWarning 两行）...
suite=batch_invariant_logp passed=True pass_rate=1.0000
candidate=ascend-batch_invariant_logp backend=ascend passed=True pass_rate=1.0000
case=batch_invariant_logp-torch.bfloat16-4x8x2048 output=0 shape=(4, 8) dtype=torch.float32 max_abs=9.53674316e-07 mean_abs=2.98023224e-08 max_rel=1.10753255e-07
tol=(atol=5.000e-02, rtol=0.000e+00) passed=True
So that means,
forward: max_abs=9.54e-07 tol=1e-05 ✓
gradient: max_abs=1.32e-07 tol=1e-05 ✓
bf16: passed=True（max_abs=9.54e-07, tol=5e-02）✓
Also the 29 unit test passed in tests/test_op_checks.py.
| from __future__ import annotations | ||
|
|
||
| import torch | ||
|
|
There was a problem hiding this comment.
Can you add these two lines to the file,
from typing import Any
and
_C_npu = None → _C_npu: Any = None
to pass a CI check called mypy.
You can check the CI lists and fix them if they are not green marks yet
| "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp." | ||
| "BatchInvariantLogpSM90Op", | ||
| "ascend": "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp." | ||
| "BatchInvariantLogpAscendOp", |
There was a problem hiding this comment.
In order to let the new logprob kernel pass gtest (it is our in-house test to measure precision of a single-card kernel), you need to update the scripts/check_operator.py file with the following 3 content,
inside the _select_device method, you need to add these two lines,
if value == "npu" and not _npu_available():
raise RuntimeError("--device npu was requested, but no Ascend NPU is available")
inside the parse_args method, you need to add the npu option,
parser.add_argument(
"--candidate",
default="pytorch",
help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, "
"ascend.",
)
i think you also need to add a new method,
def _npu_available() -> bool:
"""torch.npu only exists after torch_npu is imported; probe defensively."""
try:
import torch_npu # noqa: F401
except ImportError:
return False
return hasattr(torch, "npu") and torch.npu.is_available()
|
please resolve the CI error, other LGTM. |
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM, Thank you for work!
…rministic-attention-test Resolve conflicts with the merged PR RL-Align#297 (batch-invariant logp Ascend): - logp kernel/op/tests/docs/benchmark take the upstream reviewed versions; the kernel's PYBIND11_MODULE stays in the ops_npu.asc aggregator - registry keeps the upstream cpu-inheritance for the npu platform and adds the attention override on top - setup.py keeps the recursive .asc glob; device.py keeps the upstream broader exception probe - _C_npu.pyi exposes both operator stubs
3 participants
Summary
Adds an Ascend C (CANN) backend for the batch-invariant selected-token
log-probability operator introduced in #199 and extended to SM90 CUDA in #204.
The operator computes, from already materialized logits:
The kernel streams each vocab row through UB in fixed 4096-element tiles with a
two-pass (max, then sum-exp) fixed-order fp32 reduction, so the per-row result
depends only on the row itself — never on batch size or layout. Each row is
processed end-to-end by exactly one AI core block, and outputs are staged in UB
and written with
DataCopyPad(scalar GMSetValue/GetValueare avoided perknown hardware pitfalls). It is registered as the top-priority backend on the
new
npuplatform key and falls back cleanly to the PyTorch native op.csrc/ascend/batch_invariant_logp_ascend.asc)BatchInvariantLogpAscendOp)lse. New.rl_engine._C_npu)-x asc) extension, gated onKERNEL_ALIGN_FORCE_ASCEND=1;--npu-archdefaults todav-2201, override viaKERNEL_ALIGN_ASCEND_ARCH. New.DeviceContextnow detectstorch_npu/torch.npu; registry gains the"npu"platform key. New.batch_invariant_logpon NPU: Ascend -> PyTorch. Ascend instantiation failure (extension not built) degrades to native.benchmarks/benchmark_batch_invariant_logp.pydispatches timing/VRAM throughtorch.npuand adds anascendcolumn.Implementation
csrc/ascend/batch_invariant_logp_ascend.asc(kernel class,
__global__ __vector__entries, host entrypoint, pybind).setup.py(AscendBuildExtensioninvokingbisheng, followingthe official torch_npu
cpp_extension_ascpattern),envs.py(
KERNEL_ALIGN_FORCE_ASCEND,KERNEL_ALIGN_ASCEND_ARCH),rl_engine/_C_npu.pyi.rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py,interface aligned with
BatchInvariantLogpSM90Op.rl_engine/platforms/device.py,rl_engine/kernels/registry.py.tests/test_batch_invariant_logp.py,rl_engine/kernels/gtest/operator_specs.py,docs/operators/batch-invariant-logp.md,benchmarks/benchmark_batch_invariant_logp.py.Validation environment
KERNEL_ALIGN_FORCE_ASCEND=1 pip install -e . --no-build-isolationCorrectness / Tests
TestAscendCorrectness(fp32/bf16 vslog_softmax + gather, large vocab,single-token, 3-D, vs native op): 6/6 passed on device.
SetFlag/WaitFlagevents;AscendC::SyncAll()(cross-core barrier) is deliberately avoided because itdeadlocks when launched blocks exceed resident cores.
WIP:
TestAscendBatchInvariance/TestAscendBackward/TestAscendIgnoreIndex/TestAscendFallbackare included but still beingre-verified on device after the sync fix; will update before undrafting.
CPU-only regression (no NPU):
tests/test_batch_invariant_logp.py,tests/test_kernel_registry.py,rl_engine/tests/test_dispatch.py,tests/test_cpu_hal.py— all pass; Ascend tests skip cleanly.Files
csrc/ascend/batch_invariant_logp_ascend.asc(new)rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py(new)rl_engine/_C_npu.pyi(new)csrc/ascend/,rl_engine/kernels/ops/ascend/package init files (new)setup.py,envs.py,rl_engine/kernels/registry.py,rl_engine/platforms/device.pytests/test_batch_invariant_logp.py,rl_engine/kernels/gtest/operator_specs.py,docs/operators/batch-invariant-logp.md,benchmarks/benchmark_batch_invariant_logp.py