Add support for forcing persistent kernels.

[Chillee]

There's a variety of cases where it's useful to force persistent kernels to ensure the kernel is cudagraphable. It might be possible to do that today but having a clear way of forcing it seems useful to me.

[oulgen]

reasonable suggestion, i'll take a look. @Chillee are you thinking like an env var / setting on the decorator to force it?

[Chillee]

@oulgen setting on the kernel, similar to static shapes.

[jansel]

For persistent launch grids you could do:

@helion.kernel(autotune_config_overrides={"pid_type": "persistent_interleaved"})

I don't see how this would affect cudagraphs though. It also wouldn't let you tune between persistent_blocked/persistent_interleaved.

[oulgen]

@Chillee is jason's suggestion good for you? if so, please close the issue

[Chillee]

@jansel It's cases like https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog8_Scaling_Expert_Parallelism_in_TensorRT-LLM_part2.md#moe-auxiliary-kernels

Note that the number of valid tokens is data-dependent. To guarantee CUDA graph compatibility, we cannot rely on any data-dependent information on the host. Thus, we further modified the kernels to use persistent thread blocks, which control the loop based on the valid token number on the device.

Another simple example might be a "padded sum". For example, assume I have a tensor of size [B, 100], and I have another size scalar tensor (on GPU) that controls how many rows I should actually compute.

import torch
import helion
import helion.language as hl
@helion.kernel
def padded_sum(x: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
    out = x.new_empty(x.shape[0])
    NUM_PROGRAMS = 5
    for pid in hl.grid(NUM_PROGRAMS):
        for tile_n in hl.tile(size[0] // NUM_PROGRAMS):
            rows = NUM_PROGRAMS * tile_n.index + pid
            x_tile = x[rows, :]
            sum_tile = torch.sum(x_tile, dim=1)
            out[rows] = sum_tile
    return out

out = padded_sum(torch.ones(100, 100), torch.tensor([80]))
>>> tensor([100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100., 100.,
        100., 100., 100., 100., 100., 100., 100., 100.,   0.,   0.,   0.,   0.,
          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,
          0.,   0.,   0.,   0.], device='cuda:0')

Perhaps this most accurately describes the semantics I want (albeit this would have a D2H sync).

@helion.kernel
def padded_sum(x: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
    out = x.new_empty(x.shape[0])
    for tile_n in hl.tile(size.item()):
        x_tile = x[tile_n, :]
        sum_tile = torch.sum(x_tile, dim=1)
        out[tile_n] = sum_tile
    return out

out = padded_sum(torch.ones(100, 100), torch.tensor([80]))

but I get

helion.exc.InternalError: AssertionError: no origin found for u0
While processing:
  File "/home/horace/monorepo/personal/horace/t.py", line 12, in padded_sum
    for tile_n in hl.tile(size.item()):

However, in principle, I think I should able to write this kernel like

@helion.kernel
def padded_sum(x: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
    out = x.new_empty(x.shape[0])
    for tile_n in hl.tile(size):
        x_tile = x[tile_n, :]
        sum_tile = torch.sum(x_tile, dim=1)
        out[tile_n] = sum_tile
    return out

and this would require a persistent kernel in order for it to "work" or be cudagraphable at all.

[jansel]

I see now. You would need to pack all your tensor sizes/strides into device tensors and never read them on the host.

I don't think your proposal would work, it fixes the hl.tile/grid issue, however, x[tile_n, :] and out[tile_n] = ... are both implicitly reading x.stride()/x.size()/etc which live on the host and become kernel args.

I think you would need something like:

class CudagraphableTensor(torch.Tensor):
   def __init__(self, x):
       # setup code that must run before cudagraph
       super().__init__(x)
       self._sizes = torch.tensor(x.size(), device=original.device)
       self._strides = torch.tensor(x.stride(), device=original.device)
   ...

Then teach Helion to read the tensor metadata from that.

[Chillee]

X.stride() and x.size() are both static - the tensor itself has a "host" size that's static and works perfectly fine with cudagraphs.

In practice you would only usually only ever have one dimension that's dynamic and thus you only need a single scalar tensor that lives on the GPU.

I'd note that this isn't really a speculative/contrived example - as you can see from the Megatron example above, it's a fairly regular use case for custom kernels nowadays.

[jansel]

I see yeah. We should add support for data-dependent hl.tile(size)

[jansel]

We could also make that implicitly force persistent kernels. If the grid loop bound lives on the GPU: force persistent. No need for an extra flag.

[Chillee]

Yeah thinking about it a bit more I think that's pretty plausible - although if you're launching a grid loop based off device-side info you could just do a synchronization as well.

[jansel]

You could do that 4-byte d2h copy, but then your kernel won't cudagraph (even with static shapes). It will also flush the queue of kernel launches, which is often more expensive than running the actual kernel. Sounds like always a bad idea.

[Chillee]

@jansel yeah I agree.

One thing that we probably would still want in the case of persistent kernels is a way to limit the # of SMs we use/tune for. For example, usually for persistent matmuls we don't want them to use 132 SMs because NCCL might occupy 4 SMs, and so forcing the matmul to use 132 SMs prevents it from overlapping with NCCL.

[yf225]

@Chillee Does something like this fit the bill for forcing persistent kernels?

helion/examples/grouped_gemm.py

Lines 159 to 176 in 51580b4

	for worker_id in hl.grid(num_workers):
	# Persistent thread pattern: each worker processes tiles across all groups
	# using strided/interleaved assignment for load balancing.
	# (i.e. each worker takes every num_workers-th tile. e.g., worker 0 takes tiles 0, N, 2N, ...)
	for g in hl.grid(G):
	group_start = group_offsets[g]
	group_end = group_offsets[g + 1]
	m_size = group_end - group_start
	if m_size > 0:
	# Compute tile grid dimensions for current group
	num_m_tiles = (m_size + BLOCK_M - 1) // BLOCK_M
	# Calculate number of N tiles (shared across all groups)
	num_n_tiles = (N + BLOCK_N - 1) // BLOCK_N
	num_group_tiles = num_m_tiles * num_n_tiles
	# Distribute tiles among workers using strided access pattern
	for local_tile in hl.grid(num_group_tiles):

here num_workers can be set to # of SMs, and the num_group_tiles in hl.tile(num_group_tiles) is data-dependent.

[Chillee]

@yf225 yeah of course you can manually write out the grid, but I think that feels like it shouldn't be required in Helion.