8000 BlockMask.from_kv_blocks crashes with IndexError when kv_indices is not padded · Issue #153344 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

BlockMask.from_kv_blocks crashes with IndexError when kv_indices is not padded #153344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
kilianmandon opened this issue May 11, 2025 · 0 comments
Open

BlockMask.from_kv_blocks crashes with IndexError when kv_indices is not padded #153344

kilianmandon opened this issue May 11, 2025 · 0 comments
Labels
module: flex attention module: higher order operators torch.cond and similar module: pt2-dispatcher PT2 dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op, oncall: pt2 triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Comments

@kilianmandon
Copy link
kilianmandon commented May 11, 2025 •
edited by pytorch-bot bot
Loading

🐛 Describe the bug

In the implementation of BlockMask.from_kv_blocks, the function _ordered_to_dense constructs a dense tensor (dense_mask) using the shape of the kv_indices tensor. However, it later uses the values from kv_indices as indices into this dense tensor. If any value in kv_indices is greater than or equal to the second dimension of kv_indices, this results in an IndexError.

In the example below, the values in kv_indices go up to 2, but the tensor has shape (1, 1, 3, 1) — meaning dimension 3 has size 1, which causes out-of-bounds access when indexing with value 2.

import torch
from torch.nn.attention.flex_attention import BlockMask

kv_num_blocks = torch.tensor([1, 1, 1])
kv_indices = torch.tensor([[0], [1], [2]])
# Note: Error doesn't arise if kv_indices is sufficiently padded:
# kv_indices = torch.tensor([[0, -1, -1], [1, -1, -1], [2, -1, -1]])


kv_num_blocks = kv_num_blocks[None, None, ...]
kv_indices = kv_indices[None, None, ...]

