Skip to content

[KernelGen][Nvidia] Add in-place sign_ operator - #5184

Open
ShawnsYing wants to merge 3 commits into
flagos-ai:masterfrom
ShawnsYing:pr/sign
Open

[KernelGen][Nvidia] Add in-place sign_ operator#5184
ShawnsYing wants to merge 3 commits into
flagos-ai:masterfrom
ShawnsYing:pr/sign

Conversation

@ShawnsYing

Copy link
Copy Markdown
Contributor

Summary

Adds the in-place sign_ operator to complement the existing sign / sign.out
operators that were introduced in #5030. sign_ writes the element-wise sign
back into the input tensor.

aten::sign_ has only the default overload (sign_(Tensor(a!) self) -> Tensor(a!)),
mirroring the abs / abs_ pairing. The implementation reuses the already-merged
_sign_impl kernel with out == input, so no new kernel logic is introduced.

Changes

  • src/flag_gems/ops/sign.py: add sign_(x)_sign_impl(x, x)
  • src/flag_gems/ops/__init__.py: export sign_
  • src/flag_gems/__init__.py: register ("sign_", sign_) in _FULL_CONFIG
  • conf/operators.yaml: add id: sign_ block
  • tests/test_sign.py: test_sign_ (all real dtypes + shapes) and a
    non-contiguous case; asserts the mutated input matches torch
  • benchmark/test_sign.py: test_sign_inplace

Testing

191 pytest cases pass against this worktree (isolated launcher verifying
flag_gems resolves to the worktree). Dispatch confirmed via GEMS SIGN_
debug log; verified same storage pointer (true in-place) and NaN → 0.
sign_ supports all real dtypes including int / uint8 / bool; complex raises
NotImplementedError.

Add the in-place variant sign_ to complement the existing sign / sign.out
operators. sign_ reuses the merged _sign_impl kernel writing back into the
input tensor.

- ops/sign.py: sign_(x) = _sign_impl(x, x)
- register sign_ in ops/__init__ (__all__) and _FULL_CONFIG ("sign_")
- conf/operators.yaml: add id: sign_ block
- tests + benchmark for sign_
@ShawnsYing

Copy link
Copy Markdown
Contributor Author

KernelGen SimpleOpt result

Ran KernelGen SimpleOpt against the sign trace on NVIDIA H20 (3 rounds, all PASSED):

Round geo_mean Notes
1 1.032x Baseline: element-wise Triton kernel, tl.where for sign, BLOCK_SIZE=1024 (best)
2 regressed (-22% on large 1D) Branch-free arithmetic via bool→float conversion
3 regressed BLOCK_SIZE=4096 exceeded the 1024-thread HW limit, destroying occupancy

Best: 1.032x geo_mean speedup vs the PyTorch reference.

sign is a trivial memory-bandwidth-bound pointwise op, so the straightforward Triton
implementation is already near-optimal against PyTorch's highly-tuned reference — ~parity
is the expected result for this class of unary op. The in-place sign_ added here reuses
the same merged _sign_impl kernel (out == input), so it inherits this performance.

@bin913

bin913 commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

/test:sign_:h20

@github-actions

Copy link
Copy Markdown
Contributor

⚠️ Cannot proceed with operation

  • reviewDecision: skip_reviews
  • commitStatus: PENDING

Reviews are not required for this operation but CI checks must be passing in order to continue

@bin913

bin913 commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

/test:sign_:h20

@github-actions

github-actions Bot commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

❌ On-demand test failed

Operator: /test|sign_
Runner: h20
Backend: nvidia-cuda133

The test failed to complete. Check the workflow run for details.

@bin913

bin913 commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

/test|sign_:h20

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants