[Docs] Follow the unified kernel dispatch and the generated compile boundary

docs/api/moe.md

@@ -11,7 +11,8 @@ each token's experts, the ops that move tokens into an expert-contiguous layout
 back, and the expert GEMMs that run on it. `MoeGroupedGemmFwdOp` is one grouped
 GEMM; `MoeExpertMLPFwdOp` is the pair of them with the gated activation fused into
 the first; `FusedMoEExpertsFwdOp` is that MLP with the permutes around it, on the
-tight (no-pad) layout. The routing has to produce the layout the GEMM expects.
+tight (no-pad) layout, and `IndexedExpertMLPFwdOp` is the backend it picks instead when
+the routes are few enough to read the weights once per route rather than once per expert. The routing has to produce the layout the GEMM expects.
 
 ## Fused forward
 
@@ -66,3 +67,9 @@ tight (no-pad) layout. The routing has to produce the layout the GEMM expects.
       show_root_heading: true
       heading_level: 3
       members: ["__init__", "forward"]
+
+::: tileops.moe.IndexedExpertMLPFwdOp
+    options:
+      show_root_heading: true
+      heading_level: 3
+      members: ["__init__", "forward"]

docs/backends.md

@@ -182,11 +182,10 @@ d = op(a, b)                     # every input on one device: a.device == b.devi
 #   two or more True   → AmbiguousTargetError, asking for an explicit target=
 
 # ── op layer: the one place GemmFwdOp.forward fetches a kernel ───────
-kernel = self.get_or_build_kernel(
+kernel = self.kernel_for(
     "gemm_kernel",               # a name from kernel_map
     (a, b),                      # the tensors the kernel is about to get, in signature.inputs order
-    key=(m, n, k, a.dtype),      # in-tree only; not used on this call
-    build=lambda: GemmKernel(m, n, k, a.dtype),   # in-tree only; not used on this call
+    (m, n, k, a.dtype),          # what this call is; entry_for reads it, in-tree only
 )
 
 # ── op layer: look up the external memo table — device, then input signature ──
@@ -220,14 +219,14 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm)
 
 The op layer calls `build_gemm`; the backend never calls it itself. Importing the backend
 module only records it in the registry, and the call comes when an op call reaches
-`get_or_build_kernel` and misses the external memo table — once per device and input
+`kernel_for` and misses the external memo table — once per device and input
 signature. Whatever it returns, the op layer stores and launches.
 
 Four things follow from that:
 
-- **`key` and `build` are the op author's, not a backend's.** They serve the in-tree path
-  only: `key` decides what the in-tree kernel is looked up on, `build` how it is built.
-  Neither is used once a target serves the call.
+- **`entry_for` is the op author's, not a backend's.** It serves the in-tree path only,
+  answering with what the in-tree kernel is looked up on and how it is built. Neither
+  answer is asked for once a target serves the call.
 - **Tensors arrive positionally, params by name.** `build_kernel(*inputs, **params)`: the
   positional arguments are `TensorSpec`s (`None` for an optional input the call omitted),
   the keywords the manifest's `params` names with the values this call settled on.
@@ -236,8 +235,8 @@ Four things follow from that:
   `TensorSpec`s which kernel to return.
 - **No memoisation of its own is needed.** For the same device and input signature the op
   layer does not call again; for a finer split, or fewer rebuilds, add a cache inside
-  `build_kernel`. An op with no in-tree implementation may omit `build`, and then a call
-  with no target claiming the device raises `OpNotAvailableError`.
+  `build_kernel`. An op with no in-tree implementation may leave `entry_for` out, and then
+  a call with no target claiming the device raises `OpNotAvailableError`.
 
 ## Writing a backend that runs {#runnable}
 

docs/backends.zh.md

@@ -148,11 +148,10 @@ d = op(a, b)                     # 所有输入必须在同一设备上：a.devi
 #   两个以上返回 True    → 抛 AmbiguousTargetError，要求显式写 target=
 
 # ── 算子层：GemmFwdOp.forward 里唯一取 kernel 的那一处 ───────────────
