[Issue]: Some AITER Operators are not compatible with torch.compile

[tjtanaa]

Problem Description

Torch compile is not able to convert some of the built-in python functions into graph. e.g. highlighted in green

ERROR 03-24 02:09:11 [core.py:343]     ck_config = get_CKGEMM_config(m, n, k)
ERROR 03-24 02:09:11 [core.py:343]   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/polyfills/__init__.py", line 135, in getattr_and_trace
ERROR 03-24 02:09:11 [core.py:343]     return fn(*args[2:], **kwargs)
ERROR 03-24 02:09:11 [core.py:343]   File "/app/Quark/aiter/aiter/ops/gemm_op_a8w8.py", line 84, in get_CKGEMM_config
+ ERROR 03-24 02:09:11 [core.py:343]     if not hasattr(get_CKGEMM_config, "ckgemm_dict"):
ERROR 03-24 02:09:11 [core.py:343] 
ERROR 03-24 02:09:11 [core.py:343] Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

The full logs can be found in the attachment.

llama_fp8-int8aiterck-v1.txt

Operating System

Ubuntu 22.04.4 LTS (Jammy Jellyfish)

CPU

AMD EPYC 9654 96-Core Processor

GPU

AMD Instinct MI300X

ROCm Version

ROCm 6.3.1

ROCm Component

No response

Steps to Reproduce

Setup vLLM

git clone https://github.com/EmbeddedLLM/vllm.git --branch aiter-int8-linear  aiter-int8-linear
cd  aiter-int8-linear
# if you have installed vllm before
python3 -m pip uninstall -y vllm
export PYTORCH_ROCM_ARCH=gfx942
python3 setup.py develop

Run with V1 engine
VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=1 vllm serve neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[valarLip]

do we have any runable case for now? like with tirton moe?

[tjtanaa]

No that I am aware of.

[valarLip]

No that I am aware off.

i mean without aiter, if so i believe that's a torch compile's issue, instead of aiter's issue. with compile, we even can't have a print/log, that's not make sense to be use

[tjtanaa]

Ic. So there are a lot of restrictions on this. However, vLLM V1 architecture depends on full graph mode using torch.compile. It would be helpful moving forward if there are alternative solutions to this problem.

[tjtanaa]

I find this blogpost informative:
https://dev.to/aaronlangford31/lessons-learned-from-using-torch-inductor-for-inference-1ma7

[valarLip]

I find this blogpost informative: https://dev.to/aaronlangford31/lessons-learned-from-using-torch-inductor-for-inference-1ma7

thanks for sharing, however, it looks not good to me, for example, NV's new concept about cutile also using jit with cache, which is not allowed in this blog..... "For example, writes to self are completely disallowed by the Inductor compiler."

[tjtanaa]

I think we could refer to direct_register_custom_op from vllm.utils https://github.com/vllm-project/vllm/blob/4c41278b77a8b14fbab6dfeb95d5185c12018fbc/vllm/utils.py#L2002

torch.library.custom_op can have significant overhead because it needs to consider complicated dispatching logic.

"""
    `torch.library.custom_op` can have significant overhead because it
    needs to consider complicated dispatching logic. This function
    directly registers a custom op and dispatches it to the CUDA backend.
    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
    for more details.

    By default, the custom op is registered to the vLLM library. If you
    want to register it to a different library, you can pass the library
    object to the `target_lib` argument.

    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
    library object. If you want to bind the operator to a different library,
    make sure the library object is alive when the operator is used.
    """

[zhiding512]

Problem Description

Torch compile is not able to convert some of the built-in python functions into graph. e.g. highlighted in green

ERROR 03-24 02:09:11 [core.py:343] ck_config = get_CKGEMM_config(m, n, k)
ERROR 03-24 02:09:11 [core.py:343] File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/polyfills/init.py", line 135, in getattr_and_trace
ERROR 03-24 02:09:11 [core.py:343] return fn(*args[2:], **kwargs)
ERROR 03-24 02:09:11 [core.py:343] File "/app/Quark/aiter/aiter/ops/gemm_op_a8w8.py", line 84, in get_CKGEMM_config

• ERROR 03-24 02:09:11 [core.py:343] if not hasattr(get_CKGEMM_config, "ckgemm_dict"):
  ERROR 03-24 02:09:11 [core.py:343]
  ERROR 03-24 02:09:11 [core.py:343] Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
  The full logs can be found in the attachment.

llama_fp8-int8aiterck-v1.txt

Operating System

Ubuntu 22.04.4 LTS (Jammy Jellyfish)

CPU

AMD EPYC 9654 96-Core Processor

GPU

AMD Instinct MI300X

ROCm Version

ROCm 6.3.1

ROCm Component

No response

Steps to Reproduce

Setup vLLM

git clone https://github.com/EmbeddedLLM/vllm.git --branch aiter-int8-linear aiter-int8-linear
cd aiter-int8-linear

if you have installed vllm before

python3 -m pip uninstall -y vllm
export PYTORCH_ROCM_ARCH=gfx942
python3 setup.py develop
Run with V1 engine VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=1 vllm serve neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

Hi tjtanaa. This vllm branch seems outdated, could you provide the latest one which I can reproduce ?