More (optional) load time verification

[vasqu]

Description

Based on our slack conversation

A lot of kernels still need to be verified after loading, leading to a lot of bloated code

• https://github.com/huggingface/transformers/blob/7c8b1d4459c233295be9d25878b73a1b68f23d59/src/transformers/integrations/deepgemm.py#L80-L204
• https://github.com/huggingface/transformers/blob/7c8b1d4459c233295be9d25878b73a1b68f23d59/src/transformers/integrations/sonicmoe.py#L49-L68

It makes the development and usage way more painful and imo kernels could take a part here to allow on load verification for these certain features attributes

• NVCC
• (Cuda) Architecture
• Cute dsl versioning (so any optional dep)

Motivation

Make kernels as dumb to use as we can :D

Proposed approach

Load time verification on get kernel that checks for hw specifics and pinned versions; maybe hidden behind a flag

Alternatives considered

I mean yes, that is exactly what we currently do -> manually check ourselves

Additional context

No response

[vasqu]

cc @sayakpaul @drbh @danieldk

[sayakpaul]

Sharing some internal brainstorming messages:

these are good points but some of these problems are hard to solve with a generic solution. this is because:

1. kernel authors can add conditional logic for specific archs in their source code, which may only check arch at runtime and throw an error if some specific operation/function is not supported
2. the JIT kernels require some tool to compile/search and this information is also checked at runtime (and this may look for specific libs or nvcc or etc..)

there are only two ways for us to know if a kernel will run in a given instance - 1. we try and it succeeds or fails (and record this) 2. the kernel dev tells us (records this somewhere)

I think the suggestions above are solid and we should improve the ability to specify deps - and more validation so we can throw explicit errors from kernels.

specifically

1. I think we should be able to add versions to deps
2. maybe we can check the default location for nvcc and add a message to specify a path if in a non standard location
3. I think we can use the existing cuda-capabilities field in build.toml for sonic moe, and then we can update the _validate_variant_dependencies function to check the listed cuda caps at runtime?

[danieldk]

I think we should be able to add versions to deps

I am not sure if this is a step in the right direction if this concerns Python dependencies. Kernels were largely standalone, but we are slowly introducing more and more of our own homegrown dependency management with a subpar experience. At some point we might as well start shipping kernels as wheels, because installing kernels and then getting runtime exceptions that certain packages are missing is bad. The additional issue we'll start to run into is that if two kernels require different versions of CuTe-DSL that are not compatible, kernels are not seamless anymore, because we've brought in dependency hell.

The only direction I can see this work is if we removed Python dependencies completely and started shipping einops, etc. as kernels and provided a way for kernels to depend on other kernels (in a better way than just using get_kernel), but that has some issues for CuTe-DSL, since it contains proprietary code.

[sayakpaul]

started shipping einops, etc. as kernels

Not sure how that would be received by the community but I am keen on learning more about how you'd envision it.