-kernel = self.get_or_build_kernel(
+kernel = self.kernel_for(
     "gemm_kernel",               # kernel_map 里的名字
     (a, b),                      # 即将传给 kernel 的张量，顺序照 signature.inputs
-    key=(m, n, k, a.dtype),      # 自带实现用，这次不走
-    build=lambda: GemmKernel(m, n, k, a.dtype),   # 自带实现用，这次不走
+    (m, n, k, a.dtype),          # 本次调用是什么；由 entry_for 读，自带实现用，这次不走
 )
 
 # ── 算子层：按设备与输入签名查外部记忆表 ─────────────────────────────
@@ -184,14 +183,14 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b):
 register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm)
 ```
 
-`build_gemm` 由算子层调用，后端自己从不调它：import 后端模块时只是把它登记进注册表，真正被调是在一次调用走到 `get_or_build_kernel`、且外部记忆表未命中的时候，每个「设备 + 输入签名」一次。它返回的可调用对象随后由算子层 launch，也由算子层存进记忆表。
+`build_gemm` 由算子层调用，后端自己从不调它：import 后端模块时只是把它登记进注册表，真正被调是在一次调用走到 `kernel_for`、且外部记忆表未命中的时候，每个「设备 + 输入签名」一次。它返回的可调用对象随后由算子层 launch，也由算子层存进记忆表。
 
 四点对应关系值得记住：
 
-- **`key` 与 `build` 由算子作者写，与后端无关。** 它们只服务自带实现：`key` 决定自带 kernel 按什么查表，`build` 决定它怎么构造。target 选中后端时这两个参数整条不走。
+- **`entry_for` 由算子作者写，与后端无关。** 它只服务自带实现：给出自带 kernel 按什么查表、又怎么构造。target 选中后端时这两个答案都不会被问。
 - **张量按位置传，参数按名字传。** `build_kernel(*inputs, **params)`：位置实参是 `TensorSpec`（没传的可选输入是 `None`），关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。
 - **一个 `(算子, target)` 只注册一个 builder。** 算子内部分几种情形（GEMM 的 `gemm_kernel` 与 `gemv_kernel`）不会传进来，`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。
-- **不必自己做记忆。** 同一个设备与输入签名，算子层不会再调第二次；要更细的区分或更少的重建，在 `build_kernel` 内部另加一层缓存。算子完全没有自带实现时 `build` 可以不传，那时没有 target 认领设备，调用直接抛 `OpNotAvailableError`。
+- **不必自己做记忆。** 同一个设备与输入签名，算子层不会再调第二次；要更细的区分或更少的重建，在 `build_kernel` 内部另加一层缓存。算子完全没有自带实现时 `entry_for` 可以不写，那时没有 target 认领设备，调用直接抛 `OpNotAvailableError`。
 
 ## 实现一个可运行的后端 {#runnable}
 

docs/new-op.md

@@ -61,7 +61,7 @@ GemmFwdOp:
 ```
 
 Those names are how a kernel is asked for at runtime: `_eager_forward` picks one, passes
-the name to `get_or_build_kernel`, and the op layer looks the class up in `kernel_map` and
+the name to `kernel_for`, and the op layer looks the class up in `kernel_map` and
 builds it (see [step 2](#op-class)). An external backend registers against the same roster
 — whichever name it registers a `build_kernel` for is the kernel of the op it takes over.
 
@@ -100,56 +100,61 @@ class GemmFwdOp(Op):
         self._validate_dtypes(a, b)                  # generated by the base class
         m, n, k = self._infer_mnk(a, b)
         a, b = a.contiguous(), b.contiguous()        # handed over as the spec declares it
-        slot = "gemv_kernel" if m == 1 else "gemm_kernel"
-        kernel = self.get_or_build_kernel(
-            slot,                                    # a name from kernel_map
-            (a, b),                                  # the memo key's tensors, and what a backend receives
-            key=(m, n, k, a.dtype),                  # the cache key on the in-tree side
-            build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune),
+        role = "gemv_kernel" if m == 1 else "gemm_kernel"
+        kernel = self.kernel_for(
+            role,                                    # a name from kernel_map
+            (a, b),                                  # what a backend is described with
+            (m, n, k, a.dtype),                      # what this call is
         )
         return kernel(a, b)
+
+    def entry_for(self, role, call):                 # the in-tree recipe
+        m, n, k, dtype = call
+        return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune)
 ```
 
