[FEA] Hopper CuTe DSL FMHA example: variable-length (cu_seqlens) inputs

Open
  • Is something wrong?
#3,773 1 comment 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 4 days

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
cpp, python
Domain
machine-learning, performance

Research direction

Start with examples/python/CuTeDSL/cute/hopper/kernel/attention/fmha.py and compare its interface with cute/blackwell/kernel/attention/fmha/fmha.py. The requested work spans variable-length inputs, per-sequence tile scheduling, tail handling, and coverage for empty and length-1 sequences, causal and sliding-window modes, and GQA ratios 1/4/8. Done means those cases pass and the PR includes a TFLOPS comparison on GH200; the issue asks whether the scope is acceptable and notes it may depend on #3399.

Written by the indexing model from the issue text.

Description

Is your feature request related to a problem?
examples/python/CuTeDSL/cute/hopper/kernel/attention/fmha.py runs only fixed-length batches. The Blackwell DSL FMHA (cute/blackwell/kernel/attention/fmha/fmha.py) takes cum_seqlen_q / cum_seqlen_k, so the Hopper example can't serve packed/ragged batches, which most training and serving stacks use.

(GQA/MQA already works in the Hopper example's run(). Only the docstring still says Q and K head counts must match. #3770 fixes the docstring and adds GQA 8:2 and MQA cases, which pass the reference check on GH200.)

Describe the solution you'd like
Bring the Blackwell example's varlen interface to the Hopper example:

  • cum_seqlen_q / cum_seqlen_k;
  • a tile scheduler that issues work per sequence length;
  • tile tails handled by TMA out-of-bounds zero fill plus the existing mask.

Tests would cover empty and length-1 sequences, causal and sliding window, and GQA ratios 1/4/8. The PR would include a TFLOPS comparison on GH200.

Additional context
Before I write this: would you take a change this size to the Hopper example, or would you prefer the Hopper example stay fixed-length? I'd build on #3399 (max-reduction fix) if it's still open.

Dominant language
C++
Stars
10.7k
Forks
2.2k
Avg merge
3d 20h
Merged PRs (30d)
4

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

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.

More from NVIDIA/cutlass

All issues in NVIDIA/cutlass

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.