From 58609a0dd84923e6b9d1199cdc980b6e7bb755a4 Mon Sep 17 00:00:00 2001 From: leeyongjun Date: Thu, 3 Sep 2026 03:22:26 +0000 Subject: [PATCH 1/2] feat(gdn): adopt the CuTe-DSL disk cache for TVM-FFI GDN kernels --- .../gdn_kernels/blackwell/gdn_prefill.py | 103 +++- .../gdn_kernels/cute_dsl_cache_naming.py | 63 +++ .../gdn_kernels/gdn_decode_bf16_state.py | 235 ++++---- flashinfer/gdn_kernels/gdn_decode_mtp.py | 259 ++++++--- .../gdn_kernels/gdn_decode_nontranspose.py | 87 ++- .../gdn_kernels/gdn_decode_pretranspose.py | 100 +++- tests/gdn/test_cute_dsl_kernel_cache.py | 518 ++++++++++++++++++ 7 files changed, 1113 insertions(+), 252 deletions(-) create mode 100644 flashinfer/gdn_kernels/cute_dsl_cache_naming.py create mode 100644 tests/gdn/test_cute_dsl_kernel_cache.py diff --git a/flashinfer/gdn_kernels/blackwell/gdn_prefill.py b/flashinfer/gdn_kernels/blackwell/gdn_prefill.py index 8623d229f83..a3acbd3ca4e 100644 --- a/flashinfer/gdn_kernels/blackwell/gdn_prefill.py +++ b/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 diff --git a/flashinfer/gdn_kernels/cute_dsl_cache_naming.py b/flashinfer/gdn_kernels/cute_dsl_cache_naming.py new file mode 100644 index 00000000000..e782ee56717 --- /dev/null +++ b/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 "_" 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 diff --git a/flashinfer/gdn_kernels/gdn_decode_bf16_state.py b/flashinfer/gdn_kernels/gdn_decode_bf16_state.py index 9678211b53b..8caff454737 100644 --- a/flashinfer/gdn_kernels/gdn_decode_bf16_state.py +++ b/flashinfer/gdn_kernels/gdn_decode_bf16_state.py @@ -48,6 +48,8 @@ from cutlass.cute.runtime import from_dlpack from .dtype_compat import as_bf16 +from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel +from .cute_dsl_cache_naming import make_kernel_name def _mark_batch_dynamic(torch_t: torch.Tensor, *, assumed_align: int = 32): @@ -2843,6 +2845,18 @@ def gated_delta_rule( _compiled_kernels_mtp: dict = {} _compiled_kernels_wide_vec: dict = {} +_CUTE_DSL_MODULE = "gdn_decode_bf16_state" + + +def _bf16_state_kernel_name(variant: str, cache_key: tuple) -> str: + """Specialization name within the gdn_decode_bf16_state module. + + ``variant`` distinguishes the compiled entry points sharing this module + ("wide_vec", "wide_vec_t1", "mtp_ilp4"); ``cache_key`` is the in-process + cache tuple, which already encodes every parameter that affects codegen. + """ + return make_kernel_name(variant, *cache_key) + def _dtype_key( A_log: torch.Tensor, @@ -3225,43 +3239,48 @@ def gated_delta_rule_mtp_wide_vec( ) _compiled_kernels_wide_vec[cache_key] = { - "compiled": cute.compile( - _run_wide_vec, - h_, - inter_, - A_log_, - a_, - dt_bias_, - q_, - k_, - v_, - b_, - o_, - h0_idx_, - h0_out_idx_, - acc_steps_, - ssm_idx_, - softplus_beta, - softplus_threshold, - scale, - HV_val, - T_val, - H_val, - K_val, - V_val, - tile_v, - use_qk_l2norm_in_kernel, - effective_disable_final, - cache_intermediate_states, - use_packed_fma, - same_pool, - disable_output, - recovery_steps, - per_request_accepted_steps, - per_token_pool_scatter, - per_token_pool_scatter_flat, - stream, - options="--enable-tvm-ffi --generate-line-info --opt-level 3", + "compiled": build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _bf16_state_kernel_name("wide_vec", cache_key), + lambda: cute.compile( + _run_wide_vec, + h_, + inter_, + A_log_, + a_, + dt_bias_, + q_, + k_, + v_, + b_, + o_, + h0_idx_, + h0_out_idx_, + acc_steps_, + ssm_idx_, + softplus_beta, + softplus_threshold, + scale, + HV_val, + T_val, + H_val, + K_val, + V_val, + tile_v, + use_qk_l2norm_in_kernel, + effective_disable_final, + cache_intermediate_states, + use_packed_fma, + same_pool, + disable_output, + recovery_steps, + per_request_accepted_steps, + per_token_pool_scatter, + per_token_pool_scatter_flat, + stream, + options="--enable-tvm-ffi --generate-line-info --opt-level 3", + ), + extra_key_files=(__file__,), ), # Per-B default tensors (B-dependent shapes; can't be shared # across batch sizes — see #L bug at cache_key without B). @@ -3494,36 +3513,41 @@ def gated_delta_rule_t1_wide_vec( h0_out_idx_ = h0_idx_ _compiled_kernels_wide_vec[cache_key] = { - "compiled": cute.compile( - _run_wide_vec_t1, - h_, - inter_, - A_log_, - a_, - dt_bias_, - q_, - k_, - v_, - b_, - o_, - h0_idx_, - h0_out_idx_, - softplus_beta, - softplus_threshold, - scale, - HV_val, - T_val, - H_val, - K_val, - V_val, - tile_v, - use_qk_l2norm_in_kernel, - effective_disable_final, - cache_intermediate_states, - use_packed_fma, - same_pool, - stream, - options="--enable-tvm-ffi --generate-line-info --opt-level 3", + "compiled": build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _bf16_state_kernel_name("wide_vec_t1", cache_key), + lambda: cute.compile( + _run_wide_vec_t1, + h_, + inter_, + A_log_, + a_, + dt_bias_, + q_, + k_, + v_, + b_, + o_, + h0_idx_, + h0_out_idx_, + softplus_beta, + softplus_threshold, + scale, + HV_val, + T_val, + H_val, + K_val, + V_val, + tile_v, + use_qk_l2norm_in_kernel, + effective_disable_final, + cache_intermediate_states, + use_packed_fma, + same_pool, + stream, + options="--enable-tvm-ffi --generate-line-info --opt-level 3", + ), + extra_key_files=(__file__,), ), # Per-B default tensors (B-dependent shapes — see batch-dynamic # correctness note in gated_delta_rule_mtp_wide_vec). @@ -3885,42 +3909,47 @@ def gated_delta_rule_mtp( ) _compiled_kernels_mtp[cache_key] = { - "compiled": cute.compile( - run_gdn_decode_bf16state_mtp_ilp4, - h_, - inter_, - A_log_, - a_, - dt_bias_, - q_, - k_, - v_, - b_, - o_, - h0_idx_, - h0_out_idx_, - acc_steps_, - ssm_idx_, - softplus_beta, - softplus_threshold, - scale, - HV, - T, - H, - K, - V, - tile_v, - use_qk_l2norm_in_kernel, - disable_state_update, - cache_intermediate_states, - use_packed_fma, - same_pool, - disable_output, - per_request_accepted_steps, - per_token_pool_scatter, - per_token_pool_scatter_flat, - stream, - options="--enable-tvm-ffi --generate-line-info --opt-level 3", + "compiled": build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _bf16_state_kernel_name("mtp_ilp4", cache_key), + lambda: cute.compile( + run_gdn_decode_bf16state_mtp_ilp4, + h_, + inter_, + A_log_, + a_, + dt_bias_, + q_, + k_, + v_, + b_, + o_, + h0_idx_, + h0_out_idx_, + acc_steps_, + ssm_idx_, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + tile_v, + use_qk_l2norm_in_kernel, + disable_state_update, + cache_intermediate_states, + use_packed_fma, + same_pool, + disable_output, + per_request_accepted_steps, + per_token_pool_scatter, + per_token_pool_scatter_flat, + stream, + options="--enable-tvm-ffi --generate-line-info --opt-level 3", + ), + extra_key_files=(__file__,), ), # Per-B default tensors (B-dependent shapes — see batch-dynamic # correctness note in gated_delta_rule_mtp_wide_vec). diff --git a/flashinfer/gdn_kernels/gdn_decode_mtp.py b/flashinfer/gdn_kernels/gdn_decode_mtp.py index 08f5c4fd114..5fbe741745e 100644 --- a/flashinfer/gdn_kernels/gdn_decode_mtp.py +++ b/flashinfer/gdn_kernels/gdn_decode_mtp.py @@ -46,6 +46,9 @@ from cutlass.cute.runtime import from_dlpack import cuda.bindings.driver as cuda +from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel +from .cute_dsl_cache_naming import make_kernel_name + from .dtype_compat import as_bf16 # ============================================================================ @@ -2498,6 +2501,56 @@ def run_gdn_verify_kernel_mtp_inline( ) +_CUTE_DSL_MODULE = "gdn_decode_mtp" + + +def _mtp_kernel_name( + variant: str, + T: int, + H: int, + HV: int, + K: int, + V: int, + cache_steps: int, + disable_state_update: bool, + use_pool_indexing: bool, + pool_strides_key, + scale: float, + use_qk_l2norm: bool, + tile_v: int, + vec_size: int, + dtype_key: tuple, + ilp_rows: int = 4, + use_smem_v: bool = False, + use_packed_fma: bool = True, + per_token_pool_scatter: bool = False, +) -> str: + """Specialization name within the gdn_decode_mtp module, encoding the + kernel variant ("inline" or "warp") and every parameter that affects + codegen.""" + return make_kernel_name( + variant, + T, + H, + HV, + K, + V, + cache_steps, + disable_state_update, + use_pool_indexing, + pool_strides_key, + scale, + use_qk_l2norm, + tile_v, + vec_size, + dtype_key, + ilp_rows, + use_smem_v, + use_packed_fma, + per_token_pool_scatter, + ) + + @functools.cache def _get_compiled_mtp_kernel( T: int, @@ -2773,84 +2826,138 @@ def run_mtp_decode( ).mark_layout_dynamic() if use_inline_kernel: - compiled = cute.compile( - run_gdn_verify_kernel_mtp_inline, - h0_source_tensor, - intermediate_states_tensor, - A_log_tensor, - a_tensor, - dt_bias_tensor, - q_tensor, - k_tensor, - v_tensor, - b_tensor, - o_tensor, - h0_indices_tensor, - h0_out_indices_tensor, - cu_seqlens_tensor, - ssm_idx_tensor, - softplus_beta=1.0, - softplus_threshold=20.0, - scale=scale, - HV=HV, - T=T, - H=H, - K=K, - V=V, - tile_v=tile_v, - vec_size=vec_size, - use_initial_state=True, - use_qk_l2norm=use_qk_l2norm, - is_varlen=False, - disable_state_update=disable_state_update, - cache_intermediate_states=cutlass.Boolean(cache_intermediate_states), - use_pool_indexing=use_pool_indexing, - ilp_rows=ilp_rows, - use_smem_v=use_smem_v, - use_packed_fma=use_packed_fma, - per_token_pool_scatter=per_token_pool_scatter, - stream=stream, - options="--enable-tvm-ffi --generate-line-info", + compiled = build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _mtp_kernel_name( + "inline", + T, + H, + HV, + K, + V, + cache_steps, + disable_state_update, + use_pool_indexing, + pool_strides_key, + scale, + use_qk_l2norm, + tile_v, + vec_size, + dtype_key, + ilp_rows, + use_smem_v, + use_packed_fma, + per_token_pool_scatter, + ), + lambda: cute.compile( + run_gdn_verify_kernel_mtp_inline, + h0_source_tensor, + intermediate_states_tensor, + A_log_tensor, + a_tensor, + dt_bias_tensor, + q_tensor, + k_tensor, + v_tensor, + b_tensor, + o_tensor, + h0_indices_tensor, + h0_out_indices_tensor, + cu_seqlens_tensor, + ssm_idx_tensor, + softplus_beta=1.0, + softplus_threshold=20.0, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + is_varlen=False, + disable_state_update=disable_state_update, + cache_intermediate_states=cutlass.Boolean( + cache_intermediate_states + ), + use_pool_indexing=use_pool_indexing, + ilp_rows=ilp_rows, + use_smem_v=use_smem_v, + use_packed_fma=use_packed_fma, + per_token_pool_scatter=per_token_pool_scatter, + stream=stream, + options="--enable-tvm-ffi --generate-line-info", + ), + extra_key_files=(__file__,), ) else: - compiled = cute.compile( - run_gdn_verify_kernel_mtp, - h0_source_tensor, - intermediate_states_tensor, - A_log_tensor, - a_tensor, - dt_bias_tensor, - q_tensor, - k_tensor, - v_tensor, - b_tensor, - o_tensor, - h0_indices_tensor, - h0_out_indices_tensor, - cu_seqlens_tensor, - ssm_idx_tensor, - softplus_beta=1.0, - softplus_threshold=20.0, - scale=scale, - HV=HV, - T=T, - H=H, - K=K, - V=V, - tile_v=tile_v, - vec_size=vec_size, - use_initial_state=True, - use_qk_l2norm=use_qk_l2norm, - is_varlen=False, - disable_state_update=disable_state_update, - cache_intermediate_states=cutlass.Boolean(cache_intermediate_states), - use_pool_indexing=use_pool_indexing, - ilp_rows=ilp_rows, - use_smem_v=use_smem_v, - use_packed_fma=use_packed_fma, - per_token_pool_scatter=per_token_pool_scatter, - stream=stream, - options="--enable-tvm-ffi --generate-line-info", + compiled = build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _mtp_kernel_name( + "warp", + T, + H, + HV, + K, + V, + cache_steps, + disable_state_update, + use_pool_indexing, + pool_strides_key, + scale, + use_qk_l2norm, + tile_v, + vec_size, + dtype_key, + ilp_rows, + use_smem_v, + use_packed_fma, + per_token_pool_scatter, + ), + lambda: cute.compile( + run_gdn_verify_kernel_mtp, + h0_source_tensor, + intermediate_states_tensor, + A_log_tensor, + a_tensor, + dt_bias_tensor, + q_tensor, + k_tensor, + v_tensor, + b_tensor, + o_tensor, + h0_indices_tensor, + h0_out_indices_tensor, + cu_seqlens_tensor, + ssm_idx_tensor, + softplus_beta=1.0, + softplus_threshold=20.0, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + is_varlen=False, + disable_state_update=disable_state_update, + cache_intermediate_states=cutlass.Boolean( + cache_intermediate_states + ), + use_pool_indexing=use_pool_indexing, + ilp_rows=ilp_rows, + use_smem_v=use_smem_v, + use_packed_fma=use_packed_fma, + per_token_pool_scatter=per_token_pool_scatter, + stream=stream, + options="--enable-tvm-ffi --generate-line-info", + ), + extra_key_files=(__file__,), ) cache["compiled"] = compiled else: diff --git a/flashinfer/gdn_kernels/gdn_decode_nontranspose.py b/flashinfer/gdn_kernels/gdn_decode_nontranspose.py index 82f449f7f14..ebb2345c472 100644 --- a/flashinfer/gdn_kernels/gdn_decode_nontranspose.py +++ b/flashinfer/gdn_kernels/gdn_decode_nontranspose.py @@ -29,6 +29,9 @@ from cutlass.cute.runtime import from_dlpack import cuda.bindings.driver as cuda +from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel +from .cute_dsl_cache_naming import make_kernel_name + # ============================================================================ # Constants for NONTRANSPOSE version ([pool, HV, K, V]) # ============================================================================ @@ -675,6 +678,35 @@ def run_gdn_decode_kernel_big_batch_nontranspose( # ============================================================================ +_CUTE_DSL_MODULE = "gdn_decode_nontranspose" + + +def _nontranspose_kernel_name( + use_small_batch: bool, + T: int, + H: int, + HV: int, + K: int, + V: int, + dtype: torch.dtype, + scale: float, + use_qk_l2norm: bool, +) -> str: + """Specialization name within the gdn_decode_nontranspose module, encoding + every parameter that affects codegen.""" + return make_kernel_name( + "small" if use_small_batch else "big", + T, + H, + HV, + K, + V, + dtype, + scale, + use_qk_l2norm, + ) + + @functools.cache def _get_compiled_decode_kernel_nontranspose( use_small_batch: bool, @@ -763,31 +795,36 @@ def run_nontranspose_decode( ).mark_layout_dynamic() # Use TVM FFI to reduce runtime overhead - compiled = cute.compile( - run_func, - cu_seqlens_tensor, - q_tensor, - k_tensor, - v_tensor, - a_tensor, - b_tensor, - A_log_tensor, - dt_bias_tensor, - h0_source_tensor, - h0_indices_tensor, - o_tensor, - softplus_beta=1.0, - softplus_threshold=20.0, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - use_initial_state=True, - use_qk_l2norm=use_qk_l2norm, - stream=stream, - options="--enable-tvm-ffi", + compiled = build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _nontranspose_kernel_name(*cache_key), + lambda: cute.compile( + run_func, + cu_seqlens_tensor, + q_tensor, + k_tensor, + v_tensor, + a_tensor, + b_tensor, + A_log_tensor, + dt_bias_tensor, + h0_source_tensor, + h0_indices_tensor, + o_tensor, + softplus_beta=1.0, + softplus_threshold=20.0, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + stream=stream, + options="--enable-tvm-ffi", + ), + extra_key_files=(__file__,), ) cache["compiled"] = compiled else: diff --git a/flashinfer/gdn_kernels/gdn_decode_pretranspose.py b/flashinfer/gdn_kernels/gdn_decode_pretranspose.py index d3156beb515..e7fb27d9733 100644 --- a/flashinfer/gdn_kernels/gdn_decode_pretranspose.py +++ b/flashinfer/gdn_kernels/gdn_decode_pretranspose.py @@ -30,6 +30,9 @@ from cutlass.cute.runtime import from_dlpack import cuda.bindings.driver as cuda +from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel +from .cute_dsl_cache_naming import make_kernel_name + # ============================================================================ # Constants for PRETRANSPOSE version ([B*HV, V, K]) # ============================================================================ @@ -894,6 +897,42 @@ def run_gdn_decode_kernel_big_batch_pretranspose( # ============================================================================ +_CUTE_DSL_MODULE = "gdn_decode_pretranspose" + + +def _pretranspose_kernel_name( + T: int, + H: int, + HV: int, + K: int, + V: int, + dtype: torch.dtype, + scale: float, + use_qk_l2norm: bool, + use_pool_indexing: bool = False, + stride1: int = 0, + stride2: int = 0, + stride3: int = 0, +) -> str: + """Specialization name within the gdn_decode_pretranspose module, encoding + every parameter that affects codegen.""" + return make_kernel_name( + "decode", + T, + H, + HV, + K, + V, + dtype, + scale, + use_qk_l2norm, + use_pool_indexing, + stride1, + stride2, + stride3, + ) + + @functools.cache def _get_compiled_decode_kernel( T: int, @@ -1039,34 +1078,39 @@ def run_pretranspose_decode( run_func = run_gdn_decode_kernel_small_batch_pretranspose # Use TVM FFI to reduce runtime overhead - compiled = cute.compile( - run_func, - h0_source_tensor, - A_log_tensor, - a_tensor, - dt_bias_tensor, - q_tensor, - k_tensor, - v_tensor, - b_tensor, - o_tensor, - h0_indices_tensor, - h0_out_indices_tensor, - cu_seqlens_tensor, - softplus_beta=1.0, - softplus_threshold=20.0, - scale=scale, - HV=HV, - T=T, - H=H, - K=K, - V=V, - use_initial_state=True, - use_qk_l2norm=use_qk_l2norm, - use_pool_indexing=use_pool_indexing, - is_varlen=False, - stream=stream, - options="--enable-tvm-ffi", + compiled = build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + _pretranspose_kernel_name(*cache_key), + lambda: cute.compile( + run_func, + h0_source_tensor, + A_log_tensor, + a_tensor, + dt_bias_tensor, + q_tensor, + k_tensor, + v_tensor, + b_tensor, + o_tensor, + h0_indices_tensor, + h0_out_indices_tensor, + cu_seqlens_tensor, + softplus_beta=1.0, + softplus_threshold=20.0, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + use_pool_indexing=use_pool_indexing, + is_varlen=False, + stream=stream, + options="--enable-tvm-ffi", + ), + extra_key_files=(__file__,), ) cache["compiled"] = compiled else: diff --git a/tests/gdn/test_cute_dsl_kernel_cache.py b/tests/gdn/test_cute_dsl_kernel_cache.py new file mode 100644 index 00000000000..79bb2b2b861 --- /dev/null +++ b/tests/gdn/test_cute_dsl_kernel_cache.py @@ -0,0 +1,518 @@ +""" +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. + +Naming-contract tests for the GDN CuTe-DSL disk cache adopters. + +The specialization name is the sole per-kernel on-disk cache key (the module +meta.json guards arch / DSL version / source hash, NOT per-kernel codegen +parameters), so for every adopter the name must be a function of every +codegen argument. Pattern follows tests/jit/test_cute_dsl_cache.py: + +1. Signature coverage: every parameter of the @functools.cache'd kernel + getter appears in the name function's signature. +2. Per-argument perturbation: changing any single argument changes the name. +3. Symbol safety: names stay within [A-Za-z0-9_] (filename + TVM-FFI symbol). +""" + +import inspect +import re + +import pytest +import torch + +pytest.importorskip("cutlass") + +if not torch.cuda.is_available(): + pytest.skip( + "GDN kernel modules read device properties at import time", + allow_module_level=True, + ) + +from flashinfer.gdn_kernels.cute_dsl_cache_naming import ( # noqa: E402 + format_name_part, + make_kernel_name, +) +from flashinfer.gdn_kernels.gdn_decode_nontranspose import ( # noqa: E402 + _get_compiled_decode_kernel_nontranspose, + _nontranspose_kernel_name, +) +from flashinfer.gdn_kernels.gdn_decode_pretranspose import ( # noqa: E402 + _get_compiled_decode_kernel, + _pretranspose_kernel_name, +) +from flashinfer.gdn_kernels.gdn_decode_mtp import ( # noqa: E402 + _get_compiled_mtp_kernel, + _get_compiled_mtp_kernel_inline, + _mtp_kernel_name, +) +from flashinfer.gdn_kernels.gdn_decode_bf16_state import ( # noqa: E402 + _bf16_state_kernel_name, +) +from flashinfer.gdn_kernels.blackwell.gdn_prefill import ( # noqa: E402 + _get_compiled_cache, + _prefill_kernel_name, +) + +_SYMBOL_SAFE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") + + +# --------------------------------------------------------------------------- +# Shared name formatter +# --------------------------------------------------------------------------- + + +def test_format_name_part_symbol_safe_and_distinct(): + # One group per value kind: a cache-key slot keeps its type across calls, + # so distinctness is only required within a kind. (Deliberate aliases + # exist across kinds: bool 1/0 vs int, torch.dtype vs its str() form.) + groups = [ + [True, False], + [0, 1, -1, 128], + [1.0, 0.5, -0.5, 1e-6, 0.08838834764831845], + [torch.bfloat16, torch.float16, torch.float32, torch.int32], + ["v3_mtp_bf16_tiled_dynB", "mtp_bf16_dynB", "torch.bfloat16"], + [ + (8192, 16384, 128, 1), + (16384, 128, 1), + (torch.bfloat16, torch.float32, torch.int32), + (), + None, + ], + ] + for group in groups: + parts = [format_name_part(v) for v in group] + for v, p in zip(group, parts, strict=False): + assert re.fullmatch(r"[A-Za-z0-9_]*", p), f"{v!r} formatted to unsafe {p!r}" + assert len(set(parts)) == len(parts), (group, parts) + + +def test_make_kernel_name_caps_length_without_collision(): + long_a = make_kernel_name("v", *(range(200))) + long_b = make_kernel_name("v", *(list(range(199)) + [999])) + assert len(long_a) <= 210 + assert len(long_b) <= 210 + assert long_a != long_b + + +# --------------------------------------------------------------------------- +# Baselines (realistic decode shapes: 16 q/k heads, 32 v heads, head size 128, +# matching the shapes exercised by tests/gdn/test_decode_delta_rule.py) +# --------------------------------------------------------------------------- + +NONTRANSPOSE_BASELINE = { + "use_small_batch": True, + "T": 1, + "H": 16, + "HV": 32, + "K": 128, + "V": 128, + "dtype": torch.bfloat16, + "scale": 0.08838834764831845, + "use_qk_l2norm": True, +} + +PRETRANSPOSE_BASELINE = { + "T": 1, + "H": 16, + "HV": 32, + "K": 128, + "V": 128, + "dtype": torch.bfloat16, + "scale": 0.08838834764831845, + "use_qk_l2norm": True, + "use_pool_indexing": False, + "stride1": 0, + "stride2": 0, + "stride3": 0, +} + +MTP_BASELINE = { + "variant": "warp", + "T": 2, + "H": 16, + "HV": 32, + "K": 128, + "V": 128, + "cache_steps": 0, + "disable_state_update": False, + "use_pool_indexing": False, + "pool_strides_key": None, + "scale": 0.08838834764831845, + "use_qk_l2norm": True, + "tile_v": 64, + "vec_size": 8, + "dtype_key": (torch.float32, torch.float32, torch.int32), + "ilp_rows": 4, + "use_smem_v": False, + "use_packed_fma": True, + "per_token_pool_scatter": False, +} + +PREFILL_BASELINE = { + "io_dtype_str": "torch.bfloat16", + "state_dtype_str": "torch.float32", + "HQ": 32, + "HV": 16, + "is_GQA": True, + "use_initial_state": True, + "store_final_state": True, + "enable_checkpoints": False, + "use_state_indices": False, + "cu_seqlens_dtype_str": "torch.int32", + "state_indices_dtype_str": "none", + "cu_checkpoints_dtype_str": "none", + "initial_state_inner_strides": None, + "output_state_inner_strides": None, + "num_sm": 148, +} + +# In-process cache tuples as built at each gdn_decode_bf16_state call site. +BF16_STATE_BASELINES = { + "wide_vec": ( + "v3_mtp_bf16_tiled_dynB", + 2, # T + 16, # H + 32, # HV + 128, # K + 128, # V + -1, # pool_size_key + (-1,), # pool_slot_stride + 64, # tile_v + False, # effective_disable_final + False, # cache_intermediate_states + True, # use_qk_l2norm_in_kernel + 0.08838834764831845, # scale + 1.0, # softplus_beta + 20.0, # softplus_threshold + True, # use_packed_fma + True, # same_pool + False, # disable_output + 0, # recovery_steps + False, # per_request_accepted_steps + False, # per_token_pool_scatter + False, # per_token_pool_scatter_flat + (torch.float32, torch.float32, torch.int32), # _dtype_key + ), + "wide_vec_t1": ( + "v3_mtp_bf16_tiled_dynB", + 1, + 16, + 32, + 128, + 128, + -1, + (-1,), + 64, + False, + False, + True, + 0.08838834764831845, + 1.0, + 20.0, + True, + True, + (torch.float32, torch.float32, torch.int32), + ), + "mtp_ilp4": ( + "mtp_bf16_dynB", + 2, + 16, + 32, + 128, + 128, + -1, + (-1,), + 16, # tile_v + 4, # ilp_rows + False, # disable_state_update + False, # cache_intermediate_states + True, # use_qk_l2norm_in_kernel + 0.08838834764831845, + 1.0, + 20.0, + True, + True, + False, # disable_output + False, # per_request_accepted_steps + False, # per_token_pool_scatter + False, # per_token_pool_scatter_flat + (torch.float32, torch.float32, torch.int32), + ), +} + + +# --------------------------------------------------------------------------- +# 1. Signature coverage +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "getter,name_fn", + [ + pytest.param( + _get_compiled_decode_kernel_nontranspose, + _nontranspose_kernel_name, + id="nontranspose", + ), + pytest.param( + _get_compiled_decode_kernel, _pretranspose_kernel_name, id="pretranspose" + ), + pytest.param(_get_compiled_mtp_kernel, _mtp_kernel_name, id="mtp_warp"), + pytest.param( + _get_compiled_mtp_kernel_inline, _mtp_kernel_name, id="mtp_inline" + ), + pytest.param(_get_compiled_cache, _prefill_kernel_name, id="prefill"), + ], +) +def test_kernel_name_signature_covers_getter_params(getter, name_fn): + getter_params = set(inspect.signature(getter).parameters) + name_params = set(inspect.signature(name_fn).parameters) + missing = getter_params - name_params + assert not missing, ( + f"{getter.__name__} has codegen parameter(s) {sorted(missing)} that " + f"{name_fn.__name__} cannot encode; add them to the name function." + ) + + +# --------------------------------------------------------------------------- +# 2. Per-argument perturbation +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "param,alternate", + [ + ("use_small_batch", False), + ("T", 2), + ("H", 32), + ("HV", 64), + ("K", 64), + ("V", 64), + ("dtype", torch.float16), + ("scale", 0.0625), + ("use_qk_l2norm", False), + ], +) +def test_nontranspose_name_varies_with_every_argument(param, alternate): + baseline = _nontranspose_kernel_name(**NONTRANSPOSE_BASELINE) + kwargs = dict(NONTRANSPOSE_BASELINE) + kwargs[param] = alternate + assert _nontranspose_kernel_name(**kwargs) != baseline, ( + f"_nontranspose_kernel_name ignores {param!r}" + ) + + +@pytest.mark.parametrize( + "param,alternate", + [ + ("T", 2), + ("H", 32), + ("HV", 64), + ("K", 64), + ("V", 64), + ("dtype", torch.float16), + ("scale", 0.0625), + ("use_qk_l2norm", False), + ("use_pool_indexing", True), + ("stride1", 16384), + ("stride2", 128), + ("stride3", 1), + ], +) +def test_pretranspose_name_varies_with_every_argument(param, alternate): + baseline = _pretranspose_kernel_name(**PRETRANSPOSE_BASELINE) + kwargs = dict(PRETRANSPOSE_BASELINE) + kwargs[param] = alternate + assert _pretranspose_kernel_name(**kwargs) != baseline, ( + f"_pretranspose_kernel_name ignores {param!r}" + ) + + +@pytest.mark.parametrize( + "param,alternate", + [ + ("variant", "inline"), + ("T", 3), + ("H", 32), + ("HV", 64), + ("K", 64), + ("V", 64), + ("cache_steps", 2), + ("disable_state_update", True), + ("use_pool_indexing", True), + ("pool_strides_key", (262144, 16384, 128, 1)), + ("scale", 0.0625), + ("use_qk_l2norm", False), + ("tile_v", 128), + ("vec_size", 4), + ("dtype_key", (torch.bfloat16, torch.float32, torch.int32)), + ("ilp_rows", 8), + ("use_smem_v", True), + ("use_packed_fma", False), + ("per_token_pool_scatter", True), + ], +) +def test_mtp_name_varies_with_every_argument(param, alternate): + baseline = _mtp_kernel_name(**MTP_BASELINE) + kwargs = dict(MTP_BASELINE) + kwargs[param] = alternate + assert _mtp_kernel_name(**kwargs) != baseline, f"_mtp_kernel_name ignores {param!r}" + + +@pytest.mark.parametrize( + "param,alternate", + [ + ("io_dtype_str", "torch.float16"), + ("state_dtype_str", "torch.bfloat16"), + ("HQ", 64), + ("HV", 32), + ("is_GQA", False), + ("use_initial_state", False), + ("store_final_state", False), + ("enable_checkpoints", True), + ("use_state_indices", True), + ("cu_seqlens_dtype_str", "torch.int64"), + ("state_indices_dtype_str", "torch.int32"), + ("cu_checkpoints_dtype_str", "torch.int32"), + ("initial_state_inner_strides", (16384, 128, 1)), + ("output_state_inner_strides", (16384, 128, 1)), + ("num_sm", 132), + ], +) +def test_prefill_name_varies_with_every_argument(param, alternate): + baseline = _prefill_kernel_name(**PREFILL_BASELINE) + kwargs = dict(PREFILL_BASELINE) + kwargs[param] = alternate + assert _prefill_kernel_name(**kwargs) != baseline, ( + f"_prefill_kernel_name ignores {param!r}" + ) + + +def _perturb(value): + """Return a same-type value that must map to a different name fragment.""" + if isinstance(value, bool): + return not value + if isinstance(value, int): + return value + 1 + if isinstance(value, float): + return value * 2.0 + 1.0 + if isinstance(value, str): + return value + "_alt" + if isinstance(value, tuple): + if not value: + return (1,) + if isinstance(value[0], torch.dtype): + swapped = torch.float16 if value[0] != torch.float16 else torch.bfloat16 + return (swapped,) + value[1:] + return (_perturb(value[0]),) + value[1:] + raise TypeError(f"unhandled baseline component {value!r}") + + +@pytest.mark.parametrize("variant", sorted(BF16_STATE_BASELINES)) +def test_bf16_state_name_varies_with_every_key_component(variant): + key = BF16_STATE_BASELINES[variant] + baseline = _bf16_state_kernel_name(variant, key) + for i in range(len(key)): + perturbed = key[:i] + (_perturb(key[i]),) + key[i + 1 :] + assert _bf16_state_kernel_name(variant, perturbed) != baseline, ( + f"_bf16_state_kernel_name ignores cache_key[{i}] = {key[i]!r} " + f"for variant {variant!r}" + ) + + +def test_bf16_state_name_distinguishes_variants(): + names = { + variant: _bf16_state_kernel_name(variant, key) + for variant, key in BF16_STATE_BASELINES.items() + } + assert len(set(names.values())) == len(names), names + + +# --------------------------------------------------------------------------- +# 3. Symbol safety +# --------------------------------------------------------------------------- + + +def test_all_baseline_names_are_symbol_safe(): + names = [ + _nontranspose_kernel_name(**NONTRANSPOSE_BASELINE), + _pretranspose_kernel_name(**PRETRANSPOSE_BASELINE), + _mtp_kernel_name(**MTP_BASELINE), + _prefill_kernel_name(**PREFILL_BASELINE), + *( + _bf16_state_kernel_name(variant, key) + for variant, key in BF16_STATE_BASELINES.items() + ), + _mtp_kernel_name( + **{**MTP_BASELINE, "pool_strides_key": (262144, 16384, 128, 1)} + ), + _prefill_kernel_name( + **{**PREFILL_BASELINE, "initial_state_inner_strides": (16384, 128, 1)} + ), + ] + for name in names: + assert _SYMBOL_SAFE.fullmatch(name), f"unsafe specialization name {name!r}" + assert len(name) <= 210, f"specialization name too long: {name!r}" + + +# --------------------------------------------------------------------------- +# 4. Disk-cache round trip +# --------------------------------------------------------------------------- + + +def test_nontranspose_disk_cache_round_trip(monkeypatch, tmp_path): + """A second process must reload the exported artifact instead of recompiling. + + Simulated in-process: clear the in-process cache, then forbid cute.compile + and re-run the same specialization against the populated disk cache. + """ + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("nontranspose decode kernel requires SM90 or later") + + import flashinfer.jit.env as jit_env + import flashinfer.gdn_kernels.gdn_decode_nontranspose as nt_mod + from flashinfer.gdn_decode import gated_delta_rule_decode + + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.setattr(jit_env, "FLASHINFER_JIT_DIR", tmp_path) + nt_mod._get_compiled_decode_kernel_nontranspose.cache_clear() + + torch.manual_seed(0) + B, H, HV, D = 2, 16, 32, 128 + dev = torch.device("cuda") + q = torch.randn(B, 1, H, D, dtype=torch.bfloat16, device=dev) + k = torch.nn.functional.normalize( + torch.randn(B, 1, H, D, dtype=torch.bfloat16, device=dev), p=2.0, dim=-1 + ) + v = torch.randn(B, 1, HV, D, dtype=torch.bfloat16, device=dev) + a = torch.randn(B, 1, HV, dtype=torch.bfloat16, device=dev) * 0.1 + b = torch.randn(B, 1, HV, dtype=torch.bfloat16, device=dev) + A_log = torch.randn(HV, dtype=torch.float32, device=dev) * 0.1 + dt_bias = torch.rand(HV, dtype=torch.float32, device=dev) + state = torch.randn(B, HV, D, D, dtype=torch.float32, device=dev) + + out1, state1 = gated_delta_rule_decode(q, k, v, state.clone(), A_log, a, dt_bias, b) + artifacts = list(tmp_path.glob("gdn_decode_nontranspose_*_cute_dsl/*.o")) + assert len(artifacts) == 1, f"expected one exported artifact, got {artifacts}" + + nt_mod._get_compiled_decode_kernel_nontranspose.cache_clear() + + def _no_recompile(*args, **kwargs): + raise AssertionError("cute.compile ran despite a valid disk artifact") + + monkeypatch.setattr(nt_mod.cute, "compile", _no_recompile) + out2, state2 = gated_delta_rule_decode(q, k, v, state.clone(), A_log, a, dt_bias, b) + torch.testing.assert_close(out1, out2, atol=0, rtol=0) + torch.testing.assert_close(state1, state2, atol=0, rtol=0) From ed98c5a7661daa3fa9e6186e11452484de599ae6 Mon Sep 17 00:00:00 2001 From: leeyongjun Date: Mon, 7 Sep 2026 08:57:55 +0000 Subject: [PATCH 2/2] test(gdn): pin the disk cache off in compile-count tests --- tests/gdn/test_decode_delta_rule.py | 3 +++ tests/gdn/test_decode_pretranspose_noncontiguous_pool.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tests/gdn/test_decode_delta_rule.py b/tests/gdn/test_decode_delta_rule.py index 1e155ed24f0..1fb1d663de6 100644 --- a/tests/gdn/test_decode_delta_rule.py +++ b/tests/gdn/test_decode_delta_rule.py @@ -1282,6 +1282,9 @@ def counted_compile(*args, **kwargs): compile_count += 1 return original_compile(*args, **kwargs) + # Pin the disk cache off: a populated cache would satisfy the reuse + # property with zero compiles, breaking the count-based assertion. + monkeypatch.setenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "1") gdn_decode_mtp._get_compiled_mtp_kernel.cache_clear() gdn_decode_mtp._get_compiled_mtp_kernel_inline.cache_clear() monkeypatch.setattr(cute, "compile", counted_compile) diff --git a/tests/gdn/test_decode_pretranspose_noncontiguous_pool.py b/tests/gdn/test_decode_pretranspose_noncontiguous_pool.py index 7f40e602537..7664886061d 100644 --- a/tests/gdn/test_decode_pretranspose_noncontiguous_pool.py +++ b/tests/gdn/test_decode_pretranspose_noncontiguous_pool.py @@ -165,6 +165,9 @@ def counted_compile(*args, **kwargs): pool_compile_calls += 1 return original_compile(*args, **kwargs) + # Pin the disk cache off: a populated cache would satisfy the reuse + # property with zero compiles, breaking the count-based assertion. + monkeypatch.setenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "1") pretranspose_module._get_compiled_decode_kernel.cache_clear() monkeypatch.setattr(pretranspose_module.cute, "compile", counted_compile)