[Bug] [spec decode] [flash_attn]: CUDA illegal memory access when calling flash_attn_cuda.fwd_kvcache #5152
Description
My environment setup
1st environment (running on ec2 g6.4xlarge
)
[2024-06-01T10:14:23Z] Collecting environment information...
[2024-06-01T10:14:26Z] PyTorch version: 2.3.0+cu121
[2024-06-01T10:14:26Z] Is debug build: False
[2024-06-01T10:14:26Z] CUDA used to build PyTorch: 12.1
[2024-06-01T10:14:26Z] ROCM used to build PyTorch: N/A
[2024-06-01T10:14:26Z]
[2024-06-01T10:14:26Z] OS: Ubuntu 22.04.4 LTS (x86_64)
[2024-06-01T10:14:26Z] GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
[2024-06-01T10:14:26Z] Clang version: Could not collect
[2024-06-01T10:14:26Z] CMake version: version 3.29.3
[2024-06-01T10:14:26Z] Libc version: glibc-2.35
[2024-06-01T10:14:26Z]
[2024-06-01T10:14:26Z] Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
[2024-06-01T10:14:26Z] Python platform: Linux-6.1.90-99.173.amzn2023.x86_64-x86_64-with-glibc2.35
[2024-06-01T10:14:26Z] Is CUDA available: True
[2024-06-01T10:14:26Z] CUDA runtime version: Could not collect
[2024-06-01T10:14:26Z] CUDA_MODULE_LOADING set to: LAZY
[2024-06-01T10:14:26Z] GPU models and configuration: GPU 0: NVIDIA L4
[2024-06-01T10:14:26Z] Nvidia driver version: 525.147.05
[2024-06-01T10:14:26Z] cuDNN version: Could not collect
[2024-06-01T10:14:26Z] HIP runtime version: N/A
[2024-06-01T10:14:26Z] MIOpen runtime version: N/A
[2024-06-01T10:14:26Z] Is XNNPACK available: True
[2024-06-01T10:14:26Z]
[2024-06-01T10:14:26Z] CPU:
[2024-06-01T10:14:26Z] Architecture: x86_64
[2024-06-01T10:14:26Z] CPU op-mode(s): 32-bit, 64-bit
[2024-06-01T10:14:26Z] Address sizes: 48 bits physical, 48 bits virtual
[2024-06-01T10:14:26Z] Byte Order: Little Endian
[2024-06-01T10:14:26Z] CPU(s): 16
[2024-06-01T10:14:26Z] On-line CPU(s) list: 0-15
[2024-06-01T10:14:26Z] Vendor ID: AuthenticAMD
[2024-06-01T10:14:26Z] Model name: AMD EPYC 7R13 Processor
[2024-06-01T10:14:26Z] CPU family: 25
[2024-06-01T10:14:26Z] Model: 1
[2024-06-01T10:14:26Z] Thread(s) per core: 2
[2024-06-01T10:14:26Z] Core(s) per socket: 8
[2024-06-01T10:14:26Z] Socket(s): 1
[2024-06-01T10:14:26Z] Stepping: 1
[2024-06-01T10:14:26Z] BogoMIPS: 5299.99
[2024-06-01T10:14:26Z] Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
[2024-06-01T10:14:26Z] Hypervisor vendor: KVM
[2024-06-01T10:14:26Z] Virtualization type: full
[2024-06-01T10:14:26Z] L1d cache: 256 KiB (8 instances)
[2024-06-01T10:14:26Z] L1i cache: 256 KiB (8 instances)
[2024-06-01T10:14:26Z] L2 cache: 4 MiB (8 instances)
[2024-06-01T10:14:26Z] L3 cache: 32 MiB (1 instance)
[2024-06-01T10:14:26Z] NUMA node(s): 1
[2024-06-01T10:14:26Z] NUMA node0 CPU(s): 0-15
[2024-06-01T10:14:26Z] Vulnerability Gather data sampling: Not affected
[2024-06-01T10:14:26Z] Vulnerability Itlb multihit: Not affected
[2024-06-01T10:14:26Z] Vulnerability L1tf: Not affected
[2024-06-01T10:14:26Z] Vulnerability Mds: Not affected
[2024-06-01T10:14:26Z] Vulnerability Meltdown: Not affected
[2024-06-01T10:14:26Z] Vulnerability Mmio stale data: Not affected
[2024-06-01T10:14:26Z] Vulnerability Reg file data sampling: Not affected
[2024-06-01T10:14:26Z] Vulnerability Retbleed: Not affected
[2024-06-01T10:14:26Z] Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
[2024-06-01T10:14:26Z] Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
[2024-06-01T10:14:26Z] Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
[2024-06-01T10:14:26Z] Vulnerability Spectre v2: Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
[2024-06-01T10:14:26Z] Vulnerability Srbds: Not affected
[2024-06-01T10:14:26Z] Vulnerability Tsx async abort: Not affected
[2024-06-01T10:14:26Z]
[2024-06-01T10:14:26Z] Versions of relevant libraries:
[2024-06-01T10:14:26Z] [pip3] mypy==1.9.0
[2024-06-01T10:14:26Z] [pip3] mypy-extensions==1.0.0
[2024-06-01T10:14:26Z] [pip3] numpy==1.26.4
[2024-06-01T10:14:26Z] [pip3] nvidia-nccl-cu12==2.20.5
[2024-06-01T10:14:26Z] [pip3] torch==2.3.0
[2024-06-01T10:14:26Z] [pip3] triton==2.3.0
[2024-06-01T10:14:26Z] [conda] Could not collectROCM Version: Could not collect
[2024-06-01T10:14:26Z] Neuron SDK Version: N/A
[2024-06-01T10:14:26Z] vLLM Version: 0.4.3
[2024-06-01T10:14:26Z] vLLM Build Flags:
[2024-06-01T10:14:26Z] CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
[2024-06-01T10:14:26Z] GPU Topology:
[2024-06-01T10:14:26Z] GPU0 CPU Affinity NUMA Affinity
[2024-06-01T10:14:26Z] GPU0 X 0-15 N/A
2nd environment (running on GCP g2-standard-12
):
Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A
OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.29.3
Libc version: glibc-2.35
Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.10.0-29-cloud-amd64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L4
Nvidia driver version: 550.54.15
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 12
On-line CPU(s) list: 0-11
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family: 6
Model: 85
Thread(s) per core: 2
Core(s) per socket: 6
Socket(s): 1
Stepping: 7
BogoMIPS: 4400.45
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 192 KiB (6 instances)
L1i cache: 192 KiB (6 instances)
L2 cache: 6 MiB (6 instances)
L3 cache: 38.5 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-11
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Mitigation; Clear CPU buffers; SMT Host state unknown
Versions of relevant libraries:
[pip3] mypy==1.9.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-nccl-cu12==2.20.5
[pip3] torch==2.3.0
[pip3] triton==2.3.0
[conda] Could not collectROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.4.3
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0 CPU Affinity NUMA Affinity GPU NUMA ID
GPU0 X 0-11 0 N/A
- How to repro:
- Build vLLM Docker image from newly cloned vllm repo:
docker build --build-arg max_jobs=16 --tag vllm --target test .
- Run the test in Docker container.
docker run -it --rm --gpus all vllm bash -c "cd /vllm-workspace/tests && pytest -v -s spec_decode"
- Build vLLM Docker image from newly cloned vllm repo:
🐛 Describe the bug
-
Nothing changes in the tests/relevant code. The only difference is it's running in a different machine/environment compared to vLLM CI. I listed 2 environments which I tried and both failed.
-
The error showed when running this test in
tests/spec_decode/e2e/test_multistep_correctness.py
: -
Test name is
test_spec_decode_e2e_greedy_correctness_tiny_model_large_bs_diff_output_len[1-32-256-test_llm_kwargs0-baseline_llm_kwargs0-per_test_common_llm_kwargs1-common_llm_kwargs0]
-
kwargs={'enforce_eager': True, 'use_v2_block_manager': True, 'model': 'JackFram/llama-160m', 'speculative_model': 'JackFram/llama-68m', 'num_speculative_tokens': 5}
-
Failure message and stack trace starts here: https://buildkite.com/vllm/ci-aws/builds/82#018fcb54-3ae6-4a96-8e2a-67c66814003d/184-356
-
The error happens when
flash_attn_cuda.fwd_kvcache
is called in/attention/backends/flash_attn.py
-
Running the test with
VLLM_ATTENTION_BACKEND=XFORMERS
passes. Could this bug be related to flash attention?