diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py index 6430799305d..c8ed015cf91 100755 --- a/flashinfer/gemm/gemm_base.py +++ b/flashinfer/gemm/gemm_base.py @@ -53,6 +53,11 @@ get_hybrid_num_tokens_buckets, map_to_hybrid_bucket_uncapped, ) +from .gemm_mm_fp4_cute_dsl import ( + _compile_block_scaled_gemm, + _mm_fp4_cache_key, + precompile_mm_fp4_tactics, +) from .kernels.utils import ( _SM100_CLUSTER_SHAPE_MN_CANDIDATES, _SM100_MMA_TILER_MN_CANDIDATES, @@ -60,6 +65,7 @@ _select_sm100_mm_fp4_cute_dsl_tactic, ) from ..utils import ( + get_device_index, get_device_sm_count, get_native_fp4_dtype, is_sm100a_supported, @@ -127,11 +133,6 @@ # Error messages CUDNN_FP4_MXFP4_SM120_CUDNN_VERSION_ERROR = "cudnn FP4 GEMM with mxfp4 quantization is not supported on SM120/SM121 with cuDNN backend version < 9.14.0." -_TORCH_TO_CUTLASS_DTYPE_ATTR = { - torch.bfloat16: "BFloat16", - torch.float16: "Float16", -} - def _match_sm_version(device: torch.device, sm_version: list[str]): major, minor = get_compute_capability(device) @@ -4683,113 +4684,6 @@ def _get_sm100_block_scaled_tactics( return valid_tactics -def _compile_block_scaled_gemm( - cache, - cache_key, - make_gemm_kernel, - ab_cutlass_dtype, - sf_dtype, - c_cutlass_dtype, - ab_assumed_align, - cluster_shape_mn, - swap_ab, - sf_m, - sf_n, - sf_k, - batch_size, - cluster_shape_k=1, -): - """Compile a block-scaled GEMM kernel via CuTe DSL and cache it. - - ``make_gemm_kernel`` is a zero-arg callable that returns a kernel instance - (Sm100 or Sm103). It is only invoked on a cache miss. - - TVM-FFI compilation pattern: - - A, B, C, alpha: make_fake_compact_tensor -> torch tensors - passed directly at runtime via TVM-FFI C-level dlpack - - SF tensors: make_ptr (complex 6D BlockScaledBasicChunk - layout can't be expressed as torch tensor) -> data_ptr() at runtime - - Stream: make_fake_stream -> automatic env stream at runtime - - For FP4 runners, ``ab_cutlass_dtype`` is ``Uint8`` because FP4 data is - stored as uint8 in torch (2 FP4 values per byte); the kernel wrapper - recasts from Uint8 to Float4E2M1FN internally. - """ - if cache_key in cache: - return cache[cache_key] - - import cutlass - import cutlass.cute as cute - - from cutlass.cute.runtime import make_ptr - from flashinfer.cute_dsl.utils import get_max_active_clusters - - gemm = make_gemm_kernel() - - sym_m = cute.sym_int() - sym_k = cute.sym_int() - sym_n = cute.sym_int() - - a_fake = cute.runtime.make_fake_compact_tensor( - ab_cutlass_dtype, - (sym_m, sym_k), - stride_order=(1, 0), - assumed_align=ab_assumed_align, - ) - b_fake = cute.runtime.make_fake_compact_tensor( - ab_cutlass_dtype, - (sym_n, sym_k), - stride_order=(1, 0), - assumed_align=ab_assumed_align, - ) - if swap_ab: - c_fake = cute.runtime.make_fake_compact_tensor( - c_cutlass_dtype, - (sym_n, sym_m), - stride_order=(0, 1), - assumed_align=16, - ) - else: - c_fake = cute.runtime.make_fake_compact_tensor( - c_cutlass_dtype, - (sym_m, sym_n), - stride_order=(1, 0), - assumed_align=16, - ) - - a_sf_ptr = make_ptr(sf_dtype, 16, cute.AddressSpace.gmem, 16) - b_sf_ptr = make_ptr(sf_dtype, 16, cute.AddressSpace.gmem, 16) - alpha_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, (1,), assumed_align=4 - ) - - launch_cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] * cluster_shape_k - max_active_clusters = get_max_active_clusters(launch_cluster_size) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - - compiled_gemm = cute.compile( - gemm.wrapper, - a_fake, - b_fake, - c_fake, - sf_m, - sf_n, - sf_k, - batch_size, - a_sf_ptr, - b_sf_ptr, - alpha_fake, - max_active_clusters, - stream_fake, - swap_ab, - options="--opt-level 2 --enable-tvm-ffi", - ) - - result = (compiled_gemm, max_active_clusters) - cache[cache_key] = result - return result - - _CUTE_DSL_ALPHA_ONE_CACHE: dict = {} @@ -4846,13 +4740,14 @@ def _cute_dsl_gemm_mxfp8_runner( "Supported: torch.bfloat16, torch.float16." ) - cutlass_dtype_attr = _TORCH_TO_CUTLASS_DTYPE_ATTR.get(out_dtype) - if cutlass_dtype_attr is None: + from ..cute_dsl.utils import torch_to_cutlass_dtype + + if out_dtype not in (torch.bfloat16, torch.float16): raise ValueError( f"cute_dsl mm_mxfp8 does not support output dtype {out_dtype}. " "Supported: torch.bfloat16, torch.float16." ) - c_cutlass_dtype = getattr(cutlass, cutlass_dtype_attr) + c_cutlass_dtype = torch_to_cutlass_dtype(out_dtype) _ = sm_major, sm_minor class CuteDSLMxfp8GemmRunner(TunableRunner): @@ -5918,15 +5813,14 @@ def _cute_dsl_gemm_fp4_runner( # except ImportError: # pass - cutlass_dtype_attr = _TORCH_TO_CUTLASS_DTYPE_ATTR.get(out_dtype) - c_cutlass_dtype = ( - getattr(cutlass, cutlass_dtype_attr) if cutlass_dtype_attr is not None else None - ) - if c_cutlass_dtype is None: + from ..cute_dsl.utils import torch_to_cutlass_dtype + + if out_dtype not in (torch.bfloat16, torch.float16): raise ValueError( f"cute_dsl backend does not support output dtype {out_dtype}. " f"Supported: torch.bfloat16, torch.float16." ) + c_cutlass_dtype = torch_to_cutlass_dtype(out_dtype) class CuteDSLFp4GemmRunner(TunableRunner): """TunableRunner for CuTe DSL block-scaled FP4 dense GEMM. @@ -6078,6 +5972,26 @@ def forward( sf_dtype = cutlass.Float8E4M3FN if use_nvfp4 else cutlass.Float8E8M0FNU batch_size = 1 + if do_preparation: + try: + precompile_mm_fp4_tactics( + self.get_valid_tactics(inputs, None), + m, + n, + real_k, + use_nvfp4, + enable_pdl, + out_dtype, + _CUTE_DSL_MM_FP4_KERNEL_CACHE, + a.device, + ) + except Exception as e: # noqa: BLE001 -- serial fallback is intentional + logger.warning( + f"[mm_fp4 cute-dsl] tactic precompilation failed " + f"({type(e).__name__}: {e}); tactics will compile " + f"serially during profiling." + ) + if tactic is None or tactic == -1: # Use analytical heuristic to pick the best tactic based on # tile and wave quantization efficiency. @@ -6112,17 +6026,7 @@ def forward( sf_n = (kernel_n + 127) // 128 sf_k = (real_k // sf_vec_size + 3) // 4 - cache_key = ( - sf_vec_size, - mma_tiler_mn, - cluster_shape_mn, - swap_ab, - use_prefetch, - kernel_type, - use_tma_store, - enable_pdl, - out_dtype, - ) + cache_key = _mm_fp4_cache_key(sf_vec_size, tactic, enable_pdl, out_dtype) if kernel_type == "sm103" and Sm103Kernel is not None: make_kernel = lambda: Sm103Kernel( @@ -6155,6 +6059,8 @@ def forward( sf_n=sf_n, sf_k=sf_k, batch_size=batch_size, + cache_module_name="mm_fp4", + device_index=get_device_index(a.device), ) alpha_for_launch = _prepare_alpha_for_launch(alpha_tensor, a.device) @@ -6205,15 +6111,14 @@ def _b12x_gemm_fp4_runner( _select_default_dense_gemm_plan, ) - cutlass_dtype_attr = _TORCH_TO_CUTLASS_DTYPE_ATTR.get(out_dtype) - c_cutlass_dtype = ( - getattr(cutlass, cutlass_dtype_attr) if cutlass_dtype_attr is not None else None - ) - if c_cutlass_dtype is None: + from ..cute_dsl.utils import torch_to_cutlass_dtype + + if out_dtype not in (torch.bfloat16, torch.float16): raise ValueError( f"b12x backend does not support output dtype {out_dtype}. " f"Supported: torch.bfloat16, torch.float16." ) + c_cutlass_dtype = torch_to_cutlass_dtype(out_dtype) def _default_dense_plan(m, n, real_k, device): return _select_default_dense_gemm_plan( diff --git a/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py b/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py index 51bb135e96f..40733f673db 100644 --- a/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py +++ b/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py @@ -18,7 +18,7 @@ get_hybrid_num_tokens_buckets, map_to_hybrid_bucket_uncapped, ) -from .gemm_base import _TORCH_TO_CUTLASS_DTYPE_ATTR, _check_cute_dsl_availability +from .gemm_base import _check_cute_dsl_availability from .gemm_bf16_fp4 import _unswizzle_sf_128x4 _BF16_FP4_ALPHA_ONE_CACHE: dict = {} @@ -120,8 +120,10 @@ def _get_cute_dsl_bf16_fp4_gemm( BlackwellDenseGemmBf16Fp4Kernel, ) - a_cutlass_dtype = getattr(cutlass, _TORCH_TO_CUTLASS_DTYPE_ATTR[a_dtype]) - c_cutlass_dtype = getattr(cutlass, _TORCH_TO_CUTLASS_DTYPE_ATTR[c_dtype]) + from ..cute_dsl.utils import torch_to_cutlass_dtype + + a_cutlass_dtype = torch_to_cutlass_dtype(a_dtype) + c_cutlass_dtype = torch_to_cutlass_dtype(c_dtype) sym_m = cute.sym_int() sym_k = cute.sym_int() diff --git a/flashinfer/gemm/gemm_mm_fp4_cute_dsl.py b/flashinfer/gemm/gemm_mm_fp4_cute_dsl.py new file mode 100644 index 00000000000..f930c7167de --- /dev/null +++ b/flashinfer/gemm/gemm_mm_fp4_cute_dsl.py @@ -0,0 +1,477 @@ +""" +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 logging +import os +from typing import Optional + +import torch + +from ..utils import get_device_index + +logger = logging.getLogger(__name__) + + +def _blockscaled_gemm_cache_key_files() -> tuple: + """Source files whose content invalidates the on-disk mm_fp4 kernels.""" + from .kernels import ( + dense_blockscaled_gemm_sm100, + dense_blockscaled_gemm_sm100_common, + dense_blockscaled_gemm_sm103, + ) + + return ( + __file__, + dense_blockscaled_gemm_sm100.__file__, + dense_blockscaled_gemm_sm100_common.__file__, + dense_blockscaled_gemm_sm103.__file__, + ) + + +def _compile_block_scaled_gemm( + cache, + cache_key, + make_gemm_kernel, + ab_cutlass_dtype, + sf_dtype, + c_cutlass_dtype, + ab_assumed_align, + cluster_shape_mn, + swap_ab, + sf_m, + sf_n, + sf_k, + batch_size, + cluster_shape_k=1, + cache_module_name=None, + device_index=None, +): + """Compile a block-scaled GEMM kernel via CuTe DSL and cache it. + + ``make_gemm_kernel`` is a zero-arg callable that returns a kernel instance + (Sm100 or Sm103). It is only invoked on a cache miss. + + TVM-FFI compilation pattern: + - A, B, C, alpha: make_fake_compact_tensor -> torch tensors + passed directly at runtime via TVM-FFI C-level dlpack + - SF tensors: make_ptr (complex 6D BlockScaledBasicChunk + layout can't be expressed as torch tensor) -> data_ptr() at runtime + - Stream: make_fake_stream -> automatic env stream at runtime + + For FP4 runners, ``ab_cutlass_dtype`` is ``Uint8`` because FP4 data is + stored as uint8 in torch (2 FP4 values per byte); the kernel wrapper + recasts from Uint8 to Float4E2M1FN internally. + """ + if device_index is None: + device_index = torch.cuda.current_device() + mem_key = (device_index, cache_key) + if mem_key in cache: + return cache[mem_key] + + from flashinfer.cute_dsl.utils import get_max_active_clusters + + gemm = make_gemm_kernel() + + launch_cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] * cluster_shape_k + max_active_clusters = get_max_active_clusters(launch_cluster_size) + + compile_kernel = _make_blockscaled_gemm_compile_fn( + gemm, + ab_cutlass_dtype=ab_cutlass_dtype, + sf_dtype=sf_dtype, + c_cutlass_dtype=c_cutlass_dtype, + ab_assumed_align=ab_assumed_align, + swap_ab=swap_ab, + sf_m=sf_m, + sf_n=sf_n, + sf_k=sf_k, + batch_size=batch_size, + max_active_clusters=max_active_clusters, + ) + + if cache_module_name is None: + compiled_gemm = compile_kernel() + else: + from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel + + compiled_gemm = build_and_load_cute_dsl_kernel( + cache_module_name, + _blockscaled_kernel_disk_name(cache_key, batch_size, max_active_clusters), + compile_kernel, + extra_key_files=_blockscaled_gemm_cache_key_files(), + ) + + result = (compiled_gemm, max_active_clusters) + cache[mem_key] = result + return result + + +def _blockscaled_kernel_disk_name(cache_key, batch_size, max_active_clusters): + """On-disk kernel name encoding every mm_fp4 codegen parameter. + + Must be symbol-safe as produced (see tests/jit/test_cute_dsl_cache.py): + JitSpecCuteDsl sanitizes names, and two names differing only in + sanitized-away characters would collide on one artifact. + """ + ( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + swap_ab, + use_prefetch, + kernel_type, + use_tma_store, + enable_pdl, + out_dtype, + ) = cache_key + tma = "x" if use_tma_store is None else int(use_tma_store) + dtype = str(out_dtype).removeprefix("torch.") + return ( + f"sf{sf_vec_size}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_c{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + f"_swap{int(swap_ab)}_pf{int(use_prefetch)}_{kernel_type}" + f"_tma{tma}_pdl{int(enable_pdl)}_{dtype}" + f"_b{batch_size}_mac{max_active_clusters}" + ) + + +def _make_blockscaled_gemm_compile_fn( + gemm, + ab_cutlass_dtype, + sf_dtype, + c_cutlass_dtype, + ab_assumed_align, + swap_ab, + sf_m, + sf_n, + sf_k, + batch_size, + max_active_clusters, +): + """Build a zero-arg closure that runs ``cute.compile`` for gemm.""" + import cutlass + import cutlass.cute as cute + + from cutlass.cute.runtime import make_ptr + + def compile_kernel(): + sym_m = cute.sym_int() + sym_k = cute.sym_int() + sym_n = cute.sym_int() + + a_fake = cute.runtime.make_fake_compact_tensor( + ab_cutlass_dtype, + (sym_m, sym_k), + stride_order=(1, 0), + assumed_align=ab_assumed_align, + ) + b_fake = cute.runtime.make_fake_compact_tensor( + ab_cutlass_dtype, + (sym_n, sym_k), + stride_order=(1, 0), + assumed_align=ab_assumed_align, + ) + if swap_ab: + c_fake = cute.runtime.make_fake_compact_tensor( + c_cutlass_dtype, + (sym_n, sym_m), + stride_order=(0, 1), + assumed_align=16, + ) + else: + c_fake = cute.runtime.make_fake_compact_tensor( + c_cutlass_dtype, + (sym_m, sym_n), + stride_order=(1, 0), + assumed_align=16, + ) + + a_sf_ptr = make_ptr(sf_dtype, 16, cute.AddressSpace.gmem, 16) + b_sf_ptr = make_ptr(sf_dtype, 16, cute.AddressSpace.gmem, 16) + alpha_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (1,), assumed_align=4 + ) + + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + return cute.compile( + gemm.wrapper, + a_fake, + b_fake, + c_fake, + sf_m, + sf_n, + sf_k, + batch_size, + a_sf_ptr, + b_sf_ptr, + alpha_fake, + max_active_clusters, + stream_fake, + swap_ab, + options="--opt-level 2 --enable-tvm-ffi", + ) + + return compile_kernel + + +def _mm_fp4_precompile_worker(payload): + """Compile one mm_fp4 tactic in a spawned subprocess and persist it to + the on-disk CuTe-DSL kernel cache. + + Returns ``(kernel_name, None)`` on success or ``(kernel_name, error)``; + the parent logs failures and lets ``forward`` compile those tactics + in-process on demand. + """ + kernel_name = payload["kernel_name"] + try: + # Set the current GPU in this subprocess + torch.cuda.set_device(payload["device_index"]) + + import cutlass + + from ..cute_dsl.utils import torch_to_cutlass_dtype + from ..jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files + from .kernels.dense_blockscaled_gemm_sm100 import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + ( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + swap_ab, + use_prefetch, + _kernel_type, + _use_tma_store, + enable_pdl, + out_dtype, + ) = payload["cache_key"] + + gemm = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + use_prefetch, + enable_pdl, + ) + compile_fn = _make_blockscaled_gemm_compile_fn( + gemm, + ab_cutlass_dtype=cutlass.Uint8, + sf_dtype=cutlass.Float8E4M3FN + if sf_vec_size == 16 + else cutlass.Float8E8M0FNU, + c_cutlass_dtype=torch_to_cutlass_dtype(out_dtype), + ab_assumed_align=32, + swap_ab=swap_ab, + sf_m=payload["sf_m"], + sf_n=payload["sf_n"], + sf_k=payload["sf_k"], + batch_size=payload["batch_size"], + max_active_clusters=payload["max_active_clusters"], + ) + spec = JitSpecCuteDsl( + "mm_fp4", + kernel_name, + compile_fn, + _hash_source_files(tuple(payload["key_files"])), + ) + if not spec.is_compiled: + spec.compile_and_persist() + return (kernel_name, None) + except Exception as e: # noqa: BLE001 -- reported to the parent + return (kernel_name, f"{type(e).__name__}: {e}") + + +# Empirically measured host-RAM budget (RSS) per precompile worker -- 1 GiB +_MM_FP4_PRECOMPILE_WORKER_RAM_BYTES = 1 << 30 + + +def _cgroup_available_memory_bytes() -> Optional[int]: + """Headroom under the current cgroup memory limit, or None if + unlimited/unreadable. + + Inside containers /proc/meminfo reports *host* memory, so the cgroup + limit is what actually prevents an OOM kill. + cgroup's memory.current is profiled here to estimate the headroom. + """ + try: + with open("/sys/fs/cgroup/memory.max") as f: + limit = f.read().strip() + if limit != "max": + with open("/sys/fs/cgroup/memory.current") as f: + current = int(f.read()) + return max(0, int(limit) - current) + except (OSError, ValueError): + pass + # cgroup v1 + try: + with open("/sys/fs/cgroup/memory/memory.limit_in_bytes") as f: + limit_v1 = int(f.read()) + if limit_v1 < 1 << 60: # values near 2**63 mean "unlimited" + with open("/sys/fs/cgroup/memory/memory.usage_in_bytes") as f: + usage = int(f.read()) + return max(0, limit_v1 - usage) + except (OSError, ValueError): + pass + return None + + +def _available_host_memory_bytes() -> Optional[int]: + """Best-effort available memory: the tighter of host MemAvailable and + the cgroup limit headroom, or None if neither is readable.""" + candidates = [] + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + candidates.append(int(line.split()[1]) * 1024) + break + except OSError: + pass + cgroup = _cgroup_available_memory_bytes() + if cgroup is not None: + candidates.append(cgroup) + return min(candidates) if candidates else None + + +def _get_mm_fp4_cute_dsl_compile_workers() -> int: + """How many subprocesses to use for precompiling mm_fp4 cute-dsl tactics. + + Starts from FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS (default 4), then + lowers it to what host RAM can hold using 1 GiB per spawned worker as a safety measure. + """ + workers = int(os.environ.get("FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS", "4")) + if workers <= 1: + return workers + available = _available_host_memory_bytes() + if available is not None: + mem_cap = int(available // _MM_FP4_PRECOMPILE_WORKER_RAM_BYTES) + if mem_cap < workers: + logger.warning( + f"[mm_fp4 cute-dsl] capping tactic precompile workers " + f"{workers} -> {mem_cap} (host MemAvailable = " + f"{available / (1 << 30):.1f} GiB)." + ) + workers = mem_cap + return workers + + +def _run_mm_fp4_precompile_pool(payloads) -> None: + """Compile mm_fp4 tactics into the on-disk cache with a subprocess pool. + + Each subprocess runs _mm_fp4_precompile_worker. + """ + from multiprocessing import get_context + + num_workers = min(_get_mm_fp4_cute_dsl_compile_workers(), len(payloads)) + logger.info( + f"[mm_fp4 cute-dsl] precompiling {len(payloads)} tactics " + f"with {num_workers} workers" + ) + with get_context("spawn").Pool(num_workers) as pool: + results = pool.map(_mm_fp4_precompile_worker, payloads) + for kernel_name, err in results: + if err is not None: + logger.debug( + f"[mm_fp4 cute-dsl] precompile failed for {kernel_name}: {err}" + ) + + +def _mm_fp4_cache_key(sf_vec_size, tactic, enable_pdl, out_dtype): + """In-memory kernel-cache key for one mm_fp4 tactic tuple. + + Shared by the runner's forward path and the precompile path, which + must agree byte-for-byte: the on-disk kernel name derives from it. + """ + return (sf_vec_size, *tactic, enable_pdl, out_dtype) + + +def precompile_mm_fp4_tactics( + tactics, m, n, real_k, use_nvfp4, enable_pdl, out_dtype, kernel_cache, device +) -> None: + """Batch-compile not-yet-cached mm_fp4 tactics into the on-disk + CuTe-DSL cache with a pool of subprocesses. + + Called by autotuner's do_preparation, before the per-tactic profiling loop. + Failures are non-fatal: any tactic missing from kernel_cache and + the disk cache compiles in-process on first use. + """ + from ..jit.cute_dsl_core import ( + JitSpecCuteDsl, + _hash_source_files, + cute_dsl_cache_disabled, + ) + + if cute_dsl_cache_disabled(): + return # workers hand kernels to the parent via the disk cache + if _get_mm_fp4_cute_dsl_compile_workers() <= 1: + return + + from flashinfer.cute_dsl.utils import get_max_active_clusters + + sf_vec_size = 16 if use_nvfp4 else 32 + device_index = get_device_index(device) + key_files = _blockscaled_gemm_cache_key_files() + source_sha256 = _hash_source_files(tuple(key_files)) + + max_clusters_cache: dict = {} + payloads = [] + for tactic in tactics: + ( + mma_tiler_mn, + cluster_shape_mn, + swap_ab, + _use_prefetch, + kernel_type, + _use_tma_store, + ) = tactic + if kernel_type != "sm100": + continue + cache_key = _mm_fp4_cache_key(sf_vec_size, tactic, enable_pdl, out_dtype) + if (device_index, cache_key) in kernel_cache: + continue + + cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] + if cluster_size not in max_clusters_cache: + max_clusters_cache[cluster_size] = get_max_active_clusters(cluster_size) + mac = max_clusters_cache[cluster_size] + kernel_name = _blockscaled_kernel_disk_name(cache_key, 1, mac) + spec = JitSpecCuteDsl("mm_fp4", kernel_name, lambda: None, source_sha256) + if spec.is_compiled: + continue # already on disk; forward will JITLink it + + kernel_m, kernel_n = (n, m) if swap_ab else (m, n) + payloads.append( + { + "cache_key": cache_key, + "kernel_name": kernel_name, + "max_active_clusters": mac, + "sf_m": (kernel_m + 127) // 128, + "sf_n": (kernel_n + 127) // 128, + "sf_k": (real_k // sf_vec_size + 3) // 4, + "batch_size": 1, + "key_files": key_files, + "device_index": device_index, + } + ) + + # A single missing tactic compiles faster in-process than the + # spawn + import cost of a one-worker pool. + if len(payloads) < 2: + return + + _run_mm_fp4_precompile_pool(payloads) diff --git a/flashinfer/jit/cute_dsl_core.py b/flashinfer/jit/cute_dsl_core.py index 44cbc18691e..4398c462ee6 100644 --- a/flashinfer/jit/cute_dsl_core.py +++ b/flashinfer/jit/cute_dsl_core.py @@ -184,6 +184,36 @@ def build(self) -> None: f"{self.object_path}: {e}. The kernel will be recompiled next run." ) + def compile_and_persist(self) -> None: + """Compile outside the module lock; commit the artifact under it. + + ``build_and_load()`` holds the module lock for the whole build. + This variant runs ``compile_fn`` unlocked and takes the lock only + for the stale-module wipe and the artifact export (~ms). + + Meant for parallel precompilation workers. + """ + from filelock import FileLock + + self._compiled_kernel = self.compile_fn() + with FileLock(self.lock_path, thread_local=False): + if ( + self.module_dir.exists() + and _read_meta(self.meta_path) != self.expected_meta + ): + logger.info( + f"Invalidating stale CuTe-DSL module {self.module_dir_name}" + ) + shutil.rmtree(self.module_dir, ignore_errors=True) + try: + self._export() + except Exception as e: # noqa: BLE001 -- persistence is best-effort + logger.warning( + f"Failed to persist CuTe-DSL kernel {self.name} to " + f"{self.object_path}: {e}. The kernel will be recompiled " + f"next run." + ) + def load(self) -> Any: """The kernel compiled by build(), or the on-disk artifact.""" if self._compiled_kernel is not None: diff --git a/flashinfer/utils.py b/flashinfer/utils.py index 30160c97f06..4cc411ec804 100644 --- a/flashinfer/utils.py +++ b/flashinfer/utils.py @@ -805,6 +805,11 @@ def get_device_sm_count(device: torch.device) -> int: return torch.cuda.get_device_properties(device).multi_processor_count +def get_device_index(device: torch.device) -> int: + """Concrete CUDA device index for *device* (bare "cuda" -> current device).""" + return device.index if device.index is not None else torch.cuda.current_device() + + def get_trtllm_gen_multi_ctas_kv_counter_bytes( batch_size: int, num_qo_heads: int, sm_count: int ) -> int: diff --git a/tests/jit/test_cute_dsl_cache.py b/tests/jit/test_cute_dsl_cache.py index f260a662fa4..f510fd64275 100644 --- a/tests/jit/test_cute_dsl_cache.py +++ b/tests/jit/test_cute_dsl_cache.py @@ -156,3 +156,76 @@ def test_nvfp4_kernel_name_is_symbol_safe(): for cfg in (None, NVFP44Over6Config(err_mode="MSE")): name = _nvfp4_kernel_name(**{**NVFP4_NAME_BASELINE, "nvfp4_4over6_config": cfg}) assert re.fullmatch(r"[0-9A-Za-z_]+", name), name + + +# --------------------------------------------------------------------------- +# mm_fp4 (flashinfer/gemm/gemm_mm_fp4_cute_dsl.py) cache adopter +# --------------------------------------------------------------------------- + +import torch # noqa: E402 + +from flashinfer.gemm.gemm_mm_fp4_cute_dsl import ( # noqa: E402 + _blockscaled_kernel_disk_name, + _mm_fp4_cache_key, +) + +# A baseline argument set and, for each argument, a distinct alternative. +MM_FP4_NAME_BASELINE = { + "sf_vec_size": 16, + "mma_tiler_mn": (256, 128), + "cluster_shape_mn": (2, 1), + "swap_ab": False, + "use_prefetch": False, + "kernel_type": "sm100", + "use_tma_store": None, + "enable_pdl": False, + "out_dtype": torch.bfloat16, + "batch_size": 1, + "max_active_clusters": 74, +} +MM_FP4_NAME_PERTURBED = { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 256), # transpose of baseline: catches mixed-up axes + "cluster_shape_mn": (1, 2), + "swap_ab": True, + "use_prefetch": True, + "kernel_type": "sm103", + "use_tma_store": True, + "enable_pdl": True, + "out_dtype": torch.float16, + "batch_size": 2, + "max_active_clusters": 148, +} + + +def _mm_fp4_name(**kwargs): + tactic = ( + kwargs["mma_tiler_mn"], + kwargs["cluster_shape_mn"], + kwargs["swap_ab"], + kwargs["use_prefetch"], + kwargs["kernel_type"], + kwargs["use_tma_store"], + ) + cache_key = _mm_fp4_cache_key( + kwargs["sf_vec_size"], tactic, kwargs["enable_pdl"], kwargs["out_dtype"] + ) + return _blockscaled_kernel_disk_name( + cache_key, kwargs["batch_size"], kwargs["max_active_clusters"] + ) + + +@pytest.mark.parametrize("param", sorted(MM_FP4_NAME_BASELINE)) +def test_mm_fp4_kernel_name_varies_with_every_argument(param): + """Changing any single codegen argument must change the kernel name, and + names must be symbol-safe as produced (see the nvfp4 twins above).""" + baseline_name = _mm_fp4_name(**MM_FP4_NAME_BASELINE) + kwargs = dict(MM_FP4_NAME_BASELINE) + kwargs[param] = MM_FP4_NAME_PERTURBED[param] + perturbed_name = _mm_fp4_name(**kwargs) + assert perturbed_name != baseline_name, ( + f"_blockscaled_kernel_disk_name ignores argument {param!r}: two " + "different kernel specializations would collide on one cache artifact." + ) + for name in (baseline_name, perturbed_name): + assert re.fullmatch(r"[0-9A-Za-z_]+", name), name