submission 686072

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  """
- Monkey-patch fused_dynamic_mxfp4_quant_moe_sort with our HIP hardware quant.
- fused_moe still receives bf16 input (no KeyError), but A quant uses hardware instruction.
- Only patches for M=512 (2-stage path). M≤128 uses CKTile (no quant needed).
+ Combined optimization: HIP sort + HIP A quant, both monkey-patched into fused_moe.
+ - HIP sort: ~5µs vs AITER ~40µs → save ~35µs
+ - HIP A quant: ~5µs vs Triton ~30µs → save ~25µs (topk==1 only)
+ - Total savings: ~60µs for M=512 shapes
  """
  import os, sys
  os.environ.setdefault('PYTORCH_ROCM_ARCH', 'gfx950')
  ⋯ 4 unchanged lines
  from task import input_t, output_t
  from aiter import ActivationType, QuantType
   
- # HIP A quantization kernel
  HIP_SRC = r"""
  #include <hip/hip_runtime.h>
  __device__ __forceinline__ uint32_t f2u(float f){uint32_t u;__builtin_memcpy(&u,&f,4);return u;}
  __device__ __forceinline__ float u2f(uint32_t u){float f;__builtin_memcpy(&f,&u,4);return f;}
+  
+ // ============ Sort Kernel ============
+ __global__ void moe_sort_kernel(
+     const int* __restrict__ ti, const float* __restrict__ tw,
+     int* __restrict__ si, float* __restrict__ sw,
+     int* __restrict__ sei, int* __restrict__ nv,
+     int* __restrict__ ec, int* __restrict__ eo,
+     int M, int topk, int E, int bm
+ ) {
+     if(blockIdx.x!=0) return;
+     for(int i=threadIdx.x;i<E;i+=blockDim.x) ec[i]=0;
+     __syncthreads();
+     for(int i=threadIdx.x;i<M*topk;i+=blockDim.x) atomicAdd(&ec[ti[i]],1);
+     __syncthreads();
+     if(threadIdx.x==0){
+         int off=0;
+         for(int e=0;e<E;e++){eo[e]=off;int p=((ec[e]+bm-1)/bm)*bm;off+=p;}
+         eo[E]=off; nv[0]=M*topk; nv[1]=0;
+         int blk=0;
+         for(int e=0;e<E;e++){int p=((ec[e]+bm-1)/bm)*bm;for(int b=0;b<p/bm;b++)sei[blk++]=e;}
+     }
+     __syncthreads();
+     int tp=eo[E];
+     for(int i=threadIdx.x;i<tp;i+=blockDim.x){si[i]=M*topk;sw[i]=0.0f;}
+     __syncthreads();
+     for(int i=threadIdx.x;i<E;i+=blockDim.x)ec[i]=0;
+     __syncthreads();
+     for(int i=threadIdx.x;i<M*topk;i+=blockDim.x){
+         int eid=ti[i];int slot=atomicAdd(&ec[eid],1);
+         int pos=eo[eid]+slot;si[pos]=i;sw[pos]=tw[i];
+     }
+ }
+  
+ // ============ A Quant Kernel ============
  __global__ void quant_a_kernel(
      const uint16_t* __restrict__ A, uint8_t* __restrict__ fp4,
      uint8_t* __restrict__ scale, int M, int K
  ⋯ 19 unchanged lines
      const uint8_t*p=(const uint8_t*)pk;
      for(int i=0;i<16;i++) fp4[off+i]=p[i];
  }
- void launch_quant_a(torch::Tensor A, torch::Tensor fp4, torch::Tensor scale, int M, int K){
+  
+ void launch_sort(torch::Tensor ti, torch::Tensor tw, torch::Tensor si, torch::Tensor sw,
+     torch::Tensor sei, torch::Tensor nv, torch::Tensor ec, torch::Tensor eo,
+     int M, int topk, int E, int bm){
+     moe_sort_kernel<<<1,256>>>((const int*)ti.data_ptr(),(const float*)tw.data_ptr(),
+         (int*)si.data_ptr(),(float*)sw.data_ptr(),(int*)sei.data_ptr(),(int*)nv.data_ptr(),
+         (int*)ec.data_ptr(),(int*)eo.data_ptr(),M,topk,E,bm);
+ }
+ void launch_quant(torch::Tensor A, torch::Tensor fp4, torch::Tensor scale, int M, int K){
      quant_a_kernel<<<M, K/32>>>((const uint16_t*)A.data_ptr(),
          (uint8_t*)fp4.data_ptr(),(uint8_t*)scale.data_ptr(),M,K);
  }
  """
- CPP_SRC = "void launch_quant_a(torch::Tensor,torch::Tensor,torch::Tensor,int,int);"
   
+ CPP_SRC = """
+ void launch_sort(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
+     torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,int,int,int,int);
+ void launch_quant(torch::Tensor,torch::Tensor,torch::Tensor,int,int);
+ """
+  
  try:
-     _hip = load_inline(name='hw_quant_v3', cpp_sources=[CPP_SRC], cuda_sources=[HIP_SRC],
-         functions=['launch_quant_a'], verbose=True,
-         extra_cuda_cflags=["--offload-arch=gfx950", "-std=c++20", "-O3"])
-     _HAS_HW = True
+     _hip = load_inline(name='moe_combo_v1', cpp_sources=[CPP_SRC], cuda_sources=[HIP_SRC],
+         functions=['launch_sort','launch_quant'], verbose=True,
+         extra_cuda_cflags=["--offload-arch=gfx950","-std=c++20","-O3"])
+     _HAS = True
  except:
