_flash_attn_fwd/_flash_attn_bwd: DSLRuntimeError — window_size_left/right expects Int32 but gets plain Python int
## Summary
`interface.py` passes plain Python `int` for `window_size_left`/`window_size_right` to `cute.compile` on SM_80 and SM_12x code paths, but `cute.compile` requires `cutlass.Int32` for `Optional[Int32]` kernel parameters → `DSLRuntimeError`.
## Environment
- flash-attn-4: v4.0.0b16
- nvidia-cutlass-dsl: 4.5.2
- CUDA: 12.8
- GPU: NVIDIA GB10 (sm_121 — consumer Blackwell / DGX Spark)
## Root Cause
The kernel signatures declare `window_size_left: Optional[Int32]` and `window_size_right: Optional[Int32]`. `cute.compile` uses argument types for type-template specialization — it requires `cutlass.Int32`, not plain Python `int`.
The SM_90 path implicitly coerces (or takes a different compile route), so this issue is not visible on H100/H200. The SM_80 and SM_12x else-branches in both `_flash_attn_fwd` and `_flash_attn_bwd` pass raw Python `int` values directly into `compile_args`:
```python
# interface.py — SM_80/12x path
compile_args = (..., window_size_left, window_size_right, ...)
# ^^^^^^^^^^^^^^^^ plain int → DSLRuntimeError
```
Error:
```
DSLRuntimeError: argument #13 (window_size_left): expects (Int32, NoneType) but got int
```
## Fix
Cast to `Int32` before building `compile_args`, in both the forward and backward compile paths:
```python
# interface.py — add before compile_args in BOTH _flash_attn_fwd and _flash_attn_bwd
window_size_left_typed = Int32(window_size_left) if window_size_left is not None else None
window_size_right_typed = Int32(window_size_right) if window_size_right is not None else None
```
Then use `window_size_left_typed` / `window_size_right_typed` in the `compile_args` tuple instead of the raw values. `Int32` is already imported: `from cutlass import Int32, Float32`.
## Verification
Applied locally on GB10 (SM_121). `DSLRuntimeError` for `window_size_left` argument no longer occurs. Sliding window forward tests now run and produce results (the **forward** is numerically correct; note: backward has a separate pre-existing limitation where `AttentionMask` in `FlashAttentionBackwardSm80.kernel()` ignores `mask_local=True` — but that is a different bug, not related to this type cast).
## Affected Tests
- `test_dpa_fa4_sliding_window` (all 8 variants) — previously crashed at compile step with this error
- Any test that passes non-None `window_size_left/right` through the SM_80 or SM_12x code path
## Note
This affects all non-SM_90 hardware. Users on A100/H100 (SM_80/SM_90) without `window_size` set would not see this issue since the default `window_size=(-1, -1)` may be handled differently.
1 条评论