Skip to content

[GDN2] Add fused BT=16 inference kernels for GDN2 prefill#990

Open
Ghd-2077 wants to merge 2 commits into
fla-org:mainfrom
Ghd-2077:perf/gdn2-fused-infer
Open

[GDN2] Add fused BT=16 inference kernels for GDN2 prefill#990
Ghd-2077 wants to merge 2 commits into
fla-org:mainfrom
Ghd-2077:perf/gdn2-fused-infer

Conversation

@Ghd-2077

Copy link
Copy Markdown

[Ops] Add fused BT=16 inference kernels for GDN-2 prefill

Optimization

This implementation follows the FlashKDA two-kernel inference design for GDN-2:

  • K1 (intra-chunk): fuses optional Q/K L2 normalization, gate activation, GDN-2 vector erase/write WY auxiliary construction, and the block-local inverse for BT=16 chunks.
  • K2 (inter-chunk): propagates the recurrent state across chunks for each (batch, head, V-tile), writes the output, and optionally stores the final state.
  • The intended benefit is fewer intermediate global-memory round-trips than the baseline multi-kernel chunk path, while retaining GDN-2's independent erase gate b, write gate w, and state_v_first support.

Testing

Environment

Item Value
Git branch perf/gdn2-fused-infer
Git commit 9079232
Commit subject test
Working tree dirty yes
GPU / driver / memory NVIDIA H100 80GB HBM3, 535.183.06, 81559 MiB
GPU (PyTorch) NVIDIA H100 80GB HBM3
Python 3.12.13
PyTorch 2.10.0+cu128
CUDA (PyTorch) 12.8
Triton 3.6.0

Performance

Operation: chunk_gdn2_infer
Speedup: baseline median_ms / fused median_ms

B T H D Fused (ms) Baseline (ms) Speedup
1 8192 96 128 2.027 3.719 1.83x
2 16384 16 128 1.875 2.599 1.39x
4 2048 16 128 0.382 1.098 2.87x
4 4096 64 128 2.416 4.771 1.98x
8 1024 8 64 0.212 1.088 5.13x
8 2048 32 256 3.030 5.143 1.70x

Average speedup: 2.48x (arithmetic mean), 2.24x (geometric mean); best case: 5.13x.

Benchmark commands:

python -m benchmarks.ops.run --op chunk_gdn2_infer --modes fwd --base '' --json benchmark_baseline.json
python -m benchmarks.ops.run --op chunk_gdn2_infer --modes fwd --base '' --json benchmark_fused.json

Correctness

PASS: 5 passed, 0 skipped in 4.20s.

Test node: tests/ops/test_gdn2.py::test_chunk_gdn2_fwd_inference

Test Status o diff o ratio ht diff ht ratio Failure detail
test_chunk_gdn2_fwd_inference[small-fp32] PASSED 0.000142 0.001581 0.000526 0.001094
test_chunk_gdn2_fwd_inference[l2norm-fp32] PASSED 0.000109 0.001708 0.000681 0.001063
test_chunk_gdn2_fwd_inference[gate-dt-bias-fp32] PASSED 0.000254 0.001225 0.000615 0.001076
test_chunk_gdn2_fwd_inference[safe-gate-no-bias-bf16] PASSED 0.000802 0.003675 0.000819 0.002292
test_chunk_gdn2_fwd_inference[state-v-first-v-ne-k-fp16] PASSED 0.0001 0.000983 0.000171 0.000296

Correctness command:

python -m pytest -vv -s --log-cli-level=INFO tests/ops/test_gdn2.py::test_chunk_gdn2_fwd_inference

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant