huggingface / huggingface/candle

candle-flash-attn-v3: ragged causal varlen (different seqlens in one batch) hangs in the SingleTileScheduler mainloop

Open
#3,603 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

## Summary

Once `flash_attn_varlen(..., causal=true)` actually reaches the kernel with `is_causal=1` (a window-size clamp in the binding previously disabled causal on the varlen path silently; a PR fixing that is being submitted and will be linked here), varlen batches whose `cu_seqlens` contain **different** sequence lengths hang forever inside the kernel. Single-sequence batches and uniform-length batches work and produce correct results.

## Repro

H200 (sm90a), CUDA 12.9, with the binding-level causal fix applied:

```rust
// two ragged sequences, e.g. lengths 1743 + 1559, causal=true
let cu = Tensor::new(&[0u32, 1743, 3302], &dev)?;
let out = candle_flash_attn_v3::flash_attn_varlen(
&q, &k, &v, &cu, &cu, 1743, 1743, scale, /*causal=*/true, false,
)?; // <- never returns, GPU busy-loops
```

Uniform lengths (e.g. `[0, 1024, 2048]`) complete and match an eager f32 reference. We hit this on a real serving workload (1473 ragged prefill batches) — every ragged batch hung.

## Analysis

The varlen path hardcodes `SingleTileScheduler` (`flash_fwd_launch_template.h`). For a short sequence in a ragged batch, m_block tiles past that sequence's end take an early `continue` that skips the pipelined mainloop for that tile, but the producer/consumer warps still execute the pipeline tail (`load_tail` / `store_tail`) — an empty pipeline then deadlocks waiting on barriers that are never armed.

Switching the varlen path to the persistent scheduler used by the dense causal path does not work either: that scheduler is not varlen-aware in this kernel revision, and produces wrong results.

So as of v0.10.2 / current main, the vendored kernel cannot do *correct and terminating* ragged causal varlen — which is the common serving shape (ragged prefill). Fixing it properly seems to require the newer upstream FA3 hopper kernel (`FlashAttnFwdSm90` with the varlen-aware persistent scheduler + `prepare_varlen_num_blocks`, as used by vLLM's flash-attention fork and Dao-AILab main). We have a working port of that kernel in a downstream project (correct on the 1473-batch ragged workload, bit-identical to vLLM, and matching vLLM's speed), but it is a large change (new vendored kernel + CUTLASS 3.8); happy to discuss upstreaming it if there is interest.

## Workarounds

- Pad ragged batches to uniform lengths (wasteful), or
- split ragged batches into per-sequence calls (slow), or
- use `causal=false` shapes only (the hang requires causal + ragged).

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.