forked from ccf-ai-infra/TileOPs-Metax
22 Commits
| Author | SHA1 | Message | Date |
|---|---|---|---|
Cao Ying
|
20811a62d6
|
[Refactor][Elementwise] Extract shared base for parametric kernels (#611)
## Summary Extract a shared `ParametricUnaryKernel` base class from the 9 parametric independent kernels (LeakyRelu, Elu, Hardtanh, Softplus, Prelu, Where, Clamp, MaskedFill, NanToNum), eliminating ~800 lines of duplicated boilerplate. Each subclass is now ≤20 lines providing only `_make_kernel` and op-specific parameters. Closes #607 ## Test plan - [x] AC-1: Shared base class `ParametricUnaryKernel` extracted and importable - [x] AC-2: All 9 parametric kernels subclass the base; significant-line counts: LeakyRelu 10, Elu 10, Hardtanh 11, Softplus 11, Prelu 13, Where 10, Clamp 16, MaskedFill 16, NanToNum 17 - [x] AC-3: fp8 builder logic centralized in one place (base class lines 1912-1940) - [x] AC-4: `default_config` npt logic centralized; overrides only for Where/MaskedFill thread/npt specialization - [x] AC-5: All existing tests pass — 196 targeted tests + 682 smoke tests passed with no behavioral change - [x] AC-6: `ParametricUnaryKernel` is not exported via `__all__` — internal refactor only 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
640dc244fb
|
[Fix][Elementwise] Correct default_config npt and standardize output_dtype (#609)
## Summary Fix three kernel-level inconsistencies in `tileops/kernels/elementwise.py`: 1. **Incorrect npt defaults for fp16/bf16 in 7 independent kernels** — The 2-way check (`fp32 => 4, else => 16`) incorrectly assigned npt=16 to fp16/bf16. Replaced with the correct 3-way check: fp32 => 4, fp8 => 16, else (fp16/bf16) => 8. 2. **OUTPUT_DTYPE type inconsistency** — Comparison and logical kernels declared `OUTPUT_DTYPE = "int8"` (string) while the rest of the codebase expects `torch.dtype`. Standardized to `torch.int8` and added `dtype_to_str()` conversion at kernel construction boundaries. 3. **Missing `self.output_dtype` attribute** — `BinaryKernel.__init__` and `FusedGatedKernel.__init__` did not initialize `self.output_dtype`, causing `AttributeError` on access. Added proper initialization in both classes. Closes #605 ## Test plan - [x] All 7 independent kernels return npt=8 for fp16/bf16 and npt=16 for fp8 - [x] OUTPUT_DTYPE uses `torch.int8` consistently across all comparison/logical kernel classes - [x] `kernel.output_dtype` is a valid attribute on UnaryKernel, BinaryKernel, and FusedGatedKernel instances - [x] 200 existing tests pass with no behavioral regression (config/dtype, fp8, compile, independent fp8 suites) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Ang Gao
|
1d55896e99
|
[Perf] Add lru_cache to remaining 40 builder functions (#600)
## Summary - Add `@functools.lru_cache(maxsize=32)` to all remaining builder functions across 13 kernel families (40 files), eliminating repeated `JITImpl` construction overhead (**28–200ms → ~0.0003ms** per call) - Follow-up to #597 (which covered `gated_deltanet`, `flash_attn`, `flash_decode`) - All builder function parameters verified hashable (int, str, bool, float, torch.dtype) ### Kernel families covered | Family | Files | |--------|-------| | deepseek_mla | 5 | | deepseek_nsa | 5 | | norm | 6 | | reduction | 7 | | gla_chunkwise | 2 | | engram | 3 | | gemm | 2 | | fft | 2 | | Other (gla_recurrence, grouped_gemm, mhc, moe, rope, dropout, elementwise) | 8 | Closes #599 Ref: #571, #597 ## Test plan - [ ] Verify all 40 files have `@functools.lru_cache` on builder functions - [ ] Existing tests pass (`pytest tests/ -x`) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
937d34dd79
|
[Perf][Elementwise] Apply register_copy strategy to Elu, Hardtanh, Softplus, Clamp kernels (#561)
Closes #552 ## Summary - Apply `register_copy` strategy with `T.copy` vectorized loads to Elu, Hardtanh, Softplus, and Clamp kernels, following the same pattern used in LeakyRelu/NanToNum - Refactor `_make_elu_kernel`, `_make_hardtanh_kernel`, `_make_softplus_kernel`, `_make_clamp_kernel` to use `T.alloc_fragment` + `T.copy` for coalesced memory access - Update `default_config` npt: fp32 stays at 4, non-fp32 set to 16 (matching LeakyRelu pattern) ## Test plan - [x] pre-commit passed - [x] pytest passed: 94 tests in test_activation.py + test_special_elementwise.py (7.07s) - [x] No regression in existing activation/special elementwise test suites: 189 total tests pass (94 activation+special_elementwise + 95 compile+fp8) ## Structural Readiness All checks passed. ## Benchmark Shape: (1024, 10240) fp16 | Op | TileOPs GB/s | PyTorch GB/s | Ratio | |---|---|---|---| | elu | 2945.93 | 1754.04 | 1.680x | | hardtanh | 3150.79 | 2918.58 | 1.080x | | softplus | 1764.51 | 1415.82 | 1.246x | | clamp | 3095.03 | 2842.31 | 1.089x | All 4 ops achieve >= 1.0x PyTorch bandwidth at the target shape. Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
e021408b7a
|
[Enhancement][Elementwise] Deploy strategy-aware npt defaults for fp16/bf16 (#562)
## Summary Implement strategy-aware `default_config` so that `explicit_parallel` uses `npt=4` for fp16/bf16 while `register_copy` retains `npt=8`, unlocking a measured 58% bandwidth improvement. Closes #553 ## Changes - Add `_strategy_npt()` helper in `tileops/kernels/elementwise.py` that returns optimal `num_per_thread` per strategy+dtype pair - `explicit_parallel` fp16/bf16: npt=4 (was 8) -- 58% bandwidth improvement - `register_copy` fp16/bf16: npt=8 (unchanged) -- no regression - fp32: npt=4 for both strategies (unchanged) - fp8: npt=16 for both strategies (unchanged) - Updated `docs/perf/elementwise-evidence.md` with deployment note ## Test plan - [x] Modified files pass unit tests (124 tests pass) - [x] `explicit_parallel` strategy uses npt=4 for fp16/bf16 by default - [x] `register_copy` strategy retains npt=8 for fp16/bf16 - [x] Benchmark: 58% bandwidth improvement for explicit_parallel fp16 at (1024, 10240) vs npt=8 default (threshold: >= 30%) - [x] No regression in register_copy strategy benchmarks - [x] docs/perf/elementwise.md updated with final measured values ## Benchmark | Config | npt | BW (TB/s) | vs old | |---|---|---|---| | explicit_parallel fp16 | 4 (new) | 1.708 | +58.1% | | explicit_parallel fp16 | 8 (old) | 1.081 | baseline | | register_copy fp16 | 8 (unchanged) | 3.940 | no regression | --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
12fa73e411
|
[Fix][Elementwise] Clamp scalar literals to dtype range in kernel layer (#598)
## Summary - Add `_clamp_to_dtype_range()` helper that clamps scalar values to the target dtype's finite representable range before they reach `T.cast()`, preventing TVM `FloatImm` range-check failures for narrow dtypes - Apply clamping in `MaskedFillKernel.__init__` (`fill_value`) and `NanToNumKernel.__init__` (`nan_val`, `posinf_val`, `neginf_val`) - Fixes CI failure: `test_kernel_accepts_fp8[nan_to_num-e4m3fn]` where default `posinf_val=1e4` exceeds `float8_e4m3fn` max of 448 ## Context PR #570 added Op-layer validation (`_validate_scalar_param_repr`) that rejects out-of-range values when users go through the Op API. However, the kernel layer itself still crashes when constructed directly with overflow values — which is exactly what `test_kernel_accepts_fp8` does. This PR adds kernel-level defense so both paths are safe. ## Test plan - [x] `tests/test_elementwise_independent_fp8.py` — 22/22 passed (including previously failing `nan_to_num-e4m3fn`) - [x] `tests/test_elementwise_fp8.py` — 65/65 passed - [x] 87 total fp8 tests, 0 regressions - [x] pre-commit passed 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
fd77066565
|
[Perf][Elementwise] Generalize kernel caching and autotune to all kernel families (#560)
## Summary Extract BinaryKernel's `init_config()` caching pattern (`_compiled_fn` pre-compilation) to **UnaryKernel**, **FusedGatedKernel**, and all **11 custom kernels** (LeakyRelu, Elu, Hardtanh, Softplus, Prelu, Where, Clamp, MaskedFill, NanToNum, Alibi, Sinusoidal). Add `autotune_configs` property with serialization-fallback `autotune()` override to UnaryKernel and FusedGatedKernel. Closes #551 ## Changes - **UnaryKernel**: Added `autotune_configs` (9 configs for fp16/bf16/fp32, 6 for fp8), `autotune()` with serialization fallback, `init_config()` override that caches `_compiled_fn` - **FusedGatedKernel**: Same pattern as UnaryKernel — `autotune_configs`, `autotune()` fallback, `init_config()` caching - **11 custom kernels**: `init_config()` override to cache `_compiled_fn` (custom kernels use factory functions, so autotune is not applicable) - **docs/perf/elementwise.md**: Updated checklist to reflect all kernel families now have caching and autotune ## Test plan - [x] AC-1: Modified files pass unit tests — 432 tests pass (31 new + 401 existing), 0 failures - [x] AC-2: All UnaryKernel subclasses use cached compiled functions — `_compiled_fn` set in `init_config()`, verified by `TestUnaryCaching` (5 tests) - [x] AC-3: FusedGatedKernel subclasses use cached compiled functions — `_compiled_fn` set in `init_config()`, verified by `TestFusedGatedCaching` (4 tests) - [x] AC-4: `autotune_configs` defined for UnaryKernel and FusedGatedKernel with >= 3 configs — verified by `TestAutotuneConfigs` (8 tests) - [x] AC-5: Existing benchmark numbers show no regression — all 401 existing tests pass, BinaryKernel autotune still works - [x] AC-6: `docs/perf/elementwise.md` updated with new caching/autotune patterns ## Benchmark **Environment**: NVIDIA H200 · PyTorch 2.9.1+cu128 · CUDA 12.8 · Driver 575.57.08 173 benchmarks passed. No runtime regression. Caching eliminates per-forward JIT lookup overhead. ### Unary ops (TileOPs vs PyTorch baseline) | Op | N | dtype | TileOPs (ms) | Baseline (ms) | BW (TB/s) TileOPs | BW (TB/s) Baseline | | --- | --- | --- | --- | --- | --- | --- | | exp | 4000000 | fp16 | 0.01 | 0.01 | 2.40 | 2.42 | | exp | 4000000 | bf16 | 0.01 | 0.01 | 2.41 | 2.42 | | gelu | 4000000 | fp16 | 0.01 | 0.01 | 1.64 | 1.77 | | gelu | 4000000 | bf16 | 0.01 | 0.01 | 1.58 | 1.71 | | logical_not | 4000000 | fp16 | 0.01 | 0.01 | 0.82 | 2.07 | | bitwise_not | 4000000 | int32 | 0.02 | 0.01 | 1.94 | 3.06 | | isnan | 4000000 | fp16 | 0.01 | 0.01 | 0.82 | 2.09 | ### Binary ops (TileOPs vs PyTorch baseline, fp16) | Op | Shape | TileOPs (ms) | Baseline (ms) | BW (TB/s) TileOPs | BW (TB/s) Baseline | | --- | --- | --- | --- | --- | --- | | sub | 1024×20480 | 0.03 | 0.03 | 3.80 | 3.82 | | mul | 1024×20480 | 0.03 | 0.03 | 3.80 | 3.82 | | div | 1024×20480 | 0.03 | 0.03 | 3.73 | 3.68 | | remainder | 1024×10240 | 0.02 | 0.02 | 3.36 | 2.66 | | floor_divide | 1024×10240 | 0.02 | 0.06 | 3.36 | 1.03 | | maximum | 1024×20480 | 0.03 | 0.03 | 3.78 | 3.78 | | minimum | 1024×20480 | 0.03 | 0.03 | 3.79 | 3.78 | ### Fused gated ops (TileOPs vs PyTorch baseline, fp16) | Op | Shape | TileOPs (ms) | Baseline (ms) | Speedup | | --- | --- | --- | --- | --- | | gelu_and_mul | 1024×4096 | 0.01 | 0.04 | 4.0× | | gelu_and_mul | 1024×10240 | 0.02 | 0.08 | 4.0× | | gelu_and_mul | 1024×20480 | 0.04 | 0.15 | 3.8× | | gelu_tanh_and_mul | 1024×4096 | 0.01 | 0.03 | 3.0× | | gelu_tanh_and_mul | 1024×10240 | 0.02 | 0.07 | 3.5× | | gelu_tanh_and_mul | 1024×20480 | 0.04 | 0.14 | 3.5× | ### Broadcast ops (TileOPs vs PyTorch baseline, fp16) | Op | TileOPs (ms) | Baseline (ms) | Speedup | | --- | --- | --- | --- | | sub (bcast, largest) | 0.02 | 0.06 | 3.0× | | mul (bcast, largest) | 0.02 | 0.06 | 3.0× | | div (bcast, largest) | 0.03 | 0.07 | 2.3× | ### Strategy comparison (largest shape 1024×20480, fp16) | Op | Strategy | Latency (ms) | BW (TB/s) | | --- | --- | --- | --- | | relu (unary) | direct | 0.07 | 1.20 | | relu (unary) | explicit_parallel | 0.08 | 1.03 | | relu (unary) | register_copy | 0.02 | 3.71 | | add (binary) | direct | 0.08 | 1.65 | | add (binary) | explicit_parallel | 0.03 | 3.81 | | silu_and_mul (fused) | direct | 0.04 | 1.46 | | silu_and_mul (fused) | explicit_parallel | 0.02 | 2.84 | <details> <summary>Full benchmark report (profile_run.log)</summary> ``` 173 passed, 0 failed (516.61s) See profile_run.log for per-op latency/TFLOPS/bandwidth tables. ``` </details> 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
db6f0de738
|
[Feat][Elementwise] Extend fp8 support to independent kernels (#549)
Closes #504 ## Summary - Add `float8_e4m3fn` and `float8_e5m2` support to all 11 independent elementwise kernels (LeakyRelu, Elu, Hardtanh, Softplus, Prelu, Where, MaskedFill, Clamp, NanToNum, Alibi, Sinusoidal) - Apply fp16-accumulation strategy from PR #494: arithmetic kernels compute in fp16 intermediate, selection-only kernels (Where, MaskedFill) pass fp8 through directly - e4m3fn uses saturating `T.Cast` back to fp8 (no Inf representation); e5m2 produces fp16 output with non-saturating Op-layer cast to preserve Inf/NaN - Add `num_per_thread=16` default config for fp8 (1 byte × 16 = 128-bit alignment) **Dtype support matrix:** | Kernel | float16 | bfloat16 | float32 | float8_e4m3fn | float8_e5m2 | |--------|---------|----------|---------|---------------|-------------| | LeakyReluKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | EluKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | HardtanhKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | SoftplusKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | PreluKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | WhereKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | MaskedFillKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | ClampKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | NanToNumKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | AlibiKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | | SinusoidalKernel | ✓ | ✓ | ✓ | ✓ (new) | ✓ (new) | ## Test plan - [x] pre-commit passed - [x] pytest passed — 128/128 tests (33 new fp8 + 95 existing), 0 failures - [x] All independent kernels accept `float8_e4m3fn` and `float8_e5m2` (18 instantiation tests) - [x] Correctness tests for leaky_relu, elu, clamp with both fp8 dtypes (6 tests) - [x] Saturation/overflow behavior matches template kernel semantics (3 tests: e4m3fn saturates at 448.0, e5m2 non-saturating cast) - [x] fp8 default_config uses `num_per_thread=16` (6 tests) ## Structural Readiness - FAIL [REQUIRED] `Op.forward` validates input shape/numel: EluOp, HardtanhOp, SoftplusOp, ClampOp missing dtype/numel checks — **pre-existing gap** (LeakyReluOp has full validation, others only check `is_cuda`). Not introduced by this PR; tracked for follow-up. - SKIP [RECOMMENDED] `__init__.py` exports synchronized: Independent ops are exported via `tileops.ops.elementwise.__all__`, not re-exported from `tileops.ops.__init__`. This matches the existing pattern where users import from the elementwise submodule directly. ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 ### PyTorch native fp8 baseline: where `torch.where` is the only benchmarked op with native PyTorch fp8 support. This is a direct apples-to-apples comparison. | Op | Shape | dtype | TileOPs (ms) | PyTorch (ms) | BW (TB/s) | Speedup | |---|---|---|---|---|---|---| | where | 1024x4096 | float8_e4m3fn | 0.010 | 0.020 | 1.83 | 2.00x | | where | 1024x4096 | float8_e5m2 | 0.010 | 0.020 | 1.42 | 2.00x | | where | 1024x10240 | float8_e4m3fn | 0.050 | 0.030 | 0.85 | 0.60x | | where | 1024x10240 | float8_e5m2 | 0.050 | 0.020 | 0.91 | 0.40x | | where | 1024x20480 | float8_e4m3fn | 0.100 | 0.030 | 0.85 | 0.30x | | where | 1024x20480 | float8_e5m2 | 0.120 | 0.040 | 0.67 | 0.33x | ### No native PyTorch fp8: leaky_relu, elu, clamp, masked_fill PyTorch does not implement fp8 for these ops (`"xxx_cuda" not implemented for 'Float8_e4m3fn'`). Baseline uses `op(x.to(fp16)).to(fp8_dtype)` as a workaround — **this comparison does not have direct reference value** since the baseline includes extra dtype cast overhead that real workloads would also incur. | Op | Shape | dtype | TileOPs (ms) | Baseline\* (ms) | BW (TB/s) | Speedup | |---|---|---|---|---|---|---| | leaky_relu | 1024x4096 | float8_e4m3fn | 0.010 | 0.070 | 0.92 | 7.00x | | leaky_relu | 1024x4096 | float8_e5m2 | 0.050 | 0.060 | 0.17 | 1.20x | | leaky_relu | 1024x10240 | float8_e4m3fn | 0.050 | 0.160 | 0.39 | 3.20x | | leaky_relu | 1024x10240 | float8_e5m2 | 0.120 | 0.130 | 0.17 | 1.08x | | leaky_relu | 1024x20480 | float8_e4m3fn | 0.080 | 0.310 | 0.52 | 3.88x | | leaky_relu | 1024x20480 | float8_e5m2 | 0.230 | 0.320 | 0.18 | 1.39x | | elu | 1024x4096 | float8_e4m3fn | 0.010 | 0.080 | 0.81 | 8.00x | | elu | 1024x4096 | float8_e5m2 | 0.040 | 0.080 | 0.19 | 2.00x | | elu | 1024x10240 | float8_e4m3fn | 0.030 | 0.160 | 0.72 | 5.33x | | elu | 1024x10240 | float8_e5m2 | 0.080 | 0.160 | 0.26 | 2.00x | | elu | 1024x20480 | float8_e4m3fn | 0.060 | 0.330 | 0.75 | 5.50x | | elu | 1024x20480 | float8_e5m2 | 0.160 | 0.310 | 0.27 | 1.94x | | clamp | 1024x4096 | float8_e4m3fn | 0.020 | 0.040 | 0.46 | 2.00x | | clamp | 1024x4096 | float8_e5m2 | 0.030 | 0.070 | 0.27 | 2.33x | | clamp | 1024x10240 | float8_e4m3fn | 0.030 | 0.120 | 0.72 | 4.00x | | clamp | 1024x10240 | float8_e5m2 | 0.050 | 0.160 | 0.40 | 3.20x | | clamp | 1024x20480 | float8_e4m3fn | 0.040 | 0.300 | 1.05 | 7.50x | | clamp | 1024x20480 | float8_e5m2 | 0.120 | 0.260 | 0.36 | 2.17x | | masked_fill | 1024x4096 | float8_e4m3fn | 0.020 | 0.080 | 0.76 | 4.00x | | masked_fill | 1024x4096 | float8_e5m2 | 0.020 | 0.050 | 0.69 | 2.50x | | masked_fill | 1024x10240 | float8_e4m3fn | 0.050 | 0.150 | 0.58 | 3.00x | | masked_fill | 1024x10240 | float8_e5m2 | 0.060 | 0.190 | 0.49 | 3.17x | | masked_fill | 1024x20480 | float8_e4m3fn | 0.050 | 0.340 | 1.24 | 6.80x | | masked_fill | 1024x20480 | float8_e5m2 | 0.050 | 0.340 | 1.22 | 6.80x | \* *Baseline = `op(x.to(float16)).to(fp8_dtype)` — includes dtype cast overhead, not a fair comparison.* **Takeaways:** - **where (native baseline)**: TileOPs wins at small shapes (2x at 4M), but PyTorch's native `torch.where` is faster at medium/large shapes (0.3–0.6x) - **e4m3fn vs e5m2**: e4m3fn consistently faster — in-kernel saturating cast avoids the Op-layer non-saturating cast overhead that e5m2 requires - **Bandwidth**: fp8 kernels achieve up to 1.83 TB/s (where, small shape) and 1.24 TB/s (masked_fill, large shape) on H200 **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_independent_elementwise.py -v -k fp8 ``` --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
449250aa73
|
[Refactor][Elementwise] Extract shared _wrap_fp8_accumulation helper for fp8 kernels (#535)
## Summary - Extract a shared `_wrap_fp8_accumulation(op_func, dtype)` helper that handles both saturating (e4m3fn) and non-saturating (e5m2) fp8 accumulation paths - Replace duplicated `_get_effective_op_func()` bodies in UnaryKernel, BinaryKernel, and FusedGatedKernel with calls to the shared helper - Fix FusedGated direct strategy e5m2 output routing through fp16 buffer to preserve Inf/NaN semantics Closes #502 ## Test plan - [x] Modified files pass unit tests: `python -m pytest -q tests/test_elementwise_fp8.py tests/ops/test_special_elementwise.py` -- 77 passed - [x] No new public API -- helper is module-private (`_wrap_fp8_accumulation`, not in `__all__`) - [x] Saturating (e4m3fn) and non-saturating (e5m2) semantics preserved exactly (overflow/Inf regression coverage included) ## Changes - `tileops/kernels/elementwise.py` -- consolidated fp8 accumulation logic into `_wrap_fp8_accumulation`; fixed FusedGated direct e5m2 output dtype routing - `tests/test_elementwise_fp8.py` -- added 5 new helper tests + 1 e5m2 Inf regression test 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
9d2205b424
|
[Perf][Elementwise] Vectorize bool mask loads in where/masked_fill kernels (#538)
Closes #506 Closes #543 ## Summary - Vectorize bool mask loads in `where` and `masked_fill` kernels by packing bool tensors as `uint8` in the Op layer, enabling `T.copy` vectorized loads (TileLang does not vectorize bool tensors) - Tune kernel configs to `threads=512, npt=8` (fp16/bf16) / `npt=4` (fp32) for H200 - Write results in-place to the `x` register fragment to reduce register pressure - Fix CUPTI profiler kernel-name exclusion bug (#543): monkey-patch tilelang's `_bench_with_cupti` to only exclude `FillFunctor` cache-clearing kernels, not all `vectorized_elementwise` kernels - Preserve existing CUPTI → event fallback for global singleton contention (when `ncu` is running) ## Structural Readiness All checks passed. ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1, TileLang 0.1.8 ### where | Shape | dtype | TileOPs BW (TB/s) | Baseline BW (TB/s) | Speedup | |-------|-------|-------------------|-------------------|---------| | 1024×4096 | fp16 | 2.83 | 2.80 | 1.01x | | 1024×4096 | bf16 | 2.79 | 2.81 | 0.99x | | 1024×4096 | fp32 | 3.19 | 3.24 | 0.98x | | 1024×10240 | fp16 | 3.36 | 3.39 | 0.99x | | 1024×10240 | bf16 | 3.36 | 3.40 | 0.99x | | 1024×10240 | fp32 | 3.77 | 3.78 | 1.00x | | 1024×20480 | fp16 | 3.79 | 3.82 | 0.99x | | 1024×20480 | bf16 | 3.79 | 3.82 | 0.99x | | 1024×20480 | fp32 | 4.07 | 4.08 | 1.00x | ### masked_fill | Shape | dtype | TileOPs BW (TB/s) | Baseline BW (TB/s) | Speedup | |-------|-------|-------------------|-------------------|---------| | 1024×4096 | fp16 | 2.61 | 1.68 | **1.55x** | | 1024×4096 | bf16 | 2.62 | 1.68 | **1.56x** | | 1024×4096 | fp32 | 3.10 | 1.95 | **1.59x** | | 1024×10240 | fp16 | 3.32 | 2.14 | **1.55x** | | 1024×10240 | bf16 | 3.32 | 2.14 | **1.55x** | | 1024×10240 | fp32 | 3.68 | 2.02 | **1.82x** | | 1024×20480 | fp16 | 3.73 | 2.11 | **1.77x** | | 1024×20480 | bf16 | 3.73 | 2.12 | **1.76x** | | 1024×20480 | fp32 | 3.94 | 2.12 | **1.86x** | **Takeaways:** - `masked_fill` achieves **1.55x–1.86x** speedup over PyTorch across all dtypes/shapes, scaling better at larger sizes and fp32 (1.86x at 1024×20480 fp32) - `where` is at parity with PyTorch (0.98x–1.01x) — the original >1.15x target in #506 was based on flawed CUPTI measurements; with corrected profiling, PyTorch's `torch.where` is already highly optimized - fp32 shows the strongest masked_fill gains due to larger memory footprint amplifying the vectorization benefit **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_independent_elementwise.py -k "where or masked_fill" -v ``` ## CUPTI Profiler Fix (Closes #543) tilelang's CUPTI backend (`bench.py:198`) excludes all kernels matching `"at::native::vectorized_elementwise"` to strip `cache.zero_()` overhead. But `cache.zero_()` produces `vectorized_elementwise_kernel<FillFunctor>` — the same kernel family PyTorch uses for all standard elementwise ops (add, where, maximum, relu, etc.), so baseline latencies were reported as near-zero. **Fix**: Monkey-patch the filter to require **both** `"vectorized_elementwise"` and `"FillFunctor"` in the kernel name. Added `hasattr` guard for forward-compatibility with tilelang API changes. ## Test plan - [x] Pre-commit passed - [x] `pytest tests/ops/test_special_elementwise.py -k "masked_fill or where"` — 10 tests pass - [x] Benchmark: masked_fill ≥1.0x vs PyTorch across all shapes/dtypes (measured: 1.55x–1.86x) - [x] Benchmark: where ≥0.95x vs PyTorch across all shapes/dtypes (measured: 0.98x–1.01x) - [x] CUPTI patch: all baselines report plausible non-zero latencies - [x] `hasattr` guard: graceful degradation if tilelang API changes --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
81c0eafb99
|
[Perf][Elementwise] Optimize binary max/min kernels to close bandwidth gap with PyTorch (#539)
## Summary
- Add `register_copy` strategy to `BinaryKernel` with vectorized 128-bit
fragment loads (`uint4`) for same-shape inputs, achieving 2-3x bandwidth
improvement on complex op_funcs (maximum/minimum)
- Simplify `maximum`/`minimum` `op_func` from 5-node comparison chain
with copysign to `T.max`/`T.min` + `isnan` guards, reducing IR node
count
- Add `autotune_configs` to `BinaryKernel` with 9 configurations
(threads in {128, 256, 512} x num_per_thread in {2, 4, 8}) for
fp16/bf16/fp32 and 6 for fp8
- Cache compiled kernel to eliminate JIT lookup overhead on each
`forward()` call
Closes #512
## Benchmark
### Environment
- **GPU**: NVIDIA H200
- **PyTorch**: 2.9.1+cu128
- **CUDA**: 12.8
- **dtype**: float16
- **Measurement**: 1000 iterations, CUDA event timing, 50-iteration
warmup
### Before (main branch)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 1.18 | 1.83 | 65% |
| maximum | (1024, 10240) | 1.32 | 2.79 | 47% |
| maximum | (1024, 20480) | 1.42 | 3.37 | 42% |
| minimum | (1024, 4096) | 1.18 | 1.83 | 65% |
| minimum | (1024, 10240) | 1.31 | 2.76 | 48% |
| minimum | (1024, 20480) | 1.42 | 3.32 | 43% |
### After (this PR, independent H200 re-measurement by Reviewer)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 2.67 | 1.83 | **146%** |
| maximum | (1024, 10240) | 3.36 | 2.77 | **121%** |
| maximum | (1024, 20480) | 3.79 | 3.35 | **113%** |
| minimum | (1024, 4096) | 2.66 | 1.90 | **140%** |
| minimum | (1024, 10240) | 3.36 | 2.72 | **123%** |
| minimum | (1024, 20480) | 3.79 | 3.35 | **113%** |
All target shapes exceed the 80% acceptance threshold, with large shapes
exceeding PyTorch bandwidth.
### Kernel-level (bypassing Op layer Python overhead)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 2.64 | 4.10 | 64% |
| maximum | (1024, 10240) | 3.55 | 3.52 | **101%** |
| maximum | (1024, 20480) | 3.84 | 3.82 | **101%** |
## Test plan
- [x] 123 unit tests pass (`tests/ops/test_binary_arith.py` +
`tests/test_elementwise_strategy_bench.py`)
- [x] NaN propagation tests pass: `test_maximum_nan_propagation`,
`test_minimum_nan_propagation` (16 tests)
- [x] Signed-zero semantics tests pass: `test_maximum_signed_zero`,
`test_minimum_signed_zero` (16 tests)
- [x] Autotune config validation:
`test_binary_kernel_has_autotune_configs`,
`test_binary_kernel_autotune_configs_distinct`
- [x] 6 benchmark runs pass on H200
## Acceptance Criteria
| AC | Description | Status |
|---|---|---|
| AC-1 | Modified files pass unit tests | Pass (123/123) |
| AC-2 | maximum/minimum >= 80% PyTorch BW on target shapes | Pass
(113-146%) |
| AC-3 | Before/after benchmark comparison in PR | Pass (see above) |
| AC-4 | No regression in signed-zero/NaN tests | Pass (16/16) |
| AC-5 | autotune_configs with >= 3 configs | Pass (9 configs for
fp16/bf16/fp32) |
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
|
|
Cao Ying
|
9020a726c1
|
[Feat][Elementwise] Add strategy selection to FusedGatedKernel (#490)
Closes #441 ## Summary Add strategy selection (direct + explicit_parallel) to FusedGatedKernel, aligning with the Unary/Binary kernel patterns. - Added `_make_fused_gated_direct` kernel builder and strategy dispatch to `FusedGatedKernel` - Threaded `strategy` argument through `FusedGatedOp` in the ops layer - Added 13 new tests covering both strategies for silu_and_mul, gelu_and_mul, gelu_tanh_and_mul - Added strategy benchmark harness to `bench_binary_elementwise.py` **Supported dtypes**: `float16`, `bfloat16`, `float32` (inherited from `FusedGatedKernel.SUPPORTED_DTYPES`). ## Test plan - [x] AC-1: Decision documented with benchmark data — explicit_parallel chosen as default - [x] AC-2: FusedGatedKernel supports both strategies, all tests pass - [x] AC-3: Strategy validation uses ValueError (not assert) **Test commands:** ```bash PYTHONPATH="$PWD" python -m pytest tests/ops/test_fused_gated.py -v ``` ## Structural Compliance All checks passed. ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 Strategy comparison: `direct` vs `explicit_parallel` (before/after). ### silu_and_mul | Shape | dtype | Direct (ms) | Direct BW (TB/s) | Explicit (ms) | Explicit BW (TB/s) | Speedup | |-------|-------|------------|------------------|--------------|-------------------|---------| | 1024×4096 | fp16 | 0.020 | 1.35 | 0.010 | 2.41 | 1.8x | | 1024×4096 | bf16 | 0.020 | 1.35 | 0.010 | 2.58 | 1.9x | | 1024×4096 | fp32 | 0.020 | 2.38 | 0.020 | 3.25 | 1.4x | | 1024×10240 | fp16 | 0.040 | 1.45 | 0.020 | 2.83 | 2.0x | | 1024×10240 | bf16 | 0.040 | 1.46 | 0.020 | 3.24 | 2.2x | | 4096×4096 | fp16 | 0.070 | 1.50 | 0.030 | 3.04 | 2.0x | | 4096×4096 | bf16 | 0.070 | 1.51 | 0.030 | 3.46 | 2.3x | | 4096×4096 | fp32 | 0.070 | 2.74 | 0.050 | 4.00 | 1.5x | ### gelu_tanh_and_mul | Shape | dtype | Direct (ms) | Direct BW (TB/s) | Explicit (ms) | Explicit BW (TB/s) | Speedup | |-------|-------|------------|------------------|--------------|-------------------|---------| | 1024×4096 | fp16 | 0.020 | 1.38 | 0.010 | 2.73 | 2.0x | | 1024×10240 | fp16 | 0.040 | 1.49 | 0.020 | 3.20 | 2.1x | | 4096×4096 | fp16 | 0.070 | 1.51 | 0.030 | 3.39 | 2.2x | | 4096×4096 | bf16 | 0.070 | 1.51 | 0.030 | 3.84 | 2.5x | | 4096×4096 | fp32 | 0.070 | 2.80 | 0.050 | 4.12 | 1.5x | **Takeaways:** - `explicit_parallel` is **1.4–2.5x faster** than `direct` across all shapes/dtypes - Speedup increases with tensor size — 8 elements/thread amortizes loop overhead better at scale - bf16 ≈ fp16 latency; fp32 shows smaller strategy gap due to higher per-element bandwidth - Default choice `explicit_parallel` validated across 54 benchmark cases **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_binary_elementwise.py::test_fused_gated_strategy_bench -v # 54 passed in 157s ``` --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
Cao Ying
|
929575aaec
|
[Bench][Elementwise] Add strategy benchmarks for UnaryKernel and BinaryKernel (#499)
Closes #498 ## Summary Benchmark all strategies for UnaryKernel and BinaryKernel on DNN-realistic shapes to validate the current `DEFAULT_STRATEGY` settings. - Add `bench_unary_strategy.py`: relu across 3 shapes × 3 dtypes × 3 strategies (27 cases) - Add `bench_binary_strategy.py`: add across 3 shapes × 3 dtypes × 2 strategies (18 cases) - Add `test_elementwise_strategy_bench.py`: 12 structural validation tests (smoke tier) - Fix `assert` → `ValueError` for strategy validation in `UnaryKernel.__init__` and `BinaryKernel.__init__` - Remove issue references from source/test docstrings ## Test plan - [x] 12 smoke-tier structural tests passed - [x] 27 unary benchmark cases passed (16s) - [x] 18 binary benchmark cases passed (12s) - [x] pre-commit passed ## Structural Compliance - FAIL (fixed): `UnaryKernel.__init__` and `BinaryKernel.__init__` used `assert` for strategy validation — replaced with `ValueError` - FAIL (fixed): Source/test files contained issue references (`#498`) — removed ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 ### UnaryKernel (relu) | Shape | dtype | direct (TB/s) | explicit_parallel (TB/s) | register_copy (TB/s) | Winner | |-------|-------|---------------|--------------------------|----------------------|--------| | 1024×4096 | fp16 | 1.06 | 0.86 | **2.53** | register_copy | | 1024×4096 | bf16 | 1.06 | 0.86 | **2.53** | register_copy | | 1024×4096 | fp32 | 1.95 | 2.83 | **2.98** | register_copy | | 1024×10240 | fp16 | 1.15 | 0.97 | **3.23** | register_copy | | 1024×10240 | bf16 | 1.14 | 0.97 | **3.22** | register_copy | | 1024×10240 | fp32 | 2.12 | 3.42 | **3.69** | register_copy | | 1024×20480 | fp16 | 1.19 | 1.01 | **3.69** | register_copy | | 1024×20480 | bf16 | 1.19 | 1.01 | **3.70** | register_copy | | 1024×20480 | fp32 | 2.20 | 3.70 | **3.98** | register_copy | ### BinaryKernel (add) | Shape | dtype | direct (TB/s) | explicit_parallel (TB/s) | Winner | |-------|-------|---------------|--------------------------|--------| | 1024×4096 | fp16 | 1.43 | **2.75** | explicit_parallel | | 1024×4096 | bf16 | 1.43 | **2.75** | explicit_parallel | | 1024×4096 | fp32 | 2.47 | **3.27** | explicit_parallel | | 1024×10240 | fp16 | 1.56 | **3.43** | explicit_parallel | | 1024×10240 | bf16 | 1.55 | **3.43** | explicit_parallel | | 1024×10240 | fp32 | 2.74 | **3.77** | explicit_parallel | | 1024×20480 | fp16 | 1.63 | **3.77** | explicit_parallel | | 1024×20480 | bf16 | 1.63 | **3.78** | explicit_parallel | | 1024×20480 | fp32 | 2.89 | **4.06** | explicit_parallel | **Takeaways:** - UnaryKernel: `register_copy` wins all 9 shape×dtype pairs (2.53–3.98 TB/s vs next-best 0.86–3.70 TB/s). The advantage is largest for fp16/bf16 (~2.4–3.1× over direct) and narrows for fp32 large shapes. - BinaryKernel: `explicit_parallel` wins all 9 pairs (2.75–4.06 TB/s vs direct 1.43–2.89 TB/s), ~1.3–2.2× speedup. The gap is consistent across dtypes. - Current `DEFAULT_STRATEGY` settings (`register_copy` for unary, `explicit_parallel` for binary) are confirmed optimal. No update warranted. **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_unary_strategy.py benchmarks/ops/bench_binary_strategy.py -v ``` ## Files changed - `benchmarks/ops/bench_unary_strategy.py` — UnaryKernel strategy benchmark (relu, 27 cases) - `benchmarks/ops/bench_binary_strategy.py` — BinaryKernel strategy benchmark (add, 18 cases) - `tests/test_elementwise_strategy_bench.py` — Structural tests for benchmark modules (12 smoke tests) - `tileops/kernels/elementwise.py` — Fix `assert` → `ValueError` for strategy validation; update benchmark comment - `.claude/skills/creating-pull-request/op-compliance-checklist.md` — Add Op.forward shape/numel validation item --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
d9b83d0456
|
[Fix][Elementwise] Correct maximum/minimum signed-zero tie-break (#508)
Closes #469 ## Summary Fix `maximum` / `minimum` elementwise ops to match PyTorch / IEEE 754 signed-zero semantics: - `maximum(+0.0, -0.0)` now correctly returns `+0.0` (was returning `-0.0`) - `minimum(-0.0, +0.0)` now correctly returns `-0.0` (was returning `+0.0`) The fix adds a tie-break path using `copysign` comparison that activates only when `a == b`, preserving the existing NaN-propagation behavior from PR #463. **Dtype support matrix:** | dtype | supported | |-------|-----------| | float16 | yes | | bfloat16 | yes | | float32 | yes | ## Structural Compliance All checks passed. ## Test plan - [x] `maximum(+0.0, -0.0)` returns `+0.0` for fp16, bf16, fp32 - [x] `minimum(-0.0, +0.0)` returns `-0.0` for fp16, bf16, fp32 - [x] NaN propagation behavior from PR #463 is preserved (not regressed) - [x] 12 new signed-zero regression tests added and passing - [x] No performance regression — tie-break path only activates on `a == b` ## Benchmark **Environment**: H200 GPU, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 Before/after comparison (`main` vs this branch, fp16): | Op | Shape | Before Latency (ms) | Before BW (TB/s) | After Latency (ms) | After BW (TB/s) | Delta | |----|-------|---------------------|-------------------|---------------------|-------------------|-------| | maximum | (1024, 4096) | 0.02 | 1.18 | 0.02 | 1.18 | 0% | | maximum | (1024, 10240) | 0.05 | 1.33 | 0.05 | 1.32 | -0.8% | | maximum | (1024, 20480) | 0.09 | 1.43 | 0.09 | 1.42 | -0.7% | | minimum | (1024, 4096) | 0.02 | 1.19 | 0.02 | 1.18 | -0.8% | | minimum | (1024, 10240) | 0.05 | 1.32 | 0.05 | 1.31 | -0.8% | | minimum | (1024, 20480) | 0.09 | 1.43 | 0.09 | 1.42 | -0.7% | **Takeaways:** - No measurable regression — deltas are within noise margin (<1%) - The `copysign`-based tie-break only activates when `a == b`, which is rare in real workloads - Existing gap vs PyTorch baseline (42–65% efficiency) is a pre-existing issue tracked in #512 **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_binary_elementwise.py -v -k "maximum or minimum" ``` ## Regression - 101 tests pass (12 signed-zero + 10 NaN propagation + existing suite) - Lint clean (`pre-commit run --all-files` passes) - Signed-zero fix does not alter NaN-propagation codepath — verified by `test_maximum_signed_zero_with_nan` and `test_minimum_signed_zero_with_nan` --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
108bf5698f
|
[Feat][Elementwise] Add fp8 dtype support (e4m3fn, e5m2) (#494)
Closes #442 ## Summary Extends all elementwise kernel templates (UnaryKernel, BinaryKernel, FusedGatedKernel) to support fp8 dtypes (`float8_e4m3fn`, `float8_e5m2`). Implements fp16-accumulated compute strategy to avoid precision loss from direct fp8 arithmetic, with proper saturation semantics matching the NVIDIA spec. ### Changes - **`tileops/kernels/elementwise.py`** — fp8 default_config with `num_per_thread=16` for 128-bit alignment; fp8 accumulation wrappers that cast fp8 inputs to fp16, compute, then cast back; e4m3fn uses saturating cast, e5m2 preserves Inf/NaN - **`tileops/ops/elementwise.py`** — Register fp8 dtypes in op-level dispatch; fix `output_dtype` and `total_memory` to reflect actual post-cast dtype for e5m2 - **`tests/test_elementwise_fp8.py`** — 25 tests: template acceptance, config alignment, correctness, saturation/overflow, dtype rejection - **`benchmarks/ops/bench_elementwise_fp8.py`** — fp8 benchmarks for unary (relu, exp), binary (add), fused gated (silu_and_mul) with LLaMA-family shapes ### Dtype support matrix | Op category | e4m3fn | e5m2 | Notes | |-------------|--------|------|-------| | All template unary (relu, exp, silu, gelu, ...) | ✓ | ✓ | Via UnaryKernel fp8 accumulation | | All template binary (add, mul, div, ...) | ✓ | ✓ | Via BinaryKernel fp8 accumulation | | All fused gated (silu_and_mul, gelu_and_mul, ...) | ✓ | ✓ | Via FusedGatedKernel fp8 accumulation | | Comparison ops (eq, gt, ...) | ✓ | ✓ | Output is int8, not fp8 | | Logical ops (logical_not, ...) | ✗ | ✗ | SUPPORTED_DTYPES excludes fp8 (ValueError) | | Bitwise ops (bitwise_and, ...) | ✗ | ✗ | Integer-only (ValueError) | | Independent kernels (leaky_relu, elu, ...) | ✗ | ✗ | _FLOAT_DTYPES excludes fp8 (ValueError) | ## Test plan - [x] AC-1: fp8 accumulation strategy designed — compute in fp16, cast back to fp8 - [x] AC-2: Template base classes support fp8_e4m3fn and fp8_e5m2 (kernel instantiation tests) - [x] AC-3: Correctness tests pass for representative ops — relu, add, silu_and_mul with both fp8 dtypes - [x] AC-4: Saturation/overflow matches NVIDIA spec — e4m3fn saturates (no Inf), e5m2 produces Inf/NaN on overflow - [x] AC-5: Unsupported dtype rejection — BitwiseNotKernel, BitwiseAndKernel raise ValueError for fp8 - [x] AC-6: `output_dtype` and `total_memory` reflect actual post-cast dtype for e5m2 ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128 Shapes model real LLM workloads (LLaMA-family dimensions): - **Small**: 1×2048×4096 = 8M elements (LLaMA-7B single-batch inference) - **Medium**: 8×2048×4096 = 67M elements (LLaMA-7B multi-batch inference) - **Large**: 4×4096×8192 = 134M elements (LLaMA-70B training) - **Fused gated**: (batch×seq, intermediate_dim) — LLaMA-7B intermediate=11008, LLaMA-70B intermediate=28672 ### Unary ops (relu, exp) | Op | Shape | dtype | tileops lat(ms) | baseline lat(ms) | tflops | bandwidth(TB/s) | speedup | |----|-------|-------|-----------------|-------------------|--------|-----------------|---------| | relu | 8M | e4m3fn | 0.0153 | 0.0334 | 0.55 | 1.10 | 2.18x | | relu | 8M | e5m2 | 0.0095 | 0.0309 | 0.89 | 1.77 | 3.27x | | relu | 67M | e4m3fn | 0.0990 | 0.2415 | 0.68 | 1.36 | 2.44x | | relu | 67M | e5m2 | 0.0557 | 0.2254 | 1.20 | 2.41 | 4.04x | | relu | 134M | e4m3fn | 0.1940 | 0.4791 | 0.69 | 1.38 | 2.47x | | relu | 134M | e5m2 | 0.1086 | 0.4435 | 1.24 | 2.47 | 4.09x | | exp | 8M | e4m3fn | 0.0097 | 0.0335 | 0.86 | 1.73 | 3.45x | | exp | 8M | e5m2 | 0.0101 | 0.0310 | 0.83 | 1.66 | 3.06x | | exp | 67M | e4m3fn | 0.0597 | 0.2406 | 1.12 | 2.25 | 4.03x | | exp | 67M | e5m2 | 0.0616 | 0.2249 | 1.09 | 2.18 | 3.65x | | exp | 134M | e4m3fn | 0.1162 | 0.4756 | 1.16 | 2.31 | 4.09x | | exp | 134M | e5m2 | 0.1201 | 0.4453 | 1.12 | 2.24 | 3.71x | ### Binary op (add — residual connection) | Op | Shape | dtype | tileops lat(ms) | baseline lat(ms) | tflops | bandwidth(TB/s) | speedup | |----|-------|-------|-----------------|-------------------|--------|-----------------|---------| | add | 8M | e4m3fn | 0.0137 | 0.0665 | 0.61 | 1.84 | 4.85x | | add | 8M | e5m2 | 0.0116 | 0.0616 | 0.72 | 2.16 | 5.29x | | add | 67M | e4m3fn | 0.0851 | 0.4792 | 0.79 | 2.36 | 5.63x | | add | 67M | e5m2 | 0.0702 | 0.4472 | 0.96 | 2.87 | 6.37x | | add | 134M | e4m3fn | 0.1656 | 0.9451 | 0.81 | 2.43 | 5.71x | | add | 134M | e5m2 | 0.1372 | 0.8819 | 0.98 | 2.93 | 6.43x | ### Fused gated op (silu_and_mul — SwiGLU FFN) | Op | Shape (M×N) | dtype | tileops lat(ms) | baseline lat(ms) | tflops | bandwidth(TB/s) | speedup | |----|-------------|-------|-----------------|-------------------|--------|-----------------|---------| | silu_and_mul | 2048×11008 (LLaMA-7B) | e4m3fn | 0.0555 | 0.3135 | 2.03 | 1.22 | 5.65x | | silu_and_mul | 2048×11008 (LLaMA-7B) | e5m2 | 0.0484 | 0.3006 | 2.33 | 1.40 | 6.22x | | silu_and_mul | 16384×11008 (LLaMA-7B 8batch) | e4m3fn | 0.4072 | 2.3264 | 2.21 | 1.33 | 5.71x | | silu_and_mul | 16384×11008 (LLaMA-7B 8batch) | e5m2 | 0.3530 | 2.2319 | 2.55 | 1.53 | 6.32x | | silu_and_mul | 16384×28672 (LLaMA-70B) | e4m3fn | 1.0052 | 5.6502 | 2.34 | 1.40 | 5.62x | | silu_and_mul | 16384×28672 (LLaMA-70B) | e5m2 | 0.8441 | 5.3188 | 2.78 | 1.67 | 6.30x | **Takeaways:** - TileOPs fp8 kernels are **2.2–6.4x faster** than PyTorch fp16-compute-then-cast baselines on H200 - **Add (residual)** shows the largest gains (4.9–6.4x) — the fused cast+compute avoids 2 separate PyTorch kernel launches - **Fused gated SwiGLU** at LLaMA-70B scale: 5.6–6.3x speedup, reaching **2.78 TFLOPs** and **1.67 TB/s** - e5m2 is consistently faster than e4m3fn (~10–20%) — the non-saturating cast path avoids the `T.Cast` saturation overhead in the kernel - Peak bandwidth: 2.93 TB/s (add, 134M, e5m2) — H200 HBM3e theoretical peak is ~4.8 TB/s, so we're at ~61% utilization for this bandwidth-bound op **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_elementwise_fp8.py -v ``` --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
eb8ce2f554
|
[Perf][Elementwise] Optimize 5 elementwise ops with register_copy + bool dtype (#500)
Closes #497 ## Summary - Optimize 5 of 6 target independent elementwise ops to ≥1.0× vs PyTorch baseline at (1024, 10240) fp16 on H200 - Apply `register_copy` strategy (fragment load → compute → store) to leaky_relu, nan_to_num, prelu, where, masked_fill - Fix alibi benchmark baseline to compute correct 3D `(num_heads, seq_len, seq_len)` output - Switch where/masked_fill kernels to accept `bool` tensors directly (eliminating int8 cast overhead) - Add dtype/numel validation to 5 independent Op.forward methods (compliance fix) ## Test plan - [x] `pytest tests/ops/test_special_elementwise.py tests/ops/test_activation.py -m smoke` — 43 passed - [x] `pytest benchmarks/ops/bench_independent_elementwise.py -v` — 102 passed (125s) - [x] pre-commit passed ## Structural Compliance - FAIL (fixed): LeakyReluOp, PreluOp, WhereOp, MaskedFillOp, NanToNumOp `.forward()` missing dtype/numel validation — added ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 ### Optimized ops at target shape (1024×10240) fp16 | Op | TileOps (TB/s) | PyTorch (TB/s) | Speedup | |----|----------------|----------------|---------| | leaky_relu | 3.16 | 2.57 | **1.23×** | | nan_to_num | 3.08 | 2.51 | **1.23×** | | prelu (C=10240) | 3.16 | 1.18 | **2.68×** | | where | 3.43 | 2.99 | **1.15×** | | masked_fill | 3.22 | 4.00 | 0.81× | ### ALiBi (generative shapes, fp16) | Shape (seq_len, heads) | TileOps (ms) | PyTorch (ms) | Speedup | |------------------------|--------------|--------------|---------| | (512, 64) | 0.03 | 0.06 | **2.00×** | | (2048, 64) | 0.30 | 0.77 | **2.57×** | | (4096, 128) | 1.12 | 7.20 | **6.43×** | ### Full results across dtypes at (1024×10240) | Op | fp16 TileOps/PyTorch (TB/s) | bf16 TileOps/PyTorch (TB/s) | fp32 TileOps/PyTorch (TB/s) | |----|----|----|-----| | leaky_relu | 3.16 / 2.57 (1.23×) | 3.18 / 2.56 (1.24×) | 3.69 / 3.21 (1.15×) | | nan_to_num | 3.08 / 2.51 (1.23×) | 3.09 / 2.53 (1.22×) | 3.69 / 3.21 (1.15×) | | prelu | 3.16 / 1.18 (2.68×) | 3.16 / 1.16 (2.72×) | 3.66 / 2.10 (1.74×) | | where | 3.43 / 2.99 (1.15×) | 3.43 / 2.99 (1.15×) | 3.76 / 3.48 (1.08×) | | masked_fill | 3.22 / 4.00 (0.81×) | 3.21 / 4.00 (0.80×) | 3.60 / 4.11 (0.88×) | **Takeaways:** - leaky_relu, nan_to_num: register_copy gives 1.15–1.24× across all dtypes. Consistent wins from coalesced memory access. - prelu: Largest gain (1.74–2.72×). PyTorch PReLU is slow due to per-channel weight scatter; register_copy with flat inner loop amortizes the channel index computation. - where: 1.08–1.15× via register_copy for data tensors (x, y, out) with scalar bool cond access. T.copy does not support bool vectorization, limiting further gains. - masked_fill: 0.80–0.88× — PyTorch achieves 82% H200 peak HBM bandwidth via hand-optimized vectorized CUDA. TileLang capped at 66% due to `T.copy`/`T.alloc_fragment` not supporting bool dtype (`boolx8` vectorization error). Deferred to follow-up per issue constraint. - alibi: 2.00–6.43× — scales superlinearly with seq_len because TileLang kernel is O(N²) flat-index while PyTorch baseline uses O(N²) tensor broadcast with more memory overhead. **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_independent_elementwise.py -v ``` ## Files changed - `tileops/kernels/elementwise.py` — register_copy strategy for leaky_relu, nan_to_num, prelu, where, masked_fill; bool dtype for where/masked_fill cond/mask tensors; npt tuning - `tileops/ops/elementwise.py` — WhereOp/MaskedFillOp: skip int8 cast, use bool directly; add dtype/numel validation to 5 independent Op.forward methods - `benchmarks/ops/bench_independent_elementwise.py` — Fix alibi baseline to full (H,S,S) shape; add prelu (1024,10240) shape --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
639ee5fe74
|
[Chore][Kernels] Replace assert-based validation with ValueError and remove issue refs (#507)
## Summary Brings the codebase into full compliance with the op-compliance-checklist by: - Replacing all `assert`-based runtime input validation in `tileops/` with `if not <cond>: raise ValueError(<msg>)`, preserving original error messages - Removing all `issue #NNN` references from source/test/benchmark files - Adding compliance regression tests (`tests/test_no_assert_validation.py`) Internal-invariant asserts (post-computation sanity checks) are intentionally left untouched. Closes #505 ## Test plan - [x] AC-1: Zero assert statements used for runtime input validation in `tileops/` (grep verification) - [x] AC-2: Zero `issue #\d+` patterns in `tileops/`, `tests/`, `benchmarks/` source files - [x] AC-3: `python -m pytest tests/ -m smoke` passes (357 passed, 0 failed) - [x] AC-4: `ruff check tileops/ tests/ benchmarks/` clean - [x] AC-5: Modified files pass unit tests (92 passed via targeted test run) ## Changed files (22) **Kernels** (12 files): replaced assert-based input validation with if/raise ValueError - `tileops/kernels/deepseek_mla/deepseek_dsa_decode.py` - `tileops/kernels/deepseek_mla/deepseek_mla_decode.py` - `tileops/kernels/deepseek_nsa/gqa_sliding_window_fwd.py` - `tileops/kernels/deepseek_nsa/gqa_sliding_window_varlen_fwd.py` - `tileops/kernels/deepseek_nsa/nsa_fwd.py` - `tileops/kernels/elementwise.py` - `tileops/kernels/flash_attn/bwd.py` - `tileops/kernels/flash_attn/fwd.py` - `tileops/kernels/flash_decode/gqa_decode.py` - `tileops/kernels/flash_decode/gqa_decode_paged.py` - `tileops/kernels/flash_decode/mha_decode.py` - `tileops/kernels/flash_decode/mha_decode_paged.py` **Ops** (5 files): replaced assert-based validation + removed issue refs - `tileops/ops/deepseek_dsa_decode.py` - `tileops/ops/gated_deltanet.py` - `tileops/ops/gqa_sliding_window_fwd.py` - `tileops/ops/gqa_sliding_window_varlen_fwd.py` - `tileops/ops/op.py` **Tests** (3 files): compliance tests + removed issue refs - `tests/test_no_assert_validation.py` (new) - `tests/ops/test_unary_math.py` - `tests/test_reduction_init_files.py` **Benchmarks** (1 file): removed issue ref - `benchmarks/ops/bench_unary_elementwise.py` **Gated deltanet kernel**: `tileops/kernels/gated_delta_net/gated_deltanet_decode.py` --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
0b6c84ba7d
|
[Feat][Elementwise] Add 11 independent elementwise ops with custom signatures (#489)
Closes #439 Closes #496 ## Summary Add 11 independent elementwise operators, each inheriting the Kernel/Op ABC directly with custom signatures (no shared template base class): - **Activation ops**: `leaky_relu`, `elu`, `hardtanh`, `softplus`, `prelu` - **Special elementwise ops**: `where`, `clamp`, `masked_fill`, `nan_to_num`, `alibi`, `sinusoidal` Each operator includes a TileLang kernel implementation in `tileops/kernels/elementwise.py` and a corresponding Op wrapper in `tileops/ops/elementwise.py`. All 11 ops are registered for `torch.compile` support. **Supported dtypes**: `float16`, `bfloat16`, `float32` (enforced via `SUPPORTED_DTYPES` validation in each kernel). ## Test plan - [x] AC-1: 11 kernel + 11 op classes implemented (each inheriting Kernel/Op ABC) - [x] AC-2: `__all__` updated in both `elementwise.py` files - [x] AC-3: L1: 11 ops pass (fp16, 1M elements) - [x] AC-4: L4: 8 edge cases pass (fp32, 4K elements) - [x] AC-5: `__init__.py` not modified - [x] AC-6: Unsupported-dtype rejection tests pass (`ValueError` on `torch.int32`) **Test commands:** ```bash PYTHONPATH="$PWD" python -m pytest -q tests/ops/test_activation.py tests/ops/test_special_elementwise.py ``` ## Structural Compliance All checks passed. ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 Shapes: DNN-realistic 2-D (tokens=1024 × hidden_dim). Non-pow2 dims included (10240, 20480, d_model=300). ### Unary-like ops (fp16, TileOPs vs PyTorch baseline) | Op | Shape | Latency (ms) | TFLOPs | BW (TB/s) | Baseline (ms) | Speedup | |----|-------|-------------|--------|-----------|--------------|---------| | leaky_relu | 1024×4096 | 0.020 | 0.21 | 0.84 | 0.010 | 0.5x | | leaky_relu | 1024×10240 | 0.044 | 0.24 | 0.97 | 0.017 | 0.4x | | leaky_relu | 1024×20480 | 0.082 | 0.25 | 1.01 | 0.026 | 0.3x | | elu | 1024×4096 | 0.011 | 0.38 | 1.52 | 0.015 | 1.4x | | elu | 1024×10240 | 0.023 | 0.46 | 1.83 | 0.027 | 1.2x | | elu | 1024×20480 | 0.042 | 0.50 | 2.00 | 0.048 | 1.1x | | hardtanh | 1024×4096 | 0.007 | 0.62 | 2.50 | 0.011 | 1.6x | | hardtanh | 1024×10240 | 0.013 | 0.80 | 3.22 | 0.018 | 1.4x | | hardtanh | 1024×20480 | 0.023 | 0.92 | 3.70 | 0.029 | 1.3x | | softplus | 1024×4096 | 0.014 | 0.30 | 1.21 | 0.018 | 1.3x | | softplus | 1024×10240 | 0.029 | 0.36 | 1.42 | 0.034 | 1.2x | | softplus | 1024×20480 | 0.055 | 0.38 | 1.53 | 0.058 | 1.1x | | clamp | 1024×4096 | 0.007 | 0.63 | 2.51 | 0.011 | 1.6x | | clamp | 1024×10240 | 0.013 | 0.81 | 3.24 | 0.018 | 1.4x | | clamp | 1024×20480 | 0.023 | 0.93 | 3.71 | 0.028 | 1.2x | | nan_to_num | 1024×4096 | 0.013 | 0.32 | 1.28 | 0.011 | 0.8x | | nan_to_num | 1024×10240 | 0.029 | 0.36 | 1.45 | 0.018 | 0.6x | | nan_to_num | 1024×20480 | 0.054 | 0.39 | 1.55 | 0.028 | 0.5x | ### Multi-input ops (fp16) | Op | Shape | Latency (ms) | TFLOPs | BW (TB/s) | Baseline (ms) | Speedup | |----|-------|-------------|--------|-----------|--------------|---------| | prelu | 1024×128 | 0.003 | 0.05 | 0.19 | 0.003 | 1.0x | | prelu | 1024×4096 | 0.019 | 0.22 | 0.87 | 0.016 | 0.8x | | prelu | 1024×20480 | 0.076 | 0.28 | 1.11 | 0.066 | 0.9x | | where | 1024×4096 | 0.034 | 0.12 | 0.85 | 0.015 | 0.4x | | where | 1024×10240 | 0.073 | 0.14 | 0.99 | 0.026 | 0.4x | | where | 1024×20480 | 0.140 | 0.15 | 1.05 | 0.042 | 0.3x | | masked_fill | 1024×4096 | 0.033 | 0.13 | 0.64 | 0.007 | 0.2x | | masked_fill | 1024×10240 | 0.072 | 0.15 | 0.73 | 0.013 | 0.2x | | masked_fill | 1024×20480 | 0.138 | 0.15 | 0.76 | 0.023 | 0.2x | ### Generative ops (fp16) | Op | seq_len | dim | Latency (ms) | TFLOPs | Baseline (ms) | Speedup | |----|---------|-----|-------------|--------|--------------|---------| | alibi | 512 | 64 | 0.032 | 0.00 | 0.010 | 0.3x | | alibi | 2048 | 64 | 0.280 | 0.00 | 0.010 | 0.04x | | alibi | 4096 | 128 | 0.940 | 0.00 | 0.010 | 0.01x | | sinusoidal | 512 | 256 | 0.003 | 0.04 | 0.013 | 4.3x | | sinusoidal | 2048 | 300 | 0.005 | 0.13 | 0.016 | 3.2x | | sinusoidal | 4096 | 512 | 0.007 | 0.29 | 0.023 | 3.3x | **Takeaways:** - `hardtanh`, `clamp`: **1.2–1.6x faster** — simple compare+select maps well to TileLang - `elu`, `softplus`: **1.1–1.4x faster** — fp32 promotion overhead but still ahead - `sinusoidal`: **3.2–4.3x faster** — fused sin/cos vs PyTorch multi-pass - `leaky_relu`, `nan_to_num`: **0.3–0.8x** — simple ops where PyTorch fused kernels are hard to beat at DNN-scale - `where`, `masked_fill`: **0.2–0.4x** — multi-tensor reads not yet coalesced - `alibi`: **0.01–0.3x** — per-element kernel fundamentally wrong for broadcast pattern - bf16 ≈ fp16 across all ops; fp32 shows higher bandwidth utilization **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_independent_elementwise.py -v # 99 passed in 427s ``` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
Cao Ying
|
783fbe5f7c
|
[Feat][Elementwise] Register elementwise ops for torch.compile (#482)
## Summary Register all elementwise ops for `torch.compile` via `@torch.library.custom_op` + `register_fake`. Three factory functions handle 55 template ops; independent ops use hand-written registration. After this change, `Kernel.forward` switches from `self.kernel(x)` to `type(self)._wrapped(...)`. Closes #440 ## Changes - **`tileops/ops/elementwise.py`** — Add `_register_unary_custom_op`, `_register_binary_custom_op`, `_register_fused_gated_custom_op` factory functions; `_wrapped` binding on all kernel classes at module load time; forward path switch to use `_wrapped`; `_OP_REGISTRY` as `WeakValueDictionary` to avoid leaks - **`tileops/kernels/elementwise.py`** — Fix `RemainderKernel` and `FloorDivideKernel` fp16 precision: promote division to fp32 before `T.floor` to avoid rounding near integer quotients - **`tests/test_elementwise_compile.py`** — 60 torch.compile correctness tests covering every registered op (unary, binary, comparison, logical, bitwise, fused gated) - **`benchmarks/ops/bench_compile_elementwise.py`** — New benchmark: eager vs torch.compile latency and bandwidth ## Test plan - [x] **AC-1**: All 55 elementwise ops work under `torch.compile` — full-matrix smoke with `fullgraph=True`; 60 pytest cases passed - [x] **AC-2**: `register_fake` returns correct shape and dtype (including `torch.bool` for comparison) — verified via dtype/shape assertions - [x] **AC-3**: Eager mode L1 tests still pass — 196 passed across binary_arith, unary_math, fused_gated, special_elementwise, bitwise suites - [x] **AC-4**: Import time measured and acceptable — average 160.60ms (max 162.43ms), well below 500ms constraint ## Benchmark Environment: Torch 2.9.1+cu128, CUDA 12.8, NVIDIA H200, Driver 575.57.08 ### Unary Ops (fp16, latency us + bandwidth TB/s) | Op | N | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | relu | 262,144 | 2.3 | 2.1 | 0.92x | 0.47 | 0.51 | | relu | 1,048,576 | 3.1 | 3.0 | 0.95x | 1.34 | 1.41 | | relu | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.50 | 2.58 | | exp | 262,144 | 2.3 | 2.1 | 0.92x | 0.46 | 0.50 | | exp | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.33 | 1.40 | | exp | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.48 | 2.57 | | gelu | 262,144 | 2.4 | 2.2 | 0.93x | 0.44 | 0.47 | | gelu | 1,048,576 | 3.5 | 3.4 | 0.96x | 1.20 | 1.25 | | gelu | 4,000,000 | 8.1 | 8.0 | 0.99x | 1.99 | 2.01 | | silu | 262,144 | 2.4 | 2.2 | 0.92x | 0.44 | 0.48 | | silu | 1,048,576 | 3.4 | 3.2 | 0.95x | 1.23 | 1.29 | | silu | 4,000,000 | 7.9 | 7.7 | 0.97x | 2.01 | 2.07 | | abs | 262,144 | 2.2 | 2.1 | 0.92x | 0.47 | 0.51 | | abs | 1,048,576 | 3.1 | 3.0 | 0.94x | 1.33 | 1.41 | | abs | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.50 | 2.57 | | sqrt | 262,144 | 2.4 | 2.3 | 0.94x | 0.43 | 0.46 | | sqrt | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.31 | 1.38 | | sqrt | 4,000,000 | 6.5 | 6.2 | 0.97x | 2.48 | 2.56 | | log | 262,144 | 2.4 | 2.2 | 0.92x | 0.45 | 0.48 | | log | 1,048,576 | 3.4 | 3.2 | 0.94x | 1.24 | 1.31 | | log | 4,000,000 | 7.9 | 7.7 | 0.98x | 2.03 | 2.06 | | tanh | 262,144 | 2.3 | 2.1 | 0.92x | 0.45 | 0.49 | | tanh | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.32 | 1.40 | | tanh | 4,000,000 | 6.5 | 6.3 | 0.97x | 2.44 | 2.53 | ### Binary Ops (fp16, latency us + bandwidth TB/s) | Op | Shape (MxN) | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | add | 1024x4096 | 8.9 | 8.7 | 0.98x | 2.84 | 2.90 | | add | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.41 | 3.46 | | add | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.82 | | mul | 1024x4096 | 8.9 | 8.6 | 0.98x | 2.84 | 2.91 | | mul | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.42 | 3.46 | | mul | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.83 | | sub | 1024x4096 | 8.8 | 8.7 | 0.98x | 2.84 | 2.91 | | sub | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.42 | 3.46 | | sub | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.83 | | div | 1024x4096 | 9.1 | 8.8 | 0.98x | 2.77 | 2.84 | | div | 1024x10240 | 18.2 | 18.0 | 0.99x | 3.45 | 3.49 | | div | 1024x20480 | 33.2 | 32.8 | 0.99x | 3.80 | 3.83 | | remainder | 1024x4096 | 9.5 | 9.3 | 0.98x | 2.65 | 2.71 | | remainder | 1024x10240 | 19.3 | 19.0 | 0.99x | 3.27 | 3.31 | | remainder | 1024x20480 | 34.6 | 34.3 | 0.99x | 3.64 | 3.67 | | floor_divide | 1024x4096 | 9.5 | 9.3 | 0.98x | 2.65 | 2.71 | | floor_divide | 1024x10240 | 19.2 | 19.0 | 0.99x | 3.28 | 3.31 | | floor_divide | 1024x20480 | 34.6 | 34.2 | 0.99x | 3.64 | 3.68 | | pow | 1024x4096 | 16.2 | 16.1 | 0.99x | 1.55 | 1.56 | | pow | 1024x10240 | 35.4 | 35.4 | 1.00x | 1.78 | 1.78 | | pow | 1024x20480 | 67.2 | 67.2 | 1.00x | 1.87 | 1.87 | | maximum | 1024x4096 | 20.0 | 19.7 | 0.98x | 1.26 | 1.28 | | maximum | 1024x10240 | 44.7 | 44.2 | 0.99x | 1.41 | 1.42 | | maximum | 1024x20480 | 81.8 | 81.3 | 0.99x | 1.54 | 1.55 | | eq | 1024x4096 | 19.0 | 18.8 | 0.99x | 1.10 | 1.11 | | eq | 1024x10240 | 40.7 | 40.4 | 0.99x | 1.29 | 1.30 | | eq | 1024x20480 | 85.2 | 84.9 | 1.00x | 1.23 | 1.24 | ### Fused Gated Ops (fp16, latency us + bandwidth TB/s) | Op | Shape (MxN) | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | silu_and_mul | 1024x4096 | 9.8 | 9.5 | 0.97x | 2.57 | 2.65 | | silu_and_mul | 1024x10240 | 20.8 | 20.5 | 0.99x | 3.03 | 3.06 | | silu_and_mul | 1024x20480 | 37.6 | 37.2 | 0.99x | 3.35 | 3.38 | | gelu_and_mul | 1024x4096 | 10.6 | 10.3 | 0.97x | 2.38 | 2.45 | | gelu_and_mul | 1024x10240 | 22.2 | 22.0 | 0.99x | 2.83 | 2.86 | | gelu_and_mul | 1024x20480 | 39.4 | 39.2 | 1.00x | 3.20 | 3.21 | | gelu_tanh_and_mul | 1024x4096 | 9.1 | 8.8 | 0.97x | 2.77 | 2.85 | | gelu_tanh_and_mul | 1024x10240 | 18.1 | 17.9 | 0.99x | 3.48 | 3.52 | | gelu_tanh_and_mul | 1024x20480 | 32.9 | 32.7 | 0.99x | 3.82 | 3.84 | ### Takeaways - **Zero overhead**: `torch.compile` adds no runtime cost — compiled paths are consistently 1-8% **faster** than eager due to eliminated Python dispatch overhead - **Unary ops**: 3-8% faster under compile at small sizes (262K elements), converging to ~1-3% at large sizes (4M) - **Binary/fused ops**: 1-2% faster under compile across all DNN-realistic shapes - **Compute-bound ops** (pow, maximum, eq): negligible difference as expected — kernel time dominates 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
Cao Ying
|
d820ab34f3
|
[Feat][Elementwise] Add BinaryKernel + FusedGatedKernel template ops (22 ops) (#463)
## Summary Implements 22 elementwise operator subclasses using the template base classes from #436: - **20 BinaryKernel ops**: arithmetic (Sub, Mul, Div, Remainder, FloorDivide, Pow, Maximum, Minimum, Lerp), comparison (Eq, Ne, Gt, Lt, Ge, Le), logical (LogicalAnd, LogicalOr), bitwise (BitwiseAnd, BitwiseOr, BitwiseXor) - **2 FusedGatedKernel ops**: GeluAndMul, GeluTanhAndMul - NaN propagation fix for MaximumKernel/MinimumKernel (round 1) Closes #438 ## Dtype Support Matrix | Op Family | Supported input dtypes | Output dtype | Semantics | |---|---|---|---| | `sub`, `mul` | validated in this PR on `float16`, `bfloat16`, `float32` | same as input | PyTorch-style elementwise arithmetic with broadcast | | `div`, `remainder`, `pow`, `floor_divide`, `lerp`, `maximum`, `minimum` | `float16`, `bfloat16`, `float32` | same as input | PyTorch-aligned float arithmetic; unsupported integer dtypes are rejected at construction time | | `eq`, `ne`, `gt`, `lt`, `ge`, `le` | `float16`, `bfloat16`, `float32` | `torch.bool` | Kernel emits int8 1/0 flags, Op layer casts to bool | | `logical_and`, `logical_or` | `bool`, `uint8`, `int8`, `int16`, `int32`, `int64`, `float16`, `bfloat16`, `float32` | `torch.bool` | Non-zero truthiness, matching `torch.logical_*` semantics | | `bitwise_and`, `bitwise_or`, `bitwise_xor` | `uint8`, `int8`, `int16`, `int32`, `int64` | same as input | PyTorch-style bitwise semantics; float dtypes rejected | | `gelu_and_mul`, `gelu_tanh_and_mul` | `float16`, `bfloat16`, `float32` | same as input | Fused gated activation on `(M, 2N) -> (M, N)` | ## Test plan - [x] **AC1**: 20 BinaryKernel + 2 FusedGatedKernel subclasses implemented - [x] **AC2**: Comparison ops output torch.bool - [x] **AC3**: `__all__` updated in both elementwise.py files - [x] **AC4**: L1: 22 ops pass (fp16, same-shape) - [x] **AC5**: L3: 60 broadcast tests pass (bias-add, row, scalar patterns x 20 ops) - [x] **AC6**: L4: 6 edge cases pass (fp32, 4K) - [x] **AC7**: dtype rejection/runtime mismatch paths covered for the float-only binary and fused-gated contracts - [x] **AC8**: `__init__.py` not modified **Test results**: 196/196 passed, 0 failed ## Benchmark Environment: Torch 2.9.1+cu128, CUDA 12.8, NVIDIA H200, Driver 575.57.08 Shapes: (tokens=1024, hidden_dim) - DNN-realistic hidden dimensions ### Binary Arithmetic (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | sub | 1024x4096 | 1.97 | 1.84 | 1.07x | | sub | 1024x10240 | 2.73 | 2.77 | 0.99x | | sub | 1024x20480 | 3.37 | 3.37 | 1.00x | | mul | 1024x4096 | 1.87 | 1.83 | 1.02x | | mul | 1024x10240 | 2.82 | 2.83 | 0.99x | | mul | 1024x20480 | 3.40 | 3.40 | 1.00x | | div | 1024x4096 | 1.82 | 1.87 | 0.98x | | div | 1024x10240 | 2.72 | 2.59 | 1.05x | | div | 1024x20480 | 3.36 | 3.31 | 1.02x | | remainder | 1024x4096 | 1.81 | 1.61 | 1.13x | | remainder | 1024x10240 | 2.73 | 2.40 | 1.14x | | remainder | 1024x20480 | 3.35 | 2.82 | 1.19x | | pow | 1024x4096 | 1.05 | 1.02 | 1.02x | | pow | 1024x10240 | 1.29 | 1.26 | 1.02x | | pow | 1024x20480 | 1.42 | 1.39 | 1.02x | | floor_divide | 1024x4096 | 1.92 | 0.82 | 2.33x | | floor_divide | 1024x10240 | 2.74 | 0.98 | 2.79x | | floor_divide | 1024x20480 | 3.41 | 1.05 | 3.25x | | lerp | 1024x4096 | 1.86 | 1.85 | 1.00x | | lerp | 1024x10240 | 2.73 | 2.81 | 0.97x | | lerp | 1024x20480 | 3.36 | 3.42 | 0.98x | | maximum | 1024x4096 | 1.01 | 1.82 | 0.56x | | maximum | 1024x10240 | 1.22 | 2.71 | 0.45x | | maximum | 1024x20480 | 1.40 | 3.38 | 0.41x | | minimum | 1024x4096 | 1.02 | 1.89 | 0.54x | | minimum | 1024x10240 | 1.24 | 2.70 | 0.46x | | minimum | 1024x20480 | 1.38 | 3.34 | 0.41x | ### Comparison (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | eq | 1024x4096 | 0.78 | 1.72 | 0.45x | | eq | 1024x10240 | 1.02 | 2.49 | 0.41x | | eq | 1024x20480 | 1.05 | 3.17 | 0.33x | | ne | 1024x4096 | 0.79 | 1.65 | 0.48x | | ne | 1024x10240 | 1.02 | 2.56 | 0.40x | | gt | 1024x4096 | 0.78 | 1.72 | 0.45x | | gt | 1024x10240 | 1.02 | 2.48 | 0.41x | | lt | 1024x4096 | 0.79 | 1.72 | 0.46x | | lt | 1024x10240 | 1.02 | 2.55 | 0.40x | | ge | 1024x4096 | 0.78 | 1.72 | 0.45x | | ge | 1024x10240 | 1.02 | 2.56 | 0.40x | | le | 1024x4096 | 0.79 | 1.64 | 0.48x | | le | 1024x10240 | 1.02 | 2.55 | 0.40x | ### Logical (fp16 input, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | logical_and | 1024x4096 | 0.78 | 2.17 | 0.36x | | logical_and | 1024x10240 | 1.02 | 3.37 | 0.30x | | logical_or | 1024x4096 | 0.85 | 2.05 | 0.41x | | logical_or | 1024x10240 | 1.15 | 3.41 | 0.34x | ### Bitwise (int32, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | bitwise_and | 1024x4096 | 2.59 | 2.55 | 1.01x | | bitwise_and | 1024x10240 | 3.31 | 3.42 | 0.97x | | bitwise_or | 1024x4096 | 2.51 | 2.54 | 0.99x | | bitwise_or | 1024x10240 | 3.30 | 3.41 | 0.97x | | bitwise_xor | 1024x4096 | 2.51 | 2.61 | 0.96x | | bitwise_xor | 1024x10240 | 3.37 | 3.36 | 1.00x | ### Fused Gated (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | gelu_and_mul | 1024x4096 | 1.61 | 0.63 | 2.53x | | gelu_and_mul | 1024x10240 | 2.27 | 0.77 | 2.96x | | gelu_and_mul | 1024x20480 | 2.75 | 0.84 | 3.29x | | gelu_tanh_and_mul | 1024x4096 | 1.87 | 0.66 | 2.82x | | gelu_tanh_and_mul | 1024x10240 | 2.76 | 0.81 | 3.41x | | gelu_tanh_and_mul | 1024x20480 | 3.36 | 0.89 | 3.79x | ### Broadcast: Bias-Add Pattern (fp16, bandwidth TB/s) | Op | a_shape | b_shape | TileOPs | PyTorch | Speedup | |---|---|---|---|---|---| | sub | 1024x4096 | 1x4096 | 1.53 | 0.87 | 1.75x | | sub | 1024x10240 | 1x10240 | 2.53 | 1.14 | 2.22x | | sub | 1024x20480 | 1x20480 | 3.11 | 1.27 | 2.45x | | mul | 1024x4096 | 1x4096 | 1.54 | 0.88 | 1.75x | | mul | 1024x10240 | 1x10240 | 2.42 | 1.12 | 2.16x | | mul | 1024x20480 | 1x20480 | 3.18 | 1.28 | 2.47x | | div | 1024x4096 | 1x4096 | 1.47 | 0.82 | 1.80x | | div | 1024x10240 | 1x10240 | 2.37 | 1.06 | 2.24x | | div | 1024x20480 | 1x20480 | 2.82 | 1.16 | 2.43x | ### Takeaways - **Broadcast (bias-add)**: 1.75-2.47x faster (stride-based codegen beats PyTorch broadcast) - **Fused gated**: 2.5-3.8x faster (kernel fusion eliminates intermediate writes) - **floor_divide**: 2.3-3.3x faster (TileLang codegen advantage) - **remainder**: 1.1-1.2x faster - **Binary arith** (sub/mul/div/lerp/pow): ~parity with PyTorch at DNN-realistic sizes - **Bitwise**: ~parity (0.96-1.01x) - **Comparison**: ~0.3-0.5x slower (int8 intermediate + bool cast overhead; optimization opportunity) - **Logical**: ~0.3-0.4x slower (int8 intermediate + bool cast; optimization opportunity) - **Maximum/minimum**: ~0.4-0.6x slower (T.isnan NaN propagation overhead; optimization opportunity) ## Changes | File | Description | |---|---| | `tileops/kernels/elementwise.py` | 20 BinaryKernel + 2 FusedGatedKernel subclasses | | `tileops/ops/elementwise.py` | 22 matching Op subclasses | | `tests/ops/test_binary_arith.py` | Arithmetic smoke + broadcast + edge case tests | | `tests/ops/test_comparison.py` | Comparison smoke + broadcast + bool output tests | | `tests/ops/test_logical.py` | Logical op smoke + broadcast tests | | `tests/ops/test_bitwise.py` | Bitwise op smoke + broadcast tests | | `tests/ops/test_fused_gated.py` | Fused gated op tests | | `benchmarks/ops/bench_binary_elementwise.py` | Benchmark: same-shape + broadcast (DNN-realistic 2D shapes) | ## Follow-up Issue The next issue that should be taken immediately after this PR is: - #440: add `torch.compile` / fake-tensor support for the shared elementwise template via `custom_op` registration --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
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> |
|
|
Cao Ying
|
ec595f03ca
|
[Feat][Elementwise] Add template infrastructure: base classes, smoke tests, and strategy benchmark (#453)
## Summary Build the complete template infrastructure for 66 elementwise ops: - **3 kernel base classes** — `UnaryKernel`, `BinaryKernel`, `FusedGatedKernel` with strategy support (`direct`, `explicit_parallel`, `register_copy`) - **3 op base classes** — `UnaryOp`, `BinaryOp`, `FusedGatedOp` with `coalesce_broadcast_dims` utility for stride-based broadcast - **3 concrete example ops** — `relu`, `add`, `silu_and_mul` (one per template family) - **Strategy benchmark** — determines `DEFAULT_STRATEGY` per kernel type with bandwidth rationale Closes #436 ## Test plan - [x] AC1: UnaryKernel, BinaryKernel, FusedGatedKernel implemented with all strategies - [x] AC2: `coalesce_broadcast_dims` handles all 6+ broadcast patterns correctly (6 unit tests) - [x] AC3: relu, add, silu_and_mul pass L1 correctness tests (14 configs total) - [x] AC4: 4 broadcast patterns pass L3 tests (bias-add, row, scalar, interleaved) - [x] AC5: Strategy benchmark run; DEFAULT_STRATEGY updated with rationale - [x] AC6: `__init__.py` re-exports base classes only (not concrete ops) **29 tests pass, 14 benchmarks pass.** ## Benchmark ### Unary (relu) Strategy | strategy | bandwidth (TB/s) | |---|---| | register_copy | 2.40 | | direct | 1.04 | | explicit_parallel | 0.86 | ### Binary (add) Strategy | strategy | bandwidth (TB/s) | |---|---| | explicit_parallel | 2.62 | | direct | 1.42 | **Decision**: Unary `DEFAULT_STRATEGY = register_copy`, Binary `DEFAULT_STRATEGY = explicit_parallel`. ## Files changed (9 files, +1177 lines) | File | Description | |---|---| | `tileops/kernels/elementwise.py` | Kernel base classes + strategy factory | | `tileops/ops/elementwise.py` | Op base classes + coalesce_broadcast_dims | | `tileops/kernels/__init__.py` | Re-export base kernel classes | | `tileops/ops/__init__.py` | Re-export base op classes | | `tests/ops/test_activation.py` | Unary op tests (relu) | | `tests/ops/test_binary_arith.py` | Binary op + broadcast tests (add) | | `tests/ops/test_fused_gated.py` | Fused gated op tests (silu_and_mul) | | `benchmarks/ops/bench_activation.py` | Unary strategy benchmark | | `benchmarks/ops/bench_binary_arith.py` | Binary strategy benchmark | --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |