feat(gdn): adopt the CuTe-DSL disk cache for TVM-FFI GDN kernels

flashinfer/gdn_kernels/blackwell/gdn_prefill.py

@@ -38,6 +38,8 @@
 from flashinfer.cute_dsl.utils import get_num_sm
 
 from .gated_delta_net_chunked import GatedDeltaNetChunkedKernel
+from ...jit.cute_dsl_core import build_and_load_cute_dsl_kernel
+from ..cute_dsl_cache_naming import make_kernel_name
 
 
 # ---------------------------------------------------------------------------
@@ -47,6 +49,61 @@
 
 # Keyed on static kernel configuration. Head counts (HQ, HV) are part of
 # the key because the tile scheduler and GQA reshape logic bake them in.
+_CUTE_DSL_MODULE = "gdn_blackwell_prefill"
+
+
+def _kernel_source_files() -> tuple:
+    """Source files whose content invalidates the on-disk kernel cache."""
+    from . import gated_delta_net_chunked, gated_delta_net_tile_scheduler
+
+    return (
+        __file__,
+        gated_delta_net_chunked.__file__,
+        gated_delta_net_tile_scheduler.__file__,
+    )
+
+
+def _prefill_kernel_name(
+    io_dtype_str: str,
+    state_dtype_str: str,
+    HQ: int,
+    HV: int,
+    is_GQA: bool,
+    use_initial_state: bool,
+    store_final_state: bool,
+    enable_checkpoints: bool,
+    use_state_indices: bool,
+    cu_seqlens_dtype_str: str,
+    state_indices_dtype_str: str,
+    cu_checkpoints_dtype_str: str,
+    initial_state_inner_strides,
+    output_state_inner_strides,
+    num_sm: int,
+) -> str:
+    """Specialization name within the gdn_blackwell_prefill module.
+
+    Encodes every ``_get_compiled_cache`` key component plus ``num_sm``, which
+    the compile below bakes in as ``max_active_clusters``.
+    """
+    return make_kernel_name(
+        io_dtype_str,
+        state_dtype_str,
+        HQ,
+        HV,
+        is_GQA,
+        use_initial_state,
+        store_final_state,
+        enable_checkpoints,
+        use_state_indices,
+        cu_seqlens_dtype_str,
+        state_indices_dtype_str,
+        cu_checkpoints_dtype_str,
+        initial_state_inner_strides,
+        output_state_inner_strides,
+        num_sm,
+    )
+
+
 @functools.cache
 def _get_compiled_cache(
     io_dtype_str: str,
@@ -199,7 +256,7 @@ def chunk_gated_delta_rule_sm100(
     use_state_indices = state_indices is not None
     _state_indices = state_indices if use_state_indices else None
 
-    cache = _get_compiled_cache(
+    cache_key = (
         str(q.dtype),
         str(state_torch_dtype),
         HQ,
@@ -223,6 +280,7 @@ def chunk_gated_delta_rule_sm100(
             else None
         ),
     )
+    cache = _get_compiled_cache(*cache_key)
 
     if "compiled" not in cache:
         # --- First call: compile the kernel ---
@@ -314,25 +372,30 @@ def chunk_gated_delta_rule_sm100(
 
         stream = cuda.CUstream(torch.cuda.current_stream(device=q.device).cuda_stream)
 
-        compiled = cute.compile(
-            gdn,
-            q_cute,
-            k_cute,
-            v_cute,
-            gate_cute,
-            beta_cute,
-            o_cute,
-            cu_seqlens_cute,
-            s_in_cute,
-            s_out_cute,
-            s_indices_cute,
-            s_checkpoints_cute,
-            cu_checkpoints_cute,
-            checkpoint_every_n_tokens,
-            scale,
-            workspace_cute,
-            stream,
-            options="--enable-tvm-ffi --opt-level 3",
+        compiled = build_and_load_cute_dsl_kernel(
+            _CUTE_DSL_MODULE,
+            _prefill_kernel_name(*cache_key, num_sm),
+            lambda: cute.compile(
+                gdn,
+                q_cute,
+                k_cute,
+                v_cute,
+                gate_cute,
+                beta_cute,
+                o_cute,
+                cu_seqlens_cute,
+                s_in_cute,
+                s_out_cute,
+                s_indices_cute,
+                s_checkpoints_cute,
+                cu_checkpoints_cute,
+                checkpoint_every_n_tokens,
+                scale,
+                workspace_cute,
+                stream,
+                options="--enable-tvm-ffi --opt-level 3",
+            ),
+            extra_key_files=_kernel_source_files(),
         )
 
         cache["compiled"] = compiled

flashinfer/gdn_kernels/cute_dsl_cache_naming.py

@@ -0,0 +1,63 @@
+"""
+Copyright (c) 2026 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import hashlib
+import re
+
+import torch
+
+# The specialization name is the sole per-kernel on-disk cache key (see
+# docs/design_docs/cute_dsl_kernel_cache.md) and becomes both a filename and
+# part of the exported TVM-FFI symbol, so it must stay within [A-Za-z0-9_].
+_SANITIZE = re.compile(r"[^A-Za-z0-9_]")
+
+# ext4 caps filenames at 255 bytes; the module dir adds "<module>_" to the
+# exported symbol, so leave generous headroom before falling back to a digest.
+_MAX_NAME_LEN = 180
+
+
+def format_name_part(value) -> str:
+    """Format one cache-key component as a symbol-safe name fragment."""
+    if value is None:
+        return "none"
+    if isinstance(value, bool):
+        return "1" if value else "0"
+    if isinstance(value, float):
+        return str(value).replace(".", "_").replace("-", "m").replace("+", "p")
+    if isinstance(value, int):
+        return str(value).replace("-", "m")
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, tuple):
+        return "t" + "x".join(format_name_part(v) for v in value)
+    if isinstance(value, str):
+        return _SANITIZE.sub("_", value.removeprefix("torch."))
+    raise TypeError(
+        f"Unsupported cache-key component type {type(value).__name__}: {value!r}"
+    )
+
+
+def make_kernel_name(*parts) -> str:
+    """Join cache-key components into a specialization name.
+
+    Every codegen parameter must be passed; a component the name ignores makes
+    two different kernels collide on one on-disk artifact.
+    """
+    name = "_".join(format_name_part(p) for p in parts)
+    if len(name) > _MAX_NAME_LEN:
+        digest = hashlib.sha256(name.encode()).hexdigest()[:16]
+        name = f"{name[:_MAX_NAME_LEN]}_h{digest}"
+    return name