Spaces:
Running
Kernel Hub publishing access request for aabbdev/rwkv7-kernels
I am requesting Kernel Hub publishing access for the public v0 repository aabbdev/rwkv7-kernels.
A working edition-5 torch-noarch CUDA artifact is already built and validated:
- backend variant:
torch-cuda - CUDA capability: 8.9
- API version: v0
- source Git commit:
03cbd24173a7aa508c76099fb9559108aa73cd48 - artifact digest: SHA-256 metadata covering 7 Python files
- loader validation:
kernels==0.16.0withget_local_kernel(..., backend="cuda") - tests: 11 CPU tests and 16 CUDA tests
Purpose: optimized RWKV7 recurrent forward, prefill, decode, and native backward kernels for PyTorch models, with FP32 recurrent state and FP16/BF16 activations. The implementation uses TileLang 0.1.13 and includes an exact PyTorch fallback.
Public API:
kernel.is_available()
kernel.recurrent_forward(r, w_log, k, v, a, b, state=None, backend="auto")
kernel.rwkv7(..., return_state=False, backend="auto")
Measured on an RTX 4080 SUPER (SM89), BF16, H32, D64:
- prefill T=2048: 2.75M tokens/s
- forward + backward T=2048: best 110.16k tokens/s; final repeat 105.99k tokens/s
- decode B1: 6.53k tokens/s
- decode B8: 50.68k tokens/s
- decode B16: 101.97k tokens/s
The source tree is complete locally (build.toml, CARD.md, tests, benchmarks, and torch-ext/rwkv7_kernels). It is ready to be published publicly as soon as kernel repository creation is enabled for aabbdev. I can provide any additional code or benchmark evidence requested during review.
Update after the optimization and robustness pass:
- clean source HEAD:
f55b6b102639acc823be35df79389108edf1218c - rebuilt
torch-cudaartifact: 12 digested package files, clean Git provenance - validation: 32 CPU tests and 41 CUDA tests
- prepared-factor forward/backward, native packed resets, embedded offline schedules, PyTorch custom-op/fake/export surface, and persistent CUDA Graph decode
- BF16 B1/T2048/H32/D64 experimental sharded backward: 306.7k tokens/s at 0.94 GiB, with the stable path retained by default
- B1 CUDA Graph decode: 34.84 microseconds; B4/B8/B16 approximately 35 microseconds total
- long-context prefill validated through 256k with exact monolithic/chunked parity
- final artifact reloaded successfully with
kernels.get_local_kernel(..., backend="cuda")and the prepared/decode-graph APIs were exercised on SM89
The public repository is still ready to be created immediately once Kernel Hub publishing access is enabled for aabbdev.