Skip to content

perf(gdn): optimize MTP kernel with ILP rows and SMEM v caching #2618

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aleozlx merged 4 commits into from
Mar 7, 2026
+4,358 −2,131
Merged

perf(gdn): optimize MTP kernel with ILP rows and SMEM v caching #2618

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b1da819
perf(gdn): optimize MTP kernel with ILP rows and SMEM v caching
ameynaik-hub Feb 22, 2026
2f8d4f8
refactor(gdn): consolidate ILP register tensors into 2D layout
ameynaik-hub Feb 23, 2026
0be3d00
refactor(gdn): separate kernel implementations from API layer in gdn_…
ameynaik-hub Mar 5, 2026
71e774b
feat(gdn): port pool+indices support from eee401bb to refactored API
ameynaik-hub Mar 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions benchmarks/bench_gdn_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,6 @@
except ImportError:
GDN_DECODE_KLAST_BF16_STATE_AVAILABLE = False


# ============================================================================
# Utility Functions
# ============================================================================
Expand Down Expand Up @@ -1167,6 +1166,7 @@ def bench_gdn_mtp(
use_beta: bool = True,
use_qk_l2norm: bool = True,
cache_intermediate_states: bool = True,
disable_state_update: bool = True,
warmup_iters: int = 10,
bench_iters: int = 100,
):
Expand Down Expand Up @@ -1243,7 +1243,7 @@ def bench_gdn_mtp(
scale,
output,
intermediate_states_buffer,
disable_state_update=True,
disable_state_update=disable_state_update,
use_qk_l2norm=use_qk_l2norm,
),
enable_cupti=True,
Expand All @@ -1264,7 +1264,7 @@ def bench_gdn_mtp(
head_size,
dtype,
seq_len,
disable_state_update=True, # MTP mode: state is not written back
disable_state_update=disable_state_update,
)

kernel_tflops = flops / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0
Expand Down Expand Up @@ -1577,7 +1577,7 @@ def bench_mtp_comparison(
scale,
output_fi,
intermediate_fi,
disable_state_update=True,
disable_state_update=False,
use_qk_l2norm=use_qk_l2norm,
),
enable_cupti=True,
Expand Down Expand Up @@ -1627,7 +1627,7 @@ def bench_mtp_comparison(
scale,
output_tr,
intermediate_tr,
disable_state_update=True,
disable_state_update=False,
use_qk_l2norm=use_qk_l2norm,
),
enable_cupti=True,
Expand Down Expand Up @@ -2399,7 +2399,9 @@ def run_flashinfer_only_benchmark(args, dtype, use_qk_l2norm):
f"\nGDN MTP Benchmark "
f"(heads: q={args.num_q_heads}, k={args.num_k_heads}, "
f"v={args.num_v_heads}, d={args.head_size}, dtype={args.dtype}, "
f"qk_l2norm={'ON' if use_qk_l2norm else 'OFF'})"
f"qk_l2norm={'ON' if use_qk_l2norm else 'OFF'}, "
f"cache_intermediate={'ON' if args.cache_intermediate_states else 'OFF'}, "
f"update_state={'ON' if args.update_state else 'OFF'})"
)
print("-" * 100)
print(
Expand All @@ -2419,6 +2421,7 @@ def run_flashinfer_only_benchmark(args, dtype, use_qk_l2norm):
dtype=dtype,
use_qk_l2norm=use_qk_l2norm,
cache_intermediate_states=args.cache_intermediate_states,
disable_state_update=not args.update_state,
warmup_iters=args.warmup,
bench_iters=args.iters,
)
Expand Down Expand Up @@ -2736,6 +2739,11 @@ def main():
action="store_true",
help="Cache intermediate states for MTP benchmark",
)
parser.add_argument(
"--update-state",
action="store_true",
help="Update final state (disable_state_update=False) for MTP benchmark",
)
parser.add_argument(
"--warmup",
type=int,
Expand Down
Loading
Loading