Skip to content

Add intermediate states and exact CP transition prototype #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
0z5a wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add intermediate states and exact CP transition prototype #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
178bdb3
Add exact context-parallel KDA transitions
codex Aug 14, 2026
0b083de
Test distributed CP through FlashKDA forward
codex Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 46 additions & 16 deletions csrc/flash_kda.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ void fwd(
double lower_bound,
std::optional<torch::Tensor> initial_state = std::nullopt,
std::optional<torch::Tensor> final_state = std::nullopt,
std::optional<torch::Tensor> cu_seqlens = std::nullopt
std::optional<torch::Tensor> cu_seqlens = std::nullopt,
std::optional<torch::Tensor> intermediate_state = std::nullopt
) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && g.is_cuda() && beta.is_cuda() && out.is_cuda() && workspace.is_cuda(),
"all tensors must be on CUDA");
Expand All @@ -57,6 +58,7 @@ void fwd(
bool has_state_in = initial_state.has_value();
bool has_state_out = final_state.has_value();
bool state_fp32 = false;
bool has_intermediate_state = intermediate_state.has_value();

if (has_state_in) {
auto& is = initial_state.value();
Expand Down Expand Up @@ -150,6 +152,9 @@ void fwd(
TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided");
auto& cu_seqlens_t = cu_seqlens.value();
TORCH_CHECK(cu_seqlens_t.is_cuda(), "cu_seqlens must be on CUDA");
TORCH_CHECK(cu_seqlens_t.device() == q.device(),
"cu_seqlens must be on the same CUDA device as q");
TORCH_CHECK(cu_seqlens_t.is_contiguous(), "cu_seqlens must be contiguous");
TORCH_CHECK(cu_seqlens_t.dtype() == torch::kLong, "cu_seqlens must be int64");
TORCH_CHECK(cu_seqlens_t.dim() == 1, "cu_seqlens must be 1D");
N_val = cu_seqlens_t.numel() - 1;
Expand Down Expand Up @@ -180,36 +185,60 @@ void fwd(
total_tiles = int(N_val * ((T_seq + CHUNK - 1) / CHUNK)); // exact for batched
}

if (has_intermediate_state) {
auto& ims = intermediate_state.value();
TORCH_CHECK(ims.is_cuda() && ims.is_contiguous(),
"intermediate_state must be a contiguous CUDA tensor");
TORCH_CHECK(ims.device() == q.device(),
"intermediate_state must be on the same CUDA device as q");
TORCH_CHECK(ims.dtype() == torch::kBFloat16,
"intermediate_state must be bfloat16");
TORCH_CHECK(ims.dim() == 4 && ims.size(0) == H &&
ims.size(1) == total_tiles && ims.size(2) == D && ims.size(3) == D,
"intermediate_state must be [H, total_tiles, D, D]");
}
auto intermediate_state_ptr = has_intermediate_state
? reinterpret_cast<cutlass::bfloat16_t*>(intermediate_state->data_ptr<at::BFloat16>())
: nullptr;

// Dispatch based on state configuration and varlen
#define LAUNCH(HI, HO, FP32, VL) \
launch_fwd<128, HI, HO, FP32, VL>( \
#define LAUNCH(HI, HO, FP32, HIS, VL) \
launch_fwd<128, HI, HO, FP32, HIS, VL>( \
q_ptr, k_ptr, v_ptr, g_ptr, beta_t_ptr, \
initial_state_raw, scale_f, final_state_raw, out_ptr, \
initial_state_raw, scale_f, final_state_raw, intermediate_state_ptr, out_ptr, \
workspace_ptr, total_tiles, \
int(T_total), int(H), int(N_val), cu_seqlens_dev, \
A_log_ptr, dt_bias_ptr, gate_scale, stream)

#define DISPATCH_STATE(VL) \
#define DISPATCH_STATE(HIS, VL) \
if (!has_state_in && !has_state_out) { \
LAUNCH(false, false, false, VL); \
LAUNCH(false, false, false, HIS, VL); \
} else if (has_state_in && has_state_out && state_fp32) { \
LAUNCH(true, true, true, VL); \
LAUNCH(true, true, true, HIS, VL); \
} else if (has_state_in && has_state_out && !state_fp32) { \
LAUNCH(true, true, false, VL); \
LAUNCH(true, true, false, HIS, VL); \
} else if (!has_state_in && has_state_out && state_fp32) { \
LAUNCH(false, true, true, VL); \
LAUNCH(false, true, true, HIS, VL); \
} else if (!has_state_in && has_state_out && !state_fp32) { \
LAUNCH(false, true, false, VL); \
LAUNCH(false, true, false, HIS, VL); \
} else if (has_state_in && !has_state_out && state_fp32) { \
LAUNCH(true, false, true, VL); \
LAUNCH(true, false, true, HIS, VL); \
} else { \
LAUNCH(true, false, false, VL); \
LAUNCH(true, false, false, HIS, VL); \
}

if (is_varlen) {
DISPATCH_STATE(true);
if (has_intermediate_state) {
if (is_varlen) {
DISPATCH_STATE(true, true);
} else {
DISPATCH_STATE(true, false);
}
} else {
DISPATCH_STATE(false);
if (is_varlen) {
DISPATCH_STATE(false, true);
} else {
DISPATCH_STATE(false, false);
}
}

#undef DISPATCH_STATE
Expand All @@ -223,7 +252,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("workspace"),
py::arg("A_log"), py::arg("dt_bias"), py::arg("lower_bound"),
py::arg("initial_state") = py::none(), py::arg("final_state") = py::none(),
py::arg("cu_seqlens") = py::none());
py::arg("cu_seqlens") = py::none(),
py::arg("intermediate_state") = py::none());
m.def("get_workspace_size",
static_cast<int64_t(*)(int64_t, int64_t, int64_t)>(&get_workspace_size),
"Get workspace size in bytes",
Expand Down
3 changes: 2 additions & 1 deletion csrc/fwd.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

#include <cutlass/bfloat16.h>

template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, bool IsVarlen = true>
template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, bool HasIntermediateState = false, bool IsVarlen = true>
void launch_fwd(
cutlass::bfloat16_t const* q_ptr,
cutlass::bfloat16_t const* k_ptr,
Expand All @@ -13,6 +13,7 @@ void launch_fwd(
void const* initial_state_ptr,
float scale,
void* final_state_ptr,
cutlass::bfloat16_t* intermediate_state_ptr,
cutlass::bfloat16_t* out_ptr,
void* workspace_ptr,
int total_tiles,
Expand Down
14 changes: 14 additions & 0 deletions csrc/smxx/fwd_kernel2.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ template <
bool HasStateIn = true,
bool HasStateOut = true,
bool StateFP32 = false,
bool HasIntermediateState = false,
bool IsVarlen = true
>
__global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
Expand All @@ -143,6 +144,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
CUTE_GRID_CONSTANT TmaStoreState const tma_store_final_state,
CUTE_GRID_CONSTANT TmaStoreOut const tma_store_out,
cutlass::bfloat16_t* out_raw_ptr,
cutlass::bfloat16_t* intermediate_state_ptr,
int T_total,
int H,
int N,
Expand Down Expand Up @@ -732,6 +734,18 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
}
compute_barrier.arrive_and_wait();

if constexpr (HasIntermediateState) {
// StateSmemLayout maps the logical [D, D] state to its swizzled
// shared-memory storage. Snapshot after the whole update is
// visible so every chunk is independently reusable by a caller.
int64_t snapshot_offset =
(int64_t(head_idx) * total_tiles + tile_base + t) * D * D;
for (int i = compute_tid; i < D * D; i += kComputeThreads) {
intermediate_state_ptr[snapshot_offset + i] =
s_acc(i / D, i % D);
}
}

#ifndef TMA_DISABLE_ALL
cutlass::arch::fence_view_async_shared();
store_pipeline.producer_commit(out_write);
Expand Down
38 changes: 21 additions & 17 deletions csrc/smxx/fwd_launch.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
#include "fwd_kernel2.cuh"

// ==================== launch_fwd ====================
template <int D, bool HasStateIn, bool HasStateOut, bool StateFP32, bool IsVarlen>
template <int D, bool HasStateIn, bool HasStateOut, bool StateFP32, bool HasIntermediateState, bool IsVarlen>
void launch_fwd(
cutlass::bfloat16_t const* q_ptr,
cutlass::bfloat16_t const* k_ptr,
Expand All @@ -13,6 +13,7 @@ void launch_fwd(
void const* initial_state_ptr,
float scale,
void* final_state_ptr,
cutlass::bfloat16_t* intermediate_state_ptr,
cutlass::bfloat16_t* out_ptr,
void* workspace_ptr,
int total_tiles,
Expand Down Expand Up @@ -195,7 +196,7 @@ void launch_fwd(
decltype(tma_store_final_state),
decltype(tma_store_out),
CHUNK, D, kInputStages, kOutputStages, kK2Threads,
HasStateIn, HasStateOut, StateFP32, IsVarlen
HasStateIn, HasStateOut, StateFP32, HasIntermediateState, IsVarlen
>;

cudaFuncSetAttribute(kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k2);
Expand All @@ -210,29 +211,32 @@ void launch_fwd(
tma_load_initial_state,
tma_store_final_state,
tma_store_out,
out_ptr, T_total, H, N, cu_seqlens_ptr, total_tiles
out_ptr, intermediate_state_ptr,
T_total, H, N, cu_seqlens_ptr, total_tiles
);
}
#endif
}

// Explicit instantiations
#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, VL) \
template void launch_fwd<D, HI, HO, FP32, VL>( \
#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, HIS, VL) \
template void launch_fwd<D, HI, HO, FP32, HIS, VL>( \
cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \
cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \
cutlass::bfloat16_t const*, void const*, float, void*, \
cutlass::bfloat16_t*, void*, int, int, int, int, \
cutlass::bfloat16_t*, cutlass::bfloat16_t*, void*, int, int, int, int, \
int64_t const*, float const*, float const*, float, cudaStream_t);

#define INSTANTIATE_STATE_VARIANTS(VL) \
INSTANTIATE_LAUNCH_FWD(128, true, true, false, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, true, true, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, false, false, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, true, false, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, false, false, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, true, true, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, false, true, VL)

INSTANTIATE_STATE_VARIANTS(true) // varlen
INSTANTIATE_STATE_VARIANTS(false) // non-varlen
#define INSTANTIATE_STATE_VARIANTS(HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, true, false, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, true, true, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, false, false, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, true, false, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, false, false, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, false, true, true, HIS, VL) \
INSTANTIATE_LAUNCH_FWD(128, true, false, true, HIS, VL)

INSTANTIATE_STATE_VARIANTS(false, true) // varlen, no chunk-state snapshots
INSTANTIATE_STATE_VARIANTS(false, false) // batched, no chunk-state snapshots
INSTANTIATE_STATE_VARIANTS(true, true) // varlen, chunk-state snapshots
INSTANTIATE_STATE_VARIANTS(true, false) // batched, chunk-state snapshots
49 changes: 47 additions & 2 deletions flash_kda/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,45 @@
from flash_kda_C import fwd as _fwd_raw, get_workspace_size


def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None):
def get_intermediate_state_shape(q, cu_seqlens=None):
"""Return the required ``[H, tile_capacity, D, D]`` snapshot-buffer shape.

In ragged mode, ``tile_capacity`` is an upper bound with at most one unused
tile per sequence. Use ``get_intermediate_state_tile_prefix`` to locate the
compact valid chunk states without synchronizing to the host.
"""
B, T_seq, H, D = q.shape
if cu_seqlens is None:
total_tiles = B * ((T_seq + 15) // 16)
else:
total_tiles = (B * T_seq + 15) // 16 + cu_seqlens.numel() - 1
return H, total_tiles, D, D


def get_intermediate_state_tile_prefix(q, cu_seqlens=None):
"""Return device-resident offsets for the compact chunk-state layout.

The state after chunk ``t`` of sequence ``n`` is at
``intermediate_state[head, tile_prefix[n] + t]``.
``tile_prefix[-1]`` is the number of valid snapshot slots, so this helper
locates ragged snapshots without host synchronization. Snapshots are useful
for inspection or a serial handoff, but cannot alone implement exact CP
prefix composition: KDA's input-state dependence requires the affine
transition prototype in flash_kda.cp.
"""
B, T_seq = q.shape[:2]
if cu_seqlens is None:
tiles_per_sequence = (T_seq + 15) // 16
return torch.arange(B + 1, dtype=torch.long, device=q.device) * tiles_per_sequence
if cu_seqlens.device != q.device:
raise ValueError("cu_seqlens must be on the same device as q")
sequence_tiles = torch.div(cu_seqlens[1:] - cu_seqlens[:-1] + 15, 16,
rounding_mode="floor")
return torch.cat((torch.zeros(1, dtype=torch.long, device=cu_seqlens.device),
torch.cumsum(sequence_tiles, dim=0)))


def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None, intermediate_state=None):
"""FlashKDA forward (Flash Kimi Delta Attention).

Args:
Expand All @@ -25,6 +63,12 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state
recurrent state. Same dtype/shape rules as ``initial_state``.
cu_seqlens (torch.Tensor, optional): Cumulative sequence lengths, int64,
shape ``[N+1]``. When provided, ``B`` must be 1.
intermediate_state (torch.Tensor, optional): BF16 output buffer for the
recurrent state after each 16-token chunk. Its shape must equal
``get_intermediate_state_shape(q, cu_seqlens)``. In ragged mode, use
``get_intermediate_state_tile_prefix(q, cu_seqlens)``: chunk ``t``
of sequence ``n`` is stored at ``[head, tile_prefix[n] + t]``;
slots from ``tile_prefix[-1]`` onward are left untouched.

Notes:
* Currently requires ``K = V = 128``.
Expand All @@ -38,4 +82,5 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state
workspace = torch.empty(get_workspace_size(T_total, H, N), dtype=torch.uint8, device=q.device)

_fwd_raw(q, k, v, g, beta, float(scale), out, workspace, A_log, dt_bias, lower_bound,
initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens)
initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens,
intermediate_state=intermediate_state)
Loading