Commit Graph

5 Commits

Author SHA1 Message Date
Ang Gao 2c36c3d7bf
[Fix][MHA] Fix MHA decode auto-tuning failure for scalar params (#518)
## Summary

- Fix auto-tuning failure for MHA decode kernels caused by
`real_seqlen_kv: T.int32` scalar parameter that tilelang's autotuner
cannot auto-generate inputs for
- Add `autotune_supply_prog` property to `Kernel` base class for
subclasses with scalar kernel parameters
- Implement custom `supply_prog` in `mha_decode_kernel` that provides
`seqlen_kv` for scalar params

Closes #517

Found in nightly test:
https://github.com/tile-ai/TileOPs/actions/runs/23102471148/job/67105637462

## Test plan

- [x] `pytest
benchmarks/ops/bench_mha_decode.py::test_mha_decode_bench[bf16-long-cache]`
— previously failing, now passes
- [x] `pytest
benchmarks/ops/bench_mha_decode.py::test_mha_decode_bench[fp16-long-cache]`
— previously failing, now passes
- [x] `pytest
benchmarks/ops/bench_mha_decode.py::test_mha_decode_bench[short-kv-tail]`
— passes
- [x] All 3 mha_decode benchmarks pass

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-15 19:06:13 +08:00
Cao Ying 4d65362c58
[Feat][Elementwise] Add 30 UnaryKernel template ops (#457)
## Summary

Implement 30 `UnaryKernel` subclasses and 30 corresponding `UnaryOp`
subclasses for elementwise operations, then tighten the implementation
to match issue #437 and the follow-up review findings.

**Ops**: `exp`, `log`, `sqrt`, `rsqrt`, `abs`, `neg`, `reciprocal`,
`sign`, `sin`, `cos`, `floor`, `ceil`, `round`, `trunc`, `erf`, `log1p`,
`expm1`, `gelu`, `silu`, `sigmoid`, `tanh`, `hardswish`, `hardsigmoid`,
`mish`, `selu`, `logical_not`, `bitwise_not`, `isnan`, `isinf`,
`isfinite`

Closes #437

## Dtype Support Matrix

| Op / API | Input dtypes | Output dtype | Reference baseline |
| --- | --- | --- | --- |
| Unary math ops + activations | `fp16`, `bf16`, `fp32` | same as input
| matching `torch` / `torch.nn.functional` unary semantics |
| `logical_not` | `bool`, `uint8`, `int8`, `int16`, `int32`, `int64`,
`fp16`, `bf16`, `fp32` | `bool` | `torch.logical_not` |
| `bitwise_not` | `bool`, `uint8`, `int8`, `int16`, `int32`, `int64` |
same as input | `torch.bitwise_not` |
| `isnan`, `isinf`, `isfinite` | `fp16`, `bf16`, `fp32` | `bool` |
`torch.isnan`, `torch.isinf`, `torch.isfinite` |

## Test Plan

- [x] **AC-1**: 30 kernel subclasses and 30 op subclasses exist for the
issue #437 scope.
- [x] **AC-2**: Unary elementwise dtype contracts are explicit in code
via `SUPPORTED_DTYPES` / `OUTPUT_DTYPE` where needed.
- [x] **AC-3**: `gelu` uses the standard `erf` formulation instead of
the tanh approximation.
- [x] **AC-4**: `logical_not`, `isnan`, `isinf`, and `isfinite` produce
torch-style `bool` outputs and tests validate them with exact
comparison.
- [x] **AC-5**: `bitwise_not` support is aligned to the current torch
contract (`bool/uint8/int8/int16/int32/int64`) and reject-path tests
exist for float dtypes.
- [x] **AC-6**: L4 edge-case coverage matches issue #437 for `sqrt`,
`rsqrt`, `log`, `log1p`, `exp`, `expm1`, `erf`, `reciprocal`, `sign`,
and special predicates.
- [x] **AC-7**: Repository-level templates/docs now require future new
ops to include a dtype support matrix, acceptance checklist, and
benchmark table in the PR body.
- [x] **AC-8**: A representative unary benchmark file exists in
`benchmarks/ops/bench_unary_elementwise.py` and runs on a real GPU
machine with small, medium, and large shapes.
- [x] **AC-9**: Top-level `__init__.py` re-export files remain
unchanged.
- [x] Elementwise correctness suite passes on the real machine GPUs:
`127 passed, 2 warnings in 8.04s`.
- [x] Unary benchmark suite passes on the real machine GPUs: `21 passed,
23 warnings in 56.25s`.

## Benchmark

**Configuration**: H200, Driver `575.57.08`, CUDA `12.8` (torch), torch
`2.9.1+cu128`

| N_total | dtype | Op | TileOPs (ms) | Baseline (ms) | Ratio |
| --- | --- | --- | --- | --- | --- |
| 262,144 | fp16 | `exp` | 0.00250 | 0.00607 | 2.43x faster |
| 1,048,576 | fp16 | `exp` | 0.00334 | 0.00768 | 2.30x faster |
| 4,000,000 | fp16 | `exp` | 0.00680 | 0.01032 | 1.52x faster |
| 262,144 | bf16 | `exp` | 0.00251 | 0.00628 | 2.50x faster |
| 1,048,576 | bf16 | `exp` | 0.00334 | 0.00709 | 2.12x faster |
| 4,000,000 | bf16 | `exp` | 0.00669 | 0.01091 | 1.63x faster |
| 262,144 | fp16 | `gelu` | 0.00264 | 0.00734 | 2.78x faster |
| 1,048,576 | fp16 | `gelu` | 0.00401 | 0.00839 | 2.09x faster |
| 4,000,000 | fp16 | `gelu` | 0.00980 | 0.01323 | 1.35x faster |
| 262,144 | bf16 | `gelu` | 0.00267 | 0.00656 | 2.46x faster |
| 1,048,576 | bf16 | `gelu` | 0.00409 | 0.00845 | 2.07x faster |
| 4,000,000 | bf16 | `gelu` | 0.01009 | 0.01277 | 1.27x faster |
| 262,144 | fp16 | `logical_not` | 0.00262 | 0.00679 | 2.59x faster |
| 1,048,576 | fp16 | `logical_not` | 0.00545 | 0.00691 | 1.27x faster |
| 4,000,000 | fp16 | `logical_not` | 0.01467 | 0.00941 | 0.64x |
| 262,144 | int32 | `bitwise_not` | 0.00297 | 0.00715 | 2.41x faster |
| 1,048,576 | int32 | `bitwise_not` | 0.00604 | 0.00816 | 1.35x faster |
| 4,000,000 | int32 | `bitwise_not` | 0.01657 | 0.01404 | 0.85x |
| 262,144 | fp16 | `isnan` | 0.00263 | 0.00674 | 2.56x faster |
| 1,048,576 | fp16 | `isnan` | 0.00546 | 0.00690 | 1.26x faster |
| 4,000,000 | fp16 | `isnan` | 0.01473 | 0.01016 | 0.69x |

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-12 18:55:08 +08:00
Zheng QiHang b8707801cd
[BugFix] Add guard for missing self.kernel in Kernel.autotune() (#204)
## Summary

- Add a `hasattr` guard in `Kernel.autotune()` to check that
`self.kernel` is set before attempting to use it
- Raises a clear `AttributeError` with an actionable message instead of
a cryptic attribute error when `tune=True` is passed to a kernel that
hasn't set `self.kernel`
- Prevents crashes in all affected NSA kernels (`NSAFwdVarlenKernel`,
`NSATopkVarlenKernel`, `NSACmpFwdVarlenKernel`, `MeanPoolingFwdKernel`,
`GQAWindowSlidingKernel`) and any future kernels that omit `self.kernel`

Closes #197

## Test plan

- [ ] Verify that instantiating an affected NSA kernel with `tune=True`
now raises a clear error message instead of a cryptic `AttributeError`
- [ ] Verify that kernels which do set `self.kernel` (e.g.,
`MHAFwdKernel`) still autotune correctly when `tune=True`
- [ ] Verify that `pre-commit run --all-files` passes

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-27 14:23:22 +08:00
ChongLi 3d22d170d0
[Refactor][Kernel] Standardize kernel conventions (#210)
Closes #202

  ## Summary
  - Standardize grouped_gemm config keys to snake_case.
  - Normalize kernel_map keys to snake_case across ops.
- Normalize dtype handling to torch.dtype with consistent dtype string
conversion.
2026-02-26 21:24:57 +08:00
Zheng QiHang 5a9e54d65c
[Refactor] Rename source directory from `top` to `tileops` (#190)
## Summary

Closes #186

- Rename `top/` → `tileops/` to align the source directory name with the
package name defined in `pyproject.toml` (`name = "tileops"`)
- Update all internal imports (`from top.` → `from tileops.`) across 116
Python files
- Update documentation paths in `DEVELOPMENT.md`, `README.md`,
`SKILL.md`, and issue templates
- Update `.pre-commit-config.yaml` yapf workaround path

## Motivation

The package is named `tileops` in `pyproject.toml`, but the source code
previously lived under `top/`. This creates a mismatch between the
package name and the module import path, which is confusing for
contributors and downstream users.

## Test plan

- [x] `pre-commit run --all-files` passed
- [x] `PYTHONPATH="$PWD" python -m pytest tests/ops/test_mha.py` — 6
passed
- [x] `PYTHONPATH="$PWD" python -m pytest tests/ops/test_gemm.py
tests/ops/test_gqa.py tests/functions/test_mha_func.py` — all passed
- [x] `import tileops` works correctly

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-25 19:22:37 +08:00