[mikasenghaas]

Summary

• Don't pass topk_length to the FlashMLA sparse prefill forward on SM90 in cudnn_flashmla_forward.
• The slots past each query's valid prefix are already -1, and FlashMLA masks them. SM90 now reads the whole slot width, and the result is the same.
• SM100 and SM103 still pass topk_length. The cuDNN backward is unchanged.

Why:

• Unit (GPU) on main fails 24 cudnn_flashmla cases in tests/unit/train/models/test_deepseek_v4_kernels.py. It has failed since feat(deepseek-v4): cudnn_flashmla sparse attention backend (FlashMLA forward, cuDNN backward) #4078 landed.
• On H200, the pinned flash-mla wheel (flashmla-b7643bd) raises in sparse_prefill_fwd:

  [FlashMLA] Error: The chosen implementation does not support all required features.
  Features that are required but not supported:
    -   6: TOPK_LENGTH
  Current GPU: NVIDIA H200, SM 9.0 with 132 SMs
  

• FlashMLA's source at b7643bd lists TOPK_LENGTH as an SM90 feature. The wheel's binary prints its last two SM90 features as 0: HEAD_64, so this build does not have them.

Verification

• Before: Unit (GPU) on main at 68e5755ed (run) had 24 failed and 210 passed tests. All failures were cudnn_flashmla cases.
• After: see Unit (GPU) on this PR.

Follow-ups

• flashmla_sparse_mla_forward (GLM-5, model.dsa_backend = "cudnn_flashmla") also passes topk_length, so it likely fails the same way on H200. No CI test covers that path on H200.

🤖 Generated with Claude Code

---

Note

Low Risk
Narrow forward-path workaround for a known FlashMLA SM90 limitation; backward and non-SM90 paths unchanged, with masking preserving correctness.

Overview
Fixes DeepSeek V4 cudnn_flashmla sparse prefill on SM90 (e.g. H200) by not passing topk_length into flash_mla_sparse_fwd when the GPU’s major capability is 9.

The pinned FlashMLA wheel’s SM90 kernel errors if TOPK_LENGTH is requested; invalid gather slots are already -1 and masked, so omitting topk_length still scans the padded width with the same numerics. SM100/SM103 behavior is unchanged—they still pass topk_length. The cuDNN backward path is untouched.

^{Reviewed by Cursor Bugbot for commit 35c870b. Bugbot is set up for automated code reviews on this repo. Configure here.}