MAX 26.1: eager to compile contract, lowering pipeline, kernel selection across GPUs, and extension points for custom ops

[BradLarson]

@KrXGu

2) Symbolic dimensions and dynamic shapes: conventions and constraints

You mentioned symbolic dimensions and hinted at conventions around F.lazy().

• Is there a canonical way to declare symbolic dims in Python today (example snippet preferred), and what is the “shape contract” that the compiler assumes?
  • Example: are symbolic dims treated as runtime values with guards, or do we require bounded ranges, or both?
• How are specialization and caching keyed for symbolic shapes?
  • For example: does model.compile() generate a single artifact with runtime guards, or does it JIT/AOT multiple variants behind the scenes?

I’ll look for other examples of symbolic dimensions, but here’s a simple graph that uses them. In that case vector_width is a symbolic dimension that can be variable in length at runtime, but must match between the two input tensors when the graph is used. We can do algebraic checks on these dimensions, and I know there are more complex examples of their use in other open source code. From your description, I take these as “runtime values with guards”.

@KrXGu

4) Kernel dispatch and registration: pointers for contributors

This is the part I want to dig into in code.

• Can you point to the “main file where we register top-level kernels” in the OSS repo?
• When multiple kernel variants exist for an op (different tiling, dtype/layout, arch), how is the selection resolved?
  • Is it purely compile-time @parameter specialization, or is there also a higher-level selection layer (cost model, heuristics, autotune, per-device policy)?
• On fallbacks: is the current expected pattern “explicitly provide fallback kernels” (even if naive), and would a “generic baseline fallback” PR be welcomed for common ops on new hardware targets?

The file that I was referencing (and maybe forgot to link) is max/kernels/src/Mogg/MOGGKernelAPI/MOGGKernelAPI.mojo. That’s our primary registry for MAX kernels, and they all branch off from there.

For specialization, we largely lean on compile-time parameterization, with shapes and other compile-time values provided from the graph compiler. Unless I’m terribly mistaken or out of date with my information, we’re generally explicit with this specialization and don’t rely on a cost model / autotuning at the graph compiler or kernel level. We do autotuning ahead of time on hardware / shapes to provide optimized kernel parameters.

For general fallbacks, I think that’d be welcomed and there is at least one feature request for this. We just haven’t gotten to implementing the various paths where this would be used.