Skip to content

perf(cpu): raise GDN recurrence lane cap 4 -> 8 #694

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
chenghuaWang merged 1 commit into from
Aug 9, 2026
Merged

perf(cpu): raise GDN recurrence lane cap 4 -> 8 #694

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
9 changes: 5 additions & 4 deletions mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,11 @@ namespace mllm::cpu::gdn {
namespace {

constexpr int kMaxStackNormalizedHeadDim = 256;
// GDN state updates are bandwidth-heavy on heterogeneous mobile CPUs. A small
// fixed lane cap avoids making every efficiency core part of the per-layer
// completion barrier.
constexpr int kMaxParallelGDNLanes = 4;
// GDN state updates are bandwidth-heavy on heterogeneous mobile CPUs. The lane
// cap bounds how many tasks share the per-layer completion barrier. 8 lanes
// matches the 8-core phones this kernel targets (4B has 32 recurrence tasks,
// so all 8 cores participate); bitwise-safe because tasks are disjoint.
constexpr int kMaxParallelGDNLanes = 8;
// Scalar state elements updated across all [batch, value_head] tasks.
constexpr std::size_t kMinParallelGDNWork = 65536;

Expand Down
21 changes: 13 additions & 8 deletions mllm/engine/HpcThreadPool.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ void HpcThreadPool::wakeup() {

int HpcThreadPool::acquireTaskSlot() {
std::lock_guard<std::mutex> _l(queue_mutex_);
for (int i = 0; i < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++i) {
for (int i = 0; i < kHpcThreadPoolTaskLimit; ++i) {
if (task_available_[i]) {
task_available_[i] = false;
return i;
Expand All @@ -60,7 +60,7 @@ int HpcThreadPool::acquireTaskSlot() {
}

void HpcThreadPool::releaseTaskSlot(int task_slot_idx) {
if (task_slot_idx < 0 || task_slot_idx >= MLLM_HPC_THREAD_POOL_TASK_LIMITS) { return; }
if (task_slot_idx < 0 || task_slot_idx >= kHpcThreadPoolTaskLimit) { return; }
std::lock_guard<std::mutex> _l(queue_mutex_);
task_available_[task_slot_idx] = true;
}
Expand Down Expand Up @@ -95,9 +95,14 @@ void HpcThreadPool::splitTask(HpcThreadPoolTask&& task, int task_slot_idx) {
// e.g.: threads is 4, tiles_name is 12.
// 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
if (tiles_num > thread_cnt_) {
// Capture task and true_idx BY VALUE: splitTask's parameters (a moved-in
// rvalue task and a stack vector) die when this function returns, but the
// worker threads execute this lambda asynchronously — capturing by
// reference leaves a dangling reference (use-after-free). This was latent
// with 4 lanes and became reachable at 8 lanes.
Comment on lines +98 to +102

@coderabbitai coderabbitai Bot Aug 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Locate file:"
git ls-files | rg '(^|/)HpcThreadPool\.cpp$|(^|/)HpcThreadPool\.(hpp|h)$|(^|/)HpcThreadPool\.(cc|h)$|(^|/)HpcThreadPool\.(hh)$' || true

echo
echo "Outline:"
ast-grep outline mllm/engine/HpcThreadPool.cpp || true

echo
echo "Relevant source lines:"
cat -n mllm/engine/HpcThreadPool.cpp | sed -n '1,180p'

Repository: UbiquitousLearning/mllm

Length of output: 7671

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Inspect matching declaration and relevant code:"
cat -n mllm/engine/HpcThreadPool.hpp | sed -n '1,180p'

echo
echo "Find push/splitTask call sites:"
rg -n "splitTask|push\\(" mllm --glob '*.{cpp,hpp,c' ] 2>/dev/null || true

echo
echo "Search task lifecycle outside HpcThreadPool.cpp:"
rg -n "tasks_\\[|HpcThreadPoolTask|func\\(" mllm --glob '*.{cpp,cxx,c,hpp,h,cc,hh}' || true

Repository: UbiquitousLearning/mllm

Length of output: 10765

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Check relevant call-sites:"
cat -n mllm/core/Parallel.hpp | sed -n '1,150p'

echo
echo "Search all splitTask occurrences including tests/docs:"
rg -n "splitTask" . || true

echo
echo "Inspect repository files for splitTask declarations outside tracked text:"
python3 - <<'PY'
from pathlib import Path
p = Path('mllm/engine/HpcThreadPool')
print('exists.cpp=', p.with_suffix('.cpp').exists())
print('exists.hpp=', p.with_suffix('.hpp').exists())
for suffix in ['.cpp', '.hpp', '.h', '.cc', '.hh']:
    f = p.with_suffix(suffix)
    if f.exists():
        text = f.read_text(errors='replace')
        print(f'{suffix} len={len(text)} splitTask_count={text.count("splitTask")}')
PY

Repository: UbiquitousLearning/mllm

Length of output: 11179

Clarify the lifetime explanation.

splitTask waits until worker callbacks complete before it returns, and each worker flag is cleared after tasks_[i].first.func(thread_idx) returns. This path does not make the captured references dangling. The value capture is still useful for ownership, but update the comment to avoid saying the lifetime problem is caused by the splitTask return path.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mllm/engine/HpcThreadPool.cpp` around lines 98 - 102, Update the comment
above the worker lambda in splitTask to remove the incorrect claim that
splitTask’s return causes captured references to dangle. Explain that value
capture provides explicit ownership for asynchronous worker execution while
splitTask waits for callbacks to finish, and retain the existing value-capture
behavior.

tasks_[task_slot_idx].first = {
.func =
[tiles_num, &task, &true_idx, this](int thread_idx) {
[tiles_num, task, true_idx, this](int thread_idx) {
for (int v = thread_idx; v < tiles_num; v += thread_cnt_) { task.func(true_idx[v]); }
},
.start = 0,
Expand All @@ -107,7 +112,7 @@ void HpcThreadPool::splitTask(HpcThreadPoolTask&& task, int task_slot_idx) {
tiles_num = thread_cnt_;
} else {
tasks_[task_slot_idx].first = {
.func = [tiles_num, &task, &true_idx, this](int thread_idx) { task.func(true_idx[thread_idx]); },
.func = [tiles_num, task, true_idx, this](int thread_idx) { task.func(true_idx[thread_idx]); },
.start = 0,
.end = tiles_num,
.step = 1,
Expand Down Expand Up @@ -140,11 +145,11 @@ HpcThreadPool::HpcThreadPool(int thread_cnt) {
thread_cnt_ = thread_cnt;
available_task_slots_ = 0;
available_task_slots_old_ = 0;
task_available_.resize(MLLM_HPC_THREAD_POOL_TASK_LIMITS);
tasks_.resize(MLLM_HPC_THREAD_POOL_TASK_LIMITS);
task_available_.resize(kHpcThreadPoolTaskLimit);
tasks_.resize(kHpcThreadPoolTaskLimit);

// Each task should hold some thread ok flag that mark this thread's work is done.
for (int t = 0; t < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++t) {
for (int t = 0; t < kHpcThreadPoolTaskLimit; ++t) {
task_available_[t] = true;
for (int i = 0; i < thread_cnt_; ++i) { tasks_[t].second.emplace_back(new std::atomic_bool{false}); }
}
Expand All @@ -154,7 +159,7 @@ HpcThreadPool::HpcThreadPool(int thread_cnt) {
workers_.emplace_back([this, thread_idx]() {
while (!stop_) {
while (available_task_slots_ > 0) {
for (int i = 0; i < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++i) {
for (int i = 0; i < kHpcThreadPoolTaskLimit; ++i) {
if (*tasks_[i].second[thread_idx]) {
tasks_[i].first.func(thread_idx);
{ *tasks_[i].second[thread_idx] = false; }
Expand Down
7 changes: 5 additions & 2 deletions mllm/engine/HpcThreadPool.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,13 @@
#include <functional>
#include <condition_variable>

#define MLLM_HPC_THREAD_POOL_TASK_LIMITS 2

namespace mllm {

// One task slot per expected concurrent op. 8 covers the GDN recurrence
// 8-lane cap plus sibling parallel ops; too small a limit makes splitTask fall
// back to serial on the main thread under multi-layer concurrency.
inline constexpr int kHpcThreadPoolTaskLimit = 8;

struct HpcThreadPoolTask {
std::function<void(int)> func;
int start = 0;
Expand Down
84 changes: 82 additions & 2 deletions tests/cpu/Qwen35GDNTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,11 @@ using mllm::cpu::gdn::gatedDeltaRuleF32;
class ScopedCpuOpThreads {
public:
explicit ScopedCpuOpThreads(int32_t thread_count) : original_thread_count_(mllm::Context::instance().getCpuOpThreads()) {
// initializeContext() registers the CPU backend; SymbolTable::reg aborts on
// a duplicate key, so call it exactly once (the tests have no fixture init).
static const bool kContextInitialized = [] { mllm::initializeContext(); return true; }();
(void)kContextInitialized;
mllm::Context::instance().setCpuOpThreads(thread_count);
Comment on lines 20 to 25

@coderabbitai coderabbitai Bot Aug 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C 8 '\bsetCpuOpThreads\s*\(' mllm tests
rg -n -C 8 '\binitThreadPool\s*\(' mllm

Repository: UbiquitousLearning/mllm

Length of output: 7100

🏁 Script executed:

#!/bin/bash
set -euo pipefail

fd -a 'CPUBackend.hpp|CPUBackend.cpp|mllm.hpp|Qwen35GDNTest.cpp' . | sed 's#^\./##' | sort
printf '\n--- mllm.cpp nearby ---\n'
sed -n '1,60p' mllm/mllm.cpp
printf '\n--- CPUBackend.h/cpp pool usage ---\n'
sed -n '1,180p' mllm/backends/cpu/CPUBackend.hpp
sed -n '1,180p' mllm/backends/cpu/CPUBackend.cpp
printf '\n--- locate HpcThreadPool definitions/usages ---\n'
rg -n -C 5 'class HpcThreadPool|struct HpcThreadPool|HpcThreadPool|--set-cpu-op-threads|cpu_op_thread|CpuOpThreads|eight|8 workers|workers' .

Repository: UbiquitousLearning/mllm

Length of output: 50380

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
import re

files = {
    "CPUBackend.hpp": Path("mllm/backends/cpu/CPUBackend.hpp").read_text(),
    "CPUBackend.cpp": Path("mllm/backends/cpu/CPUBackend.cpp").read_text(),
    "mllm.hpp": Path("mllm/mllm.hpp").read_text(),
    "Context.cpp": Path("mllm/engine/Context.cpp").read_text(),
    "Qwen35GDNTest.cpp": Path("tests/cpu/Qwen35GDNTest.cpp").read_text(),
}
for name, text in files.items():
    print(f"\n--- {name} relevant lines ---")
    for n,line in enumerate(text.splitlines(), 1):
        if re.search(r'\b(setCpuOpThreads|getCpuOpThreads|initThreadPool|initializeContext|MLLM_CONDITIONAL_PARALLEL_FOR|scoped_thread_pool|thread_pool_|__threadPoolDestroy)', line):
            print(f"{name}:{n}: {line}")

for n,line in enumerate(files["Qwen35GDNTest.cpp"].splitlines(), 1):
    if "ScopedCpuOpThreads" in line or "kEightThreads" in line or "kCpuOpThreads" in line or "ASSERT_FLOAT_EQ" in line:
        print(f"{name}:{n}: {line}")
PY

Repository: UbiquitousLearning/mllm

Length of output: 2950

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- CPU backend thread-count API definitions/usages ---'
rg -n -C 4 'thread_count|thread_cnt|num_threads|cpu_op_thread|CpuOpThreads|ConditionalParallel|conditional|parallel|thread_pool_->|getThreadPool\(\)' mllm/backends mllm/common mllm/engine mllm/mllm.hpp

printf '%s\n' '--- HpcThreadPool definitions/usages ---'
rg -n -C 4 'class HpcThreadPool|struct HpcThreadPool|HpcThreadPool' mllm

printf '%s\n' '--- relevant Qwen35GDNTest setup ---'
sed -n '1,60p' tests/cpu/Qwen35GDNTest.cpp
sed -n '250,330p' tests/cpu/Qwen35GDNTest.cpp

Repository: UbiquitousLearning/mllm

Length of output: 50380

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- test thread count macros/constants ---'
sed -n '1,40p' tests/cpu/Qwen35GDNTest.cpp
printf '%s\n' '--- MLLM_CONDITIONAL_PARALLEL_FOR definition/usages (top-level files) ---'
rg -n '\bMLLM_(CONDITIONAL_)?PARALLEL_FOR|MLLM_AUTO_PARALLEL_FOR|MLLM_KERNEL_USE_THREADS_VENDOR_MLLM|struct\s+HpcThreadPool|class\s+HpcThreadPool' \
  -g '*.{h,hpp,cpp,h.in,cmake,txt,md}' \
  mllm tests | head -n 200

printf '%s\n' '--- CPUBackend exact implementation ---'
sed -n '88,102p' mllm/backends/cpu/CPUBackend.cpp

printf '%s\n' '--- Context setCpuOpThreads exact implementation ---'
sed -n '158,164p' mllm/engine/Context.cpp

printf '%s\n' '--- initializeContext exact implementation ---'
sed -n '145,158p' mllm/mllm.hpp

printf '%s\n' '--- Qwen35GDNTest relevant constructor and scoped-threads calls ---'
sed -n '18,28p' tests/cpu/Qwen35GDNTest.cpp
sed -n '256,269p' tests/cpu/Qwen35GDNTest.cpp
sed -n '312,346p' tests/cpu/Qwen35GDNTest.cpp

Repository: UbiquitousLearning/mllm

Length of output: 22515

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
import re

text = Path("mllm/core/Parallel.hpp").read_text()
needle = "MLLM_KERNEL_USE_THREADS_VENDOR_MLLM"
idx = text.find(needle)
print("MLLM_KERNEL_USE_THREADS_VENDOR_MLLM at", idx)
if idx != -1:
    print(text[max(0, idx-400):idx+700])
print("\nmacro branch with MLLM_KERNEL_USE_THREADS_VENDOR_MLLM:")
for i,line in enumerate(text.splitlines(), 1):
    if re.search(r'#(?:ifndef|define) ' + re.escape(needle)|MLLM_CONDITIONAL_PARALLEL_FOR|`#define` MLLM_KERNEL_USE_THREADS_VENDOR_MLLM|`#endif`\s*(//)?\s*?{}'.format(re.escape(needle).replace("_", r"\_")), line):
        print(f"{i}: {line}")

for i,line in enumerate(Path("tests/cpu/Qwen35GDNTest.cpp").read_text().splitlines(), 1):
    if re.match(r'\s*constexpr int kThreadCount\s*=', line) or re.search(r'mllm\.cmake|MLLM_KERNEL_USE_THREADS_VENDOR_MLLM|kThreadCount', line):
        print(f"tests/cpu/Qwen35GDNTest.cpp:{i}: {line}")
PY

printf '%s\n' '--- GDN implementation around parallel macro ---'
sed -n '340,380p' mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp

printf '%s\n' '--- exact hpc threadpool definition excerpt ---'
sed -n '33,120p' mllm/engine/HpcThreadPool.hpp

Repository: UbiquitousLearning/mllm

Length of output: 403

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- Parallel.hpp around MLLM_KERNEL_USE_THREADS_VENDOR_MLLM ---'
rg -n -C 20 'MLLM_KERNEL_USE_THREADS_VENDOR_MLLM|MLLM_CONDITIONAL_PARALLEL_FOR' mllm/core/Parallel.hpp

printf '%s\n' '--- GDN parallel loop macro usage ---'
sed -n '340,380p' mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp

printf '%s\n' '--- HpcThreadPool definition excerpt ---'
sed -n '20,120p' mllm/engine/HpcThreadPool.hpp

printf '%s\n' '--- Qwen35GDNTest constants and flag references ---'
rg -n -C 6 'kThreadCount|kThread|MLLM_KERNEL_USE_THREADS_VENDOR_MLLM|CPU backend|eight' tests/cpu/Qwen35GDNTest.cpp

Repository: UbiquitousLearning/mllm

Length of output: 15390

Set the CPU pool count before creating the context.

initializeContext() creates the CPU backend pool from Context::getCpuOpThreads(), and Context::setCpuOpThreads() only updates the scalar. The current code calls initializeContext() first, so this 8-lane path can keep the prior default worker count. Set thread_count before initializing the context; reset the count in the destructor as an additional guard.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/cpu/Qwen35GDNTest.cpp` around lines 20 - 25, Update ScopedCpuOpThreads
so thread_count is applied with Context::setCpuOpThreads() before the one-time
initializeContext() call, ensuring the CPU pool is created with the requested
count. In the destructor, restore the saved original_thread_count_ after the
test scope completes.

mllm::initializeContext();
}

~ScopedCpuOpThreads() { mllm::Context::instance().setCpuOpThreads(original_thread_count_); }
Expand Down Expand Up @@ -218,7 +221,10 @@ TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) {
constexpr int kValueHeads = 32;
constexpr int kKeyDim = 128;
constexpr int kValueDim = 128;
constexpr int kThreadCount = 4;
// Exercises the parallel lane partition up to the 8-lane cap
// (kMaxParallelGDNLanes); tasks are disjoint so output must be bitwise
// identical regardless of how many lanes the scheduler picks.
constexpr int kThreadCount = 8;

std::vector<float> q(kBatch * kSequence * kKeyHeads * kKeyDim);
std::vector<float> k(q.size());
Expand Down Expand Up @@ -266,4 +272,78 @@ TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) {
}
}

// 4B real geometry (B=1, S=69, 16 key heads, 32 value heads, 128 dims) at the
// 8-lane cap — exercises the full task fan-out (32 tasks) that the small
// geometry above does not. Guards against the device crash observed on
// OnePlus with the 8-lane product build.
TEST(Qwen35GDNTest, FourBGeometry8LaneDoesNotCrash) {
constexpr int kBatch = 1;
constexpr int kSequence = 69;
constexpr int kKeyHeads = 16;
constexpr int kValueHeads = 32;
constexpr int kKeyDim = 128;
constexpr int kValueDim = 128;
constexpr int kThreadCount = 8;

std::vector<float> q(kBatch * kSequence * kKeyHeads * kKeyDim);
std::vector<float> k(q.size());
std::vector<float> v(kBatch * kSequence * kValueHeads * kValueDim);
std::vector<float> a(kBatch * kSequence * kValueHeads);
std::vector<float> b(a.size());
std::vector<float> a_log(kValueHeads);
std::vector<float> dt_bias(kValueHeads);

for (std::size_t i = 0; i < q.size(); ++i) {
q[i] = 0.01F * static_cast<float>(static_cast<int>(i % 7) - 3);
k[i] = 0.01F * static_cast<float>(static_cast<int>(i % 5) - 2);
}
for (std::size_t i = 0; i < v.size(); ++i) { v[i] = 0.01F * static_cast<float>(static_cast<int>(i % 11) - 5); }
for (std::size_t i = 0; i < a.size(); ++i) {
a[i] = 0.001F * static_cast<float>(static_cast<int>(i % 3));
b[i] = 0.001F * static_cast<float>(static_cast<int>(i % 9));
}
for (int i = 0; i < kValueHeads; ++i) { a_log[i] = -1.0F; dt_bias[i] = 0.0F; }

std::vector<float> state(kBatch * kValueHeads * kValueDim * kKeyDim, 0.0F);
std::vector<float> output(v.size());
std::vector<float> ref_output(v.size());
std::vector<float> ref_state = state;

// Serial reference, then 8-lane parallel — must be bitwise identical.
gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), ref_state.data(),
ref_output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim,
/*thread_count=*/1);
const ScopedCpuOpThreads scoped_threads(kThreadCount);
gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), state.data(),
output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount);

for (std::size_t i = 0; i < output.size(); ++i) {
ASSERT_EQ(ref_output[i], output[i]) << "output index " << i;
}
for (std::size_t i = 0; i < state.size(); ++i) {
ASSERT_EQ(ref_state[i], state[i]) << "state index " << i;
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// Repeat the full 4B GDN pass 24 times (one per layer) to mimic the real
// model's layer loop, which interleaves the recurrence with other parallel
// ops on the shared thread pool. Context init is now once-only (see
// ScopedCpuOpThreads), so this exercises multi-call thread-pool reuse.
// Run the recurrence 24 times on a FRESH copy of the initial state each
// time (mirroring one GDN layer per model layer from the same prefill input),
// and compare each run's output to the serial reference for that same input.
// This exercises repeated thread-pool push/acquire/release cycles — the
// multi-call reuse pattern that crashed on device.
for (int layer = 0; layer < 24; ++layer) {
std::vector<float> layer_state(state.size(), 0.0F);
std::vector<float> layer_ref_state(state.size(), 0.0F);
gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), layer_ref_state.data(),
ref_output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, /*thread_count=*/1);
gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), layer_state.data(),
output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount);
for (std::size_t i = 0; i < output.size(); ++i) {
ASSERT_EQ(ref_output[i], output[i]) << "layer " << layer << " output index " << i;
}
}
}

} // namespace
Loading