Skip to content
KernelIndex
Search⌘K
Sign in

submission 553926

bloomberg9383 · python · License unknown
Kernel source · 95 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-gated-deltanet-chunk-fwd-h-553926?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 chunk_fwd_hsuite of 3 cases
NVIDIA B200
12.6µs
#9 of 28
2026-03-14
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:0fb062f1ab7100962547b49fdbf9a8f58ace8f8bcfff6d0425865e5803e4ecbb
license declaredunknown
license concludedunknown
authorsbloomberg9383
imported2026-08-15

Techniques

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

num-warps = 16(1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
stages = 1(1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),

Kernel source

submission.py95 lines
from task import input_t, output_t

import torch
import helion
import helion.language as hl


# Per-shape configs: map (B, T, H, K, V) to optimized helion.Config objects.
# Autotune locally for each shape, then paste the best config here.
SHAPE_CONFIGS: dict[tuple, helion.Config] = {
    # Test shapes
    (1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
    (2, 128, 4, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=2, pid_type='flat'),
    (1, 256, 4, 64, 128): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
    # Benchmark shapes
    (1, 64, 1, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
    (2, 512, 3, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
    (2, 1024, 3, 64, 64): helion.Config(block_sizes=[], num_stages=3, num_warps=4, pid_type='flat'),
    (3, 1024, 4, 100, 100): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
    (4, 1024, 4, 128, 128): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
    (2, 1536, 4, 128, 128): helion.Config(block_sizes=[], num_stages=2, num_warps=8, pid_type='flat'),
    (4, 2048, 8, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
}


# Optional: add advanced_controls_file to your Config for extra performance (see docs).
# Autotune with autotune_search_acf to find the best ACF, then hardcode it:
#     helion.Config(..., advanced_controls_file="/opt/booster_pack/chunk_fwd_h_0.acf")


def _make_kernel(config: helion.Config):
    @helion.kernel(static_shapes=True, dot_precision="ieee", config=config)
    def kernel(
        k: torch.Tensor,   # [B, T, H, K]
        w: torch.Tensor,   # [B, T, H, K]
        u: torch.Tensor,   # [B, T, H, V]
        g: torch.Tensor,   # [B, T, H]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        B, T, H, K = k.shape
        V = u.shape[-1]
        C = 64
        K = hl.specialize(K)
        V = hl.specialize(V)

        NT = (T + C - 1) // C
        h_out = torch.empty(B, NT, H, K, V, dtype=k.dtype, device=k.device)
        v_out = torch.empty_like(u)

        BH = B * H

        for flat, tv in hl.tile([BH, V], block_size=[1, 8]):
            b_idx = flat.begin // H
            h_idx = flat.begin % H
            state = hl.zeros([K, tv], dtype=torch.float32)

            for tc in hl.tile(T, block_size=C):
                chunk_idx = tc.begin // C
                t_end = min(tc.begin + C, T) - 1

                h_out[b_idx, chunk_idx, h_idx, :, tv] = state.to(k.dtype)

                proj = hl.dot(
                    w[b_idx, tc, h_idx, :], state, out_dtype=torch.float32
                )
                diff = u[b_idx, tc, h_idx, tv].to(torch.float32) - proj
                v_out[b_idx, tc, h_idx, tv] = diff.to(u.dtype)

                g_end = g[b_idx, t_end, h_idx].to(torch.float32)
                g_t = g[b_idx, tc, h_idx].to(torch.float32)
                valid = tc.index < T
                alpha = torch.where(valid, torch.exp(g_end - g_t), 0.0)
                k_adj = k[b_idx, tc, h_idx, :] * alpha[:, None]

                state = state * torch.exp(g_end)
                upd = hl.dot(k_adj.T, diff, out_dtype=torch.float32)
                state = state + upd

        return h_out, v_out

    return kernel


_KERNEL_CACHE: dict[tuple, object] = {}


def custom_kernel(data: input_t) -> output_t:
    k, w, u, g = data
    B, T, H, K = k.shape
    V = u.shape[-1]
    key = (B, T, H, K, V)
    if key not in _KERNEL_CACHE:
        _KERNEL_CACHE[key] = _make_kernel(SHAPE_CONFIGS[key])
    kernel = _KERNEL_CACHE[key]
    return kernel(k, w, u, g)
scrolls · 95 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 552580.

  ⋯ 8 unchanged lines
  # Autotune locally for each shape, then paste the best config here.
  SHAPE_CONFIGS: dict[tuple, helion.Config] = {
      # Test shapes
-     (1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
-     (2, 128, 4, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
-     (1, 256, 4, 64, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
+     (1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
+     (2, 128, 4, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=2, pid_type='flat'),
+     (1, 256, 4, 64, 128): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
      # Benchmark shapes
-     (1, 64, 1, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 512, 3, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 1024, 3, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (3, 1024, 4, 100, 100): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (4, 1024, 4, 128, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 1536, 4, 128, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (4, 2048, 8, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
+     (1, 64, 1, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
+     (2, 512, 3, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
+     (2, 1024, 3, 64, 64): helion.Config(block_sizes=[], num_stages=3, num_warps=4, pid_type='flat'),
+     (3, 1024, 4, 100, 100): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
+     (4, 1024, 4, 128, 128): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
+     (2, 1536, 4, 128, 128): helion.Config(block_sizes=[], num_stages=2, num_warps=8, pid_type='flat'),
+     (4, 2048, 8, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=4, pid_type='flat'),
  }
   
   
  ⋯ 2 unchanged lines
  #     helion.Config(..., advanced_controls_file="/opt/booster_pack/chunk_fwd_h_0.acf")
   
   
- # NOTE: This is an intentionally inefficient baseline implementation.
  def _make_kernel(config: helion.Config):
      @helion.kernel(static_shapes=True, dot_precision="ieee", config=config)
      def kernel(
  ⋯ 25 unchanged lines
   
                  h_out[b_idx, chunk_idx, h_idx, :, tv] = state.to(k.dtype)
   
-                 proj1 = hl.dot(
+                 proj = hl.dot(
                      w[b_idx, tc, h_idx, :], state, out_dtype=torch.float32
                  )
-                 proj2 = hl.dot(
-                     w[b_idx, tc, h_idx, :], state, out_dtype=torch.float32
-                 )
-                 proj = (proj1 + proj2) * 0.5
                  diff = u[b_idx, tc, h_idx, tv].to(torch.float32) - proj
                  v_out[b_idx, tc, h_idx, tv] = diff.to(u.dtype)
   
  ⋯ 4 unchanged lines
                  k_adj = k[b_idx, tc, h_idx, :] * alpha[:, None]
   
                  state = state * torch.exp(g_end)
-                 upd1 = hl.dot(k_adj.T, diff, out_dtype=torch.float32)
-                 upd2 = hl.dot(k_adj.T, diff, out_dtype=torch.float32)
-                 state = state + (upd1 + upd2) * 0.5
+                 upd = hl.dot(k_adj.T, diff, out_dtype=torch.float32)
+                 state = state + upd
   
          return h_out, v_out
   
      return kernel
   
   
- _KERNELS = {shape: _make_kernel(cfg) for shape, cfg in SHAPE_CONFIGS.items()}
+ _KERNEL_CACHE: dict[tuple, object] = {}
   
   
  def custom_kernel(data: input_t) -> output_t:
      k, w, u, g = data
      B, T, H, K = k.shape
      V = u.shape[-1]
-     kernel = _KERNELS[(B, T, H, K, V)]
+     key = (B, T, H, K, V)
+     if key not in _KERNEL_CACHE:
+         _KERNEL_CACHE[key] = _make_kernel(SHAPE_CONFIGS[key])
+     kernel = _KERNEL_CACHE[key]
      return kernel(k, w, u, g)
scrolls · 80 diff lines total

Best evidence level for this revision: reported

JSON