Conversation
The BF16-state and FP32-state MTP kernels are bf16-only internally, so documented fp16 inputs were silently reinterpreted. Convert those operands at the boundary, key remaining polymorphic dtypes in the compile caches, and allocate fresh default outputs instead of returning a cached buffer.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughGDN decode paths now validate tensor dtypes, normalize BF16 kernel operands, isolate dtype-specific compilation caches, allocate fresh default outputs, and preserve caller-provided output dtypes. Regression tests cover conversions, cache interleaving, output staging, and padding behavior. ChangesGDN dtype hardening
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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. Comment |
Use output.to(bfloat16) instead of empty_like so negative-index padding rows keep caller-initialized values; add regression test.
Keep only short why-comments; move design narrative out of the wrappers.
Member
Author
|
/bot run tests/gdn |
Collaborator
Reject unsupported dtypes early for dt_bias (bf16/fp32) and initial_state_indices (int32/int64); convert WY indices to int32.
kahyunnam
marked this pull request as ready for review
July 28, 2026 19:59
kahyunnam
requested review from
bkryu,
jiahanc,
yongwww and
yzh119
as code owners
July 28, 2026 19:59
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tests/gdn/test_decode_delta_rule.py (1)
4484-4513: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd int64-index coverage for the WY int32 conversion fix.
The PR explicitly fixes WY slot-index conversion to int32, but this test only passes
torch.int32indices (line 4491), so the conversion path itself is untested here. Consider adding an int64-indices variant to actually exercise the fix.✅ Suggested addition
def test_gdn_decode_wy_output_only_fp16_inputs_are_converted(batch_size=4): """The WY output-only kernel is bf16-only too (`io = cutlass.BFloat16`).""" _skip_if_not_sm90_or_later() if not GDN_DECODE_BF16_WY_OUTPUT_ONLY_AVAILABLE: pytest.skip("gdn_decode_bf16_wy_output_only kernel not available") x = _dtype_case_inputs("bf16_state_mtp", torch.bfloat16, batch_size) - indices = torch.arange(batch_size, dtype=torch.int32, device=x["q"].device) + indices = torch.arange(batch_size, dtype=torch.int64, device=x["q"].device)🤖 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 `@tests/gdn/test_decode_delta_rule.py` around lines 4484 - 4513, Extend test_gdn_decode_wy_output_only_fp16_inputs_are_converted to also run with the same indices converted to torch.int64, ensuring the WY slot-index int32 conversion path is exercised. Keep the existing int32 case and assertions, and verify the int64-index output matches the established bf16 reference.flashinfer/gdn_kernels/gdn_decode_bf16_state.py (1)
3628-3633: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSynchronize the MTP dtype documentation.
as_bf16(...)now accepts non-BF16q/k/v/a/band converts them, but thegated_delta_rule_mtpdocstring still describes those operands as BF16-only. Document the accepted FP16/BF16 inputs and the BF16 normalization.As per coding guidelines, keep documentation synchronized with code changes.
🤖 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 `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py` around lines 3628 - 3633, Update the gated_delta_rule_mtp docstring to state that q, k, v, a, and b accept FP16 or BF16 inputs and are normalized to BF16 by as_bf16; retain BF16-only wording only for outputs or other operands that still require it.Source: Coding guidelines
🤖 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 `@flashinfer/gdn_kernels/gdn_decode_bf16_wy_output_only.py`:
- Around line 2117-2123: Validate initial_state_indices before the narrowing
conversion in the wrapper containing the dtype assertion, rejecting any values
outside the signed INT32 range (or enforce the equivalent
initial_state_source.shape[0] limit) instead of allowing wraparound. Preserve
valid INT64 indices by converting them to INT32 only after validation, and add
coverage for the out-of-range boundary case.
In `@flashinfer/gdn_kernels/gdn_decode_mtp.py`:
- Line 2515: Update _get_compiled_mtp_kernel and _get_compiled_mtp_kernel_inline
so they use a process-level cache keyed by the complete inline/warp cache key,
including dtype_key, instead of returning a new empty dictionary on each call.
Ensure run_mtp_decode reuses cached compiled kernels and only compiles missing
entries.
---
Nitpick comments:
In `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py`:
- Around line 3628-3633: Update the gated_delta_rule_mtp docstring to state that
q, k, v, a, and b accept FP16 or BF16 inputs and are normalized to BF16 by
as_bf16; retain BF16-only wording only for outputs or other operands that still
require it.
In `@tests/gdn/test_decode_delta_rule.py`:
- Around line 4484-4513: Extend
test_gdn_decode_wy_output_only_fp16_inputs_are_converted to also run with the
same indices converted to torch.int64, ensuring the WY slot-index int32
conversion path is exercised. Keep the existing int32 case and assertions, and
verify the int64-index output matches the established bf16 reference.
🪄 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 Plus
Run ID: 0cb3ef2e-1f26-471a-a7e5-53656365e6e8
📒 Files selected for processing (6)
flashinfer/gdn_decode.pyflashinfer/gdn_kernels/dtype_compat.pyflashinfer/gdn_kernels/gdn_decode_bf16_state.pyflashinfer/gdn_kernels/gdn_decode_bf16_wy_output_only.pyflashinfer/gdn_kernels/gdn_decode_mtp.pytests/gdn/test_decode_delta_rule.py
Collaborator
|
[SUCCESS] Pipeline #60011072: 18/18 executed test jobs passed |
kahyunnam
added
run-ci
op: linear attention
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Avoid silent wraparound when the WY path converts slot indices to int32.
Member
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
@flashinfer-bot run |
bkryu
approved these changes
Aug 11, 2026
Keep one representative case per bug (8 tests) instead of the 22-case grid, to cut GDN suite wall time without dropping C1/C2/C3/H1 coverage.
Collaborator
|
@flashinfer-bot run |
5 tasks
5 tasks
kahyunnam
added a commit
that referenced
this pull request
Aug 15, 2026
…ions (#4513) ## 📌 Description `tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins many times over. Two independent reasons, both verified against the cache keys rather than assumed: **1. Batch size is not a compile key.** It has been dynamic since #3649, so it reaches the cache only through coarse buckets: | Path | How `B` enters the key | |------|------------------------| | pretranspose (`gdn_decode_pretranspose.py:964`) | not at all | | bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all — the tests build a contiguous pool, so `pool_size_key = -1` and `pool_slot_stride = (-1,)` are B-independent sentinels | | nontranspose (`gdn_decode_nontranspose.py:725`) | only via `use_small_batch = B < 32` | | fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`, `gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` / `_get_bf16_mtp_config` | The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`: 378 of the file's 817 parametrized cases (46%) but only **42** compile keys, because `tile_v` is an explicit monkeypatched axis and the 9 batch sizes contribute nothing. The first commit keeps one batch size per bucket. I verified each kept set reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a real mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel` silently dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}). **2. Intermediate `seq_len` values only re-specialize on `T`.** `get_mtp_config` returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for every `T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile config that T=4 or T=8 does not already cover. `T=2` is kept as the one structurally distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2` branches. ### Effect | | cases | compile keys (retuned tests) | |---|---|---| | before | 817 | 102 | | after batch collapse | 501 | 102 | | after `seq_len` trim | 409 | 62 | Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized). ### Coverage cost The first commit costs nothing in specialization coverage — the same cubins still run, just at fewer runtime batch sizes. The second commit is a deliberate reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a T-specific off-by-one would no longer be caught. It is a separate commit so it can be dropped if reviewers would rather keep the full sweep. ## 🔍 Related Issues Refs #4110 (GDN cold-compile CI time). Complements #4128 and #4444, which remove key entries that provably do not reach codegen; this removes test cases that map onto keys already covered. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416 passed in 32m25s** - [ ] GPU CI for the timing comparison ## Reviewer Notes - I did not measure a clean before/after wall clock: the ~38 min baseline I was working from comes from #4219's description rather than the same machine, so I'd rather let CI provide the comparison than quote a number I can't stand behind. The case and key counts above are exact and static. - Worth noting for #4110 more broadly: cutting 50% of the cases bought substantially less than 50% of the wall clock, which suggests the remaining cost is dominated by compilation and fixed overhead rather than per-case execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in #4214) as the larger lever. - The batch sizes kept per test are load-bearing, not arbitrary — each set is one representative per config bucket. I left a one-line comment at each site so they don't get "restored" later. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Streamlined decode and MTP test coverage to use representative batch sizes and sequence lengths. * Preserved coverage for key thresholds, tile configurations, transposition modes, precision variants, and sequence-length scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…nfer-ai#4219) ## Issue flashinfer-ai#4214 Addresses GDN-C1 / GDN-C2 / GDN-C3 / GDN-H1 from the GDN CuTe-DSL cache audit (PR flashinfer-ai#1 of the suggested sequence). ## Summary - Documented FP16 `q/k/v/a/b` on the BF16-state and FP32-state MTP decode paths were silently reinterpreted as BF16 because the kernels hard-code `cutlass.BFloat16` fragments. Convert those operands to BF16 at the kernel boundary (and stage non-BF16 caller `output=` on MTP). - Add polymorphic operand dtypes (`A_log`, `dt_bias`, slot indices) to the compile-cache identities so mixed-dtype sequences no longer collide. - Stop returning a cached per-batch default `output` buffer from BF16-state paths when `output=None`. - Apply the same FP16 conversion on the WY output-only kernel; convert slot indices to int32 when needed. - Follow-up: non-BF16 MTP `output=` staging uses `output.to(bfloat16)` (not `empty_like`) so negative-index padding rows keep caller-initialized values. - Assert documented `dt_bias` (bf16/fp32) and `initial_state_indices` (int32/int64) dtypes at the public API. ## Test plan - [x] New regressions in `tests/gdn/test_decode_delta_rule.py` (8 parametrizations): FP16 conversion, dtype / `dt_bias` interleaving, default-output non-aliasing, non-BF16 `output=`, padding-slot preservation, WY FP16 - [x] Those regressions fail on unmodified `main` (most on main) and pass on this branch - [x] Full `tests/gdn/test_decode_delta_rule.py`: **838 passed** on H100 NVL (`CUDA_VISIBLE_DEVICES=1`, `-x -vv`, ~38 min) - [ ] GPU CI: `@flashinfer-bot run` ## Review Independent re-review at `024e7c4f`: **approve-with-nits** (padding critical fixed). Follow-up commits add comment trim + API dtype asserts. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved validation for tensor data types, shapes, and index values. * Fixed mixed-precision decode workflows, including proper BF16 conversion and preservation of requested output types. * Prevented stale output reuse and preserved padding in partially filled output buffers. * Improved compilation behavior when switching between supported input data types. * **Tests** * Added coverage for FP16/BF16 conversion, output handling, cache isolation, index validation, and mixed-precision correctness. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…ions (flashinfer-ai#4513) ## 📌 Description `tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins many times over. Two independent reasons, both verified against the cache keys rather than assumed: **1. Batch size is not a compile key.** It has been dynamic since flashinfer-ai#3649, so it reaches the cache only through coarse buckets: | Path | How `B` enters the key | |------|------------------------| | pretranspose (`gdn_decode_pretranspose.py:964`) | not at all | | bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all — the tests build a contiguous pool, so `pool_size_key = -1` and `pool_slot_stride = (-1,)` are B-independent sentinels | | nontranspose (`gdn_decode_nontranspose.py:725`) | only via `use_small_batch = B < 32` | | fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`, `gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` / `_get_bf16_mtp_config` | The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`: 378 of the file's 817 parametrized cases (46%) but only **42** compile keys, because `tile_v` is an explicit monkeypatched axis and the 9 batch sizes contribute nothing. The first commit keeps one batch size per bucket. I verified each kept set reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a real mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel` silently dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}). **2. Intermediate `seq_len` values only re-specialize on `T`.** `get_mtp_config` returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for every `T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile config that T=4 or T=8 does not already cover. `T=2` is kept as the one structurally distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2` branches. ### Effect | | cases | compile keys (retuned tests) | |---|---|---| | before | 817 | 102 | | after batch collapse | 501 | 102 | | after `seq_len` trim | 409 | 62 | Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized). ### Coverage cost The first commit costs nothing in specialization coverage — the same cubins still run, just at fewer runtime batch sizes. The second commit is a deliberate reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a T-specific off-by-one would no longer be caught. It is a separate commit so it can be dropped if reviewers would rather keep the full sweep. ## 🔍 Related Issues Refs flashinfer-ai#4110 (GDN cold-compile CI time). Complements flashinfer-ai#4128 and flashinfer-ai#4444, which remove key entries that provably do not reach codegen; this removes test cases that map onto keys already covered. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416 passed in 32m25s** - [ ] GPU CI for the timing comparison ## Reviewer Notes - I did not measure a clean before/after wall clock: the ~38 min baseline I was working from comes from flashinfer-ai#4219's description rather than the same machine, so I'd rather let CI provide the comparison than quote a number I can't stand behind. The case and key counts above are exact and static. - Worth noting for flashinfer-ai#4110 more broadly: cutting 50% of the cases bought substantially less than 50% of the wall clock, which suggests the remaining cost is dominated by compilation and fixed overhead rather than per-case execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in flashinfer-ai#4214) as the larger lever. - The batch sizes kept per test are load-bearing, not arbitrary — each set is one representative per config bucket. I left a one-line comment at each site so they don't get "restored" later. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Streamlined decode and MTP test coverage to use representative batch sizes and sequence lengths. * Preserved coverage for key thresholds, tile configurations, transposition modes, precision variants, and sequence-length scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>
5 tasks
bkryu
pushed a commit
that referenced
this pull request
Sep 8, 2026
## 📌 Description Migrate the GDN CuTe-DSL call sites that already compile with `--enable-tvm-ffi` to `build_and_load_cute_dsl_kernel`, so each specialization is exported to the on-disk kernel cache and later processes JITLink it instead of recompiling. Covered call sites (8 compile sites in 5 files): | File | Sites | Module | | --- | --- | --- | | `gdn_decode_nontranspose.py` | small/big batch | `gdn_decode_nontranspose` | | `gdn_decode_pretranspose.py` | 1 | `gdn_decode_pretranspose` | | `gdn_decode_mtp.py` | inline, warp | `gdn_decode_mtp` | | `gdn_decode_bf16_state.py` | wide_vec, wide_vec_t1, mtp_ilp4 | `gdn_decode_bf16_state` | | `blackwell/gdn_prefill.py` | 1 | `gdn_blackwell_prefill` | How it works: - Only the `cute.compile(...)` call is wrapped in a `compile_fn` closure. Every in-process cache structure (getter dicts, `defaults_by_B`, aux buffers) is unchanged, and on a disk hit the closure never runs. - Specialization names mirror the in-process cache keys, which #4219 and #4436 already audited to cover every codegen parameter. - Where several entry points share one module, the name adds a variant tag. The Blackwell prefill name also adds `num_sm`, which the compile bakes in as `max_active_clusters`. - A small shared formatter (`gdn_kernels/cute_dsl_cache_naming.py`) keeps names symbol-safe and caps their length. Measurements with fresh processes and an isolated cache dir: | GPU | Selection | Cold (compile + export) | Warm (disk hit) | | --- | --- | --- | --- | | RTX PRO 6000 (SM120) | bf16_state + fp32 MTP verify subsets (56 tests) | 604s | **2.3s** | | RTX PRO 6000 (SM120) | pretranspose + nontranspose basics (7 tests) | 2.9s | **1.4s** | | B200 (SM100) | prefill suites (2835 tests, includes GPU test execution) | 217s | **110s** | ## 🔍 Related Issues #4214 (GDN-P1, PR 4 of the suggested implementation sequence) ## 🚀 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Ran on RTX PRO 6000 (SM120) and B200 (SM100): - `tests/gdn/test_cute_dsl_kernel_cache.py` (new, modeled on `tests/jit/test_cute_dsl_cache.py`): 68 passed on both GPUs. Naming contracts plus a disk-cache round trip (reload with `cute.compile` forbidden, bit-identical results). - `tests/jit/test_cute_dsl_cache.py`: 42 passed - decode parity subsets from `tests/gdn/test_decode_delta_rule.py` (SM120): passed cold, warm, and with `FLASHINFER_CUTE_DSL_DISABLE_CACHE=1` - `tests/gdn/test_multistream_overlap.py`: 4 passed on B200 (1 of them skipped on SM120) - prefill suites on B200: 2835 passed, 864 skipped, 24 failed, identical cold/warm/cache-disabled. The 24 failures (fp8-state `test_prefill_kernel_state_dtype` cases) fail identically on unmodified main on the same machine, i.e. pre-existing in my environment. ## Reviewer Notes Out of scope, as follow-ups: - WY kernels: off the `gdn_decode.py` dispatch and not TVM-FFI yet (persisting them needs a calling-convention change first). - `delta_rule_dsl` kernels: use explicit `cute.compile[gpu_arch]` targets, which the disk cache's arch keying does not reflect yet (documented limitation). - GDN-H2/H3 (explicit compile targets, PR 2 of the sequence) is unchanged. Open to adjusting the split if a different scope was intended for GDN-P1. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Performance** - Added persistent caching for GDN GPU kernel builds, reducing compilation overhead for repeated workloads. - Cached kernels are reused across sessions when the relevant configuration and source remain unchanged. - Kernel builds automatically refresh when source code or compilation settings change, helping ensure correct results across updates. - Improved cache handling across multiple GDN decode and prefill kernel variants for more consistent build reuse. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue
#4214
Addresses GDN-C1 / GDN-C2 / GDN-C3 / GDN-H1 from the GDN CuTe-DSL cache audit (PR #1 of the suggested sequence).
Summary
q/k/v/a/bon the BF16-state and FP32-state MTP decode paths were silently reinterpreted as BF16 because the kernels hard-codecutlass.BFloat16fragments. Convert those operands to BF16 at the kernel boundary (and stage non-BF16 calleroutput=on MTP).A_log,dt_bias, slot indices) to the compile-cache identities so mixed-dtype sequences no longer collide.outputbuffer from BF16-state paths whenoutput=None.output=staging usesoutput.to(bfloat16)(notempty_like) so negative-index padding rows keep caller-initialized values.dt_bias(bf16/fp32) andinitial_state_indices(int32/int64) dtypes at the public API.Test plan
tests/gdn/test_decode_delta_rule.py(8 parametrizations): FP16 conversion, dtype /dt_biasinterleaving, default-output non-aliasing, non-BF16output=, padding-slot preservation, WY FP16main(most on main) and pass on this branchtests/gdn/test_decode_delta_rule.py: 838 passed on H100 NVL (CUDA_VISIBLE_DEVICES=1,-x -vv, ~38 min)@flashinfer-bot runReview
Independent re-review at
024e7c4f: approve-with-nits (padding critical fixed). Follow-up commits add comment trim + API dtype asserts.Summary by CodeRabbit
Bug Fixes
Tests