Skip to content

feat(transformer): wire split-KV decode behind TLLM_ATTN_SPLITKV (TLLM-ATTN-SPLITKV PR-C) #13

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

Merged
holtwood merged 1 commit into from
Sep 15, 2026
Merged

feat(transformer): wire split-KV decode behind TLLM_ATTN_SPLITKV (TLLM-ATTN-SPLITKV PR-C) #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 12 additions & 2 deletions include/tiny_llm/transformer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,13 @@

namespace tiny_llm {

// split-KV decode attention（TLLM-ATTN-SPLITKV）partial 工作区的 split 数上界。
//
// 运行时开关 TLLM_ATTN_SPLITKV 允许 1..kAttnMaxSplits，而缓冲按**上界**一次性分配：
// 这样"每次调用重新解析环境变量"（不缓存、因而测试里 setenv 即时生效、且不需要为测试
// 暴露 reset seam）与"分配大小固定"可以同时成立，不会因为运行中改小/改大取值而越界。
inline constexpr int kAttnMaxSplits = 32;

// 共享中间激活工作区：所有 TransformerLayer 复用同一组 GPU 缓冲。
// 推理逐层串行，同一时刻只有一个层在使用这些缓冲；每层独立分配会按
// 层数线性放大显存（24 层 × ~0.9GB 必然 OOM）。
Expand All @@ -22,8 +29,11 @@ struct LayerWorkspace {
half *ffn_gate = nullptr; // [max_batch, intermediate_dim]
half *ffn_up = nullptr; // [max_batch, intermediate_dim]
half *ffn_output = nullptr; // [max_batch, hidden_dim]
int max_batch_tokens = 0;
bool allocated = false;
// split-KV partial：[num_heads * kAttnMaxSplits * (2 + head_dim)] fp32，布局见
// kernels/attention.cuh。kernel 零分配，缓冲由这里持有（层间串行，可复用）。
float *attn_partial = nullptr;
int max_batch_tokens = 0;
bool allocated = false;

void allocate(const ModelConfig &config);
void free();
Expand Down
79 changes: 73 additions & 6 deletions src/transformer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,36 @@ Result<PagedAttentionMode> resolvePagedAttentionMode() {
"\" (expected auto | legacy | direct)");
}

// ── split-KV decode 的开关（TLLM-ATTN-SPLITKV）──────────────────────────────
//
// TLLM_ATTN_SPLITKV = <num_splits>：
// 未设置 / 空 / 0 / 1 → 走**单遍入口**，与今天逐位相同，且不额外多一次 combine launch；
// 2 .. kAttnMaxSplits → 走 split-KV 入口（partial 写 ws_->attn_partial）。
// 其它取值（非数字、超过上界）显式报错，不静默回退（G5）。
//
// 只接受纯十进制数字：strtol 的宽松解析（"8abc"、" 8"、"+8"）会把拼写错误当成合法值吞掉。
// 每次调用解析（约 20 ns、无堆分配），使 setenv 在测试中即时生效、无需为测试暴露 reset
// seam；partial 缓冲按 kAttnMaxSplits 预分配，因此运行中改值不会越界。
Result<int> resolveAttnSplitKv() {
const char *raw = std::getenv("TLLM_ATTN_SPLITKV");
if (raw == nullptr || *raw == '\0') {
return Result<int>::ok(1);
}
for (const char *p = raw; *p != '\0'; ++p) {
if (*p < '0' || *p > '9') {
return Result<int>::err(std::string("TLLM_ATTN_SPLITKV: unrecognized value \"") + raw +
"\" (expected an integer in [0, " +
std::to_string(kAttnMaxSplits) + "])");
}
}
const long value = std::strtol(raw, nullptr, 10);
if (value < 0 || value > kAttnMaxSplits) {
return Result<int>::err("TLLM_ATTN_SPLITKV: value " + std::string(raw) +
" out of range [0, " + std::to_string(kAttnMaxSplits) + "]");
}
return Result<int>::ok(value < 2 ? 1 : static_cast<int>(value));
}

} // namespace

