[Feature][Mamba] Batched two-phase grouped prefill execution for hybrid GDN/Mamba (+7.58x)

[nicholaskh-ai]

Purpose

Implements RFC #55697.

Part 3 of 3 for Application-Directed Mamba Prefix Checkpointing (depends on PR #55873 and PR #55875).

In hybrid GDN/Mamba models (such as Qwen3.5 35B-A3B), prefill batches containing shared prefixes (1 common Producer prompt + N distinct Consumer queries) previously required either sequential per-request execution or full prompt recomputation due to state recurrence.

This PR introduces batched two-phase grouped prefill execution in the worker and GDN layer:

• Two-Phase Kernel Execution: Splits the grouped prefill forward into batched Producer and Consumer phases. In Phase 1, all Producers run in a single packed kernel invocation and populate shared checkpoint state rows. In Phase 2, Consumers read the shared checkpoint rows (index_select), copy them to private destination slots (index_copy_), and compute suffixes concurrently in a second packed kernel invocation.
• Attention Metadata Builder: Automatically constructs prefix_producer_ranges, consumer_ranges, consumer_shared_state_sources, and private_final_state_destination to dispatch requests into packed phases.
• Microbenchmark & Unit Tests: Validates numerical equivalence against the sequential baseline and provides a reproducible benchmark demonstrating up to 7.58x speedup.

Microbenchmark Results (NVIDIA L40S, PyTorch 2.11.0+cu130)

Measured on an NVIDIA L40S GPU using benchmarks/kernels/benchmark_grouped_gdn_prefill.py:

Workload (Producers, Consumers)	Tokens	Legacy Serial (us)	Batched Phased (us)	Speedup	P10 ~ P90 (us)
1P + 1C (single candidate)	576	2,146.41	1,548.13	1.39x	1,534 ~ 1,562
1P + 2C	640	2,876.00	1,546.33	1.86x	1,532 ~ 1,569
1P + 4C	768	4,245.61	1,545.72	2.75x	1,532 ~ 1,565
1P + 8C (typical batch)	1,024	7,055.63	1,582.38	4.46x	1,569 ~ 1,594
1P + 16C (wide candidate batch)	1,536	12,698.95	1,674.81	7.58x	1,660 ~ 1,716

Performance Highlights:

• Near-Flat Latency ($O(N) \to O(1)$): As consumers increase from 1 to 16, latency only increases from 1.54ms to 1.67ms (+8.4%), whereas serial execution scales linearly up to 12.7ms.
• Substantial Throughput Gains: At 16 consumers, latency drops by 86.8%, yielding a 7.58x speedup.

Duplicate-work Check

• PR [Perf][Model Runner V2] GDN all-mode decode-path optimizations: packed QKV, fused index prep, weight de-interleave + GEMM fold #55260 and [Perf] Qwen3-Next (GDN): decode-path kernel optimizations for all-mode prefix caching + MTP #54612 optimize GDN decode-path kernels for MRV2/MTP, not prefill grouping.
• PR Conv1d optimised kernel for aarch64 #54093 targets generic conv1d refactor.
• PR [PD][Feature] Add KV consumer partial-group caching for hybrid Mamba models #42524 targets partial-group caching for PD disaggregation, not batched phased prefill execution.

Test Plan

• Kernel split equivalence test:
  pytest tests/kernels/mamba/test_gdn_forward_core_split.py -v
• Metadata builder test:
  pytest tests/v1/attention/test_gdn_metadata_builder.py -k "test_grouped_metadata" -v
• CPU GDN ops alignment test:
  pytest tests/kernels/mamba/cpu/test_cpu_gdn_ops.py -v
• Benchmark verification:
  python -m benchmarks.kernels.benchmark_grouped_gdn_prefill --num-producers 1 --num-consumers 1 2 4 8 16

AI assistance was used to prepare this change. The human submitter is responsible for reviewing the changed code and test results.

Generated with Devin

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: nizhang1 nizhang1@coupang.com

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nicholaskh-ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nicholaskh-ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[nicholaskh-ai]

CC @mgoin @jeejeelee

Hi reviewers, this is Part 3 of the implementation stack for RFC #55697. It implements the batched two-phase grouped prefill execution in the GDN layer, cutting kernel launch overhead and delivering up to 7.58x speedup on NVIDIA L40S.

Whenever you have a moment, we would love your review and feedback. Thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nicholaskh-ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nicholaskh-ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nicholaskh-ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork