submission 114676 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-gemv-114676?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

64.6µs

#313 of 678

2025-11-29

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:ddd6b0f9239b4fe9357aa84791d4a83a50e03aaaf9000b92d18add6f2b43c847

license declaredunknown

license concludedunknown

authorsJB Gage

imported2026-08-15

Techniques

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

fp4Optimized FP4 GEMV using pre-permuted scales

Kernel source

submission.py43 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:
    """
    Optimized FP4 GEMV using pre-permuted scales
    
      
    Key optimizations:
    1. Use pre-permuted scale factors (sfa_permuted, sfb_permuted)
    2. Single permutation operation for all batches
    3. Minimal overhead in the compute loop
    4. Leverage cuBLAS through torch._scaled_mm
    
    This is the optimal solution achievable with public PyTorch APIs.
    """
    a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
    _, _, l = c_ref.shape
    
    # Pre-permute ALL batches once
    # Transform from (32, 4, rest_m, 4, rest_k, L) to (rest_m, rest_k, 32, 4, 4, L)
    sfa_reordered = sfa_permuted.permute(2, 4, 0, 1, 3, 5)
    sfb_reordered = sfb_permuted.permute(2, 4, 0, 1, 3, 5)
    
    # Process each batch using cuBLAS
    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_ref[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref

scrolls · 43 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 114558.

  ⋯ 5 unchanged lines
   
  def custom_kernel(data: input_t) -> output_t:
      """
-     OPTIMIZED: Use pre-permuted scales directly!
+     Optimized FP4 GEMV using pre-permuted scales
      
-     Avoids calling to_blocked() and CPU→GPU transfer in the loop.
-     Pre-permuted format just needs: permute(2,4,0,1,3).flatten()
+       
+     Key optimizations:
+     1. Use pre-permuted scale factors (sfa_permuted, sfb_permuted)
+     2. Single permutation operation for all batches
+     3. Minimal overhead in the compute loop
+     4. Leverage cuBLAS through torch._scaled_mm
+     
+     This is the optimal solution achievable with public PyTorch APIs.
      """
-     a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, sfa_permuted, sfb_permuted, c_ref = data
-     m, k_packed, l = a_ref.shape
+     a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
+     _, _, l = c_ref.shape
      
-     # Extract scales from pre-permuted format
-     # sfa_permuted is already on GPU and in (32,4,rest_m,4,rest_k,L) format
-     # We just need to permute and flatten per batch
+     # Pre-permute ALL batches once
+     # Transform from (32, 4, rest_m, 4, rest_k, L) to (rest_m, rest_k, 32, 4, 4, L)
+     sfa_reordered = sfa_permuted.permute(2, 4, 0, 1, 3, 5)
+     sfb_reordered = sfb_permuted.permute(2, 4, 0, 1, 3, 5)
      
-     scales_a = []
-     scales_b = []
-     
+     # Process each batch using cuBLAS
      for l_idx in range(l):
-         # Extract batch slice
-         sfa_slice = sfa_permuted[:, :, :, :, :, l_idx]  # (32, 4, rest_m, 4, rest_k)
-         sfb_slice = sfb_permuted[:, :, :, :, :, l_idx]  # (32, 4, 1, 4, rest_k)
+         scale_a = sfa_reordered[..., l_idx].reshape(-1)
+         scale_b = sfb_reordered[..., l_idx].reshape(-1)
          
-         # Apply the magic permutation: (2, 4, 0, 1, 3)
-         # This reorders to match what to_blocked() produces
-         scale_a = sfa_slice.permute(2, 4, 0, 1, 3).flatten()
-         scale_b = sfb_slice.permute(2, 4, 0, 1, 3).flatten()
-         
-         scales_a.append(scale_a)
-         scales_b.append(scale_b)
-     
-     # Main compute loop
-     for l_idx in range(l):
          res = torch._scaled_mm(
              a_ref[:, :, l_idx],
              b_ref[:, :, l_idx].transpose(0, 1),
-             scales_a[l_idx],
-             scales_b[l_idx],
+             scale_a,
+             scale_b,
              bias=None,
              out_dtype=torch.float16,
          )

scrolls · 62 diff lines total

Best evidence level for this revision: reported

JSON