Problem
CHUNK=16 and D=128 appear as magic numbers in multiple locations:
csrc/flash_kda.cpp — static_assert and workspace size calculationcsrc/smxx/fwd_kernel1.cuh — kernel constantscsrc/smxx/fwd_kernel2.cuh — kernel constantstests/torch_ref.py — reference implementationtests/test_fwd.py — test parametersdocs/20260420-flashkda-v1-deep-dive.md — documentation
There is no single source of truth. If someone wanted to support D=64 or D=256, they'd need to change dozens of locations.
Recommendation
Define CHUNK and D as compile-time constants in one place (e.g., setup.py as compiler defines, or a shared header), and reference them everywhere else. Add a test that verifies the error message when D != 128.
Impact
Extensibility. Makes it possible to support different D values without a multi-file hunt for magic numbers.
Problem
CHUNK=16andD=128appear as magic numbers in multiple locations:csrc/flash_kda.cpp—static_assertand workspace size calculationcsrc/smxx/fwd_kernel1.cuh— kernel constantscsrc/smxx/fwd_kernel2.cuh— kernel constantstests/torch_ref.py— reference implementationtests/test_fwd.py— test parametersdocs/20260420-flashkda-v1-deep-dive.md— documentationThere is no single source of truth. If someone wanted to support D=64 or D=256, they'd need to change dozens of locations.
Recommendation
Define CHUNK and D as compile-time constants in one place (e.g.,
setup.pyas compiler defines, or a shared header), and reference them everywhere else. Add a test that verifies the error message when D != 128.Impact
Extensibility. Makes it possible to support different D values without a multi-file hunt for magic numbers.