-Four members to write, each of them from the spec:
+Five members to write, the first four of them from the spec:
 
 | # | Member | Written from |
 | --- | --- | --- |
 | 1 | `__init__` | the names and defaults in `signature.params`, plus `kernel_map` and `tune`, closing with `self.dispatch_kernel(kernel_map)` to establish this instance's kernel_map |
 | 2 | `default_kernel_map` | `source.kernel_map`: the same names, against the Kernel classes themselves |
 | 3 | `_infer_output_shapes` | the rules in `signature.shape_rules` that derive an output's shape |
 | 4 | `forward` | `signature.inputs` — its order and defaults, optional inputs last — plus the validation, the contiguity, fetching the kernel and launching it |
+| 5 | `entry_for` | what two calls must share to reuse one kernel, and how that kernel is built |
 
 Two more members arrive on their own. When the subclass is defined, the base class
 synthesises `_validate_dtypes` and `eval_roofline` from the spec's dtype declarations and
 its `roofline`, so they are there to call — and worth overriding only where the op needs
 something the spec cannot say.
 
-### `get_or_build_kernel`
+### `kernel_for` and `entry_for`
 
 A kernel is a compiled artefact, hundreds of milliseconds to seconds to build, while an op
 instance is called over and over at different shapes and dtypes. The op layer therefore
 keeps a memo table: a kernel this call needs and has built before comes straight back,
-and only otherwise is one built and stored. `get_or_build_kernel` is that table's only
-entrance, and the point where the in-tree implementation and an external backend part
-ways — the second layer of selection in [the backend protocol](backends.md).
+and only otherwise is one built and stored. `kernel_for` is that table's only entrance,
+and the point where the in-tree implementation and an external backend part ways — the
+second layer of selection in [the backend protocol](backends.md).
 
-Its four arguments:
+Its three arguments:
 
-**`name`** — which kernel this call wants, as a name from `kernel_map`.
+**`role`** — which of this op's kernels the call wants, as a name from `kernel_map`.
 
 ```python
-slot = "gemv_kernel" if m == 1 else "gemm_kernel"
+role = "gemv_kernel" if m == 1 else "gemm_kernel"
 ```
 
 The in-tree side looks up the Kernel class under that name; a backend looks up the
-`build_kernel` it registered under it. An op has as many names as it has cases.
+`build_kernel` it registered under it. A role is a memoization bucket, one per kernel the
+op runs — never the name of the implementation selection picked for this call.
 
 **`inputs`** — the tensors the kernel is about to be handed, in `signature.inputs` order,
 one slot per input.
 
 ```python
-self.get_or_build_kernel(slot, (a, b), ...)                     # GEMM: two required inputs
-self.get_or_build_kernel("group_norm", (x, weight, bias), ...)  # an absent optional input is None
+self.kernel_for(role, (a, b), ...)                     # GEMM: two required inputs
+self.kernel_for("group_norm", (x, weight, bias), ...)  # an absent optional input is None
 ```
 
 The external path keys on it — the device, plus each slot's `(dtype, shape)`. The device
@@ -161,45 +166,44 @@ An optional input that was not passed keeps its slot, as `None`; that is what a
 reads presence off, rather than counting slots. Squeeze the empty slots out, and a clamp
 with only a lower bound looks exactly like one with only an upper bound.
 
