[Improvement] Re-structure sonic MoE kernel interface to use a Handle object

[remi-or]

In #47126 and offline we talked about how we want to take care of loading kernels like DeepGEMM.

Notably, we can put everything related to loading and calling the kernel inside a singleton Handle object. This object would only keep track of a few state variable and provide a single entry point to access the kernel. This has the following advantages:

• we can keep track of the cached state without using @lru_cache which can be opaque and cause memory leaks
• failures are actually cached: @lru_cache does not cache exceptions, so the cache never worked in a machine / env that did not support sonicmoe
• we can reset the cached state of the handle whenever we want through a simple reset method
• we can check the availability of the kernel easily through a sonicmoe_is_available property, that accesses the same underlying cached kernel as the function call itself (quite convenient, as is shown in the reworked _test_eager_matches_batched_and_grouped_inference)

The downside of the Handle API are that transient failures are also cached (hence the reset method, but it is not called automatically) and it's beefier than @lru_cache. IMO this is an ok tradeoff, but I would love some feedback on this.

Since I was fixing a failure in sonicmoe, I thought this was a good occasion to try out the Handle API. The fail was an image failure, which is now guarded against.

I also added a test for sonicmoe + compile and it passes on B200. Also tested the sonicmoe regular tests on a qwen and they pass.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 29320447311
Result: success | Grafana metrics are not available yet.

[IlyasMoutawwakil]

can you please also add the specific versions that worked for you of tvm ffi and cutedsl, i didn't expect it to start failing so fast with newer versions, when we tell users to install them they'll just do "pip install nvidia-cutlass-dsl" which would then result in version mismatch between the vendored quack and the installed cutedsl
cc @vasqu

[vasqu]

💯 we just had discussions around that for the past few days so let's guard it on our side properly

[remi-or]

Here you go:

nvidia-cutlass-dsl                       4.5.2
apache-tvm-ffi                           0.1.9

[IlyasMoutawwakil]

sorry but i am once again totally against syncs and tensor creation in the loader (also does this not break cuda graph capture ?) when a more direct alternative exist, in this case checking cutedsl exists and explicitly verifying its version. here the user sees import error no package named cutlass.

[remi-or]

Yeah the synchronize should be guarded against, mb. As for tensor creation, sure I don't mind avoiding them, I will look into the version check.

[IlyasMoutawwakil]

the scope of the handle is a bit fuzzy, it doesn't handle the full loading (lets _load_sonicmoe_kernel do part of it) and here it implements something very implementation specific.. With FinegrainedFP8 for example which implements multiple implementations, it will be a very big handle. I think the implementation part should stay outside of it.

[vasqu]

Yea I think the goal here was to only handle the import / cache structure; in that case we really should be strict and dont mix functionality

[remi-or]

Makes a lot of sense, will change it up

[IlyasMoutawwakil]

love the idea of the handle and standardizing the interaction with the kernel package (loading, checking for availability, completeness and usability), i just think it can be scoped better and also it should generalize to other kernel integrations. for example we can make it so that it has a class attr of required packages and min/max versions etc, a bit like we do in exporters' validate_environment. the list of required attributes as well can be programatically checked. ofc some custom checks will be necessary for deepgemm and its cuda toolkit for example.

[vasqu]

I feel like the pain points are valid but the direction is not ideal

I think we already have existing patterns we could apply to solve all these problems, notably

• the lru cache (with the loading structure of deepgemm for example)
• the availability check can very well be moved to import utils (and it makes things a lot clearer, not relying on try catch)

[vasqu]

Yea I think the goal here was to only handle the import / cache structure; in that case we really should be strict and dont mix functionality

[vasqu]

Tbh I'm not a fan of singleton objects (but that's more of a personal opinion) - it feels like OOP for the sake of OOP 😅

I'm still a bit confused about the actual benefits and why we may not adopt a generalized pattern similar to flash attention

transformers/src/transformers/modeling_flash_attention_utils.py

