diff --git a/include/tiny_llm/transformer.h b/include/tiny_llm/transformer.h index c13ba25..650fb14 100644 --- a/include/tiny_llm/transformer.h +++ b/include/tiny_llm/transformer.h @@ -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）。 @@ -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(); diff --git a/src/transformer.cpp b/src/transformer.cpp index ed2b668..7a4ed26 100644 --- a/src/transformer.cpp +++ b/src/transformer.cpp @@ -70,6 +70,36 @@ Result resolvePagedAttentionMode() { "\" (expected auto | legacy | direct)"); } +// ── split-KV decode 的开关（TLLM-ATTN-SPLITKV）────────────────────────────── +// +// TLLM_ATTN_SPLITKV = ： +// 未设置 / 空 / 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 resolveAttnSplitKv() { + const char *raw = std::getenv("TLLM_ATTN_SPLITKV"); + if (raw == nullptr || *raw == '\0') { + return Result::ok(1); + } + for (const char *p = raw; *p != '\0'; ++p) { + if (*p < '0' || *p > '9') { + return Result::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::err("TLLM_ATTN_SPLITKV: value " + std::string(raw) + + " out of range [0, " + std::to_string(kAttnMaxSplits) + "]"); + } + return Result::ok(value < 2 ? 1 : static_cast(value)); +} + } // namespace void LayerWorkspace::allocate(const ModelConfig &config) { @@ -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() 会因早退检查跳过，必须在此手动清理， @@ -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; } @@ -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; @@ -468,16 +511,32 @@ Result 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::err(split_result.error()); + } + const bool use_split = is_decode && split_result.value() > 1; + // Attention float scale = 1.0f / std::sqrt(static_cast(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, @@ -486,8 +545,16 @@ Result 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, diff --git a/tests/test_paged_dispatch.cpp b/tests/test_paged_dispatch.cpp index f4635a4..561bd97 100644 --- a/tests/test_paged_dispatch.cpp +++ b/tests/test_paged_dispatch.cpp @@ -275,6 +275,43 @@ class PagedDispatchTest : public ::testing::Test { ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); } + // 构造"刚好能进入 attentionPaged"的最小视图并跑一次 decode。 + // 供只验证**入口错误语义**（开关非法取值）的用例使用——这些用例关心的是路由前的 + // 校验分支，不需要完整的 fixture。 + static Result 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(kMaxBlocks) * static_cast(kBlockSize) * f.kv_dim; + f.k_pool = DeviceBuffer(pool_elems); + f.v_pool = DeviceBuffer(pool_elems); + f.k_scratch = DeviceBuffer(pool_elems); + f.v_scratch = DeviceBuffer(pool_elems); + f.hidden_a = DeviceBuffer(static_cast(f.config.max_seq_len) * f.hidden); + f.table = DeviceBuffer(static_cast(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 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(); @@ -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(kMaxBlocks) * static_cast(kBlockSize) * f.kv_dim; - f.k_pool = DeviceBuffer(pool_elems); - f.v_pool = DeviceBuffer(pool_elems); - f.k_scratch = DeviceBuffer(pool_elems); - f.v_scratch = DeviceBuffer(pool_elems); - f.hidden_a = DeviceBuffer(static_cast(f.config.max_seq_len) * f.hidden); - f.table = DeviceBuffer(static_cast(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 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(); @@ -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 +std::vector 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 h(static_cast(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 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 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); +}