Ameyn/wide vec t1

[ameynaik-hub]

Summary

Replaces the legacy gdn_decode_bf16state_cooprow_kernel and the
gdn_decode_bf16state_mtp_kernel (ILP=8) with a new
gdn_wide_vec_kernel (LDG.E.128 / STG.E.128 fast path) plus a
small-batch mtp_ilp4 fallback. Drops ~1900 LOC of dead/unused code,
adds split-pool support (#2905-compatible) to both surviving BF16
kernels, and ships the OOB fix mirroring upstream PR #3145 — for the
BF16 kernels that survived the cleanup.

Supersedes #3118. That PR's perf delta (T=1 per-call overhead +
pool+padding for the ILP kernel) is the first commit on this branch
(8a6e9819).

What changes

• New kernel: gdn_wide_vec_kernel — 128 threads/CTA = 8 groups
  × 16 threads, vec=8 BF16 → LDG.E.128 / STG.E.128, ILP=4 V-rows per
  thread. Configurable tile_v ∈ {32, 64, 128} so the kernel covers
  small/medium/large B*HV work-unit sizes uniformly.
• Pool-only: BF16 GDN dispatch is strictly pool-mode (matches
  the production serving contract). Wrapper
  gated_delta_rule_decode_pretranspose auto-promotes legacy non-pool
  callers internally — public API unchanged.
