submission 107135

- """NVFP4 GEMV - Target 12μs: All Optimizations Combined
+ """NVFP4 GEMV - Direct Assignment without Intermediate
   
- Building on submission_final_ultra.py (24.4μs on server), this adds:
+ Based on submission_final_ultra.py (24.4μs baseline on server).
   
- 1. Scale factor caching (by tensor data_ptr) - eliminates ~L×2 permute kernels
- 2. Optimized stream wait (wait_stream instead of synchronize per stream)
- 3. Direct view operations (select().copy_() instead of slice assignment)
-  
- Target: 24.4 μs → 12 μs on competition server
+ Try to minimize intermediate tensor creation by restructuring the assignment.
  """
   
  import torch
   
- # Global caches
- _streams = {}
- _scale_cache = {}
+ _s = [torch.cuda.Stream() for _ in range(8)]
+ _mm = torch._scaled_mm
   
   
- def get_stream(idx):
-     """Get or create CUDA stream."""
-     if idx not in _streams:
-         _streams[idx] = torch.cuda.Stream()
-     return _streams[idx]
-  
-  
- def permuted_to_blocked(permuted_tensor, l_idx):
-     """Convert pre-permuted scale factor to blocked format."""
-     tensor = permuted_tensor.select(-1, l_idx)
-     return tensor.permute(2, 4, 0, 1, 3).reshape(-1, 32, 16).flatten().contiguous()
-  
-  
- def get_all_scales(sfa_permuted, sfb_permuted, L):
-     """Get cached scales or compute and cache them."""
-     # Use data pointer as cache key
-     key = (sfa_permuted.data_ptr(), sfb_permuted.data_ptr(), L)
-  
-     if key not in _scale_cache:
-         # Compute all scales at once and cache
-         scales_a = []
-         scales_b = []
-         for l in range(L):
-             scales_a.append(permuted_to_blocked(sfa_permuted, l))
-             scales_b.append(permuted_to_blocked(sfb_permuted, l))
-         _scale_cache[key] = (scales_a, scales_b)
-  
-     return _scale_cache[key]
-  
-  
  def custom_kernel(data):
-     """Optimized NVFP4 GEMV kernel targeting 12μs."""
-     a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data
-     M, K_half, L = a.shape
+     a, b, _, _, sfa_p, sfb_p, c = data
+     L = a.shape[2]
   
-     # Get cached scale factors
-     scales_a, scales_b = get_all_scales(sfa_permuted, sfb_permuted, L)
+     for l in range(L):
+         with torch.cuda.stream(_s[l]):
+             # Get result first
+             result = _mm(
+                 a.select(2, l),
+                 b.select(2, l).T,
+                 sfa_p.select(-1, l).permute(2, 4, 0, 1, 3).flatten(),
+                 sfb_p.select(-1, l).permute(2, 4, 0, 1, 3).flatten(),
+                 bias=None,
+                 out_dtype=torch.float16
+             )
   
-     if L == 1:
-         # L=1: Single GEMM, no streams needed
-         result = torch._scaled_mm(
-             a[:, :, 0],
-             b[:, :, 0].transpose(0, 1),
-             scales_a[0],
-             scales_b[0],
-             bias=None,
-             out_dtype=torch.float16
-         )
-         # Direct copy using select (potentially faster than slice assignment)
-         c.select(2, 0).select(1, 0).copy_(result.select(1, 0))
+             # Direct select on both sides - may avoid some indexing overhead
+             c.select(2, l).select(1, 0).copy_(result.select(1, 0))
   
-     else:
-         # L>1: Parallel streams with pre-cached scales
-         for l_idx in range(L):
-             stream = get_stream(l_idx)
-             with torch.cuda.stream(stream):
-                 result = torch._scaled_mm(
-                     a[:, :, l_idx],
-                     b[:, :, l_idx].transpose(0, 1),
-                     scales_a[l_idx],
-                     scales_b[l_idx],
-                     bias=None,
-                     out_dtype=torch.float16
-                 )
-                 c.select(2, l_idx).select(1, 0).copy_(result.select(1, 0))
-  
-         # Efficient stream synchronization
-         for l_idx in range(L):
-             torch.cuda.current_stream().wait_stream(get_stream(l_idx))
-  
      torch.cuda.synchronize()
      return c