[KernelGen][Nvidia] Add special_round operator with Triton kernel (#3969)

repair: rebase special_round onto origin/master, resolve registration conflicts

fix: correct import/registration alphabetical order for special_round

fix: apply isort+black formatting for special_round registrations

chore: retrigger CI for special_round

fix: apply black+isort formatting to test_special_round

fix: replace nearbyint with rint for CI compatibility

fix: use manual round-half-to-even without tl_extra_shim externs

fix: apply black formatting to manual round-half-to-even implementation

fix: preserve signed zero in round-half-to-even and fix device assertion in midpoint test

fix: add decimals support and remove unsupported special_round_ in-place variant



Add decimals parameter to special_round kernel, following the same
scaling approach used by the existing round operator
Remove special_round_ in-place variant: PyTorch does not provide
torch.special.round_ as a native function
Remove unavailable special_round_ from operators.yaml and
registrations
Fix duplicate special.round.out in operators.yaml: separate
special_round (only special.round) and special_round_out
(only special.round.out) entries, following the special_i1 pattern
Add decimals test coverage (positive, negative, zero, out variant)
Remove benchmark for nonexistent special_round_

Co-Authored-By: Claude Opus 4.7 noreply@anthropic.com

fix: reorder special_round entries alphabetically in operators.yaml

Update operators.yaml to modify special.log_softmax


Removed KernelGen label and stages from special.log_softmax.
Signed-off-by: Dongxu-H dxhan@baai.ac.cn

Add nn.functional label to special.log_softmax

Signed-off-by: Dongxu-H dxhan@baai.ac.cn

Remove special_round_out operator

Removed special_round_out operator definition from YAML.
Signed-off-by: Dongxu-H dxhan@baai.ac.cn

Add special_round_out operator definition

Added a new operator ‘special_round_out’ with its details.
Signed-off-by: Dongxu-H dxhan@baai.ac.cn

Apply suggestion from @Caeruleann

Co-authored-by: wanghairui wanghairui0927@gmail.com
Signed-off-by: Dongxu-H dxhan@baai.ac.cn

fix: inline _round_half_to_even logic in special_round kernel

The helper function pattern (_round_half_to_even called from @triton.jit)
is unique in this codebase and may cause Triton compilation failures.
Inline the banker’s rounding logic directly into special_round_func,
and remove the incorrect -0.0 preservation (torch.special.round(-0.0)
returns 0.0, not -0.0).
Co-Authored-By: Claude noreply@anthropic.com

fix: cast to fp32 before multiply in special_round kernel

The original (x * scale).to(tl.float32) performs multiplication in the
input dtype (e.g. fp16), causing potential overflow for large scale values.
PyTorch’s torch.special.round does all computation in fp32 internally.
Cast to fp32 first to match PyTorch behavior: x.to(tl.float32) * scale.
Co-Authored-By: Claude noreply@anthropic.com

fix: rewrite special_round kernel with libentry, use fp32 mul + fp64 div


Replace pointwise_dynamic with manual @libentry kernel to avoid
segfault caused by pointwise_dynamic tensor+scalar out variant bug
Use fp32 for multiplication to match PyTorch’s truncation behavior
(1.2345 * 1000 = 1234.5 exactly in fp32)
Use fp64 for division to avoid 1-ULP errors from fp32 reciprocal
approximation (26.0 / 10.0 should be 0x40266666, not 0x40266667)
Fix out variant signature: (self, out, *, decimals=0)

All 49 unit tests + 2 benchmark tests pass on NVIDIA H20.
Co-Authored-By: Claude noreply@anthropic.com

Update init.py

Signed-off-by: Dongxu-H dxhan@baai.ac.cn

Fix tengqm review comments for special_round PR


benchmark/test_special_round.py: Use torch.ops.aten.special_round as
torch_op for the out variant test, instead of torch.special.round,
matching the pattern used by special_i0e benchmark.

conf/operators.yaml: Restore accidentally removed kind, stages, and
KernelGen label from special_log_softmax entry.


Co-Authored-By: Claude noreply@anthropic.com

Signed-off-by: Dongxu-H dxhan@baai.ac.cn
Co-authored-by: Claude Opus 4.7 noreply@anthropic.com
Co-authored-by: chx7514 hxchen7514@gmail.com
Co-authored-by: Dongxu-H dxhan@baai.ac.cn
Co-authored-by: wanghairui wanghairui0927@gmail.com
Co-authored-by: yzw1128 2508161484@qq.com
Co-authored-by: yzw1128 yzw1128@users.noreply.github.com