diff --git a/src/tilegym/backend/__init__.py b/src/tilegym/backend/__init__.py index 53324d3e..013265af 100644 --- a/src/tilegym/backend/__init__.py +++ b/src/tilegym/backend/__init__.py @@ -9,6 +9,7 @@ from .dispatcher import dispatch from .dispatcher import get_available_backends_for_op from .dispatcher import get_registry_info +from .dispatcher import has_backend_impl from .dispatcher import print_registry_info from .dispatcher import register_impl from .selector import assert_backend_available @@ -56,6 +57,7 @@ def __getattr__(name): "dispatch", "register_impl", "get_available_backends_for_op", + "has_backend_impl", "get_registry_info", "print_registry_info", # Cutile utilities diff --git a/src/tilegym/backend/dispatcher.py b/src/tilegym/backend/dispatcher.py index da8365b9..4c55ce65 100644 --- a/src/tilegym/backend/dispatcher.py +++ b/src/tilegym/backend/dispatcher.py @@ -131,6 +131,11 @@ def wrapper(*args, **kwargs): _REGISTRY[name]["default"] = default_impl + # Expose the op name on the dispatched callable so tooling (e.g. test + # collection) can map a public op back to its registry entry and query + # which backends implement it. + wrapper._tilegym_op_name = name + return wrapper return decorator @@ -152,6 +157,24 @@ def get_available_backends_for_op(name: str) -> list: return list(_REGISTRY[name].keys()) +def has_backend_impl(name: str, backend: str) -> bool: + """ + Check whether a specific backend implementation is registered for an op. + + Unlike a plain membership test, this ignores the synthetic ``"default"`` + entry (the NotImplementedError stub), so it answers "is there a real + ``backend`` kernel for this op?". + + Args: + name: Operation name (registry key, e.g. ``"softmax"``) + backend: Backend name (e.g. ``"cutile"``, ``"triton"``) + + Returns: + True if a real implementation for ``backend`` is registered. + """ + return name in _REGISTRY and backend in _REGISTRY[name] + + def get_registry_info() -> Dict[str, Dict[str, str]]: """ Get information about all registered implementations diff --git a/src/tilegym/ops/cutile/softmax.py b/src/tilegym/ops/cutile/softmax.py index 1b04ce5b..ea7b5978 100644 --- a/src/tilegym/ops/cutile/softmax.py +++ b/src/tilegym/ops/cutile/softmax.py @@ -6,6 +6,7 @@ import cuda.tile as ct import torch +from cuda.tile import RoundingMode as RMd from tilegym.backend import register_impl from tilegym.experimental import experimental_kernel @@ -39,13 +40,13 @@ def _softmax_kernel( row_minus_max = row - row_max # Compute exponential - numerator = ct.exp(row_minus_max) + numerator = ct.exp(row_minus_max, rounding_mode=RMd.APPROX) # Compute sum for normalization denominator = ct.sum(numerator, 0, keepdims=True) # Final softmax computation - softmax_output = numerator / denominator + softmax_output = ct.truediv(numerator, denominator, rounding_mode=RMd.APPROX, flush_to_zero=True) softmax_output = ct.astype(softmax_output, input.dtype) ct.scatter(output, (row_idx, offsets), softmax_output, check_bounds=True) @@ -69,10 +70,11 @@ def _softmax_kernel_multi_wave_full_row_reg_cached_ldg( row = ct.astype(row, ct.float32) row_max = ct.max(row, 0, keepdims=True) - numerator = ct.exp(row - row_max) + numerator = ct.exp(row - row_max, rounding_mode=RMd.APPROX) denominator = ct.sum(numerator, 0, keepdims=True) - softmax_output = ct.astype(numerator / denominator, input.dtype) + softmax_output = ct.truediv(numerator, denominator, rounding_mode=RMd.APPROX, flush_to_zero=True) + softmax_output = ct.astype(softmax_output, input.dtype) ct.scatter(output, (row_idx, offsets), softmax_output, check_bounds=check_bound) @@ -101,13 +103,13 @@ def _softmax_kernel_tma( row_minus_max = row - row_max # Compute exponential - numerator = ct.exp(row_minus_max) + numerator = ct.exp(row_minus_max, rounding_mode=RMd.APPROX) # Compute sum for normalization denominator = ct.sum(numerator, 1, keepdims=True) # Final softmax computation - softmax_output = numerator / denominator + softmax_output = ct.truediv(numerator, denominator, rounding_mode=RMd.APPROX, flush_to_zero=True) # Convert back to original dtype and store softmax_output = ct.astype(softmax_output, input.dtype) @@ -150,10 +152,13 @@ def _softmax_kernel_chunked( chunk = ct.gather(input, (row_idx, col_indices), check_bounds=True, padding_value=-math.inf) chunk = ct.astype(chunk, ct.float32) row_minus_max = chunk - row_max - numerator = ct.exp(row_minus_max) + numerator = ct.exp(row_minus_max, rounding_mode=RMd.APPROX) exponentials_sum = ct.sum(numerator, 0, keepdims=True) denominator = denominator + exponentials_sum + # Reciprocal once per row, multiplied inside the pass-3 chunk loop. + inv_denominator = ct.truediv(1.0, denominator, rounding_mode=RMd.APPROX, flush_to_zero=True) + # Pass 3: Compute final softmax for chunk_idx in range(num_chunks): chunk_start = chunk_idx * TILE_SIZE @@ -162,8 +167,8 @@ def _softmax_kernel_chunked( chunk = ct.gather(input, (row_idx, col_indices), check_bounds=True, padding_value=-math.inf) chunk = ct.astype(chunk, ct.float32) row_minus_max = chunk - row_max - numerator = ct.exp(row_minus_max) - softmax_output = numerator / denominator + numerator = ct.exp(row_minus_max, rounding_mode=RMd.APPROX) + softmax_output = numerator * inv_denominator softmax_output = ct.astype(softmax_output, input.dtype) # Use scatter with bounds checking to avoid writing padded zeros ct.scatter(output, (row_idx, col_indices), softmax_output, check_bounds=True) diff --git a/src/tilegym/ops/cutile/splitk_reduce.py b/src/tilegym/ops/cutile/splitk_reduce.py index 5b4063c2..b73000d7 100644 --- a/src/tilegym/ops/cutile/splitk_reduce.py +++ b/src/tilegym/ops/cutile/splitk_reduce.py @@ -45,6 +45,7 @@ def _splitk_reduce_kernel( order=(0, 1, 2, 3), allow_tma=True, latency=2, + padding_mode=ct.PaddingMode.ZERO, ) out_splitk = ct.reshape(out_splitk, (NUM_KV_SPLITS_POW2, TILE_D)) diff --git a/src/tilegym/ops/tilecpp/attention.cuh b/src/tilegym/ops/tilecpp/attention.cuh index 20249cf8..efa44720 100644 --- a/src/tilegym/ops/tilecpp/attention.cuh +++ b/src/tilegym/ops/tilecpp/attention.cuh @@ -251,7 +251,11 @@ __tile_global__ void prefill_fmha_fwd_kernel( // Convert back to input type and store output auto acc_T = ct::element_cast(acc); auto acc_4d = ct::reshape(acc_T, ct::shape<1, 1, BLOCK_M, BLOCK_D>{}); - Out_view.store(acc_4d, batch_idx, head_idx, pid_x, 0); + if constexpr (EVEN_Q) { + Out_view.store(acc_4d, batch_idx, head_idx, pid_x, 0); + } else { + Out_view.store_masked(acc_4d, batch_idx, head_idx, pid_x, 0); + } if constexpr (HAS_BACKWARD) { auto L_span = ct::tensor_span{L_ptr, ct::extents{}}; @@ -260,7 +264,11 @@ __tile_global__ void prefill_fmha_fwd_kernel( auto lse_2d = m_i + ct::log2(l_i); // (TILE_M, 1) auto lse_1d = ct::reshape(lse_2d, ct::shape{}); auto lse_3d = ct::reshape(lse_1d, ct::shape<1, 1, BLOCK_M>{}); - L_view.store(lse_3d, batch_idx, head_idx, pid_x); + if constexpr (EVEN_Q) { + L_view.store(lse_3d, batch_idx, head_idx, pid_x); + } else { + L_view.store_masked(lse_3d, batch_idx, head_idx, pid_x); + } } } diff --git a/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.cuh b/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.cuh index 74efa0a6..ed947f8d 100644 --- a/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.cuh +++ b/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.cuh @@ -74,6 +74,8 @@ __tile__ inline TileType cgdr_solve_tril(TileType A) { // Grid: (B * NUM_HEADS, num_chunks, 1). // ============================================================================ template(ct::reshape>(v_4d)); auto vb = v * beta_col; auto vc = ct::matmul(attn, vb); - pVcorr.store(ct::reshape>(vc), - b, h, pid_chunk, 0, vt); + pVcorr.store_masked(ct::reshape>(vc), + b, h, pid_chunk, 0, vt); } } diff --git a/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.py b/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.py index 22ccf531..fae133fa 100644 --- a/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.py +++ b/src/tilegym/ops/tilecpp/chunk_gated_delta_rule.py @@ -19,6 +19,7 @@ from tilegym.backend import register_impl from tilegym.ops.tilecpp.utils._cuda_utils import TileCppKernel +from tilegym.ops.tilecpp.utils._cuda_utils import get_cpp_type from tilegym.ops.tilecpp.utils._dump_types import dump_kernel_types _intra_kernel = TileCppKernel( @@ -64,7 +65,11 @@ def _launch_intra( occupancy = 1 bool_to_str = lambda b: "true" if b else "false" + beta_cpp_type = get_cpp_type(Beta.dtype) + g_cpp_type = get_cpp_type(G.dtype) template_params = [ + beta_cpp_type, + g_cpp_type, chunk_size, block_k, bool_to_str(use_qk_l2norm), @@ -74,7 +79,7 @@ def _launch_intra( dtype=dtype, template_params=template_params, signature=( - "const {T}*, const {T}*, const {T}*, const {T}*, const {T}*, " + f"const {{T}}*, const {{T}}*, const {{T}}*, const {beta_cpp_type}*, const {g_cpp_type}*, " "float*, float*, float*, float*, float*, " "float, int, int, int, int, int, int" ), diff --git a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py index 5119a5b2..5c7dcf81 100644 --- a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py +++ b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py @@ -13,9 +13,9 @@ from tilegym.kernel_utils import get_kernel_configs from tilegym.ops.cutile.utils import cached_replace_hints -# Module-level tune caches for standard and swap_ab ragged BMM -_ragged_bmm_standard_tune_cache: dict = {} -_ragged_bmm_swap_ab_tune_cache: dict = {} +# Module-level tune cache for ragged BMM: shape key -> (config, tuned kernel). +# The tuned kernel is whichever of the standard / swap_ab variants measured faster. +_ragged_bmm_tune_cache: dict = {} @ct.kernel @@ -425,11 +425,17 @@ def _get_default_kernel_configs(): } -def _ragged_bmm_autotune_standard( - stream, a, b, c, m_indptr, Q, max_m, max_m_device, N, total_m, transpose_a, transpose_b -): +def _ragged_bmm_autotune(stream, a, b, c, m_indptr, Q, max_m, max_m_device, N, total_m, transpose_a, transpose_b): """ - Autotuned launch for standard ragged BMM kernel. + Autotuned launch for ragged BMM. + + Tunes the standard and the swap_ab kernel over their respective config spaces + and launches whichever measured faster. The two variants only differ in the + accumulator layout, so they are interchangeable for any shape; which one wins + is not predictable from the shape alone (it depends on the per-arch config + space and is non-monotonic in the per-batch M), so it is measured rather than + guessed. Both variants read their grid bound from max_m_device, so the choice + never affects correctness. """ NUM_SMS = torch.cuda.get_device_properties(a.device).multi_processor_count @@ -473,87 +479,28 @@ def hints_fn(cfg): return {"num_ctas": cfg.num_ctas, "occupancy": cfg.occupancy} cache_key = (Q, max_m, N, total_m, transpose_a_int, transpose_b_int, a.dtype, str(a.device)) - if cache_key not in _ragged_bmm_standard_tune_cache: - result = exhaustive_search( - list(_ragged_bmm_autotune_configs_standard()), - stream, - grid_fn, - _ragged_bmm_kernel, - args_fn, - hints_fn, - ) - best_cfg = result.best.config - _ragged_bmm_standard_tune_cache[cache_key] = ( + if cache_key not in _ragged_bmm_tune_cache: + best = None + for kernel, configs in ( + (_ragged_bmm_kernel, list(_ragged_bmm_autotune_configs_standard())), + (_ragged_bmm_swap_ab_kernel, list(_ragged_bmm_autotune_configs_swap_ab())), + ): + try: + result = exhaustive_search(configs, stream, grid_fn, kernel, args_fn, hints_fn) + except Exception: + # A whole config space can fail to build on a given arch (e.g. smem + # limits); fall back to whichever variant did tune successfully. + continue + if best is None or result.best.mean_us < best[0]: + best = (result.best.mean_us, kernel, result.best.config) + if best is None: + raise RuntimeError("ragged_bmm autotune found no working configuration") + _, best_kernel, best_cfg = best + _ragged_bmm_tune_cache[cache_key] = ( best_cfg, - _ragged_bmm_kernel.replace_hints(**hints_fn(best_cfg)), + best_kernel.replace_hints(**hints_fn(best_cfg)), ) - best_cfg, tuned_kernel = _ragged_bmm_standard_tune_cache[cache_key] - ct.launch(stream, grid_fn(best_cfg), tuned_kernel, args_fn(best_cfg)) - - -def _ragged_bmm_autotune_swap_ab( - stream, a, b, c, m_indptr, Q, max_m, max_m_device, N, total_m, transpose_a, transpose_b -): - """ - Autotuned launch for swap_ab ragged BMM kernel. - """ - NUM_SMS = torch.cuda.get_device_properties(a.device).multi_processor_count - - transpose_a_int = 1 if transpose_a else 0 - transpose_b_int = 1 if transpose_b else 0 - - def args_fn(cfg): - BM = cfg.BLOCK_M - BN = cfg.BLOCK_N - BK = cfg.BLOCK_K - GSM = cfg.GROUP_SIZE_M - - return ( - a, - b, - c, - m_indptr, - Q, - max_m, - max_m_device, - N, - transpose_a_int, - transpose_b_int, - BM, - BN, - BK, - GSM, - ) - - def grid_fn(cfg): - BM = cfg.BLOCK_M - BN = cfg.BLOCK_N - num_pid_m = ct.cdiv(max_m, BM) - num_pid_n = ct.cdiv(N, BN) - tiles_per_batch = num_pid_m * num_pid_n - total_tiles = tiles_per_batch * Q - num_programs = min(NUM_SMS // cfg.num_ctas, total_tiles) * cfg.occupancy - return (num_programs, 1, 1) - - def hints_fn(cfg): - return {"num_ctas": cfg.num_ctas, "occupancy": cfg.occupancy} - - swap_cache_key = (Q, max_m, N, total_m, transpose_a_int, transpose_b_int, a.dtype, str(a.device)) - if swap_cache_key not in _ragged_bmm_swap_ab_tune_cache: - result = exhaustive_search( - list(_ragged_bmm_autotune_configs_swap_ab()), - stream, - grid_fn, - _ragged_bmm_swap_ab_kernel, - args_fn, - hints_fn, - ) - best_cfg = result.best.config - _ragged_bmm_swap_ab_tune_cache[swap_cache_key] = ( - best_cfg, - _ragged_bmm_swap_ab_kernel.replace_hints(**hints_fn(best_cfg)), - ) - best_cfg, tuned_kernel = _ragged_bmm_swap_ab_tune_cache[swap_cache_key] + best_cfg, tuned_kernel = _ragged_bmm_tune_cache[cache_key] ct.launch(stream, grid_fn(best_cfg), tuned_kernel, args_fn(best_cfg)) @@ -627,49 +574,25 @@ def ragged_bmm( # Check if autotune is enabled enable_autotune = is_autotune_enabled() - # Decide whether to use swap_ab based on M vs N ratio. - # swap_ab (small BLOCK_M) is beneficial when the per-batch M is small relative - # to N. Use the per-batch average M (total_m / Q) rather than the host `max_m` - # hint: `max_m` is only a grid/cache-key upper bound and callers may pass a - # coarse over-estimate (e.g. the fused-MoE path passes total tokens_in_chunk, - # not the per-expert max), which would wrongly route small-per-expert MoE - # GEMMs to the large-M standard kernel (BLOCK_M=128) instead of swap_ab - # (BLOCK_M<=64). The grid bound stays exact via the device-side max_m_device, - # so this only affects config selection, never correctness. - avg_m = total_m // Q if Q > 0 else max_m - use_swap_ab = avg_m <= 128 and N >= 256 - if enable_autotune: - if use_swap_ab: - _ragged_bmm_autotune_swap_ab( - torch.cuda.current_stream(), - a, - b, - c, - m_indptr, - Q, - max_m, - max_m_device, - N, - total_m, - transpose_a, - transpose_b, - ) - else: - _ragged_bmm_autotune_standard( - torch.cuda.current_stream(), - a, - b, - c, - m_indptr, - Q, - max_m, - max_m_device, - N, - total_m, - transpose_a, - transpose_b, - ) + # The standard and swap_ab kernels are both tuned and the faster one wins. + # There is no host-side shape heuristic here on purpose: a threshold on the + # per-batch M mis-routes MoE GEMMs, because which variant is faster is not + # monotonic in M and differs per arch (measured on sm90 / sm103 / sm120). + _ragged_bmm_autotune( + torch.cuda.current_stream(), + a, + b, + c, + m_indptr, + Q, + max_m, + max_m_device, + N, + total_m, + transpose_a, + transpose_b, + ) else: # Use fixed default configs default_configs = _get_default_kernel_configs() diff --git a/tests/ops/test_chunk_gated_delta_rule.py b/tests/ops/test_chunk_gated_delta_rule.py index 818265bb..771f6601 100644 --- a/tests/ops/test_chunk_gated_delta_rule.py +++ b/tests/ops/test_chunk_gated_delta_rule.py @@ -176,9 +176,6 @@ def test_op(self, B, T, H, K, V, CS, use_init, out_final, use_l2, dtype, backend if dtype == torch.float32: pytest.skip("Skipping fp32 tests due to known failures; under investigation") - if backend == "tilecpp" and use_l2: - pytest.skip("Skipping tilecpp l2norm case due to known failure; under investigation") - self.setUp() from tilegym.ops import chunk_gated_delta_rule @@ -284,8 +281,6 @@ def test_op_correlated_keys_stable_triangular_solve(self, use_l2, backend, arch, that reduced-precision Neumann products amplified catastrophically. The pre-normalized case also exercises the non-L2 infinity-norm guard. """ - if backend == "tilecpp": - pytest.skip("Skipping tilecpp case due to known failure; under investigation") monkeypatch.setenv("TILEGYM_DISABLE_AUTOTUNE", "1") if not tilegym.is_backend_available(backend): pytest.skip(f"Backend {backend} is not available")