• Split-pool support (PR feat(gdn): separate input and output pool indices #2905 contract): both surviving BF16
  kernels (gdn_wide_vec_kernel, gdn_decode_bf16state_mtp_ilp4_kernel)
  natively support output_state_indices != initial_state_indices,
  with bit-equivalent single-pool behavior selected at compile time
  via Constexpr[bool] same_pool for zero-overhead dispatch.
• OOB fix (PR Fix OOB crash in intermediate_states indexing for GDN decode MTP kernel #3145 equivalent): intermediate_states is indexed
  by the per-call batch index i_n (not the pool-scoped cache_idx),
  so the buffer can be sized [B, T, HV, V, K] as production callers
  expect. Regression test catches the bug; pre-fix triggers
  cudaErrorIllegalAddress in <2 s.

Removed (~1900 LOC of dead code)

Kernel	Why removed
gdn_decode_bf16state_cooprow_kernel (~280 LOC)	Replaced by wide_vec + ILP=4 MTP; had known correctness issues at small batch
gdn_decode_bf16state_ilp_kernel (~740 LOC)	Only reachable at HV<32 with B≥16 — not a Qwen3.5 shape; MTP path covers it
gdn_decode_bf16state_mtp_kernel (ILP=8) (~940 LOC)	After wide_vec extension to split-pool + tile_v=32, mtp_kernel was unreachable

End-state BF16 surface = 2 @cute.kernels in one file:

• gdn_wide_vec_kernel — production hot path
• gdn_decode_bf16state_mtp_ilp4_kernel — small-batch fallback

Both pool-only, both split-pool capable, both indexed batch-scoped.

Speedup vs previous baseline

Baseline = pre-wide_vec dispatch (the mtp_kernel ILP=8 path, captured
on this same branch by monkey-patching _select_wide_vec_tile_v to
return None for every shape). Same harness, same hardware, same
config — so the comparison isolates the kernel-level speedup that
wide_vec + the cleanup deliver.

Setup: B200, HV=64, K=V=128, BF16, qk_l2norm=ON, warmup=5, iters=50,
T=1 invoked with --update-state, T≥2 invoked with
--cache-intermediate-states. Kernel time in microseconds (CUPTI).

Speedup (×, baseline / post-PR)

B	T=1	T=2	T=3	T=4	T=5	T=6	T=7	T=8
1	1.03×	1.04×	1.03×	1.03×	1.00×	1.02×	1.02×	1.01×
4	0.97×	1.23×	1.10×	1.12×	1.11×	1.11×	1.14×	1.14×
8	1.08×	1.11×	1.11×	1.12×	1.13×	1.15×	1.15×	1.14×
16	1.04×	1.09×	1.11×	1.13×	1.13×	1.12×	1.11×	1.10×
32	1.06×	1.12×	1.10×	1.11×	1.10×	1.09×	1.09×	1.09×
64	1.04×	1.11×	1.08×	1.09×	1.06×	1.07×	1.08×	1.06×
128	1.04×	1.11×	1.07×	1.09×	1.06×	1.06×	1.07×	1.07×
256	1.04×	1.11×	1.07×	1.09×	1.07×	1.06×	1.07×	1.07×

Headline

• T=1 production decode (B≥16): 4–6 % time reduction across the full batch sweep — the Qwen3.5 hot path.
• T≥2 with cache=ON (B≥4): 6–18 % time reduction at every shape. Best at small-T / mid-batch (B=4 T=2: 1.23×; B=8 T=6: 1.15×).
• Tiny shapes (B=1): within ±3 % of baseline (kernel isn't DRAM-bound; small fixed-cost overheads dominate; the ILP=4 fallback was already efficient there).

Sustained DRAM bandwidth post-PR (TB/s, 8 TB/s peak on B200)

B	T=1	T=2	T=3	T=4	T=5	T=6	T=7	T=8
1	1.25	1.21	1.49	1.63	1.45	1.55	1.64	1.70
4	2.83	3.24	3.54	3.78	3.53	3.68	3.84	3.95
8	3.97	4.09	4.40	4.54	4.42	4.55	4.59	4.61
16	4.73	4.73	5.02	5.03	4.95	4.91	4.92	4.87
32	5.39	5.36	5.44	5.46	5.27	5.23	5.21	5.17
64	5.83	5.76	5.80	5.77	5.45	5.44	5.45	5.33
128	6.31	6.05	6.03	6.01	5.68	5.61	5.57	5.54
256	6.57	6.23	6.20	6.17	5.85	5.74	5.72	5.66

Post-PR peaks at 6.57 TB/s = 82 % of B200 peak DRAM (T=1 B=256 production decode shape).

Split-pool

With wide_vec now supporting split-pool natively, split-pool
matches single-pool to within ±1 % at every measured shape.

Tests

513 passed, 0 failed in 18m18s on B200.

Including: 477 existing BF16/wide_vec/pool tests, 12 new split-pool MTP tests, 12 new OOB regression tests covering pool_size_multiplier ∈ {1, 4} × B ∈ {1, 8, 32} × T ∈ {2, 4}, 12 wrapper-level split-pool tests.

Files changed (4)

• flashinfer/gdn_decode.py — wrapper auto-promotes BF16 non-pool → pool
• flashinfer/gdn_kernels/gdn_decode_bf16_state.py — wide_vec inlined; dead kernels removed; split-pool plumbing; OOB fix; same_pool DCE
• tests/gdn/test_decode_delta_rule.py — split-pool + OOB regression tests
• benchmarks/bench_gdn_decode.py — --pool-mode {single,split} flag

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• New Features

  • Added --pool-mode option to benchmark tool for configuring state pool allocation (single or split modes).

• Tests

  • Expanded BF16 test coverage with regression tests for split-pool semantics and out-of-bounds scenarios; improved batch-dimension handling for intermediate-state comparisons.

[coderabbitai]

📝 Walkthrough

Walkthrough

Implements support for separate input and output state indices in BF16-accelerated gated-delta-rule decoding, enabling split-pool dispatch where outputs are written to different pool slots than inputs are read from. Unifies both T==1 and T>1 code paths to route through consistent pool/index handling, updates the benchmark wrapper with a pool_mode parameter, and adds comprehensive test coverage including split-pool semantics and edge-case validation.

Changes

BF16 Split-Pool State Dispatch

Layer / File(s)	Summary
Core Implementation flashinfer/gdn_decode.py	Unified pool/index handling for both T==1 and T>1 BF16 paths; gated_delta_rule_decode_pretranspose now routes non-pool state as a synthetic pool and forwards both initial_state_indices and output_state_indices into BF16 state kernel calls, enabling split-pool dispatch.
Benchmark Wrapper benchmarks/bench_gdn_decode.py	gdn_decode_bf16_state_wrapper adds output_state_indices argument; bench_gdn_decode_bf16_state adds pool_mode parameter ("single" or "split"); CLI --pool-mode flag propagates pool mode from command line through benchmark function.
Test Coverage tests/gdn/test_decode_delta_rule.py	State kernel tests now explicitly pass initial_state_indices=torch.arange(batch_size) for pool-only semantics; BF16 T==1 test expanded to include num_v_heads=64 variant; BF16 MTP intermediate-state comparison chunked to prevent OOM; new wide-vector BF16 MTP test uses monkeypatch to exercise varied tile_v parameters; two new BF16 MTP regression tests validate split-pool semantics and OOB index handling.

Estimated Code Review Effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly Related Issues

• [RFC] Unified GDN Decode/Prefill API #2687: Proposes unified gated-delta decode API and pool+indices/BF16 dispatch semantics that align directly with this PR's implementation of split-pool and index forwarding.

Possibly Related PRs

• feat: add pool+indices support to gated_delta_rule_decode_pretranspose (bf16 path)  #2619: Modifies BF16 pool/index handling in gdn_decode by adding pool-based initial_state/indices and threading indices through BF16 state kernels.
• perf(gdn): optimize MTP kernel with ILP rows and SMEM v caching #2618: Updates the GDN decode API/implementation in gated_delta_rule_decode_pretranspose and related MTP/BF16 state dispatch and benchmark wiring.
• feat(gdn): separate input and output pool indices #2905: Adds separate output_state_indices/write-index parameter to the same BF16 decode codepath and kernel call sites.

Suggested Labels

ready, run-ci

Suggested Reviewers

• yzh119
• bkryu
• nvmbreughe
• jimmyzho
• yongwww
• kahyunnam

Poem

🐰 Split the pool in two, we say,
Read from here, write far away,
Indices dance both high and low,
BF16 kernels steal the show!
Tests now catch the edge-case flaws,
Split-pool state obeys new laws! ✨

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title 'Ameyn/wide vec t1' is vague and uses non-descriptive terms that do not clearly convey the main purpose of this substantial refactoring.	Use a clear, descriptive title such as 'Replace legacy BF16 GDN kernels with wide_vec path and add split-pool support' or 'Add gdn_wide_vec_kernel with split-pool support and OOB fix'.
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description is comprehensive and detailed, covering all required sections with extensive technical documentation, performance data, and implementation details.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get your free trial and get 200 agent minutes per Slack user (a $50 value).

---

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 and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces a wide-vector BF16 GDN MTP decode kernel and an ILP=4 variant to optimize performance across various batch sizes and work units. It also refactors the dispatch logic to support pool-based state management with indices and includes updated tests. Feedback from the review focuses on the significant code duplication across the new kernels and the repetitive nature of the memory write-back logic, suggesting that these be refactored into parameterized implementations or loops to improve maintainability.

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #49720223 is currently running. I'll report back once the pipeline job completes.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/conftest.py`:
- Around line 143-149: The traversal that checks for CUDA OOM only follows one
exception link with "exc = exc.__cause__ or exc.__context__", so wrapped OOMs on
the other branch can be missed; update the logic in tests/conftest.py (the loop
that inspects "exc" and calls is_cuda_oom_error_str and checks isinstance(...,
torch.cuda.OutOfMemoryError)) to traverse both __cause__ and __context__
branches (e.g., use a stack/queue or set to visit exceptions
breadth/depth-first, pushing exc.__cause__ and exc.__context__ when present)
until all linked exceptions are inspected and return True if any is an OOM.

🪄 Autofix (Beta)

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

Run ID: 34a6f2f1-b813-434d-bf58-e5d66218c68a

📥 Commits

Reviewing files that changed from the base of the PR and between 7266b5e and ec76b2aac8ec8b2c7428af8ac1c68bd0453bf85e.

📒 Files selected for processing (1)

• tests/conftest.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
python - <<'PY'
def old_walk(exc):
    while exc is not None:
        if "out of memory" in str(exc).lower() and "cuda" in str(exc).lower():
            return True
        exc = exc.__cause__ or exc.__context__
    return False

def new_walk(exc):
    seen, stack = set(), [exc]
    while stack:
        cur = stack.pop()
        if id(cur) in seen:
            continue
        seen.add(id(cur))
        s = str(cur).lower()
        if "cuda" in s and "out of memory" in s:
            return True
        if cur.__cause__ is not None:
            stack.append(cur.__cause__)
        if cur.__context__ is not None:
            stack.append(cur.__context__)
    return False

ctx = RuntimeError("CUDA out of memory while formatting tensor diff")
cause = RuntimeError("wrapper cause without oom text")
top = RuntimeError("top-level runtime")
top.__cause__ = cause
top.__context__ = ctx

print("old_walk:", old_walk(top))  # expected False (missed)
print("new_walk:", new_walk(top))  # expected True (detected)
PY

Repository: flashinfer-ai/flashinfer

Length of output: 99

---

Traverse both __cause__ and __context__ branches to detect OOM in wrapped exceptions.

At line 148, the code uses exc = exc.__cause__ or exc.__context__, which follows only one link. When both are present, an OOM error in the non-selected branch will be missed, causing a test to fail instead of skip. The exception chain from torch.testing.assert_close can wrap inner OOM errors in multiple branches, so both must be examined.

Proposed fix

 def _is_oom_in_chain(exc: BaseException) -> bool:
-    # torch.testing.assert_close wraps the inner OOM in a RuntimeError("Comparing ...")
-    # whose str() does not mention "out of memory" — only __cause__ does.
-    while exc is not None:
-        if isinstance(exc, torch.cuda.OutOfMemoryError) or is_cuda_oom_error_str(
-            str(exc)
-        ):
-            return True
-        exc = exc.__cause__ or exc.__context__
+    # torch.testing.assert_close may wrap inner exceptions; inspect the full chain.
+    seen: set[int] = set()
+    stack: list[BaseException] = [exc]
+    while stack:
+        cur = stack.pop()
+        cur_id = id(cur)
+        if cur_id in seen:
+            continue
+        seen.add(cur_id)
+        if isinstance(cur, torch.cuda.OutOfMemoryError) or is_cuda_oom_error_str(
+            str(cur)
+        ):
+            return True
+        if cur.__cause__ is not None:
+            stack.append(cur.__cause__)
+        if cur.__context__ is not None:
+            stack.append(cur.__context__)
     return False

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/conftest.py` around lines 143 - 149, The traversal that checks for CUDA
OOM only follows one exception link with "exc = exc.__cause__ or
exc.__context__", so wrapped OOMs on the other branch can be missed; update the
logic in tests/conftest.py (the loop that inspects "exc" and calls
is_cuda_oom_error_str and checks isinstance(..., torch.cuda.OutOfMemoryError))
to traverse both __cause__ and __context__ branches (e.g., use a stack/queue or
set to visit exceptions breadth/depth-first, pushing exc.__cause__ and
exc.__context__ when present) until all linked exceptions are inspected and
return True if any is an OOM.

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #49751839 is currently running. I'll report back once the pipeline job completes.

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #49882216 is currently running. I'll report back once the pipeline job completes.

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #49978944 is currently running. I'll report back once the pipeline job completes.

[kahyunnam]

LGTM, thanks!

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #50042869 is currently running. I'll report back once the pipeline job completes.

[ameynaik-hub]

/bot run

[flashinfer-bot]

GitLab MR !613 has been updated with latest changes, and the CI pipeline #50118003 is currently running. I'll report back once the pipeline job completes.