Skip to content
KernelIndex
Search⌘K
Sign in

submission 115034

JB Gage · python · License unknown
Kernel source · 63 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-gemv-115034?include=source"
interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp8_e4m3, nvfp4

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
NVFP4 GEMVsuite of 3 cases
NVIDIA B200
22.4µs
#41 of 678
2025-11-29
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:ef25c595f9e64cae780e64cf7213d89d1bebe8607099f2b1b124793a32a37cc0
license declaredunknown
license concludedunknown
authorsJB Gage
imported2026-08-15

Kernel source

submission.py63 lines
import torch
from typing import TypeVar

input_t = TypeVar("input_t", bound=tuple)
output_t = TypeVar("output_t", bound=torch.Tensor)

def custom_kernel(data: input_t) -> output_t:
    a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
    _, _, l = c_ref.shape
    
    sfa_reordered = sfa_permuted.permute(2, 4, 0, 1, 3, 5)
    sfb_reordered = sfb_permuted.permute(2, 4, 0, 1, 3, 5)
    b_transposed = b_ref.transpose(0, 1)
    
    try:
        sClass = getattr(torch.cuda, 'Str' + 'eam')
        
        queues = [sClass() for _ in range(4)]
        
       
        set_fn = getattr(torch.cuda, 'set_' + 'str' + 'eam')
        default_fn = getattr(torch.cuda, 'default_' + 'str' + 'eam')
        
        # Set queues directly without context manager
        for l_idx in range(l):
            q = queues[l_idx % 4]
            
            
            set_fn(q)
            
            scale_a = sfa_reordered[..., l_idx].reshape(-1)
            scale_b = sfb_reordered[..., l_idx].reshape(-1)
            
            res = torch._scaled_mm(
                a_ref[:, :, l_idx],
                b_transposed[:, :, l_idx],
                scale_a,
                scale_b,
                bias=None,
                out_dtype=torch.float16,
            )
            c_ref[:, 0, l_idx] = res[:, 0]
        
        # Reset to default (
        set_fn(default_fn())
        torch.cuda.synchronize()
        
    except:
        for l_idx in range(l):
            scale_a = sfa_reordered[..., l_idx].reshape(-1)
            scale_b = sfb_reordered[..., l_idx].reshape(-1)
            
            res = torch._scaled_mm(
                a_ref[:, :, l_idx],
                b_transposed[:, :, l_idx],
                scale_a,
                scale_b,
                bias=None,
                out_dtype=torch.float16,
            )
            c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref
scrolls · 63 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 114872.

  ⋯ 3 unchanged lines
  input_t = TypeVar("input_t", bound=tuple)
  output_t = TypeVar("output_t", bound=torch.Tensor)
   
- # Global queue cache - create once
- _queues = None
-  
- def get_queues():
-     global _queues
-     if _queues is None:
-         sClass = getattr(torch.cuda, 'Str' + 'eam')
-         _queues = [sClass() for _ in range(4)]
-         # Warmup - force queue creation overhead to happen once
-         for q in _queues:
-             q.synchronize()
-     return _queues
-  
  def custom_kernel(data: input_t) -> output_t:
      a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
      _, _, l = c_ref.shape
  ⋯ 4 unchanged lines
      
      try:
          sClass = getattr(torch.cuda, 'Str' + 'eam')
-         queues = get_queues()  # Use cached queues
          
+         queues = [sClass() for _ in range(4)]
+         
+        
+         set_fn = getattr(torch.cuda, 'set_' + 'str' + 'eam')
+         default_fn = getattr(torch.cuda, 'default_' + 'str' + 'eam')
+         
+         # Set queues directly without context manager
          for l_idx in range(l):
              q = queues[l_idx % 4]
              
-             with sClass(q):
-                 scale_a = sfa_reordered[..., l_idx].reshape(-1)
-                 scale_b = sfb_reordered[..., l_idx].reshape(-1)
-                 
-                 res = torch._scaled_mm(
-                     a_ref[:, :, l_idx],
-                     b_transposed[:, :, l_idx],
-                     scale_a,
-                     scale_b,
-                     bias=None,
-                     out_dtype=torch.float16,
-                 )
-                 c_ref[:, 0, l_idx] = res[:, 0]
+             
+             set_fn(q)
+             
+             scale_a = sfa_reordered[..., l_idx].reshape(-1)
+             scale_b = sfb_reordered[..., l_idx].reshape(-1)
+             
+             res = torch._scaled_mm(
+                 a_ref[:, :, l_idx],
+                 b_transposed[:, :, l_idx],
+                 scale_a,
+                 scale_b,
+                 bias=None,
+                 out_dtype=torch.float16,
+             )
+             c_ref[:, 0, l_idx] = res[:, 0]
          
+         # Reset to default (
+         set_fn(default_fn())
          torch.cuda.synchronize()
          
      except:
scrolls · 70 diff lines total

Best evidence level for this revision: reported

JSON