submission 545257 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

matmul_v6.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-matmul-v2-545257?include=source"

interfacepython

Compatibility

measured onNVIDIA B200

declared hardwareNVIDIA B200

architecturessm_100

dtypesfp16

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload

Hardware

Latency

Rank

Observed

FP16 matmulsuite of 8 cases

NVIDIA B200

136.9µs

#28 of 53

2026-03-13

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:db535e902cfd32541b31c6257d6af9291f50894c78c1f6a2376965fc82d86d2c

license declaredunknown

license concludedunknown

authorsrajesh0042

imported2026-08-15

Kernel source

matmul_v6.py33 lines

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
from task import input_t, output_t

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True

# Extensive warmup for cuBLAS on B200
def _warmup():
    sizes = [
        (128, 128, 128), (256, 256, 256), (512, 512, 512),
        (1024, 1024, 1024), (2048, 2048, 2048),
        (4096, 4096, 4096), (4096, 5120, 4096),
        (1024, 1536, 1024), (2048, 3072, 2048),
    ]
    for m, n, k in sizes:
        a = torch.randn(m, k, device='cuda', dtype=torch.float16)
        b = torch.randn(k, n, device='cuda', dtype=torch.float16)
        c = torch.empty(m, n, device='cuda', dtype=torch.float16)
        for _ in range(5):
            torch.mm(a, b, out=c)
    torch.cuda.synchronize()

_warmup()

def custom_kernel(data: input_t) -> output_t:
    a, b, c = data
    torch.mm(a, b, out=c)
    return c

scrolls · 33 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 545198.

  ⋯ 5 unchanged lines
   
  torch.backends.cuda.matmul.allow_tf32 = True
  torch.backends.cudnn.allow_tf32 = True
-  
- # Try FP32 accumulation with cuBLAS - this uses different tensor core paths
- # torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction controls this
  torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
   
- # Pre-allocate and warm cuBLAS for benchmark shapes
- _cache = {}
-  
+ # Extensive warmup for cuBLAS on B200
  def _warmup():
      sizes = [
          (128, 128, 128), (256, 256, 256), (512, 512, 512),
          (1024, 1024, 1024), (2048, 2048, 2048),
-         (1024, 1536, 1024), (2048, 3072, 2048), (4096, 5120, 4096),
+         (4096, 4096, 4096), (4096, 5120, 4096),
+         (1024, 1536, 1024), (2048, 3072, 2048),
      ]
      for m, n, k in sizes:
          a = torch.randn(m, k, device='cuda', dtype=torch.float16)
          b = torch.randn(k, n, device='cuda', dtype=torch.float16)
          c = torch.empty(m, n, device='cuda', dtype=torch.float16)
-         # Run twice to warm cuBLAS handle caching
-         torch.mm(a, b, out=c)
-         torch.mm(a, b, out=c)
+         for _ in range(5):
+             torch.mm(a, b, out=c)
      torch.cuda.synchronize()
   
  _warmup()

scrolls · 33 diff lines total

Best evidence level for this revision: reported

JSON