can helion support in-tile reshape for quantization casting?

[vkuzo]

Describe the bug

I'm looking at representing quantization casting with tile invariant functions, so they can be easily tiled. This doesn't seem to work today with helion. I'm wondering if this is a bug or a limitation of the programming model.

For example, given the following pytorch code

def deepseek_1x128_f(x):
    fp8_max = torch.finfo(torch.float8_e4m3fn).max  # 448.0
    *lead, last = x.shape
    x_b = x.reshape(*lead, last // 128, 128)
    amax = x_b.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12).to(torch.float32)
    scale = (amax / fp8_max).to(torch.float32)  # forward scale
    qdata = (x_b.to(torch.float32) * (1.0 / scale)).to(torch.float8_e4m3fn)
    return qdata.reshape(*lead, last), scale.squeeze(-1)

I want, but currently cannot, write the following helion code

@helion.kernel(config=helion.Config(block_sizes=[32, BLOCK_N], num_warps=4))
def _tile_map_broken_kernel(x: torch.Tensor, f):
    M, N = x.size()
    out = torch.empty((M, N), dtype=torch.float8_e4m3fn, device=x.device)
    aux = torch.empty((M, N // BLOCK_N), dtype=torch.float32, device=x.device)
    for tile_m, tile_n in hl.tile([M, N]):
        acc = x[tile_m, tile_n]  # (BM, BLOCK_N)
        # Pass the reference recipe straight through. It does
        #   x.reshape(*lead, last // 128, 128)
        # on a symbolic tile -> Helion cannot prove the tile extent factors as
        # (n_groups, 128), tracing fails here.
        out_local, aux_local = f(acc)
        out[tile_m, tile_n] = out_local
        aux[tile_m, tile_n.begin // BLOCK_N] = aux_local.squeeze(-1)
    return out, aux


def run_helion_does_not_work(x):
    return _tile_map_broken_kernel(x, deepseek_1x128_f)

error:

helion.exc.TorchOpTracingError: RuntimeError: shape '[u1, (u2//128), 128]' is invalid for input of size u1*u2

To Reproduce

Full repro: https://gist.github.com/vkuzo/f2e2bf10e9526146ff836023f33b735c

Stdout of running the repro:
https://gist.github.com/vkuzo/afe6f2064fa44e58775c71da849ab4c5

Versions

torch 2.14.0.dev20260720+cu130, helion 1.2.0

[ethche]

@vkuzo Helion can be a bit picky about tile reshapes. But you can still squeeze unsqueeze. Could you try the following:

Requires N % 256 == 0

  def helion_deepseek_1x128_f(block):                                                                                                                                                                                                                                        
      fp8_max = torch.finfo(torch.float8_e4m3fn).max                                                                                                                                                                                                                         
                                                                                                                                                                                                                                                                             
      # [BM, 128] -> [BM, 1, 128]                                                                                                                                                                                                                                            
      grouped = block.unsqueeze(-2)                                                                                                                                                                                                                                          
      amax = grouped.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)                                                                                                                                                                                                       
      scale = amax.float() / fp8_max                                                                                                                                                                                                                                         
      quantized = (grouped.float() / scale).to(torch.float8_e4m3fn)                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                             
      return quantized.squeeze(-2), scale.squeeze(-1)                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                             
  @helion.kernel(config=helion.Config(block_sizes=[32], num_warps=4))                                                                                                                                                                                                        
  def kernel(x):                                                                                                                                                                                                                                                             
      M, N = x.shape                                                                                                                                                                                                                                                         
      out = torch.empty_like(x, dtype=torch.float8_e4m3fn)                                                                                                                                                                                                                   
      aux = torch.empty((M, N // 128), dtype=torch.float32, device=x.device)                                                                                                                                                                                                 
                                                                                                                                                                                                                                                                             
      # N is fixed; only M consumes a Config.block_sizes entry.                                                                                                                                                                                                              
      for tile_m, tile_n in hl.tile([M, N], block_size=[None, 128]):                                                                                                                                                                                                         
          out_local, aux_local = helion_deepseek_1x128_f(x[tile_m, tile_n])                                                                                                                                                                                                  
          out[tile_m, tile_n] = out_local                                                                                                                                                                                                                                    
          aux[tile_m, tile_n.begin // 128] = aux_local.squeeze(-1)                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                             
      return out, aux