FrRay
/
TileOPs-Metax
forked from ccf-ai-infra/TileOPs-Metax
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
429 Commits 6 Branches 0 Tags 15 MiB
Commit Graph

4 Commits

Author SHA1 Message Date
Cao Ying 978615d1b3
[Refactor][Benchmark] Detach benchmarks from tests via workloads layer (#787) 
## Summary

Repo-wide refactor to fully detach `benchmarks/` from `tests/` by
introducing a shared `workloads/` layer.

Resolves #783. Case study documented in #789.

### Problem

`benchmarks/` had 81 `tests.*` imports across 67 files, plus 42
`self.test.*` accesses. Benchmark code broke when test-only code changed
— the wrong dependency direction.

### Solution

Introduce `workloads/` as a neutral shared layer owning only workload
parameters, `gen_inputs()`, and fixture metadata. Tests and benchmarks
both import from `workloads/` but never from each other.

### Architecture

```
workloads/base.py       → WorkloadBase (gen_inputs only), FixtureMeta, FixtureBase
workloads/ops/*.py      → 55 concrete workload classes (params + gen_inputs)

tests/test_base.py      → TestBase(WorkloadBase) + @abstractmethod ref_program + check()
tests/ops/*.py          → Test classes inherit (Workload, TestBase), define ref_program locally

benchmarks/benchmark.py → BenchmarkBase(workload: WorkloadBase), self.workload
benchmarks/ops/*.py     → Import from workloads.ops/*, define baselines locally
```

### Trust model

- `workloads/` contains NO reference implementations, NO correctness
logic, NO assertion/tolerance code
- `ref_program()` is `@abstractmethod` on `TestBase`, defined
independently in each test class
- Benchmark baselines are independent local copies — no shared oracle
surface between tests and benchmarks
- `workloads/` ships in the wheel with lazy `pytest` import (no
undeclared runtime dependency)

## Changes

**189 files changed** (+5329 / -3732)

| Category | Count | Description |
|----------|-------|-------------|
| `workloads/` (new) | 58 files | `WorkloadBase`, `FixtureMeta`,
`FixtureBase`, 55 workload classes |
| `benchmarks/` | 73 files | All imports migrated from `tests.*` →
`workloads.*`; `self.test` → `self.workload`; local baseline copies |
| `tests/` | 58 files | Workload params extracted to `workloads/`;
`ref_program` kept local; `TestBase` enforces abstract contract |
| `pyproject.toml` | 1 file | `workloads` included in wheel package
discovery |

**Zero changes** to `tileops/ops/`, `tileops/kernels/`, or
`tileops/ops_manifest.yaml`.

## Acceptance criteria

- [x] `rg -n "from tests\.|import tests\." benchmarks` → no matches
- [x] `rg -n "self\.test\." benchmarks` → no matches
- [x] `benchmarks/benchmark.py` does not import or reference `TestBase`
- [x] `BenchmarkBase` stores `self.workload`, not `self.test`
- [x] `workloads/base.py` exports `WorkloadBase`, `FixtureMeta`,
`FixtureBase`
- [x] `tests/test_base.py` no longer defines `FixtureMeta` or
`FixtureBase`
- [x] `workloads/` contains no correctness-only methods or
assertion/tolerance logic
- [x] No shared oracle surface — reference functions duplicated
independently
- [x] Op implementation files identical to upstream/main (pure refactor)
- [x] Manifest unchanged from upstream/main
- [x] Test coverage preserved (fp32/fp16/bf16, 1D-4D, all original
params)
- [x] Representative benchmark smoke runs produce numeric output
- [x] Representative correctness smoke tests pass

## Test plan

- [x] Structural audit: zero `tests.*` imports in `benchmarks/`
- [x] Structural audit: zero `self.test` in `benchmarks/`
- [x] `pytest --collect-only benchmarks` — 1004 tests collected, 0
import errors
- [x] `pytest tests/ops/test_softmax.py -m smoke` — 15 passed
- [x] `pytest tests/ops/test_activation.py -m smoke` — 20 passed
- [x] `pytest benchmarks/ops/bench_softmax.py -m smoke` — passed with
numeric output
- [x] `pytest benchmarks/ops/bench_activation.py -m smoke` — passed with
numeric output
- [x] Wheel build: `workloads` importable without pytest installed

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

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
 2026-04-05 11:38:51 +08:00
Ang Gao fce542af21
[Fix][Bench] Standardize benchmark baseline tags to match actual implementations (#624) (#657) 
## Summary

- Replace generic `tag="baseline"` across 52 benchmark files with
precise tags that reflect the actual backend implementation
- Nightly report Baseline column now shows exactly what was used for
comparison instead of uninformative "baseline"

## Changes

1. **Generic → specific**: `"baseline"` → descriptive names across all
benchmark files
2. **Backend precision**: `torch.matmul` → `"torch-cublas"`,
`torch.fft.fft` → `"torch-cufft"`, `F.scaled_dot_product_attention` →
`"torch-sdpa"`
3. **Ref classification**: hand-written multi-step references unified as
`"torch-ref"`
4. **Naming consistency**: `"FA3"` → `"fa3"`, `"pytorch-*"` →
`"torch-*"`, underscores → hyphens, `"fla_bwd_with_recompute"` → `"fla"`
5. **Rules update**: `.claude/rules/benchmark.md` tag registry expanded

## Tag taxonomy

| Tag | Meaning |
|-----|---------|
| `"torch"` | PyTorch built-in API with its own CUDA kernel (softmax,
layer_norm, sum, etc.) |
| `"torch-cublas"` | PyTorch API → cuBLAS (matmul) |
| `"torch-cufft"` | PyTorch API → cuFFT (fft) |
| `"torch-sdpa"` | `F.scaled_dot_product_attention` |
| `"torch-cudnn"` | PyTorch → cuDNN (batch_norm fwd) |
| `"torch-autograd"` | PyTorch built-in fwd + autograd bwd (batch_norm
bwd) |
| `"torch-ref"` | Hand-written multi-step PyTorch reference |
| `"fa3"` / `"fla"` / `"triton"` / `"sgl-kernel"` / `"vllm"` | External
libraries |

## Test plan

- [ ] `pytest benchmarks/ --collect-only` passes (no import errors)
- [ ] Nightly report baseline column shows specific tags instead of
"baseline"

Closes #624
 2026-03-25 13:00:58 +08:00
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>
 2026-03-16 12:34:20 +08:00
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>
 2026-03-13 13:21:27 +08:00