NVIDIA / NVIDIA/cutlass

[BUG] MLA example is broken for split-kv and larger Q

Open
#2,274 3 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
Currently example initializes Q to quite small values (mean -1, stddev 1). If I initialize Q to a bit bigger values (e.g. stddev 100), split-kv stops working. Larger Qs are typical for LLMs.
Looks like is some overflow in LSE compute.

Steps/Code to reproduce bug
Apply the following patch https://gist.github.com/divchenko/10f1991a7a197b706b5c46aaca1a9bd2 to
commit f535c33634b640a4c0bee131f2f6e9f81877a18c (HEAD, tag: v3.9.1)
Then run w/o split-kv

./77_blackwell_mla_2sm_fp16 --verify --split_kv=1
###### B 64 MLA H 128 D_rope 64 D_latent 512 Q 1 K 256 Gen None Split 1 Gen None #SM 148
 [OK]  128x128 fp16 persistent          : 156.768 TFLOPS/s 1.26077 TB/s
 [OK]  128x128 fp16 individual          : 162.176 TFLOPS/s 1.30426 TB/s

But when split-kv is enabled, it fails:


./77_blackwell_mla_2sm_fp16 --verify --split_kv=2
###### B 64 MLA H 128 D_rope 64 D_latent 512 Q 1 K 256 Gen None Split 2 Gen None #SM 148
failed O: max diff 6.26562 mean 1.15391
failed LSE: max diff inf mean inf
Reference check failed
[FAIL] 128x128 fp16 persistent          : 70.6905 TFLOPS/s 0.568513 TB/s
failed O: max diff 6.26562 mean 1.15391
failed LSE: max diff inf mean inf
Reference check failed
[FAIL] 128x128 fp16 individual          : 65.94 TFLOPS/s 0.530308 TB/s

Expected behavior
Verification should pass for larger Q

Environment details (please complete the following information):
B200
NVIDIA-SMI 570.124.06 Driver Version: 570.124.06 CUDA Version: 12.8

Additional context

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

Apply the patch from the linked gist to commit f535c33634b640a4c0bee131f2f6e9f81877a18c, then run 77_blackwell_mla_2sm_fp16 with --verify and --split_kv=1 and 2. Compare the split-kv path with the non-split path, focusing on the LSE computation that produces infinite differences for larger Q values. Done means verification passes with larger Q and the LSE and output checks remain finite and accurate.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
machine-learning, performance, testing
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.