feat(dispatcher): add has_backend_impl and expose op name on dispatch wrapper & other updates

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

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

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)

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))
 

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<T>(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<uint32_t, B, H, S_QO>{}};
@@ -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<BLOCK_M>{});
         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);
+        }
     }
 }
 

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<typename T,
+         typename BetaT,
+         typename GT,
          int CHUNK_SIZE,
          int BLOCK_K,
          bool USE_QK_L2NORM,
@@ -83,8 +85,8 @@ __tile_global__ void chunk_gated_delta_rule_intra_kernel(
     const T* __restrict__ Q,           // (B, T, H, K)
     const T* __restrict__ K,           // (B, T, H, K)
     const T* __restrict__ V,           // (B, T, H, V)
-    const T* __restrict__ Beta,        // (B, T, H)
-    const T* __restrict__ G,           // (B, T, H)
+    const BetaT* __restrict__ Beta,    // (B, T, H)
+    const GT* __restrict__ G,          // (B, T, H)
     float* __restrict__ Q_out,         // (B, H, num_chunks, CHUNK_SIZE, K)
     float* __restrict__ K_out,         // (B, H, num_chunks, CHUNK_SIZE, K)
     float* __restrict__ V_corr,        // (B, H, num_chunks, CHUNK_SIZE, V)
@@ -249,12 +251,12 @@ __tile_global__ void chunk_gated_delta_rule_intra_kernel(
     // --- v_corrected = attn @ (v * beta[:, None]) : iterate V tiles ---
     int num_v_tiles = (V_dim + BLOCK_K - 1) / BLOCK_K;
     for (auto vt : ct::irange(0, num_v_tiles)) {
-        auto v_4d = pV.load(b, pid_chunk, h, vt);
+        auto v_4d = pV.load_masked(b, pid_chunk, h, vt);
         auto v    = ct::element_cast<float>(ct::reshape<ct::shape<CHUNK_SIZE, BLOCK_K>>(v_4d));
         auto vb   = v * beta_col;
         auto vc   = ct::matmul(attn, vb);
-        pVcorr.store(ct::reshape<ct::shape<1, 1, 1, CHUNK_SIZE, BLOCK_K>>(vc),
-                     b, h, pid_chunk, 0, vt);
+        pVcorr.store_masked(ct::reshape<ct::shape<1, 1, 1, CHUNK_SIZE, BLOCK_K>>(vc),
+                            b, h, pid_chunk, 0, vt);
     }
 }
 

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"
         ),