diff --git a/.codespellrc b/.codespellrc index d8c810ec..ec6dbc3a 100644 --- a/.codespellrc +++ b/.codespellrc @@ -1,3 +1,3 @@ [codespell] -skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/* +skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/*,src/xorl/ops/quack/*,src/xorl/ops/linear_attention/flashqla/* ignore-words-list = dout,te,subtile,parm,mot,numer,notin diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4fad13c9..aea412a2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,10 @@ default_stages: [pre-commit, pre-push, manual] -exclude: ^(submodules/|\.venv/|src/xorl/ops/quack/) - +# Vendored trees and byte-contract-gated files are never rewritten by any +# hook: quack/ and flashqla/ are third-party snapshots (see their VENDORED.md), +# and bi_families_v2.py is vendored byte-identical into the serving engine +# with both copies sha256-gated, so it keeps that engine's formatting +# (black, 88 columns). See src/xorl/ops/README.md. +exclude: ^(submodules/|\.venv/|src/xorl/ops/quack/|src/xorl/ops/linear_attention/flashqla/|src/xorl/ops/bi_families_v2\.py$) repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 @@ -28,11 +32,6 @@ repos: - id: ruff args: [--fix-only] - id: ruff-format - # bi_families_v2.py is vendored byte-identical into the serving engine - # and both copies are sha256-gated, so it carries that engine's - # formatting (black, 88 columns). Reformatting it here at 120 would - # break the byte-equality the contract depends on. Lint still runs. - exclude: ^src/xorl/ops/bi_families_v2\.py$ - repo: https://github.com/codespell-project/codespell rev: v2.4.1 hooks: diff --git a/pyproject.toml b/pyproject.toml index cb86f301..fefd626e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -142,6 +142,16 @@ where = ["src"] [tool.ruff] target-version = "py312" line-length = 120 +# Vendored trees and the byte-contract-gated serving copy are never linted or +# formatted by first-party tooling (see src/xorl/ops/README.md). +# force-exclude keeps this true even when paths are passed explicitly, which +# is how pre-commit invokes ruff. +force-exclude = true +extend-exclude = [ + "src/xorl/ops/quack", + "src/xorl/ops/linear_attention/flashqla", + "src/xorl/ops/bi_families_v2.py", +] [tool.ruff.lint] ignore = ["C901", "E501", "E741", "W605", "C408"] diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md new file mode 100644 index 00000000..94559c56 --- /dev/null +++ b/src/xorl/ops/README.md @@ -0,0 +1,45 @@ +# xorl.ops + +Compute kernels and autograd ops. This README is the map: what each +subtree is, what may be edited, and where things are headed +(reorganization plan: [issue #78](https://github.com/togethercomputer/xorl/issues/78)). + +## Edit policies + +Three kinds of code live here with three different rules: + +| kind | rule | +| --- | --- | +| **Vendored** — `quack/`, `linear_attention/flashqla/` | Never hand-edit, lint, or reformat. Each tree carries a `VENDORED.md` with provenance and the local-patch ledger. First-party tooling skips them (`[tool.ruff]` excludes in `pyproject.toml`; top-level `exclude:` in `.pre-commit-config.yaml`). | +| **Byte-contract-gated** — `bi_families_v2.py` | Vendored byte-identical into the serving engine; both copies are sha256-gated. Any edit here without the paired serving-side edit breaks the gate. It keeps the engine's formatting (black, 88 columns) and is excluded from all rewriting hooks. | +| **First-party** — everything else | Normal rules. | + +## Map (current) + +- `bi_families_v2.py`, `batch_invariant_ops.py`, `bi_gemm_configs.py`, + `exact_sampling_transforms.py`, `rope_class_b.py`, `canonical_moe_leaf.py`, + `canonical_moe_cast.py`, `fused_silu_and_mul.py`, `kernel_config_pin.py`, + `block_fp8_native.py` — the **serving-parity (exact) contract family**: + byte-pinned programs shared with the serving engine. The `bi_` / `exact_` / + `canonical_` / `class_b` prefixes are historical names for the same + concept. Planned home: `ops/exact/` (#78 phase 3). +- `loss/` — the CE/selected-logprob kernel stack **plus** (for now) the RL + objective functions; the objectives move to `xorl/objectives/` + (#78 phase 2). +- `moe/` — MoE expert compute backends (triton/quack/native, LoRA variants). + `ep_kernels/` (DeepEP sort/scatter) merges in here (#78 phase 5). +- `linear_attention/` — GDN/linear-attention kernels; also currently hosts + the `GatedDeltaNet` layer class and the vendored `flashqla/` tree + (layer moves to `models/layers/`, vendored tree to `ops/_vendored/`, #78). +- `ssm/` — Mamba-2 kernels; also currently hosts the `Mamba2Mixer` layer + class (same plan as above). +- `quantize/` — NF4/INT4/FP4/FP8 quantization codecs and fake-quant ops. +- `glm5_kernels/`, `dsv4/` — model-family-specific kernels; planned home: + `models/transformers/{glm5,deepseek_v4}/kernels/` (#78 phase 4). +- `quack/` — vendored (see above). + +## What does NOT belong here + +New `nn.Module` layer classes (→ `models/layers/`), RL objectives +(→ `xorl/objectives/` once it exists), and orchestration logic. `ops/` is +for kernels and the autograd boundaries directly over them. diff --git a/src/xorl/ops/quack/VENDORED.md b/src/xorl/ops/quack/VENDORED.md new file mode 100644 index 00000000..99088f29 --- /dev/null +++ b/src/xorl/ops/quack/VENDORED.md @@ -0,0 +1,26 @@ +# Vendored: QuACK (partial fork) + +This directory is a vendored, locally-patched subset of +[QuACK](https://github.com/Dao-AILab/quack) (Wentao Guo, Ted Zadouri, +Tri Dao — CuTeDSL kernels), snapshotted at upstream `0.4.1` +(see `__init__.py:__version__`). + +It is a *partial* fork: modules in this tree still import helper modules +(`copy_utils`, `layout_utils`, …) from the separate PyPI dependency +`quack-kernels` pinned in `pyproject.toml`, so the pinned PyPI version and +this tree must stay compatible. + +Known local additions/patches (not upstream): + +- `cute_dsl_elf_fix.py` — works around cutlass#3161 (duplicate `.text` + section flags break MCJIT in multi-process loads). +- `cute_dsl_mlir_threading.py` — works around cutlass#3062 (leaked LLVM + thread pools exhaust pthreads across compiles). + +## Edit policy + +Do not hand-edit, lint, or reformat files in this tree. First-party tooling +skips it (see `pyproject.toml` `[tool.ruff]` excludes and the top-level +`exclude:` in `.pre-commit-config.yaml`). Fixes should go upstream first; +a local patch that cannot wait must be listed above with its upstream +issue/PR so the next re-vendor can reconcile it.