Skip to content

Commit 23b5013

Browse files
authored
add multi_tensor_compute_scale_inv_e8m0 and change call erro (#74)
# Description Added the binding and invocation for the `multi_tensor_compute_scale_inv_e8m0` operator for the Kunlunxin vendor, and corrected a syntax error in the invocation of the `multi_tensor_compute_scale_and_scale_inv` operator. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [1] Bug fix (non-breaking change which fixes an issue) - [1] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Added code for the binding and invocation of the `multi_tensor_compute_scale_inv_e8m0` operator. Modified the code calling `multi_tensor_compute_scale_and_scale_inv`. # Checklist: - [ 1] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ 1] The functionality is complete - [ 1] I have commented my code, particularly in hard-to-understand areas - [ 1] I have made corresponding changes to the documentation - [ 1] My changes generate no new warnings - [ 1] I have added tests that prove my fix is effective or that my feature works - [ 1] New and existing unit tests pass locally with my changes
1 parent 25e80e9 commit 23b5013

2 files changed

Lines changed: 24 additions & 1 deletion

File tree

transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,21 @@ def multi_tensor_compute_scale_and_scale_inv(
357357
epsilon: float,
358358
) -> None:
359359
tex = self._get_tex()
360-
return self.multi_tensor_compute_scale_and_scale_inv(
360+
return tex.multi_tensor_compute_scale_and_scale_inv(
361361
chunk_size, noop_flag, tensor_lists, max_fp8, force_pow_2_scales, epsilon
362362
)
363+
364+
def multi_tensor_compute_scale_inv_e8m0(
365+
self,
366+
chunk_size: int,
367+
noop_flag: torch.Tensor,
368+
tensor_lists: List[List[torch.Tensor]],
369+
block_len: int,
370+
) -> None:
371+
tex = self._get_tex()
372+
return tex.multi_tensor_compute_scale_inv_e8m0(
373+
chunk_size,
374+
noop_flag,
375+
tensor_lists,
376+
block_len,
377+
)

transformer_engine/plugin/core/backends/vendor/kunlunxin/register_ops.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,14 @@ def register_builtins(registry) -> None:
173173
vendor="KUNLUNXIN",
174174
priority=100,
175175
),
176+
OpImpl(
177+
op_name="multi_tensor_compute_scale_inv_e8m0",
178+
impl_id="vendor.kunlunxin",
179+
kind=BackendImplKind.VENDOR,
180+
fn=_bind_is_available(backend.multi_tensor_compute_scale_inv_e8m0, is_avail),
181+
vendor="KUNLUNXIN",
182+
priority=100,
183+
),
176184
]
177185

178186
registry.register_many(impls)

0 commit comments

Comments
 (0)