Lines 124 to 272 in a689565

	# `globals()` is not compatible with dynamo, hence we have do define them in global scope ourselves
	_loaded_implementation = None
	_flash_fn = None
	_flash_varlen_fn = None
	_flash_with_kvcache_fn = None
	_pad_fn = None
	_unpad_fn = None
	# function that processes kwargs, generalized to handle any supported kwarg within the function
	_process_flash_kwargs_fn = None
	# exceptions where hf API doesn't match the original flash attention API
	_hf_api_to_flash_mapping = {
	"dropout": "dropout_p",
	"sliding_window": "window_size",
	}
	# alternative names within the different flash attention APIs, e.g. for attention sinks
	_flash_api_alternative_names = {"s_aux": "learnable_sink"}
	def _lazy_imports(
	implementation: str | None, attention_wrapper: Callable | None = None, allow_all_kernels: bool = False
	):
	"""
	Lazy loads the respective flash attention implementations.
	Return:
	flash_attn_func: The base flash attention function.
	flash_attn_varlen_func: The flash attention function supporting variable sequence lengths,
	e.g. for padding-free training.
	pad_input: The function to pad inputs into one sequence and returning the respective kwargs.
	unpad_input: The function to unpad outputs based on the kwargs (from pad_input).
	"""
	is_fa2 = is_flash_attn_2_available()
	is_fa3 = is_flash_attn_3_available()
	is_fa4 = is_flash_attn_4_available()
	pad_input, unpad_input = _pad_input, _unpad_input
	is_paged, implementation = split_attention_implementation(implementation)
	if (implementation == "flash_attention_2" and is_fa2) or (
	implementation is None and is_fa2 and not is_fa3 and not is_fa4
	):
	from flash_attn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache
	from flash_attn.bert_padding import pad_input, unpad_input
	elif is_torch_npu_available():
	# Package `flash-attn` is unavailable on Ascend NPU, which will cause ImportError
	# Flash-Attention2 related apis for Ascend NPU must be imported from `.integrations.npu_flash_attention` module
	from .integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func
	from .integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func
	from .integrations.npu_flash_attention import npu_flash_attn_with_kvcache as flash_attn_with_kvcache
	else:
	if implementation == "flash_attention_3" or (implementation is None and is_fa3 and not is_fa4):
	from flash_attn_interface import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache
	elif implementation == "flash_attention_4" or (implementation is None and is_fa4):
	from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
	flash_attn_with_kvcache = None # not supported yet
	# Kernels fallback
	else:
	from .integrations.hub_kernels import load_and_register_attn_kernel
	# Map standard attention names to hub kernel repos
	kernel_repo = FLASH_ATTN_KERNEL_FALLBACK.get(implementation, implementation)
	# We want to explicitly register the name with `paged|` if found
	kernel_implementation = f"paged|{implementation}" if is_paged else kernel_repo
	kernel = load_and_register_attn_kernel(
	kernel_implementation, attention_wrapper, allow_all_kernels=allow_all_kernels
	)
	flash_attn_func = getattr(kernel, "flash_attn_func", None)
	flash_attn_varlen_func = getattr(kernel, "flash_attn_varlen_func", None)
	flash_attn_with_kvcache = getattr(kernel, "flash_attn_with_kvcache", None)
	# Block-sparse kernels (e.g. ``kernels-staging/msa``) expose ``sparse_atten_func`` rather than
	# ``flash_attn_varlen_func``. ``load_and_register_attn_kernel`` already registered their dedicated
	# wrapper into ``ALL_ATTENTION_FUNCTIONS``, so they dispatch through the attention interface and
	# never touch the flash varlen globals -- preloading them here is a no-op, not an error.
	if flash_attn_varlen_func is None and hasattr(kernel, "sparse_atten_func"):
	return flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache, pad_input, unpad_input
	if flash_attn_varlen_func is None:
	raise ValueError(
	f"Could not find the currently requested flash attention implementation at `{implementation}`."
	"Make sure that you request a valid kernel from the hub, e.g. `kernels-community/flash-attn2`."
	)
	if flash_attn_func is None:
	logger.warning(
	f"The loaded flash attention implementation at `{implementation}` only supports varlen, i.e. "
	"it can only be used with continuous batching and does not support the full functionality for "
	"the base transformers generation methods."
	)
	if flash_attn_with_kvcache is None:
	logger.warning(
	f"The loaded flash attention implementation at `{implementation}` does not support block tables, so"
	" the full performances of continuous batching will not be achieved, only the varlen path will be "
	"used."
	)
	return flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache, pad_input, unpad_input
	def _lazy_define_process_function(flash_function):
	"""
	Depending on the version and kernel some features are not supported. Due to limitations in
	`torch.compile`, we opt to statically type which (optional) kwarg parameters are supported
	within `_process_flash_attention_kwargs`.
	NOTE: While all supported kwargs are marked as `True`, everything else is marked as `False`.
	This might be confusing for kwargs that we use in any case, e.g. `is_causal`.
	"""
	flash_parameters = inspect.signature(flash_function).parameters
	process_parameters = inspect.signature(_process_flash_attention_kwargs).parameters
	supports_mapping = {}
	for param in process_parameters:
	fa_param = _hf_api_to_flash_mapping.get(param, param)
	supports_mapping[fa_param] = fa_param in flash_parameters
	if (fa_alternative_name := _flash_api_alternative_names.get(param, param)) != fa_param:
	supports_mapping[fa_alternative_name] = fa_alternative_name in flash_parameters
	return partial(_process_flash_attention_kwargs, supports_mapping=supports_mapping)
	def lazy_import_flash_attention(
	implementation: str | None, attention_wrapper: Callable | None = None, allow_all_kernels: bool = False
	):
	"""
	Lazily import flash attention and return the respective functions + flags.
	NOTE: For fullgraph, this needs to be called before compile, while no fullgraph can
	work without preloading. See `load_and_register_attn_kernel` in `integrations.hub_kernels`.
	"""
	global _loaded_implementation
	if implementation is None and _loaded_implementation is None:
	raise ValueError("Could not find any flash attn implementation based on your environment.")
	global _flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn, _process_flash_kwargs_fn
	if implementation is not None and _loaded_implementation != implementation:
	_loaded_implementation = implementation
	_flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn = _lazy_imports(
	implementation, attention_wrapper, allow_all_kernels=allow_all_kernels
	)
	# Block-sparse kernels register their own attention interface and expose no varlen fn to introspect;
	# skip building the kwargs-support map (it is only consumed by the flash varlen path they never take).
	_process_flash_kwargs_fn = _lazy_define_process_function(_flash_varlen_fn) if _flash_varlen_fn else None
	return (_flash_fn, _flash_varlen_fn, _flash_with_kvcache_fn, _pad_fn, _unpad_fn), _process_flash_kwargs_fn

(which is more complicated but I mean the lazy import x globals structure)

I think the biggest point is is_sonicmoe_available and the avoidance of the lru cache

• is_sonicmoe_available could be properly moved to import utils and we make proper versioning guards, e.g. is kernels available and the respective cute version.
• Re lru cache, I feel like this is valid but we can restructure to an FA like structure as well (I think deepgemm does something similar)

[remi-or]

Haha my personal would be a singleton object over globals, but then again it comes back to the same results (pretty much) and FA already uses globals so might as well do the same. Will reformat into the same pattern and we can re-evaluate from there.

[remi-or]

Lots of feedback, thanks! This was more of a rough PoC so lots to take in, and I did not plan to spend too much time on this, it was more because tests were crashing on main that I took the opportunity to look into it. Will update when I have the bandwidth!

[vasqu]

No worries, we also nudged the kernels team so maybe we can simplify a bit more in the future. It's a bit awkward to handle these kernels 😢