BlockMask.from_kv_blocks(kv_num_blocks, kv_indices, BLOCK_SIZE=32)
Traceback (most recent call last):
  File "test.py", line 10, in <module>
    BlockMask.from_kv_blocks(kv_num_blocks, kv_indices, BLOCK_SIZE=32)
  File "lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 352, in from_kv_blocks
    q_num_blocks, q_indices = _transpose_ordered(kv_num_blocks, kv_indices)
  File "lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 186, in _transpose_ordered
    dense = _ordered_to_dense(num_blocks_in_row, col_indices)
  File "lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 171, in _ordered_to_dense
    out = create_dense_batched(num_blocks_in_row, col_indices)
  File "lib/python3.10/site-packages/torch/_functorch/apis.py", line 202, in wrapped
    return vmap_impl(
  File "lib/python3.10/site-packages/torch/_functorch/vmap.py", line 334, in vmap_impl
    return _flat_vmap(
  File "lib/python3.10/site-packages/torch/_functorch/vmap.py", line 484, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "lib/python3.10/site-packages/torch/_functorch/apis.py", line 202, in wrapped
    return vmap_impl(
  File "lib/python3.10/site-packages/torch/_functorch/vmap.py", line 334, in vmap_impl
    return _flat_vmap(
  File "lib/python3.10/site-packages/torch/_functorch/vmap.py", line 484, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 164, in create_dense_one
    dense_mask[row_indices, valid_indices] = dense_mask.new_ones(())
IndexError: index 2 is out of bounds for dimension 3 with size 2

Versions

PyTorch version: 2.7.0
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Rocky Linux release 8.10 (Green Obsidian) (x86_64)
GCC version: (GCC) 11.3.0
Clang version: Could not collect
CMake version: version 3.26.5
Libc version: glibc-2.28

Python version: 3.10.17 | packaged by conda-forge | (main, Apr 10 2025, 22:19:12) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-553.36.1.el8_10.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 12.9.41
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA H100
GPU 1: NVIDIA H100
GPU 2: NVIDIA H100
MIG 1g.12gb Device 0:
MIG 1g.12gb Device 1:
MIG 1g.12gb Device 2:
MIG 1g.12gb Device 3:
MIG 1g.12gb Device 4:
MIG 1g.12gb Device 5:
MIG 1g.12gb Device 6:
GPU 3: NVIDIA H100

Nvidia driver version: 565.57.01
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
Byte Order: Little Endian
CPU(s): 96
On-line CPU(s) list: 0-95
Thread(s) per core: 1
Core(s) per socket: 48
Socket(s): 2
NUMA node(s): 8
Vendor ID: GenuineIntel
CPU family: 6
Model: 143
Model name: Intel(R) Xeon(R) Platinum 8468
Stepping: 8
CPU MHz: 2100.000
CPU max MHz: 3800.0000
CPU min MHz: 800.0000
BogoMIPS: 4200.00
L1d cache: 48K
L1i cache: 32K
L2 cache: 2048K
L3 cache: 107520K
NUMA node0 CPU(s): 0-11
NUMA node1 CPU(s): 12-23
NUMA node2 CPU(s): 24-35
NUMA node3 CPU(s): 36-47
NUMA node4 CPU(s): 48-59
NUMA node5 CPU(s): 60-71
NUMA node6 CPU(s): 72-83
NUMA node7 CPU(s): 84-95
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single intel_ppin cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] optree==0.15.0
[pip3] torch==2.7.0
[pip3] triton==3.3.0+gitd9a1100d
[conda] cuda-cudart 12.9.37 h5888daf_0 conda-forge
[conda] cuda-cudart-dev_linux-64 12.9.37 h3f2d84a_0 conda-forge
[conda] cuda-cudart-static_linux-64 12.9.37 h3f2d84a_0 conda-forge
[conda] cuda-cudart_linux-64 12.9.37 h3f2d84a_0 conda-forge
[conda] cuda-cupti 12.9.19 h9ab20c4_0 conda-forge
[conda] cuda-cupti-dev 12.9.19 h9ab20c4_0 conda-forge
[conda] cuda-nvrtc 12.9.41 h5888daf_0 conda-forge
[conda] cuda-nvtx 12.9.19 h5888daf_0 conda-forge
[conda] cudnn 9.8.0.87 h81d5506_1 conda-forge
[conda] libblas 3.9.0 31_hfdb39a5_mkl conda-forge
[conda] libcblas 3.9.0 31_h372d94f_mkl conda-forge
[conda] libcublas 12.9.0.13 h9ab20c4_0 conda-forge
[conda] libcublas-dev 12.9.0.13 h9ab20c4_0 conda-forge
[conda] libcufft 11.4.0.6 h5888daf_0 conda-forge
[conda] libcufft-dev 11.4.0.6 h5888daf_0 conda-forge
[conda] libcurand 10.3.10.19 h9ab20c4_0 conda-forge
[conda] libcurand-dev 10.3.10.19 h9ab20c4_0 conda-forge
[conda] libcusolver 11.7.4.40 h9ab20c4_0 conda-forge
[conda] libcusolver-dev 11.7.4.40 h9ab20c4_0 conda-forge
[conda] libcusparse 12.5.9.5 h5888daf_0 conda-forge
[conda] libcusparse-dev 12.5.9.5 h5888daf_0 conda-forge
[conda] liblapack 3.9.0 31_hc41d3b0_mkl conda-forge
[conda] libmagma 2.9.0 h19665d7_1 conda-forge
[conda] libnvjitlink 12.9.41 h5888daf_0 conda-forge
[conda] libtorch 2.7.0 cuda126_mkl_h99b69db_300 conda-forge
[conda] mkl 2024.2.2 ha957f24_16 conda-forge
[conda] nccl 2.26.5.1 ha44e49d_0 conda-forge
[conda] numpy 1.26.4 pypi_0 pypi
[conda] optree 0.15.0 py310h3788b33_0 conda-forge
[conda] pytorch 2.7.0 cuda126_mkl_py310_h5ee0071_300 conda-forge
[conda] triton 3.3.0 cuda126py310h05ca3d0_1 conda-forge

cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @Chillee @drisspg @yanboliang @BoyuanFeng
@colesbury colesbury added triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module module: flex attention labels May 13, 2025
@pytorch-bot pytorch-bot bot added module: higher order operators torch.cond and similar oncall: pt2 module: pt2-dispatcher PT2 dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op, labels May 13, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
module: flex attention module: higher order operators torch.cond and similar module: pt2-dispatcher PT2 dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op, oncall: pt2 triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@colesbury @kilianmandon
0