Reorganize xorl.ops: vendored quarantine, objectives split, ops/exact, ops/families (#78 phases 1-5)

.codespellrc

@@ -1,3 +1,3 @@
 [codespell]
-skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/*,src/xorl/ops/quack/*,src/xorl/ops/linear_attention/flashqla/*
+skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/*,src/xorl/ops/_vendored/*
 ignore-words-list = dout,te,subtile,parm,mot,numer,notin

.github/workflows/pr-test-cpu.yml

@@ -62,7 +62,7 @@ jobs:
               tests/distributed/test_lm_head_tp_*.py
               tests/distributed/test_tensor_parallel.py
               tests/distributed/test_vocab_parallel_ce.py
-              tests/distributed/test_bi_fused_lm_head_tp.py
+              tests/distributed/test_batch_invariant_lm_head_tp.py
               tests/distributed/test_fused_linear_logprob_tp.py
               tests/distributed/test_olmo2_*.py
 
@@ -75,7 +75,7 @@ jobs:
               tests/distributed/test_native_fp8_fsdp2_materialization.py
               tests/distributed/test_bi_trunk_linear_fsdp.py
               tests/distributed/test_qwen35_lora_projection_fsdp2.py
-              tests/distributed/test_rope_class_b_fsdp_transport.py
+              tests/distributed/test_rope_fp32_single_round_fsdp_transport.py
               tests/distributed/test_torch_parallelize_policies.py
               tests/distributed/test_parallel_plan_meta_slice.py
 

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_stages: [pre-commit, pre-push, manual]
 # 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$)
+exclude: ^(submodules/|\.venv/|src/xorl/ops/_vendored/|src/xorl/ops/sglang/bi_families_v2\.py$)
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0

certification/glm52/benchmark_sparse_mla_backward.py

@@ -36,7 +36,7 @@ def main() -> None:
     except ImportError as exc:
         raise SystemExit("This certification requires TileLang") from exc
 
-    from xorl.ops.glm5_kernels.sparse_mla import SparseMLA  # noqa: PLC0415
+    from xorl.ops.families.glm5.sparse_mla import SparseMLA  # noqa: PLC0415
 
     sequence, kv_sequence, heads, rank, tail, topk = 2048, 32768, 64, 512, 64, 2048
     scale = (rank + tail) ** -0.5

certification/opd/vocab_parallel_kl_gathered.py

@@ -24,9 +24,9 @@
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import Shard, distribute_tensor
 
-from xorl.ops.loss.opd_loss import opd_vocab_parallel_loss_function
-from xorl.ops.loss.opd_streaming_kl import streaming_reverse_kl_function
-from xorl.ops.loss.reducers import TokenPartial
+from xorl.objectives.opd_loss import opd_vocab_parallel_loss_function
+from xorl.objectives.opd_streaming_kl import streaming_reverse_kl_function
+from xorl.objectives.reducers import TokenPartial
 from xorl.ops.loss.vocab_parallel_reverse_kl import vocab_parallel_reverse_kl_gathered
 from xorl.server.runner.model_runner import ModelRunner
 

docs/k3/GEMM_CONTRACT.md

@@ -15,7 +15,7 @@ The output tile, group size, pipeline depth, and warp count do not split the K
 reduction, so they remain performance-tuning axes after passing the bitwise
 gate.
 
-`bi_gemm_configs.py` contains the shared, shape-keyed table. Each entry keeps
+`batch_invariant_gemm_configs.py` contains the shared, shape-keyed table. Each entry keeps
 the dtype's pinned K tile, compares bitwise with the baseline configuration,
 and checks that an identical row keeps identical output bits across row-count
 buckets. Set `XORL_BI_GEMM_CONFIG_TABLE=0` to use the pinned baseline table.

docs/k3/LM_HEAD_CONTRACT.md

@@ -27,7 +27,7 @@ vocabulary chunk at a time, records
 the same maximum, exponential sum, and selected logit, then merges chunks in
 pinned order. This rollback is exact but uses more launches.
 
-The trainer exposes this as `ce_mode="bi_fused"`. Its custom autograd function
+The trainer exposes this as `ce_mode="batch_invariant"`. Its custom autograd function
 saves the exact forward LSE and recomputes the conventional closed-form CE
 gradient by chunks. Only forward values enter the train/serve equality
 contract; backward remains ordinary checked training numerics.
@@ -42,7 +42,7 @@ second fp32 copy of the LM-head weight.
 Run on a Hopper GPU:
 
 ```bash
-pytest tests/ops/test_bi_fused_lm_head.py -v
+pytest tests/ops/test_batch_invariant_lm_head.py -v
 ```
 
 The gate compares forward probabilities and loss with an eager fp32 reference,

docs/src/content/docs/config-reference/local.md

@@ -175,7 +175,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict:
 | `activation_gpu_limit` | `0.0` | GB of activations to keep on GPU when offloading. `0.0` = offload all. |
 | `enable_compile` | `false` | `torch.compile` for model forward pass. |
 | `compile_dynamic_shapes` | `false` | Pass `dynamic=True` to `torch.compile`; keep disabled unless a workload has benchmarked a dynamic-shape win. |
-| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `bi_fused`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. |
+| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `batch_invariant`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. |
 | `ce_num_chunks` | `8` | Number of token chunks for chunked/compiled cross-entropy. |
 | `enable_fp8_training` | `false` | Experimental full-weight block-FP8 compute with BF16/FP32 master parameters. Mutually exclusive with QARL. |
 | `enable_qarl` | `false` | Experimental dynamic fake-quant training with full-precision masters and STE gradients. E4M3 applies to dense `nn.Linear` modules; NVFP4 also supports MoE expert containers. Mutually exclusive with full-weight FP8 training. |

docs/src/content/docs/config-reference/server.md

@@ -53,7 +53,7 @@ These stored defaults are resolved after the model architecture is known. Ordina
 | `rmsnorm_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `native`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs require `sglang_fused`. Other explicit diagnostic modes are also accepted by the argument type. |
 | `activation_native` | `false` (resolved) | Use native SiLU instead of the fused Triton kernel. Exact Qwen3.5-family programs resolve this to `true`; the other exact programs retain their architecture-owned fused arithmetic. |
 | `rope_native` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `false`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `true`. |
-| `rope_class_b` | `null` (resolved) | Select the compiled Class-B RoPE FP32-chain path. It is enabled for exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; DSV4 owns a separate RoPE program. |
+| `rope_fp32_single_round` | `null` (resolved) | Select the compiled Class-B RoPE FP32-chain path. It is enabled for exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; DSV4 owns a separate RoPE program. |
 | `attention_cast_bf16` | `false` (resolved) | Explicitly cast Q/K to BF16 after RoPE. Exact Qwen3.5-family programs resolve this to `true`; dense Qwen3, GLM-5.2, and DSV4-Flash exact programs require `false`. |
 | `qwen35_rmsnorm_family` | `null` (resolved) | Exact Qwen3.5/3.6 programs require the qualified `v2` arithmetic; other architectures reject an override. |
 | `sparse_mla_enabled` | `null` (resolved) | Canonical GLM-5.2 enables the sparse-MLA path; ordinary models resolve to `false`. |
@@ -96,7 +96,7 @@ These stored defaults are resolved after the model architecture is known. Ordina
 | `enable_forward_prefetch` | `false` | FSDP forward prefetch. |
 | `init_device` | `meta` | Model initialization device: `cpu`, `meta`, `cuda`. |
 | `load_weights_mode` | `grouped` | Weight loading mode: `grouped` (default, with rank-0 fallback), `all_ranks`, or `skip`. |
-| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `bi_fused`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. |
+| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `batch_invariant`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. |
 | `enable_fp8_training` | `false` | Experimental full-weight block-FP8 compute. Mutually exclusive with LoRA/QLoRA and QARL. |
 | `enable_qarl` | `false` | Experimental dynamic fake-quant training with full-precision masters and STE gradients. E4M3 applies to dense `nn.Linear` modules; NVFP4 also supports MoE expert containers. Mutually exclusive with LoRA/QLoRA and full-weight FP8 training. |
 | `qarl_quant_cfg` | `null` | QARL alias or dictionary. `null`/`FP8_DEFAULT_CFG` resolves to dynamic E4M3 W8A8 with `[128, 128]` weight blocks. `nvfp4` resolves to dynamic, weight-only W4 with `group_size: 16`; set `activation: true` for W4A4. NVFP4 covers dense linears and MoE expert containers, while E4M3 is dense-only. |

docs/src/content/docs/loss-functions.mdx

@@ -71,13 +71,13 @@ Per-token output behavior depends on the selected loss and the request's `return
 
 ### Compiled cross-entropy (`ce_mode`)
 
-The `ce_mode` setting controls the LM-head and cross-entropy implementation. When omitted, XoRL resolves it from the model: ordinary models and exact DSV4-Flash use `compiled`, while exact dense Qwen3, Qwen3.5-family, and GLM-5.2 numerical programs require `bi_fused`.
+The `ce_mode` setting controls the LM-head and cross-entropy implementation. When omitted, XoRL resolves it from the model: ordinary models and exact DSV4-Flash use `compiled`, while exact dense Qwen3, Qwen3.5-family, and GLM-5.2 numerical programs require `batch_invariant`.
 
 | Value | Description | Important constraints |
 |---|---|---|
 | `compiled` | `torch.compile`-compiled chunked cross-entropy | General production default |
 | `eager` | Standard eager cross-entropy | Debugging; may materialize the full logits tensor |
-| `bi_fused` | Batch-invariant selected-token logprob/CE path | Required by the current exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; topology and dtype restrictions are checked at runtime |
+| `batch_invariant` | Batch-invariant selected-token logprob/CE path | Required by the current exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; topology and dtype restrictions are checked at runtime |
 | `quack_linear` | Quack chunked linear plus scalar cross-entropy | Causal-LM path; pipeline parallelism supports this mode on its last stage |
 | `fused_quack` | Chunked matmul plus fused selected-token CE | Used by supported per-token loss paths; not supported by every loss/topology combination |
 

docs/src/content/docs/moe/kernels.mdx

@@ -236,6 +236,6 @@ The retained tensors and avoided work differ by model and backend. Measure peak
 | [`src/xorl/ops/group_gemm/kernel/group_gemm.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/group_gemm/kernel/group_gemm.py) | `group_gemm_same_nk_kernel` — Triton autotuned grouped GEMM with optional STORE_ACTIVATIONS |
 | [`src/xorl/ops/__init__.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/__init__.py) | `triton_moe_forward`, `triton_moe_lora_forward`, backend registry |
 | [`src/xorl/ops/group_gemm/kernel/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/group_gemm/kernel/quack.py) | `quack_group_gemm_same_nk` — CuTe GEMM wrapper; `XORL_QUACK_TUNED` env var |
-| [`src/xorl/ops/quack/gemm_interface.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/quack/gemm_interface.py) | Quack GEMM interface — SM90/SM100 configs, autotuner, gated activation kernels |
+| [`src/xorl/ops/_vendored/quack/gemm_interface.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/_vendored/quack/gemm_interface.py) | Quack GEMM interface — SM90/SM100 configs, autotuner, gated activation kernels |
 | [`src/xorl/models/layers/moe/backend/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/models/layers/moe/backend/quack.py) | `quack_expert_forward` — non-EP quack MoE forward |
 | [`src/xorl/ops/moe/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/moe/quack.py) | `QuackEPGroupGemm`, `QuackEPGroupGemmMoeAct` — EP + moe_act variants |

examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml

@@ -56,6 +56,6 @@ enable_packing: true
 
 # Full-weights training
 enable_lora: false
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 freeze_router: false

examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml

@@ -30,8 +30,8 @@ sample_packing_sequence_len: 128000
 enable_packing: true
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 optimizer: muon
 optimizer_dtype: bf16

examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml

@@ -31,8 +31,8 @@ sample_packing_sequence_len: 128000
 enable_packing: true
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 optimizer: muon
 optimizer_dtype: bf16

examples/server/configs/full/qwen3_5_35b_a3b_full.yaml

@@ -36,8 +36,8 @@ sample_packing_sequence_len: 128000
 enable_packing: true
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 optimizer: muon
 optimizer_dtype: bf16

examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml

@@ -36,8 +36,8 @@ sample_packing_sequence_len: 128000
 enable_packing: true
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 optimizer: muon
 optimizer_dtype: bf16

examples/server/configs/full/qwen3_5_397b_a17b_full.yaml

@@ -45,8 +45,8 @@ sample_packing_sequence_len: 65536
 enable_packing: true
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 optimizer: muon
 optimizer_dtype: bf16

examples/server/configs/full/qwen3_8b_full.yaml

@@ -36,5 +36,5 @@ enable_packing: true
 
 # Full weight training
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant

examples/server/configs/full/qwen3_coder_30b_a3b_full.yaml

@@ -41,8 +41,8 @@ enable_packing: true
 
 # Full weight training
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant
 
 # Optimizer: Muon with bf16 states on GPU (no offload with 8 GPUs).
 # 8 GPUs: weights ~7.5GB/GPU + Muon momentum ~7.5GB/GPU = ~15GB/GPU. Fits easily.

examples/server/configs/lora/qwen3_30b_a3b_lora.yaml

@@ -74,5 +74,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p
 
 # Skip initial checkpoint
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss).
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss).
+ce_mode: batch_invariant

examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml

@@ -48,5 +48,5 @@ lora_alpha: 16
 lora_target_modules: ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant

examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml

@@ -46,5 +46,5 @@ lora_alpha: 32
 lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "g_proj", "down_proj"]
 
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant

examples/server/configs/lora/qwen3_8b_lora.yaml

@@ -42,5 +42,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p
 
 # Skip initial checkpoint
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant

examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml

@@ -65,5 +65,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p
 
 # Skip initial checkpoint
 skip_initial_checkpoint: true
-# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused)
-ce_mode: bi_fused
+# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant)
+ce_mode: batch_invariant

pyproject.toml

@@ -65,7 +65,7 @@ dependencies = [
   # PR #2303, merged post-v0.1.10 so absent from every PyPI wheel); we pin
   # the blessed prebuilt wheel of stock tile-ai/tilelang@a8d93798 (includes
   # #2303, cp38-abi3, CUDA 13.1) from togethercomputer/xorl-wheels — see
-  # src/xorl/ops/linear_attention/flashqla/VENDORED.md. Switch to PyPI
+  # src/xorl/ops/_vendored/flashqla/VENDORED.md. Switch to PyPI
   # tilelang>=0.1.11 once a release carrying #2303 ships; bump deliberately
   # + revalidate the GPU smokes (sparse-MLA bwd NaN history and the FlashQLA
   # contract-pin tests) when picking up a new build.
@@ -148,9 +148,8 @@ line-length = 120
 # 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",
+  "src/xorl/ops/_vendored",
+  "src/xorl/ops/sglang/bi_families_v2.py",
 ]
 
 [tool.ruff.lint]