[Refactor][Ops] read dtype at forward instead of taking it in op constructors - #1826
Open
lcy-seso wants to merge 11 commits into
Open
[Refactor][Ops] read dtype at forward instead of taking it in op constructors#1826lcy-seso wants to merge 11 commits into
lcy-seso wants to merge 11 commits into
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This was referenced Aug 3, 2026
lcy-seso
changed the title
An op today takes a `dtype` kwarg and the caller then hands it tensors that must agree. Two sources for one fact, and the constructor's copy is the one that can be wrong. Make the tensors the only source. Consequences recorded in the slot rules: - S12/S13: `dtype` is not a kwarg, and no kernel is built in `__init__` — a kernel is dtype-specialized and there is no dtype yet. `dispatch_kernel` stays, so an unsupported arch still fails at construction. - S16: the kernel cache is keyed by `(_cache_key(*shapes), dtype)`, so a second dtype builds a second kernel instead of reusing the first. - S19: `eval_roofline` reads attributes `forward()` binds, so it is post-forward only. - The fixed-rank / arbitrary-rank split now governs when shape inference runs, not when the kernel is built. Spec only; the op layer does not conform yet.
…lies 20 op constructors stop taking `dtype`; each reads it from its input in `forward()` and keys its kernel cache by shape and dtype. Also removes the `_committed_dtype` dual path (an optional ctor dtype, validated when supplied) that cumulative, softmax and the four ada/fused-add norm ops carried. `LayerNormFwdOp` gains a real kernel cache — it kept one kernel and rebuilt whenever the leading-dims product changed, which with a second dtype in play would rebuild every call. `_dtype_codegen` now emits the validator body unrolled per input instead of looping over `input_names` through `locals()`. Dynamo cannot trace `locals()`, and this body runs inside `forward()`, so an op that validates dtypes would lose `fullgraph`. The unreachable `same_as(ref)`-not-supplied branch goes with it: synthesis already rejects a ref that names no sibling input. `_SoftmaxBaseOp` keeps its inline dtype-union check rather than delegating to `_validate_dtypes`: `LogSumExpFwdOp`'s manifest entry declares only float16|bfloat16 while the kernel and its tests use float32. Narrowing the code to match is the wrong direction, and widening the manifest is a separate change. Two tests asserted that a ctor dtype disagreeing with the input raises. That coupling is gone, but the invariant underneath is not — the `dim=[]` short-circuit must still gate dtype — so they now feed a dtype outside the manifest union instead of one outside the ctor's.
…ated The library targets H200, so `is_hopper()` was always true and every branch behind its false arm was unreachable. Removing the helper removes one axis from kernel selection entirely: what remains depends on dtype, shape and the semantic flags. Two consequences that had to be handled rather than left implicit: - The manifest named the fallbacks the ops no longer dispatch to. `source.kernel_map` now names the Wgmma/Ws kernels actually used, and `gqa_bwd_postprocess_kernel` is gone — that slot existed only on the non-Hopper path, where it held `None` on Hopper. - `test_gqa_fwd_dispatch_falls_back_off_h200` asserted the non-Hopper fallback. Its premise is gone, but the other half of what it covered is not, so it now checks that a shape outside the H200 warp-specialized contract still lands on the WGMMA kernel. The six fallback kernel classes are left in place and recorded for a separate decision: each is still exported and still passes the arch check on H200, so a caller can reach one through a `kernel_map` override. `supported_archs == [80, 89, 90]` does not identify them — several kernels carry that declaration while being the only implementation of their slot.
…ng undeclared `GroupedQueryAttentionPrefillFwdOp` takes a `dtype` constructor argument that selects the fp8 path's output element type, but no manifest param declared it — the entry's `dtype_combos` already recorded that fp8 inputs admit either fp16 or bf16 output, with nothing saying who chooses. Also drops the writer identity from the trust-model diagram and narrows `implement-op`'s blanket "do not modify manifest or design docs" to what it was protecting: the spec is not rewritten to match code that does not conform to it. Known advisory: the validator compares parameter defaults with `!=`, so a `torch.dtype` default can never match its manifest spelling. No entry could declare one before this change either.
All nine MoE constructors stop taking `dtype`. The four leaves build their kernel on first forward and cache it per dtype; the five composites had nothing left to forward once the leaves stopped accepting it, so the parameter simply disappears from them. `FusedMoEExpertsNopadPersistent3WGFwdOp` anchors its shared dtype validator on `hidden_states.dtype` — the helper already required output, `w_gate_up` and `w_down` to agree with it, so the constructor's copy was the redundant one. `SharedFusedMoE` builds its shared-expert MLP kernel lazily per dtype; a `Kernel` still takes `dtype`, only the op stops deciding it. `build_activation_op` drops its `dtype` parameter: the `FusedGatedOp` it returns already accepted `None` and deferred the build to forward, so the MoE pipeline never needed to name a dtype for its activation. Call sites needed three forms handled beyond a keyword argument: a positional fourth argument to `MoeUnpermuteFwdOp` that would otherwise have bound silently to `padded_batch_sum`, and a `dict(...)` of kwargs unpacked at the call.
`CBProducerOp`, `EngramGateConvFwdOp`, `EngramGateConvBwdOp`, `EngramDecodeOp` and `MeanPoolingForwardOp` build their kernel on first forward, cached per dtype. `DaCumsumFwdOp` keeps its `dtype` argument: its inputs are float32 while `dt_out` is float16/bfloat16/float32, so the output element type is a choice its inputs cannot express. The engram ops route their existing dtype check through the synthesized `_validate_dtypes`, which needs every declared input rather than just the anchor tensor. `CBProducerOp` keeps an inline check — its `_validate_dtypes` is still the L1 stub — and now compares its two inputs against each other instead of against a constructor argument. `MeanPoolingForwardOp` collects its kernel arguments through `locals()`; that dict becomes every argument except the element type, with dtype supplied at the cache lookup. Four call sites passed dtype positionally, where the constructor took it before `eps`. Removing the parameter would have bound a dtype to `eps` and carried it into the kernel — the failure surfaced as a TileLang type error about argument 3, far from the cause.
…de ops `GroupedQueryAttentionBwdOp`, `MultiHeadAttentionBwdOp` and the two MHA decode ops build their kernels on first forward, cached per dtype. The GQA backward op caches the preprocess and backward kernels as one entry so a dtype switch cannot pair one dtype's preprocess kernel with another's backward kernel. `MultiHeadAttentionBwdOp` stops re-exporting `prep_kernel` / `kernel` from the GQA op it wraps: those attributes no longer exist once the kernels live in a cache, and nothing outside the op read them. Seven call sites passed dtype as the last positional argument; they are handled by position rather than by keyword, which a name-based sweep would have missed. Deferred to a follow-up commit: the four ops whose `default_kernel_map` selects a kernel class from `self.dtype` (GQA forward, GQA prefill, GQA prefill-paged, MHA forward). Those need the candidate set installed as separate map keys with the choice made at forward, which also renames manifest `source.kernel_map` entries.
…indow and DeepSeek ops Eight more ops build their kernel on first forward, cached per dtype: the three NSA ops, the DeepSeek sparse-MLA and MLA decode ops, GQA prefill varlen, and both GQA sliding-window ops. `GroupedQueryAttentionSlidingWindowFwdOp.total_memory` computed `torch.tensor([], dtype=self.dtype).element_size()`. With no constructor dtype, `dtype=None` makes that float32, so it silently returned twice the real byte count. It now refuses to answer before the first forward rather than guessing. Three tests changed because the coupling they asserted is gone, not the invariant underneath: - an unsupported element type is still rejected, but the check moved with the dtype to first use, so the test feeds float32 tensors instead of naming float32 at construction; - `total_memory` is asserted after a forward, and before one the test now requires the refusal; - the dtype disagreement that still matters is between q, k and v, so the mismatch test makes k differ from q rather than making all three differ from a constructor argument. Call sites needed all four forms: keyword, last-positional, an explicit `params` dict inside the op, and a `**params` dict at the call.
lcy-seso
force-pushed
the
design/ops/ctor-dtype
branch
from
August 4, 2026 12:39
8ae6063 to
61222aa
Compare
…at construction Both GQA decode ops chose their kernel *slot* from `self.dtype` inside `__init__`. The slot table already held both candidates, so only the timing was wrong: `_select_decode_kernel_key` and `_uses_bs1_fast_path` now take the element type, and the kernel is built and cached on first forward. `_uses_bs1_fast_path` stops being a property of the op. Whether a request takes the batch=1 warp-specialized kernel depends on the element type of the call, so one instance answers differently for float16 and bfloat16 — which is what the dispatch test now asserts against a single instance instead of constructing a second op it could no longer distinguish. Two paged tests moved their expectation to where the work now happens: an unsupported page layout is still rejected, but by the kernel the op builds on first use rather than by the constructor.
…istry
`elementwise/_base.py`, `rope.py` and `dropout.py` each kept a private
`WeakValueDictionary` alongside the shared one in `compile_boundary.py`,
and thirteen constructors did
self.dispatch_kernel(kernel_map) # sets a str key, registers
self._instance_key = id(self) # overwrites it with an int
_OP_REGISTRY[self._instance_key] = self
`compile_boundary` documents why the key must be a string: dynamo
generalizes an int custom-op argument to an unhashable `SymInt` once a
second instance compiles through the same frame. Every one of the 23
wrappers was annotated `instance_key: int`, so that hazard was live
rather than hypothetical. The private dictionaries are gone, the manual
registrations with them, and the wrappers resolve through `get_instance`.
`_IntIdentityUnaryOp`'s integer path called `_install_kernel_map`
directly, which skips the registration entirely — those instances were
never registered at all. It now routes through `dispatch_kernel`.
Also picks up two more call sites that arrived with the rebase: a
`dtype=` construction and a shape-only kernel-cache key in the autotune
tests added upstream.
…efaults as dtypes Three sequence-modeling entries still declared a `dtype` param after the engram ops stopped taking one. Every output of all three is `same_as(<input>)`, so the element type was never theirs to choose. This was the only real CI failure: `validate-manifest`, `compile-contract-gate` and `gpu-smoke` all assert the validator exits 0, so one signature divergence failed three checks. The validator also compared a `torch.dtype` default against its YAML spelling with `!=`, so `torch.float16 != "float16"` made every dtype default a mismatch — advisory locally, blocking under `--strict` in CI. No entry could declare a dtype default at all before this. Where the declared type is `torch.dtype`, the default is now resolved before comparison, with a test covering both the matching and the genuinely mismatched case.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
An op takes a
dtypekwarg and the caller then hands it tensors that must agree. Two sources for one fact, and the constructor's copy is the one that can be wrong. This makes the tensors the only source: dtype is read inforward(), never taken by an op constructor.Consequences, recorded in the slot rules:
dtypeis not a kwarg, and no kernel is built in__init__: a kernel is dtype-specialized and there is no dtype yet.dispatch_kernel()stays, so an unsupported GPU still fails at construction.(_cache_key(*shapes), dtype), so a second dtype builds a second kernel instead of silently reusing the first.eval_rooflinereads attributes thatforward()binds, so it is post-forward only.Kernelis unaffected and keeps itsdtypector argument — it compiles one program for one element type. This PR only moves who supplies it.Scope
Spec only. 79 op constructors still take
dtypeand do not conform; two follow-up issues bring them over, split by whether the constructor currently builds a kernel eagerly (40) or lazily (39).Test plan
pre-commit runon both files cleanmainshows one addition,#parameter-design, which resolves)