pytorch / pytorch/pytorch

varlen_attn with a paged KV cache returns incorrect results when split-KV is used

Open
#194,167 0 comments 0 reactions 1 assignee Claimed by @liangel-02 View on GitHub
bot-triaged high priority module: correctness (silent) module: sdpa triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

On a paged varlen_attn call, split-KV fails to write the output rows for sequences after the first. In the example below, the output buffer is initialized to NaN. With num_splits=2, the output row for sequence 1 remains NaN, indicating that the kernel does not write a result for that row.

```python
import torch
from torch.nn.attention.varlen import varlen_attn_out

torch.manual_seed(0)
b, page, pages, h, d = 2, 256, 4, 4, 64
ctx = page * pages

q = torch.randn(b, h, d, dtype=torch.bfloat16, device="cuda")
k = torch.randn(b * pages, page, h, d, dtype=torch.bfloat16, device="cuda")
v = torch.randn(b * pages, page, h, d, dtype=torch.bfloat16, device="cuda")
cu = torch.arange(b + 1, dtype=torch.int32, device="cuda")

kw = dict(
seqused_k=torch.full((b,), ctx, dtype=torch.int32, device="cuda"),
block_table=torch.arange(b * pages, dtype=torch.int32, device="cuda").view(b, pages),
)

# Prefill the output with NaN, so anything the kernel fails to write stays NaN.
for num_splits in (1, 2):
out = torch.full((b, h, d), float("nan"), dtype=torch.bfloat16, device="cuda")
varlen_attn_out(out, q, k, v, cu, cu * ctx, 1, ctx, num_splits=num_splits, **kw)
print(f"num_splits={num_splits}: sequence 1 -> {out[1, 0, :4].tolist()}")
```

Output on an RTX 5090, torch 2.13.0+cu130:

```
num_splits=1: sequence 1 -> [0.09716796875, -0.01531982421875, 0.0208740234375, 0.031005859375]
num_splits=2: sequence 1 -> [nan, nan, nan, nan]
```

### Versions

Collecting environment information...
PyTorch version: 2.13.0+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.13 (main, Jun 23 2026, 15:18:55) [Clang 22.1.3 ] (64-bit runtime)
Python platform: Linux-6.8.0-137-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.3.73
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 5090
Nvidia driver version: Could not collect
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 48 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 24
On-line CPU(s) list: 0-23
Vendor ID: AuthenticAMD
Model name: AMD Ryzen 9 7900X 12-Core Processor
CPU family: 25
Model: 97
Thread(s) per core: 2
Core(s) per socket: 12
Socket(s): 1
Stepping: 2
CPU(s) scaling MHz: 72%
CPU max MHz: 5733.0000
CPU min MHz: 400.0000
BogoMIPS: 9400.72
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 amd_lbr_v2 nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tcetopoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor smca fsrm flush_l1d ibpb_exit_to_user
Virtualization: AMD-V
L1d cache: 384 KiB (12 instances)
L1i cache: 384 KiB (12 instances)
L2 cache: 12 MiB (12 instances)
L3 cache: 64 MiB (2 instances)
NUMA node(s): 1
NUMA node0 CPU(s): 0-23
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Mitigation; Safe RET
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsa: Mitigation; Clear CPU buffers
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Mitigation; IBPB before exit to userspace

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect

cc @ezyang @gchanan @kadeng @msaroufim @drisspg @liangel-02 @howardzhang-cv

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.