NVIDIA / NVIDIA/cutlass

[BUG] FMHA fwd kernel causal misbehaves when qlen != klen

Open
#2,336 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug inactive-30d inactive-90d
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Describe the bug
The Blackwell FMHA fwd kernel causal mask implementation doesn't behave correctly when qlen != klen.

Steps/Code to reproduce bug
Unfortunately the FMHA verifier is also buggy because it reuses the same wrong Mask implementation.
But something like ./77_blackwell_fmha_fp16 --b=1 --h=1 --d=128 --q=3 --k=128 --b=1 --verify --mask=causal --verify should reveal the issue if the reference implementation is correct.

Expected behavior
Referenced from FA repo README (https://github.com/Dao-AILab/flash-attention):

    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
        1 1 1 1 0
        1 1 1 1 1
    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
        0 0
        0 0
        0 0
        1 0
        1 1

Environment details (please complete the following information):

  • Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]

Additional context
Add any other context about the problem here.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running ./77_blackwell_fmha_fp16 with q=3, k=128, --verify, and --mask=causal. Inspect the Blackwell FMHA forward kernel and the verifier's Mask implementation, noting that the verifier currently reuses the same suspected-wrong mask. Done means causal masking follows the bottom-right alignment shown in the issue for unequal query and key lengths, with verification reporting the expected result.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.