fix: add topk parameter to MLA profiler to enable sparse dsv4 decode (closes #4336)

[botbikamordehai2-sketch]

What

This PR fixes the MLA profiler script to support sparse (top-k) decode, which is needed for DeepSeek-V4-Flash-0731 on SM120. The script previously hardcoded dense attention, so the (32, 256) top-k kernel instance was never invoked, causing a crash in DSPark drafts due to a missing kernel instance.

Fix

• Added an optional --topk CLI argument (default: None for dense).
• Passed topk to BatchMLAPagedAttentionWrapper and its plan() method.
• Updated the profile_deepseek_mla_decode function signature and CLI parser.

Closes #4336

Summary by CodeRabbit

• New Features
  • Added optional top-k configuration to MLA profiling.
  • Profiling can now evaluate top-k attention scenarios when specified.

[coderabbitai]

📝 Walkthrough

Walkthrough

The MLA decode profiler now accepts an optional topk parameter and forwards it to BatchMLAPagedAttentionWrapper and its plan call.

Changes

MLA top-k profiling

Layer / File(s)	Summary
Thread top-k through MLA profiling profiler/mla.py	profile_deepseek_mla_decode accepts topk and passes it to the wrapper and planning call.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#4108: Adds related top-k configuration for sparse MLA decode.
• flashinfer-ai/flashinfer#4309: Adds related DeepSeek MLA top-k support.
• flashinfer-ai/flashinfer#4372: Introduces MLA paths that this profiler can exercise with topk=256.

Suggested labels: op: attention

Suggested reviewers: saltyminty, bkryu, cherichy

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the change and links the issue, but it omits the required checklist and test status sections.	Add the template sections for pre-commit checks, tests, and reviewer notes, and report the applicable completion status.
Linked Issues check	⚠️ Warning	The change enables profiler top-k configuration but does not add the missing SM120 kernel instance or runtime support required by issue #4336.	Add and compile the (num_heads=32, topk=256) SM120 DSV4 kernel instance and verify sparse decode support for num_tokens <= 64.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the MLA profiler fix and its top-k purpose, and it references the linked issue.
Out of Scope Changes check	✅ Passed	The changes are limited to MLA profiler top-k support, which is related to the sparse MLA decode problem in issue #4336.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@profiler/mla.py`:
- Line 46: Update BatchMLAPagedAttentionWrapper.__init__ and plan() to accept
topk, then propagate it through the decode planner and launch path so
profile_deepseek_mla_decode() can pass topk without raising TypeError; otherwise
defer the profiler topk arguments until the wrapper API supports them.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d6f92617-7ce0-44c5-90e9-8f6ea6a3b868

📥 Commits

Reviewing files that changed from the base of the PR and between 2ab910c and dc92bb5.

📒 Files selected for processing (1)

• profiler/mla.py

[coderabbitai]

🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Expect: both wrapper signatures declare topk before this profiler change is merged.
rg -n -A18 -B3 'def __init__\(|def plan\(' flashinfer/mla/_core.py

# Trace all topk forwarding and wrapper API declarations.
rg -n -C4 '\btopk\b|BatchMLAPagedAttentionWrapper|\.plan\(' \
  profiler/mla.py flashinfer/mla --glob '*.py'

Repository: flashinfer-ai/flashinfer

Length of output: 22270

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== profiler/mla.py =="
cat -n profiler/mla.py

echo
echo "== BatchMLAPagedAttentionWrapper __init__ around signature =="
sed -n '1505,1525p' flashinfer/mla/_core.py | cat -n

echo
echo "== BatchMLAPagedAttentionWrapper plan signature and body =="
sed -n '1585,1720p' flashinfer/mla/_core.py | cat -n

echo
echo "== cached_module.plan references surrounding wrapper =="
rg -n -C20 'self\._cached_module\.plan\(' flashinfer/mla/_core.py

Repository: flashinfer-ai/flashinfer

Length of output: 14305

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast
from pathlib import Path

files = {"profiler/mla.py": "profiler/mla.py", "wrapper init": "flashinfer/mla/_core.py", "wrapper plan": "flashinfer/mla/_core.py"}
for label, path in files.items():
    source = Path(path).read_text()
    tree = ast.parse(source)
    if label == "wrapper init":
        cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "BatchMLAPagedAttentionWrapper")
        func = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "__init__")
    elif label == "wrapper plan":
        cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "BatchMLAPagedAttentionWrapper")
        func = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "plan")
    else:
        func = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "profile_deepseek_mla_decode")
    args = {a.arg: a.value is None for a in func.args.args if hasattr(a, "arg")}
    kwargs = {}
    for node in ast.walk(func):
        if isinstance(node, (ast.Call, ast.keyword)) and hasattr(node, "keyword"):
            continue
        if isinstance(node, ast.Call):
            for kw in node.keywords:
                kwargs[kw.arg] = True
    print(label, path, {
        "func_args": list(args),
        "accept_defaults": {k for k,v in args.items() if v},
        "call_topk_kwargs": list(kwargs),
    })
PY

Repository: flashinfer-ai/flashinfer

Length of output: 331

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast
from pathlib import Path

path = "profiler/mla.py"
source = Path(path).read_text()
tree = ast.parse(source)
func = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "profile_deepseek_mla_decode")

kwarg_calls_by_line = {}
for node in ast.walk(func):
    if not isinstance(node, ast.Call):
        continue
    line_no = node.lineno
    accepted_args = [a.arg for a in node.args] if hasattr(node.args, "args") else []
    accepted_keywords = [kw.arg for kw in node.keywords]
    for kw in node.keywords:
        arg_name = kw.arg
        arg_value_is_none = isinstance(kw.value, ast.Constant) and kw.value.value is None
        kwarg_calls_by_line.setdefault(line_no, []).append({
            "func": getattr(node.func, "id", None) or getattr(node.func, "attr", "<call>"),
            "args": accepted_args,
            "keywords": accepted_keywords,
            "topk_kwarg": arg_name == "topk",
            "topk_value_is_none": arg_value_is_none,
        })

print("topk arguments in profile_deepseek_mla_decode calls:")
for line_no, items in sorted(kwarg_calls_by_line.items()):
    print(line_no, items)

print("\nPython static argument count for the two changed call sites:")
for lineno in [45, 52]:
    func = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "profile_deepseek_mla_decode")
    for node in ast.walk(func):
        if isinstance(node, ast.Call) and node.lineno == lineno:
            print(f"line {lineno}: args={len(node.args)}, keywords={[kw.arg for kw in node.keywords]}")
PY

Repository: flashinfer-ai/flashinfer

Length of output: 2780

---

Block the PR until the MLA wrapper API supports topk.

BatchMLAPagedAttentionWrapper.__init__ and plan() do not declare topk, so profile_deepseek_mla_decode()’s topk arguments at profiler/mla.py:46 and profiler/mla.py:66 will raise TypeError whenever the profiler runs. Add topk to both wrapper signatures and forward it into the decode planner/launch path, or defer this profiler change until that API exists.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@profiler/mla.py` at line 46, Update BatchMLAPagedAttentionWrapper.__init__
and plan() to accept topk, then propagate it through the decode planner and
launch path so profile_deepseek_mla_decode() can pass topk without raising
TypeError; otherwise defer the profiler topk arguments until the wrapper API
supports them.