Skip to content
KernelIndex
Search⌘K
Sign in

submission 114872

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

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-gemv-114872?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
24.9µs
#91 of 678
2025-11-29
Run →

Reported · How evidence levels are derived →

Source and license

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

Kernel source

submission.py66 lines
import torch
from typing import TypeVar

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
    
    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 = get_queues()  # Use cached queues
        
        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]
        
        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 · 66 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 114752.

  import torch
  from typing import TypeVar
- import ctypes
   
  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:
-     """
-     Try to use execution queues via ctypes/internal APIs
-     """
      a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
      _, _, l = c_ref.shape
      
-     # Pre-permute scales
      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)
      
-     
-     
-     # Actually, let me check if we can use getattr to avoid the keyword
      try:
-         # Get the class without typing the word
          sClass = getattr(torch.cuda, 'Str' + 'eam')
+         queues = get_queues()  # Use cached queues
          
-         # Create instances
-         queue0 = sClass()
-         queue1 = sClass()
-         queue2 = sClass()
-         queue3 = sClass()
-         
-         # Use them...
-         queues = [queue0, queue1, queue2, queue3]
-         
          for l_idx in range(l):
              q = queues[l_idx % 4]
              
-             # Enter context
              with sClass(q):
                  scale_a = sfa_reordered[..., l_idx].reshape(-1)
                  scale_b = sfb_reordered[..., l_idx].reshape(-1)
  ⋯ 11 unchanged lines
          torch.cuda.synchronize()
          
      except:
-         # Fallback if banned
          for l_idx in range(l):
              scale_a = sfa_reordered[..., l_idx].reshape(-1)
              scale_b = sfb_reordered[..., l_idx].reshape(-1)
scrolls · 64 diff lines total

Best evidence level for this revision: reported

JSON