`top_k_top_p_sampling_from_logits` incompatible with torch.compile + CUDAGraph

[sharvil]

Python version: 3.10
PyTorch version: 2.6.0
FlashInfer version: 0.2.3+cu124torch2.6

Minimal repro:

import torch

from flashinfer.sampling import top_k_top_p_sampling_from_logits


def sample(x):
  return top_k_top_p_sampling_from_logits(
    logits=x,
    top_k=0,
    top_p=1.0,
  )

sample = torch.compile(sample, mode='reduce-overhead', fullgraph=True)
sample(torch.randn([1, 32]).cuda())

Produces the following error:

torch._dynamo.exc.Unsupported: call_function UserDefinedClassVariable(<class 'types.SimpleNamespace'>) [] {'sampling_from_probs': NestedUserFunctionVariable(), 'top_p_sampling_from_probs': NestedUserFunctionVariable(), 'top_k_sampling_from_probs': NestedUserFunctionVariable(), 'min_p_sampling_from_probs': NestedUserFunctionVariable(), 'top_k_top_p_sampling_from_probs': NestedUserFunctionVariable(), 'top_p_renorm_probs': NestedUserFunctionVariable(), 'top_k_renorm_probs': NestedUserFunctionVariable(), 'top_k_mask_logits': NestedUserFunctionVariable(), 'chain_speculative_sampling': NestedUserFunctionVariable()}

from user code:
   File ".../compile.py", line 7, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 810, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1129, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 393, in get_sampling_module
    _sampling_module = SimpleNamespace(

Setting fullgraph=False succeeds.

[sharvil]

Fails on nightly as well:

torch._dynamo.exc.Unsupported: builtin: add [<class 'torch._dynamo.variables.misc.GetAttrVariable'>, <class 'torch._dynamo.variables.lists.ListVariable'>] False

from user code:
   File ".../compile.py", line 9, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 810, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1129, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 39, in get_sampling_module
    FLASHINFER_CSRC_DIR / "sampling.cu",
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 855, in __truediv__
    return self._make_child((key,))
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 617, in _make_child
    drv, root, parts = self._flavour.join_parsed_parts(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 108, in join_parsed_parts
    return drv, root, parts + parts2

[sharvil]

Calling the kernel first to get around the lazy module load followed by torch.compile on nightly leads to:

torch._dynamo.exc.Unsupported: SETUP_WITH ConstantVariable(device: device(type='cuda', index=0))

from user code:
   File ".../compile.py", line 9, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 810, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1129, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 317, in top_k_mask_logits
    with logits.device as device:

[yzh119]

Fails on nightly as well:

torch._dynamo.exc.Unsupported: builtin: add [<class 'torch._dynamo.variables.misc.GetAttrVariable'>, <class 'torch._dynamo.variables.lists.ListVariable'>] False

from user code:
   File ".../compile.py", line 9, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 810, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1129, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 39, in get_sampling_module
    FLASHINFER_CSRC_DIR / "sampling.cu",
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 855, in __truediv__
    return self._make_child((key,))
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 617, in _make_child
    drv, root, parts = self._flavour.join_parsed_parts(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/pathlib.py", line 108, in join_parsed_parts
    return drv, root, parts + parts2

Seems dynamo is trying to capture the JIT compilation process in flashinfer. Would you mind JIT compile and cache the kernels first (or using the AOT version of flashinfer) and try it again?

[yzh119]

Calling the kernel first to get around the lazy module load followed by torch.compile on nightly leads to:

torch._dynamo.exc.Unsupported: SETUP_WITH ConstantVariable(device: device(type='cuda', index=0))

from user code:
   File ".../compile.py", line 9, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 810, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1129, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 317, in top_k_mask_logits
    with logits.device as device:

The device guard should have been removed in #969

[sharvil]

Here's what I'm seeing on the 0.2.4 release:

torch._dynamo.exc.TorchRuntimeError: Failed running call_function sampling.top_k_mask_logits.default(*(FakeTensor(..., device='cuda:0', size=(1, 32)), FakeTensor(..., device='cuda:0', size=(1, 32)), None, 0), **{}):
Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor). If you're using torch.compile/export/fx, it is likely that we are erroneously tracing into a custom kernel. To fix this, please wrap the custom kernel into an opaque custom op. Please see the following for details: https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html

from user code:
   File ".../sample.py", line 7, in sample
    return top_k_top_p_sampling_from_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 803, in top_k_top_p_sampling_from_logits
    masked_logits = top_k_mask_logits(logits, top_k)
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 1123, in top_k_mask_logits
    return get_sampling_module().top_k_mask_logits(
  File "/home/sharvil/.conda/envs/lmnt/lib/python3.10/site-packages/flashinfer/sampling.py", line 316, in top_k_mask_logits
    module.top_k_mask_logits.default(

Code again for reference:

import torch

from flashinfer.sampling import top_k_top_p_sampling_from_logits


def sample(x):
  return top_k_top_p_sampling_from_logits(
    logits=x,
    top_k=0,
    top_p=1.0,
  )

sample(torch.randn([1, 32]).cuda())
sample = torch.compile(sample, mode='reduce-overhead', fullgraph=True)
sample(torch.randn([1, 32]).cuda())

[yzh119]

Hi @sharvil , oh yes it was because we removed the custom op registration (it incurs significant cpu overhead, when kernel is not captured by cudagraph), feel free to add me back:

flashinfer/flashinfer/utils.py

Lines 239 to 249 in 5751fc6

	# NOTE(Zihao): torch.library.custom_op has significant overhead as mentioned in the following link
	# https://github.com/vllm-project/vllm/blob/36e76700453924c8d421db99af70a88a1df835cd/vllm/utils.py#L1660-L1674
	# return torch.library.custom_op(
	# name,
	# fn,
	# mutates_args=mutates_args,
	# device_types=device_types,
	# schema=schema,
	# )
	return lambda x: x

[guilhermeleobas]

Hi @yzh119, I’ve resumed work on adding torch.compile support to flashinfer. One issue I ran into is the optional generator argument (gen_) introduced in #912. The type torch._C.Generator isn't supported as a schema type. Do you know how widely this argument is actually used? If it’s rare, I can remove it, as the CUDA code has an if-else that creates a generator if the gen_ argument is not given:

flashinfer/csrc/sampling.cu

Lines 217 to 218 in 018b551

	auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
	gen_, at::cuda::detail::getDefaultCUDAGenerator());