-Omitting `inputs` raises nothing until a backend is installed, and then
-`OpNotAvailableError`: the op stays in-tree only, out of reach of any target (see [after
-install: two states](backends.md#three-states)).
+**`call`** — what this call is, in whatever form the op's own `entry_for` reads. An op
+that selects among several implementations passes the record its family defines; an op
+with one implementation passes the few values its kernel is built from.
 
-**`key`** — what the in-tree kernel specializes on; the in-tree path only. What becomes of
-these last two once a backend serves the op is in [how one call reaches
+`entry_for(role, call)` answers with the pair the memo table needs. What becomes of it
+once a backend serves the op is in [how one call reaches
 `build_kernel`](backends.md#from-op-layer).
 
 ```python
-key=(m, n, k, a.dtype)                             # GEMM: three dimensions and the dtype
-key=(self._cache_key(*input_shapes), x.dtype)      # the general form
+def entry_for(self, role, call):
+    m, n, k, dtype = call
+    return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune)
 ```
 
-The default `_cache_key` takes the sizes of every non-static axis across the inputs —
-always correct, but it can over-fragment: one compile per distinct shape. Where the kernel
-depends on fewer quantities, override it to project the shape onto those, flattening the
-leading dims to one product when the kernel treats its input as 2-D.
-
-**`build`** — how that in-tree kernel is constructed; the in-tree path only.
-
-```python
-build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune)
-```
+The **identity** it returns first is what two calls must share to be one entry: the
+construction arguments, plus the device wherever the constructor could produce a different
+object on another card. Carry too little and a second dtype reuses the first dtype's
+kernel; carry the whole shape where the kernel depends on fewer quantities and you compile
+once per distinct shape.
 
-Called once per `key`, which is why compiling belongs here. It may return one Kernel, a
-sequence of Kernels built together, or a dataclass carrying them — the last two suit an op
-that launches several kernels per call.
+The **builder** it returns second runs once per identity, which is why compiling belongs
+inside it. It may return one Kernel, a sequence of Kernels built together, or a dataclass
+carrying them — the last two suit an op that launches several kernels per call.
 
-An op with no in-tree implementation at all, one written to depend on a backend, may leave
-`build` out; a call on a device no target claims then raises `OpNotAvailableError`.
+An op that selects among candidate Kernel classes writes no `entry_for` at all: the
+default asks the class selection chose, which states its own identity and builder. An op
+with no in-tree implementation, written to depend on a backend, leaves both out; a call on
+a device no target claims then raises `OpNotAvailableError`.
 
 ### Finishing: the compile boundary, and registering
 
 Two things to finish, a few lines each:
 
 - **To support `torch.compile`**, declare a compile boundary as well: `forward` only calls
   the opaque operator, and the validation, the kernel lookup and the launch move into
-  `_eager_forward`. The op above declares none, so its `forward` holds all the work. How to
-  declare it is in [bringing an op into torch.compile](torch-compile.md).
+  `_eager_forward`. The operator itself is generated from the spec, so declaring the
+  boundary is one class attribute. The op above declares none, so its `forward` holds all
+  the work. How to declare it is in [bringing an op into torch.compile](torch-compile.md).
 - **Add the op's name** to the imports and `__all__` in two places: its family's
   [`src/tileops/ops/<family>/__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops), where the class
   is implemented, and [`src/tileops/<family>.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops), the public
@@ -247,8 +251,8 @@ kernel = AttnKernel(num_heads, head_dim, dtype)
 out = kernel(q, k, v)                       # seq_len is read off the tensor shapes
 ```
 
-With the first form, `seq_len` ends up in `get_or_build_kernel`'s `key`, every step misses,
-every step compiles, and decode goes nowhere.
+With the first form, `seq_len` ends up in the identity `entry_for` returns, every step
+misses, every step compiles, and decode goes nowhere.
 
 ## Step 4: write the test