Skip to content
KernelIndex
Search⌘K
Sign in

submission 555027

Ayush10 · python · License unknown
Kernel source · 82 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 82 lines, June 9 Researcher Reciprocity License v1.0.

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-gated-deltanet-recompute-w-u-555027?include=source"
interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp32

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
Gated DeltaNet recompute_w_usuite of 3 cases
NVIDIA B200
10.6µs
#12 of 28
2026-03-14
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:3670d5416c9e631b1e3e7bdab48a0d390eb36840e30882b92bba68b73b04f8bc
license declaredunknown
license concludedunknown
authorsAyush10
imported2026-08-15

Techniques

Extracted from the mirrored source by pattern, never inferred. Each row cites its line.

num-warps = 2(64, 64): helion.Config(num_warps=2, num_stages=3, l2_groupings=[8]),
stages = 3(64, 64): helion.Config(num_warps=2, num_stages=3, l2_groupings=[8]),

Kernel source

submission.py82 lines
#!POPCORN leaderboard gated_deltanet_recompute_w_u
#!POPCORN gpu B200_Nebius
# Team: KernalForge
# TF32 dots (default, no ieee or bf16) + lazy compilation + tuned configs
from task import input_t, output_t

import torch
import helion
import helion.language as hl


KV_CONFIGS: dict[tuple[int, int], helion.Config] = {
    (64, 64): helion.Config(num_warps=2, num_stages=3, l2_groupings=[8]),
    (64, 128): helion.Config(num_warps=2, num_stages=2),
    (100, 100): helion.Config(num_warps=2, num_stages=2, l2_groupings=[4]),
    (128, 128): helion.Config(num_warps=4, num_stages=1),
}


def _make_kernel(config: helion.Config):
    @helion.kernel(config=config)
    def kernel(
        k: torch.Tensor,
        v: torch.Tensor,
        beta: torch.Tensor,
        A: torch.Tensor,
        beta_g: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        B, T, H, K = k.shape
        V = v.shape[-1]
        C = hl.specialize(A.shape[-1])
        K = hl.specialize(K)
        V = hl.specialize(V)

        w_out = torch.empty_like(k)
        u_out = torch.empty_like(v)
        BH = B * H

        for flat_bh, rt in hl.tile([BH, T], block_size=[1, C]):
            b_idx = flat_bh.begin // H
            h_idx = flat_bh.begin % H

            a_chunk = A[b_idx, rt, h_idx, :].to(torch.float32)
            beta_chunk = beta[b_idx, rt, h_idx].to(torch.float32)
            beta_g_chunk = beta_g[b_idx, rt, h_idx].to(torch.float32)
            rhs_k = k[b_idx, rt, h_idx, :].to(torch.float32) * beta_g_chunk[:, None]
            rhs_v = v[b_idx, rt, h_idx, :].to(torch.float32) * beta_chunk[:, None]

            w_out[b_idx, rt, h_idx, :] = hl.dot(
                a_chunk,
                rhs_k,
                out_dtype=torch.float32,
            ).to(w_out.dtype)
            u_out[b_idx, rt, h_idx, :] = hl.dot(
                a_chunk,
                rhs_v,
                out_dtype=torch.float32,
            ).to(u_out.dtype)

        return w_out, u_out

    return kernel


# Lazy compilation — only compile the (K,V) variant actually needed
_KERNEL_CACHE: dict[tuple[int, int], object] = {}


def _get_kernel(kv: tuple[int, int]):
    if kv not in _KERNEL_CACHE:
        _KERNEL_CACHE[kv] = _make_kernel(KV_CONFIGS[kv])
    return _KERNEL_CACHE[kv]


def custom_kernel(data: input_t) -> output_t:
    k, v, beta, A, g = data
    K = k.shape[-1]
    V = v.shape[-1]
    beta_g = beta * torch.exp(g)
    kernel = _get_kernel((K, V))
    return kernel(k, v, beta, A, beta_g)
scrolls · 82 lines total

Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0

Changes from previous submission

Against this author's previous submission submission 553231.

  #!POPCORN leaderboard gated_deltanet_recompute_w_u
  #!POPCORN gpu B200_Nebius
  # Team: KernalForge
- # Fix: Group configs by (K,V) to reduce JIT compilations from 10 to 4
- # Remove static_shapes so B,T,H are runtime values — fits within 420s timeout
+ # TF32 dots (default, no ieee or bf16) + lazy compilation + tuned configs
  from task import input_t, output_t
   
  import torch
  ⋯ 1 unchanged lines
  import helion.language as hl
   
   
- # 4 unique (K,V) pairs across all test+benchmark shapes:
- #   (64,64)   — 7 shapes
- #   (64,128)  — 1 shape
- #   (100,100) — 1 shape
- #   (128,128) — 2 shapes
- # Each group compiles ONE Triton kernel. 4 compilations × ~60s = ~240s < 420s timeout.
  KV_CONFIGS: dict[tuple[int, int], helion.Config] = {
-     (64, 64): helion.Config(num_warps=4, num_stages=3, l2_groupings=[4]),
-     (64, 128): helion.Config(num_warps=8, num_stages=3),
-     (100, 100): helion.Config(num_warps=8, num_stages=4, l2_groupings=[4]),
-     (128, 128): helion.Config(num_warps=8, num_stages=4, l2_groupings=[8]),
+     (64, 64): helion.Config(num_warps=2, num_stages=3, l2_groupings=[8]),
+     (64, 128): helion.Config(num_warps=2, num_stages=2),
+     (100, 100): helion.Config(num_warps=2, num_stages=2, l2_groupings=[4]),
+     (128, 128): helion.Config(num_warps=4, num_stages=1),
  }
   
   
  def _make_kernel(config: helion.Config):
-     @helion.kernel(dot_precision="ieee", config=config)
+     @helion.kernel(config=config)
      def kernel(
          k: torch.Tensor,
          v: torch.Tensor,
  ⋯ 37 unchanged lines
      return kernel
   
   
- _KERNELS = {kv: _make_kernel(cfg) for kv, cfg in KV_CONFIGS.items()}
+ # Lazy compilation — only compile the (K,V) variant actually needed
+ _KERNEL_CACHE: dict[tuple[int, int], object] = {}
   
   
+ def _get_kernel(kv: tuple[int, int]):
+     if kv not in _KERNEL_CACHE:
+         _KERNEL_CACHE[kv] = _make_kernel(KV_CONFIGS[kv])
+     return _KERNEL_CACHE[kv]
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      k, v, beta, A, g = data
      K = k.shape[-1]
      V = v.shape[-1]
      beta_g = beta * torch.exp(g)
-     kernel = _KERNELS[(K, V)]
+     kernel = _get_kernel((K, V))
      return kernel(k, v, beta, A, beta_g)
scrolls · 60 diff lines total

Best evidence level for this revision: reported

JSON