-     _HAS_HW = False
+     _HAS = False
   
- # Monkey-patch: replace Triton quant with HIP hw quant
+ # ============ Monkey-patch sorting ============
+ import aiter
  import aiter.fused_moe as _fm
+  
+ _sort_bufs = {}
+  
+ def _fast_sort(topk_ids, topk_weights, num_experts, model_dim,
+                moebuf_dtype, block_size, expert_mask, num_local_tokens,
+                dispatch_policy, use_opus):
+     if not _HAS:
+         return _fm._moe_sorting_impl.__wrapped__(topk_ids, topk_weights, num_experts,
+             model_dim, moebuf_dtype, block_size, expert_mask, num_local_tokens,
+             dispatch_policy, use_opus)
+  
+     device = topk_ids.device
+     M, topk = topk_ids.shape
+     bm = int(block_size)
+     E = num_experts
+     max_pad = M*topk + E*bm - topk
+     max_blk = (max_pad + bm - 1) // bm
+  
+     key = (M, E, model_dim, bm)
+     if key not in _sort_bufs:
+         _sort_bufs[key] = {
+             'si': torch.empty(max_pad, dtype=torch.int32, device=device),
+             'sw': torch.empty(max_pad, dtype=torch.float32, device=device),
+             'sei': torch.empty(max_blk, dtype=torch.int32, device=device),
+             'nv': torch.empty(2, dtype=torch.int32, device=device),
+             'ec': torch.zeros(E, dtype=torch.int32, device=device),
+             'eo': torch.zeros(E+1, dtype=torch.int32, device=device),
+             'buf': torch.empty((M, model_dim), dtype=moebuf_dtype, device=device),
+         }
+     b = _sort_bufs[key]
+     _hip.launch_sort(topk_ids, topk_weights, b['si'], b['sw'], b['sei'], b['nv'],
+                      b['ec'], b['eo'], M, topk, E, bm)
+     return b['si'], b['sw'], b['sei'], b['nv'], b['buf']
+  
+ _fm._moe_sorting_impl = _fast_sort
+  
+ # ============ Monkey-patch A quant ============
  from aiter.ops.triton.quant.fused_mxfp4_quant import fused_dynamic_mxfp4_quant_moe_sort as _orig_quant
  from aiter.utility import fp4_utils as _fp4u
   
  _qcache = {}
   
- _patch_call_count = 0
- def _fast_quant_moe_sort(input, sorted_ids, num_valid_ids, token_num, topk, block_size):
-     """Replace Triton fused quant+sort with HIP hw quant + separate scale sort."""
-     global _patch_call_count
-     _patch_call_count += 1
-     if _patch_call_count <= 2:
-         print(f"[PATCH] Called #{_patch_call_count}: M={token_num} topk={topk} shape={input.shape} dtype={input.dtype}", file=sys.stderr)
-     # ONLY replace A quant (topk==1). Keep Triton for ALL intermediate quant.
-     if not _HAS_HW or input.dtype != torch.bfloat16 or topk != 1:
+ def _fast_quant(input, sorted_ids, num_valid_ids, token_num, topk, block_size):
+     if not _HAS or input.dtype != torch.bfloat16 or topk != 1:
          return _orig_quant(input, sorted_ids, num_valid_ids, token_num, topk, block_size)
   
-     M = input.shape[0]
-     K = input.shape[1]
+     M, K = input.shape
      dev = input.device
      key = (M, K)
      if key not in _qcache:
  ⋯ 2 unchanged lines
              torch.empty(M, K//32, dtype=torch.uint8, device=dev),
          )
      fp4_buf, scale_buf = _qcache[key]
+     _hip.launch_quant(input, fp4_buf, scale_buf, M, K)
   
-     # Fast HIP quantization (~5µs vs Triton ~72µs)
-     _hip.launch_quant_a(input, fp4_buf, scale_buf, M, K)
-  
-     # Reinterpret as typed tensors
      fp4 = torch.tensor([], dtype=torch.float4_e2m1fn_x2, device=dev).set_(
          fp4_buf.untyped_storage(), fp4_buf.storage_offset(), (M, K//2), (K//2, 1))
      scale = torch.tensor([], dtype=torch.float8_e8m0fnu, device=dev).set_(
          scale_buf.untyped_storage(), scale_buf.storage_offset(), (M, K//32), (K//32, 1))
   
-     # Sort only the scale (like M>1024 separate path)
      scale_sorted = _fp4u.moe_mxfp4_sort(
-         scale.view(token_num, topk, -1) if topk > 1 else scale,
-         sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
+         scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
          token_num=token_num, block_size=block_size)
-  
      return fp4, scale_sorted
   
- # Apply the patch
  import aiter.ops.triton.quant.fused_mxfp4_quant as _qmod
- _qmod.fused_dynamic_mxfp4_quant_moe_sort = _fast_quant_moe_sort
- # Also patch the import in fused_moe module
- _fm.fused_dynamic_mxfp4_quant_moe_sort = _fast_quant_moe_sort
+ _qmod.fused_dynamic_mxfp4_quant_moe_sort = _fast_quant
+ _fm.fused_dynamic_mxfp4_quant_moe_sort = _fast_quant
   
+ # ============ Main ============
  from aiter.fused_moe import fused_moe
  _ACT = ActivationType.Silu
  _QT = QuantType.per_1x32