Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
15 lines
339 B
Python
15 lines
339 B
Python
from .cumsum import (
|
|
chunk_local_cumsum,
|
|
chunk_local_cumsum_scalar,
|
|
chunk_local_cumsum_vector,
|
|
)
|
|
from .index import prepare_chunk_indices, prepare_chunk_offsets
|
|
|
|
__all__ = [
|
|
"chunk_local_cumsum",
|
|
"chunk_local_cumsum_scalar",
|
|
"chunk_local_cumsum_vector",
|
|
"prepare_chunk_indices",
|
|
"prepare_chunk_offsets",
|
|
]
|