[NPU]: Add NPU support for the tvd operator #998
Merged
+227
−0
Conversation
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
TianHao324
force-pushed
the
tvd
branch
3 times, most recently
from
1a31a76 to
d8740a3
Compare
Contributor
Author
|
Hi @Tcc0403 @zheliuyu @noemotiovon |
noemotiovon
reviewed
Jan 7, 2026
| # Fallback to desired block size if no best practice found (no tiling needed) | ||
| BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) | ||
|
|
||
| MAX_BATCH_PER_KERNEL = 65535 # 每个kernel最大处理量 |
Contributor
There was a problem hiding this comment.
Please use English comments.
| def tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label): | ||
| BT, V = p.shape | ||
| # NPU does not support bfloat16 type | ||
| p = p.to(torch.float32) |
Contributor
There was a problem hiding this comment.
In some previous operators, I recall that there are test cases using the bf16 dtype that can run correctly on the Ascend Triton backend. It would be helpful to explicitly document the reason why this operator requires conversion to float32.
Additionally, once the NPU backend adds native support for bf16, this cast should be removed. Consider adding a TODO comment here, for example:
# TODO: Remove float32 conversion after Ascend NPU supports bf16 natively
| grads_chunk = grads[start:end] | ||
| labels_chunk = shift_labels[start:end] if has_label else torch.empty(1, device=p.device) | ||
|
|
||
| _tv_distance_kernel[grid]( |
Contributor
There was a problem hiding this comment.
It is possible to redesign the kernel to process multiple rows per program using a kernel-side loop with a fixed grid (e.g., 65535 programs).
Given that the current kernel is a simple per-row implementation, the overhead of multiple launches is likely negligible.
TianHao324
force-pushed
the
tvd
branch
7 times, most recently
from
5ae19a8 to
d7b5254
Compare
Tcc0403
reviewed
Jan 8, 2026
Comment on lines
13
to
25
| REDUCTION_LITERAL = Literal["none", "sum", "mean", "batchmean"] | ||
|
|
||
| _REDUCTION_MODE_NONE = tl.constexpr(0) | ||
| _REDUCTION_MODE_SUM = tl.constexpr(1) | ||
| _REDUCTION_MODE_MEAN = tl.constexpr(2) | ||
| _REDUCTION_MODE_BATCHMEAN = tl.constexpr(3) | ||
|
|
||
| _str_to_reduction_mode = { | ||
| "none": _REDUCTION_MODE_NONE.value, | ||
| "sum": _REDUCTION_MODE_SUM.value, | ||
| "mean": _REDUCTION_MODE_MEAN.value, | ||
| "batchmean": _REDUCTION_MODE_BATCHMEAN.value, | ||
| } |
Collaborator
There was a problem hiding this comment.
You should be able to assign str to tl.constexpr variable directly without having to manually convert it to int. For instance, reduction in liger's cross_entropy is just a str. It doesn't need extra mapping onto int.
This mapping exists in the original tvd implementation, but I think we should remove it for readibility as well.
Contributor
Author
There was a problem hiding this comment.
It has been modified and passed the test.
Tcc0403
reviewed
Jan 8, 2026
Comment on lines
162
to
170
| if reduction == _REDUCTION_MODE_BATCHMEAN.value: | ||
| # TODO: Remove float32 conversion after Ascend NPU supports bf16 natively | ||
| return output_tensor.sum().to(torch.float32) / n_non_ignore, grads.to(torch.float32) / n_non_ignore | ||
| elif reduction == _REDUCTION_MODE_SUM.value: | ||
| return output_tensor.sum(dim=0), grads | ||
| elif reduction == _REDUCTION_MODE_MEAN.value: | ||
| return output_tensor.sum().to(torch.float32) / (n_non_ignore * V), grads.to(torch.float32) / (n_non_ignore * V) | ||
| else: | ||
| return output_tensor, grads |
Collaborator
There was a problem hiding this comment.
Without str->int mapping, this part will be more readable.
Tcc0403
reviewed
Jan 8, 2026
|
|
||
| tl.store(grads_row_ptr + offsets, grad_res, mask=mask) | ||
|
|
||
| if reduction == _REDUCTION_MODE_NONE: |
Collaborator
There was a problem hiding this comment.
And it will be a simple if reduction == "none":
TianHao324
changed the title
Collaborator
|
Feel free to re-request review when it's ready |
TianHao324
force-pushed
the
tvd
branch
2 times, most recently
from
cdae03e to
7b7f071
Compare
Contributor
Author
|
On NPU, bfloat16 execution of the tvd operator involves low-precision accumulation in the underlying matmul, and the stricter tolerance may therefore lead to false negatives in correctness tests. |
Contributor
Author
|
Meanwhile, in the checkstyle check, I did not modify any related files. This might be caused by other commits. |
Contributor
Author
|
@Tcc0403 It's ready for review, if you have time. |
Tcc0403
reviewed
Jan 12, 2026
3 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 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.
Summary
Mainly to complete the adaptation of the tvd operator on the NPU：
1、Solving the operator ub overflow problem.
2、Use the chunking strategy to solve the problem where the grid maximum limit of trirton-ascend is 65535.
3、The data type is not supported for bf16, so all of them have been converted to f32.
Testing Done
tvc forward and backward pass tests
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence