submission 69575

+ from task import input_t, output_t
  import torch
  import triton
  import triton.language as tl
- from task import input_t, output_t
   
- """
- NVFP4 Batched GEMV Optimization - Research-Based Approach
+ # =============================================================================
+ # OPTIMIZED SCALE FACTOR BLOCKING - From your original kernel
+ # =============================================================================
   
- KEY INSIGHTS:
- 1. NVFP4 has dual scaling: FP8 per 16 elements + global FP32
- 2. Dequantization happens IN HARDWARE on Blackwell Tensor Cores
- 3. torch._scaled_mm is optimized for GEMM not GEMV
- 4. Sequential batch loop kills parallelism
- 5. Scale blocking transformation can be overlapped
-  
- OPTIMIZATION TARGETS:
- - Eliminate sequential batch processing
- - Use async streams for computation-communication overlap
- - Optimize scale factor transformation with better memory access
- - Leverage torch's native operations where possible
- """
-  
  def ceil_div(a, b):
      return (a + b - 1) // b
   
- # ============================================================================
- # APPROACH 1: Stream-based Overlap (Quick Win)
- # Hide scale transformation behind previous batch's computation
- # ============================================================================
-  
  @triton.jit
- def blocked_transform_optimized(
+ def blocked_transform_kernel(
      inp, out, M, K, L,
      s_im, s_ik, s_il, s_ol, s_oe,
      BLK: tl.constexpr
  ):
-     """Optimized scale factor blocking with coalesced access"""
+     """Optimized blocking transformation for scale factors"""
      pid_l = tl.program_id(0)
      pid_b = tl.program_id(1)
      
  ⋯ 1 unchanged lines
      offs = pid_b * BLK + tl.arange(0, BLK)
      mask = offs < mk
      
-     # Decompose to (i, j)
      i = offs // K
      j = offs % K
      
-     # Blocking parameters
      nrb = (M + 127) // 128
      ncb = (K + 3) // 4
      
-     # Block indices
-     rb, ri = i // 128, i % 128
-     cb, ci = j // 4, j % 4
+     rb = i // 128
+     ri = i % 128
+     cb = j // 4
+     ci = j % 4
      
-     # Permute: view(nrb,128,ncb,4) -> permute(0,2,1,3)
-     perm = rb * ncb * 512 + cb * 512 + ri * 4 + ci
+     # Blocking logic from reference
+     perm = rb * ncb * 128 * 4 + cb * 128 * 4 + ri * 4 + ci
      
-     # Reshape and transpose
      chunk = perm // 512
-     rest = perm % 512
-     d1, d2, d3 = rest // 128, (rest % 128) // 4, rest % 4
+     in_chunk = perm % 512
+     d1 = in_chunk // 128
+     rest = in_chunk % 128
+     d2 = rest // 4
+     d3 = rest % 4
      
      out_idx = chunk * 512 + d2 * 16 + d1 * 4 + d3
      
-     # Coalesced load/store
+     # Load and store
      inp_idx = pid_l * s_il + i * s_im + j * s_ik
      out_idx_final = pid_l * s_ol + out_idx * s_oe
      
      val = tl.load(inp + inp_idx, mask=mask)
      tl.store(out + out_idx_final, val, mask=mask)
   
- def transform_scales_async(tensor, stream):
-     """Transform scales asynchronously on given stream"""
+ def transform_scales_gpu(tensor):
+     """GPU-based scale transformation"""
      M, K, L = tensor.shape
      mk = M * K
      result = torch.empty((L, mk), dtype=tensor.dtype, device='cuda')
      
-     t_gpu = tensor if tensor.is_cuda else tensor.cuda()
+     t = tensor.cuda() if not tensor.is_cuda else tensor
      
-     with torch.cuda.stream(stream):
-         BLK = 256
-         grid = (L, (mk + BLK - 1) // BLK)
-         blocked_transform_optimized[grid](
-             t_gpu, result, M, K, L,
-             t_gpu.stride(0), t_gpu.stride(1), t_gpu.stride(2),
-             result.stride(0), result.stride(1),
-             BLK=BLK
-         )
+     BLK = 256
+     grid = (L, (mk + BLK - 1) // BLK)
      
+     blocked_transform_kernel[grid](
+         t, result, M, K, L,
+         t.stride(0), t.stride(1), t.stride(2),
+         result.stride(0), result.stride(1),
+         BLK=BLK
+     )
+     
      return [result[i] for i in range(L)]
   
- def custom_kernel_streams(data: input_t) -> output_t:
-     """
-     Stream-based overlap: Transform scales for batch N+1 
-     while computing batch N
-     """
-     a, b, sfa_cpu, sfb_cpu, _, _, c = data
-     m, _, l = c.shape
+ # =============================================================================
+ # CUDA GRAPH OPTIMIZATION
+ # =============================================================================
+  
+ _graph_cache = {}
+  
+ class CUDAGraphExecutor:
+     """Captures CUDA graph to eliminate kernel launch overhead"""
+     def __init__(self, M, K, L):
+         self.M = M
+         self.K = K
+         self.L = L
+         self.graph = None
+         self.static_a = None
+         self.static_b = None
+         self.static_sfa_list = None
+         self.static_sfb_list = None
+         self.static_c = None
+         
+     def capture(self, a, b, sfa_list, sfb_list, c):
+         """Capture CUDA graph"""
+         # Warmup
+         for _ in range(3):
+             for i in range(self.L):
+                 res = torch._scaled_mm(
+                     a[:, :, i],
+                     b[:, :, i].transpose(0, 1),
+                     sfa_list[i],
+                     sfb_list[i],
+                     bias=None,
+                     out_dtype=torch.float16,
+                 )
+                 c[:, 0, i] = res[:, 0]
+         torch.cuda.synchronize()
+         
+         # Create static tensors
+         self.static_a = a.clone()
+         self.static_b = b.clone()
+         self.static_sfa_list = [s.clone() for s in sfa_list]
+         self.static_sfb_list = [s.clone() for s in sfb_list]
+         self.static_c = c.clone()
+         
+         # Capture
+         self.graph = torch.cuda.CUDAGraph()
+         with torch.cuda.graph(self.graph):
+             for i in range(self.L):
+                 res = torch._scaled_mm(
+                     self.static_a[:, :, i],
+                     self.static_b[:, :, i].transpose(0, 1),
+                     self.static_sfa_list[i],
+                     self.static_sfb_list[i],
+                     bias=None,
+                     out_dtype=torch.float16,
+                 )
+                 self.static_c[:, 0, i] = res[:, 0]
+         
+         return self
      
-     # Create streams for overlap
-     compute_stream = torch.cuda.Stream()
-     transform_stream = torch.cuda.Stream()
-     
-     # Transform first batch
-     with torch.cuda.stream(transform_stream):
-         sfa_list = transform_scales_async(sfa_cpu, transform_stream)
-         sfb_list = transform_scales_async(sfb_cpu, transform_stream)
-     
-     # Wait for first transformation
-     transform_stream.synchronize()
-     
-     # Process all batches
-     for i in range(l):
-         with torch.cuda.stream(compute_stream):
-             res = torch._scaled_mm(
-                 a[:, :, i],
-                 b[:, :, i].transpose(0, 1),
-                 sfa_list[i],
-                 sfb_list[i],
-                 bias=None,
-                 out_dtype=torch.float16,
-             )
-             c[:, 0, i] = res[:, 0]
-     
-     compute_stream.synchronize()
-     return c
+     def execute(self, a, b, sfa_list, sfb_list, c):
+         """Execute with graph"""
+         self.static_a.copy_(a, non_blocking=True)
+         self.static_b.copy_(b, non_blocking=True)
+         for i in range(self.L):
+             self.static_sfa_list[i].copy_(sfa_list[i], non_blocking=True)
+             self.static_sfb_list[i].copy_(sfb_list[i], non_blocking=True)
+         
+         self.graph.replay()
+         
+         c.copy_(self.static_c, non_blocking=True)
+         return c
   
- # ============================================================================
- # APPROACH 2: Pre-transform + Optimized Loop (Better)
- # Transform all scales once, then tight loop
- # ============================================================================
+ # =============================================================================
+ # MAIN KERNEL
+ # =============================================================================
   
- def custom_kernel_pretransform(data: input_t) -> output_t:
+ def custom_kernel(data: input_t) -> output_t:
      """
-     Pre-transform all scales, then run tight loop.
-     Reduces per-batch overhead.
-     """
-     a, b, sfa_cpu, sfb_cpu, _, _, c = data
-     m, _, l = c.shape
+     Optimized NVFP4 batched GEMV kernel.
      
-     # Transform ALL scales at once (single kernel launch overhead)
-     M_sfa, K_sfa, _ = sfa_cpu.shape
-     M_sfb, K_sfb, _ = sfb_cpu.shape
+     KEY OPTIMIZATIONS:
+     1. GPU-based scale transformation (parallel across L batches)
+     2. CUDA graphs (eliminates kernel launch overhead)
+     3. Optimized memory access patterns
      
-     sfa_all = torch.empty((l, M_sfa * K_sfa), dtype=sfa_cpu.dtype, device='cuda')
-     sfb_all = torch.empty((l, M_sfb * K_sfb), dtype=sfb_cpu.dtype, device='cuda')
+     LIMITATIONS:
+     - Still uses torch._scaled_mm which doesn't use Blackwell tensor cores
+     - To go faster: Need CUTLASS with tcgen05 instructions (multi-file setup)
      
-     # Single transformation for all batches
-     sfa_gpu = sfa_cpu if sfa_cpu.is_cuda else sfa_cpu.cuda()
-     sfb_gpu = sfb_cpu if sfb_cpu.is_cuda else sfb_cpu.cuda()
-     
-     BLK = 256
-     grid_sfa = (l, ceil_div(M_sfa * K_sfa, BLK))
-     grid_sfb = (l, ceil_div(M_sfb * K_sfb, BLK))
-     
-     blocked_transform_optimized[grid_sfa](
-         sfa_gpu, sfa_all, M_sfa, K_sfa, l,
-         sfa_gpu.stride(0), sfa_gpu.stride(1), sfa_gpu.stride(2),
-         sfa_all.stride(0), sfa_all.stride(1),
-         BLK=BLK
-     )
-     
-     blocked_transform_optimized[grid_sfb](
-         sfb_gpu, sfb_all, M_sfb, K_sfb, l,
-         sfb_gpu.stride(0), sfb_gpu.stride(1), sfb_gpu.stride(2),
-         sfb_all.stride(0), sfb_all.stride(1),
-         BLK=BLK
-     )
-     
-     # Tight loop over pre-transformed scales
-     for i in range(l):
-         res = torch._scaled_mm(
-             a[:, :, i],
-             b[:, :, i].transpose(0, 1),
-             sfa_all[i],
-             sfb_all[i],
-             bias=None,
-             out_dtype=torch.float16,
-         )
-         c[:, 0, i] = res[:, 0]
-     
-     return c
-  
- # ============================================================================
- # APPROACH 3: Fused Transformation (Best for large L)
- # Combine all transformations into single grid
- # ============================================================================
-  
- def custom_kernel_fused(data: input_t) -> output_t:
+     Expected speedup: 1.5-2.5x over reference
      """
-     Fuse both scale transformations into one kernel grid.
-     Best when L is large.
-     """
-     a, b, sfa_cpu, sfb_cpu, _, _, c = data
-     m, _, l = c.shape
+     a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data
      
-     M_sfa, K_sfa, _ = sfa_cpu.shape
-     M_sfb, K_sfb, _ = sfb_cpu.shape
+     M = a.size(0)
+     K = a.size(1)
+     L = a.size(2)
      
-     # Allocate output
-     sfa_blocked = torch.empty((l, M_sfa * K_sfa), dtype=sfa_cpu.dtype, device='cuda')
-     sfb_blocked = torch.empty((l, M_sfb * K_sfb), dtype=sfb_cpu.dtype, device='cuda')
+     # Transform scales using the CORRECT format for torch._scaled_mm
+     # (Not the permuted format - that's for CUTLASS)
+     sfa_list = transform_scales_gpu(sfa)
+     sfb_list = transform_scales_gpu(sfb)
      
-     # Move to GPU
-     sfa_gpu = sfa_cpu.cuda() if not sfa_cpu.is_cuda else sfa_cpu
-     sfb_gpu = sfb_cpu.cuda() if not sfb_cpu.is_cuda else sfb_cpu
+     # Use CUDA graphs for repeated calls
+     cache_key = (M, K, L)
      
-     # Launch both transformations (they'll run concurrently)
-     BLK = 256
-     blocked_transform_optimized[(l, ceil_div(M_sfa * K_sfa, BLK))](
-         sfa_gpu, sfa_blocked, M_sfa, K_sfa, l,
-         sfa_gpu.stride(0), sfa_gpu.stride(1), sfa_gpu.stride(2),
-         sfa_blocked.stride(0), sfa_blocked.stride(1),
-         BLK=BLK
-     )
-     blocked_transform_optimized[(l, ceil_div(M_sfb * K_sfb, BLK))](
-         sfb_gpu, sfb_blocked, M_sfb, K_sfb, l,
-         sfb_gpu.stride(0), sfb_gpu.stride(1), sfb_gpu.stride(2),
-         sfb_blocked.stride(0), sfb_blocked.stride(1),
-         BLK=BLK
-     )
+     if cache_key not in _graph_cache:
+         # First time - capture graph
+         executor = CUDAGraphExecutor(M, K, L)
+         executor.capture(a, b, sfa_list, sfb_list, c)
+         _graph_cache[cache_key] = executor
+     else:
+         # Reuse cached graph
+         executor = _graph_cache[cache_key]
      
-     # Main compute loop (still sequential, but tight)
-     for i in range(l):
-         res = torch._scaled_mm(
-             a[:, :, i],
-             b[:, :, i].transpose(0, 1),
-             sfa_blocked[i],
-             sfb_blocked[i],
-             bias=None,
-             out_dtype=torch.float16,
-         )
-         c[:, 0, i] = res[:, 0]
+     result = executor.execute(a, b, sfa_list, sfb_list, c)
      
-     return c
-  
- # Entry point - use best approach
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Main entry - chooses best approach based on problem size.
-     For now, use fused transformation (Approach 3).
-     """
-     return custom_kernel_fused(data)
   No newline at end of file
+     return result
   No newline at end of file