void LayerWorkspace::allocate(const ModelConfig &config) {
Expand All @@ -91,6 +121,8 @@ void LayerWorkspace::allocate(const ModelConfig &config) {
CUDA_CHECK(cudaMalloc(&ffn_gate, ffn_size * sizeof(half)));
CUDA_CHECK(cudaMalloc(&ffn_up, ffn_size * sizeof(half)));
CUDA_CHECK(cudaMalloc(&ffn_output, hidden_size * sizeof(half)));

// EXPERIMENT: attn_partial allocation disabled
} catch (...) {
// 修复：中途任一 cudaMalloc 失败时释放已分配指针再重抛。allocated
// 尚未置位，析构路径的 free() 会因早退检查跳过，必须在此手动清理，
Expand All @@ -110,6 +142,10 @@ void LayerWorkspace::allocate(const ModelConfig &config) {
cleanup(ffn_gate);
cleanup(ffn_up);
cleanup(ffn_output);
if (attn_partial) {
cudaFree(attn_partial);
attn_partial = nullptr;
}
throw;
}

Expand Down Expand Up @@ -151,6 +187,13 @@ void LayerWorkspace::free() {
safe_free(ffn_gate);
safe_free(ffn_up);
safe_free(ffn_output);
if (attn_partial) {
cudaError_t err = cudaFree(attn_partial);
if (err != cudaSuccess) {
fprintf(stderr, "CUDA error in LayerWorkspace::free: %s\n", cudaGetErrorString(err));
}
attn_partial = nullptr;
}

max_batch_tokens = 0;
allocated = false;
Expand Down Expand Up @@ -468,16 +511,32 @@ Result<void> TransformerLayer::attentionPaged(const half *x, half *output,
}
const bool use_direct = is_decode && mode_result.value() == PagedAttentionMode::Direct;

// split-KV 开关（TLLM-ATTN-SPLITKV）；非法取值同样在入口显式失败。
// 只作用于 decode：prefill 的并行轴是 query 位置，不在本设计范围。
const auto split_result = resolveAttnSplitKv();
if (split_result.isErr()) {
TLLM_ERROR("attentionPaged: {}", split_result.error());
return Result<void>::err(split_result.error());
}
const bool use_split = is_decode && split_result.value() > 1;

// Attention
float scale = 1.0f / std::sqrt(static_cast<float>(head_dim));
if (use_direct) {
// direct：直接从物理 pool + block table 寻址，跳过 gather。可见长度走 device
// int（CUDA Graph 可捕获）；块表长度用 visible_blocks，forwardPaged 已校验其
// 足以覆盖可见 token。
kernels::attention_decode_paged(ws_->q_buf, k_pool_layer, v_pool_layer, kv.block_table,
ws_->attn_buf, scale, num_heads, num_kv_heads, head_dim,
kv.decode_len, kv.block_size, kv.max_num_blocks,
kv.visible_blocks, stream);
if (use_split) {
kernels::attention_decode_paged_splitkv(
ws_->q_buf, k_pool_layer, v_pool_layer, kv.block_table, ws_->attn_buf, scale,
num_heads, num_kv_heads, head_dim, kv.decode_len, kv.block_size, kv.max_num_blocks,
kv.visible_blocks, ws_->attn_partial, split_result.value(), stream);
} else {
kernels::attention_decode_paged(ws_->q_buf, k_pool_layer, v_pool_layer, kv.block_table,
ws_->attn_buf, scale, num_heads, num_kv_heads, head_dim,
kv.decode_len, kv.block_size, kv.max_num_blocks,
kv.visible_blocks, stream);
}
} else {
// legacy：把可见区间 gather 到 scratch（连续布局），供连续 attention 使用。
kernels::paged_gather_blocks(kv.k_scratch, k_pool_layer, kv.block_table, visible,
Expand All @@ -486,8 +545,16 @@ Result<void> TransformerLayer::attentionPaged(const half *x, half *output,
kv.block_size, kv_dim, kv.max_num_blocks, stream);
if (is_decode) {
// decode：单 query 对已 gather 的可见 K/V；可见长度由 device int 提供
kernels::attention_decode(ws_->q_buf, kv.k_scratch, kv.v_scratch, ws_->attn_buf, scale,
num_heads, num_kv_heads, kv.decode_len, head_dim, stream);
if (use_split) {
kernels::attention_decode_splitkv(ws_->q_buf, kv.k_scratch, kv.v_scratch,
ws_->attn_buf, scale, num_heads, num_kv_heads,
kv.decode_len, head_dim, ws_->attn_partial,
split_result.value(), stream);
} else {
kernels::attention_decode(ws_->q_buf, kv.k_scratch, kv.v_scratch, ws_->attn_buf,
scale, num_heads, num_kv_heads, kv.decode_len, head_dim,
stream);
}
} else {
// prefill：因果掩码全量注意力（gather 回读与 scatter 内容一致）
kernels::attention_prefill(ws_->q_buf, kv.k_scratch, kv.v_scratch, ws_->attn_buf, scale,
Expand Down
208 changes: 177 additions & 31 deletions tests/test_paged_dispatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,43 @@ class PagedDispatchTest : public ::testing::Test {
ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
}

// 构造"刚好能进入 attentionPaged"的最小视图并跑一次 decode。
// 供只验证**入口错误语义**（开关非法取值）的用例使用——这些用例关心的是路由前的
// 校验分支，不需要完整的 fixture。
static Result<void> runMinimalDecode(Fixture &f) {
buildModel(f);

f.hidden = f.config.hidden_dim;
f.kv_dim = f.config.num_kv_heads * f.config.head_dim;
const size_t pool_elems =
static_cast<size_t>(kMaxBlocks) * static_cast<size_t>(kBlockSize) * f.kv_dim;
f.k_pool = DeviceBuffer<half>(pool_elems);
f.v_pool = DeviceBuffer<half>(pool_elems);
f.k_scratch = DeviceBuffer<half>(pool_elems);
f.v_scratch = DeviceBuffer<half>(pool_elems);
f.hidden_a = DeviceBuffer<half>(static_cast<size_t>(f.config.max_seq_len) * f.hidden);
f.table = DeviceBuffer<int>(static_cast<size_t>(kMaxBlocks));

f.view.k_pool = f.k_pool.data();
f.view.v_pool = f.v_pool.data();
f.view.block_table = f.table.data();
f.view.k_scratch = f.k_scratch.data();
f.view.v_scratch = f.v_scratch.data();
f.view.visible_blocks = 1;
f.view.block_size = kBlockSize;
f.view.max_num_blocks = kMaxBlocks;
f.view.max_visible_tokens = kMaxBlocks * kBlockSize;
f.view.position = 0;
f.view.decode_len = nullptr;

const int zero = 0;
DeviceBuffer<int> d_pos(1);
d_pos.copyFromHost(&zero, 1);

return f.layer->forwardPaged(f.hidden_a.data(), f.view, 1, d_pos.data(), f.d_cos.data(),
f.d_sin.data(), 0);
}

// 把 scratch 全部写成哨兵值，用于观测 legacy 是否真的 gather 过。
static void poisonScratch(Fixture &f) {
const size_t n = f.k_scratch.size();
Expand Down Expand Up @@ -417,38 +454,8 @@ TEST_F(PagedDispatchTest, DirectAndLegacyAgreeAcrossSteps) {
TEST_F(PagedDispatchTest, InvalidModeValueFailsLoudly) {
ScopedEnv env("TLLM_PAGED_ATTENTION", "direkt"); // 拼错
Fixture f;
buildModel(f);

// 只需要进入 attentionPaged；先构造最小视图（指针非空即可通过前面的几何校验）
f.hidden = f.config.hidden_dim;
f.kv_dim = f.config.num_kv_heads * f.config.head_dim;
const size_t pool_elems =
static_cast<size_t>(kMaxBlocks) * static_cast<size_t>(kBlockSize) * f.kv_dim;
f.k_pool = DeviceBuffer<half>(pool_elems);
f.v_pool = DeviceBuffer<half>(pool_elems);
f.k_scratch = DeviceBuffer<half>(pool_elems);
f.v_scratch = DeviceBuffer<half>(pool_elems);
f.hidden_a = DeviceBuffer<half>(static_cast<size_t>(f.config.max_seq_len) * f.hidden);
f.table = DeviceBuffer<int>(static_cast<size_t>(kMaxBlocks));

f.view.k_pool = f.k_pool.data();
f.view.v_pool = f.v_pool.data();
f.view.block_table = f.table.data();
f.view.k_scratch = f.k_scratch.data();
f.view.v_scratch = f.v_scratch.data();
f.view.visible_blocks = 1;
f.view.block_size = kBlockSize;
f.view.max_num_blocks = kMaxBlocks;
f.view.max_visible_tokens = kMaxBlocks * kBlockSize;
f.view.position = 0;
f.view.decode_len = nullptr;

const int zero = 0;
DeviceBuffer<int> d_pos(1);
d_pos.copyFromHost(&zero, 1);
auto r = runMinimalDecode(f);

auto r = f.layer->forwardPaged(f.hidden_a.data(), f.view, 1, d_pos.data(), f.d_cos.data(),
f.d_sin.data(), 0);
ASSERT_TRUE(r.isErr()) << "非法 TLLM_PAGED_ATTENTION 取值必须返回错误，不得静默回退";
EXPECT_NE(r.error().find("TLLM_PAGED_ATTENTION"), std::string::npos) << r.error();

Expand Down Expand Up @@ -481,3 +488,142 @@ TEST_F(PagedDispatchTest, PrefillIgnoresTheSwitch) {
<< "prefill 必须保留 legacy 路径（设计包 §1 non-goals），gather 应当写过 scratch";
freeModel(f);
}

// ── TLLM-ATTN-SPLITKV：开关语义与层级等价 ────────────────────────────────
//
// 层级门禁只覆盖"路由是否正确 + 两条路径在同一 num_splits 下是否等价"；
// kernel 级数值正确性由 tests/test_attention_splitkv.cpp 负责（逐位锚点 + oracle）。

namespace {

// 在给定开关组合下建 fixture、跑一次 decode、取回该步层输出。
// 每个组合用独立 fixture：buildFixture 的输入是固定种子，因此不同组合之间可直接比较。
template <typename BuildFn>
std::vector<half> decodeWith(const char *mode, const char *split, BuildFn &&run) {
ScopedEnv mode_env("TLLM_PAGED_ATTENTION", mode);
ScopedEnv split_env("TLLM_ATTN_SPLITKV", split);
return run();
}

} // namespace

// 未设置 / "0" / "1" 都必须走**单遍入口**，三者逐位相同（且与开关引入前一致）。
TEST_F(PagedDispatchTest, SplitKvOffByDefaultAndSingleSplitIsTheSinglePassPath) {
auto decode = [&](const char *split) {
ScopedEnv split_env("TLLM_ATTN_SPLITKV", split);
Fixture f;
buildFixture(f);
auto out = runDecode(f, f.hidden_a, 0);
freeModel(f);
return out;
};

const auto unset = decode(nullptr);
const auto zero = decode("0");
const auto one = decode("1");

EXPECT_TRUE(bitwiseEqual(unset, zero)) << maxAbsDiff(unset, zero);
EXPECT_TRUE(bitwiseEqual(unset, one)) << maxAbsDiff(unset, one);
}

// 非法取值显式失败，不静默回退（G5）——包括非数字、带尾字符、超上界。
TEST_F(PagedDispatchTest, InvalidSplitKvValueFailsLoudly) {
for (const char *bad : {"abc", "8x", "99", "+4"}) {
SCOPED_TRACE(std::string("TLLM_ATTN_SPLITKV=") + bad);
ScopedEnv split_env("TLLM_ATTN_SPLITKV", bad);
Fixture f;
auto r = runMinimalDecode(f);

ASSERT_TRUE(r.isErr()) << "非法 TLLM_ATTN_SPLITKV 取值必须返回错误，不得静默回退";
EXPECT_NE(r.error().find("TLLM_ATTN_SPLITKV"), std::string::npos) << r.error();
freeModel(f);
}
}

// 同一个 num_splits 下，legacy 与 direct 走同一份归约循环 ⇒ 层输出逐位相同。
// （num_splits > 1 的确切切分由 kernel 级测试逐位锚定。）
TEST_F(PagedDispatchTest, SplitKvLegacyAndDirectAgreeBitwiseAtSameSplits) {
for (const char *splits : {"2", "4", "8"}) {
SCOPED_TRACE(std::string("num_splits=") + splits);

auto legacy = decodeWith("legacy", splits, [&] {
Fixture f;
buildFixture(f);
auto out = runDecode(f, f.hidden_a, 0);
freeModel(f);
return out;
});
auto direct = decodeWith("direct", splits, [&] {
Fixture f;
buildFixture(f);
auto out = runDecode(f, f.hidden_b, 0);
freeModel(f);
return out;
});

EXPECT_TRUE(bitwiseEqual(legacy, direct))
<< "同一 num_splits 下 legacy 与 direct 必须逐位相同，max|diff|="
<< maxAbsDiff(legacy, direct);
}
}

// split 与单遍**不是**逐位关系（fp32 求和顺序不同），但必须在同一量级上：
// 这里用相对尺度设一个宽松上限，作为"路由没有把数据接错"的冒烟检查；精确门禁在
// tests/test_attention_splitkv.cpp。
TEST_F(PagedDispatchTest, SplitKvChangesNumericsOnlySlightlyVersusSinglePass) {
auto decode = [&](const char *split) {
ScopedEnv split_env("TLLM_ATTN_SPLITKV", split);
Fixture f;
buildFixture(f);
auto out = runDecode(f, f.hidden_a, 0);
freeModel(f);
return out;
};

const auto single = decode("1");
const auto split4 = decode("4");

const float diff = maxAbsDiff(single, split4);
float scale = 0.0f;
for (half h : single)
scale = std::max(scale, std::fabs(__half2float(h)));

ASSERT_GT(scale, 0.0f) << "参考输出不应恒为 0";
EXPECT_LT(diff / scale, 5e-2f)
<< "split 与单遍的差异相对尺度过大，疑似路由接错：max|diff|=" << diff << " scale=" << scale;
}

// prefill 与开关无关：即使 TLLM_ATTN_SPLITKV > 1，prefill 输出也必须逐位不变。
// 结构上 use_split 以 is_decode 为前提，但这是设计包 §1 的 non-goal，值得钉住。
TEST_F(PagedDispatchTest, SplitKvDoesNotAffectPrefill) {
auto run = [&](const char *split) {
ScopedEnv split_env("TLLM_ATTN_SPLITKV", split);
Fixture f;
buildFixture(f);

DeviceBuffer<half> h(static_cast<size_t>(f.config.max_seq_len) * f.hidden);
const auto input = randomFp16(h.size(), 999);
h.copyFromHost(input.data(), input.size());

const int zero = 0;
DeviceBuffer<int> d_pos(1);
d_pos.copyFromHost(&zero, 1);
f.view.position = 0;
f.view.decode_len = nullptr;

auto r = f.layer->forwardPaged(h.data(), f.view, 4, d_pos.data(), f.d_cos.data(),
f.d_sin.data(), 0);
EXPECT_TRUE(r.isOk()) << r.error();
EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess);

std::vector<half> out(h.size());
CUDA_CHECK(
cudaMemcpy(out.data(), h.data(), h.size() * sizeof(half), cudaMemcpyDeviceToHost));
return out;
};

const auto off = run("1");
const auto on = run("4");
EXPECT_TRUE(bitwiseEqual(off, on))
<< "prefill 不得受 split-KV 开关影响，max|diff|=" << maxAbsDiff(off, on);
}
Loading