ee18916abf
Removed the PROFILER_ZONE from the most highly called functions to reduce the overhead.
...
PiperOrigin-RevId: 819739402
2025-10-15 07:10:04 -07:00
e3e8511e79
Initialization of profiler zones.
...
PiperOrigin-RevId: 819662587
2025-10-15 03:05:58 -07:00
fb6fa793f4
Added a global (to gemma) zones list to enable most call sites to PROFILER_ZONE3 to avoid the sychronization required for the static const initialization of the zone handle.
...
Improved flash_attention to enable profiling using the new zones.
PiperOrigin-RevId: 819235421
2025-10-14 08:30:58 -07:00
035273c184
tune pool kSpin mode in threading_context
...
Previously, this happened concurrently with the matmul autotune, which could lead to incorrect outcomes.
threading: de-singleton Pinning (no longer stores affinity); pass PoolWorkerMapping; fix Pool dtor order
Also enable SPR target (Zen4 is AMD-only),
update Highway version for renamed Thread()->GlobalIdx().
PiperOrigin-RevId: 816223017
2025-10-07 08:36:26 -07:00
9dc802c7aa
Add logging to io.cc on failed write and read.
...
This should provide insights into any failures.
PiperOrigin-RevId: 815784482
2025-10-06 10:25:41 -07:00
684a0444e9
Reduced parallelism for TransposeQ, making each thread read and write within its own cache lines
...
PiperOrigin-RevId: 814241032
2025-10-02 08:15:16 -07:00
14244664c8
Avoid transposing Q when it isn't needed
...
PiperOrigin-RevId: 814187984
2025-10-02 05:16:35 -07:00
fe5a39990e
Improve FlashAttention threading:
...
kFlat for RMSNorm (hierarchical is excessive),
profiler zone naming improvements.
PiperOrigin-RevId: 814144012
2025-10-02 02:37:05 -07:00
6098a022b3
Increased parallelism for RMSNormAndPositionalEncoding
...
PiperOrigin-RevId: 813738994
2025-10-01 07:11:14 -07:00
2f6cbde8ff
Added a smaller tile size to flash attention for smaller batch sizes
...
PiperOrigin-RevId: 813226193
2025-09-30 05:49:20 -07:00
4974f24832
Fixed bug with softcap in single flash attention
...
PiperOrigin-RevId: 813164938
2025-09-30 02:17:58 -07:00
16536996d1
Remove less useful spammy log lines.
...
PiperOrigin-RevId: 812694572
2025-09-29 02:28:41 -07:00
667a3f117a
Utilize multiple cores to read weight batches.
...
PiperOrigin-RevId: 811893059
2025-09-26 11:28:33 -07:00
d15731d201
Used hn::BroadcastLane instead of Set(..., x.raw)
...
PiperOrigin-RevId: 811386295
2025-09-25 09:42:03 -07:00
4f0c633248
(1) Added QueryResultAndMetrics and BatchQueryModelWithMetrics to also return TimingInfo besides query results.
...
PiperOrigin-RevId: 810634261
2025-09-23 17:02:29 -07:00
fac8aac4cb
Internal change
...
PiperOrigin-RevId: 809975026
2025-09-22 05:37:03 -07:00
501fdf000e
Remove no longer used MatVec
...
PiperOrigin-RevId: 809059409
2025-09-19 09:03:22 -07:00
b603425bf3
Fix batch inference: dangling reference
...
Also add more detailed asserts/error messages.
PiperOrigin-RevId: 807695421
2025-09-16 08:01:56 -07:00
f3bc1c17da
1.03x speedup: fused FFN
...
matmul-inl: support CView=StridedView or RowPtrs; rename to C_MC_NC
matmul.cc: Allow 1 more rep for MC/NC to allow half-sized tiles, which helps.
PiperOrigin-RevId: 807291701
2025-09-15 10:26:37 -07:00
59db30e209
add const restriction for benchmark_helper.cc, and paligemma_helper.cc to remove a few uncessary copies.
...
PiperOrigin-RevId: 807004597
2025-09-14 16:27:26 -07:00
c9b8479f7d
Added zero-initialization to att_out.
...
Re-enabled flash attention when HWY_NATIVE_DOT_BF16 is not available.
PiperOrigin-RevId: 806284756
2025-09-12 07:48:23 -07:00
2695aab5d2
Temporarily disable flash pending msan fix
...
PiperOrigin-RevId: 805350234
2025-09-10 07:25:41 -07:00
ba6131311a
Fix gemma_batch_bench for flash attention
...
q_T rows do not change.
Also repeat prefill to reflect perf after autotuning.
PiperOrigin-RevId: 805319377
2025-09-10 05:32:34 -07:00
9457258330
Refactor MatMul to accept views in the kernel functions
...
Make arg order consistent.
Move StridedView into mat.h.
Add view support to RowPtrs.
PiperOrigin-RevId: 805197381
2025-09-09 22:09:47 -07:00
f10ac41a20
Added flash attention, with both a single-q function, and a register-tiled function.
...
The register-tiled version achieves a speed-up by a factor of about 9.7 over the previous attention function on an AVX3-enabled machine.
PiperOrigin-RevId: 804913784
2025-09-09 08:05:26 -07:00
24b1760f03
Refactor: move Worker to ThreadingContext, factor out MMDecompress
...
PiperOrigin-RevId: 804909921
2025-09-09 07:56:12 -07:00
461a9c7d1b
Matmul refactoring towards fusion
...
MMLoops: move dispatch code out, use overloads
split build target into matmul_env (for MatMulEnv/MMOptions)
weights: no longer call BindB
Fix potential out of bounds in gemma_batch_bench
PiperOrigin-RevId: 804895985
2025-09-09 07:13:38 -07:00
34ceee6c30
Update MatMul comments, removing mention of partial.
...
PiperOrigin-RevId: 804872289
2025-09-09 05:57:33 -07:00
a5ab99e4ba
Memory use reduction: smaller/single MMStorage
...
PiperOrigin-RevId: 804865029
2025-09-09 05:32:46 -07:00
06e5da1e22
Cleanup: split CacheInfo from Allocator, MatMul helper functions
...
Lift DecompressA out of main autotuner to prevent interference
Also use kMaxNR / kNR constants instead of extra args
Fix: only require vector alignment, not cache alignment
PiperOrigin-RevId: 804333769
2025-09-08 02:23:58 -07:00
6e52a835c6
Faster startup on tsan: use hierarchical parallelism for BF16 conversion
...
Also re-enable profiler zones
PiperOrigin-RevId: 804273899
2025-09-07 22:50:31 -07:00
cbe24eac51
1.15x speedup: parallel sampling, enabled by new RNG
...
Also pass pos to SampleFunc, for seeding the RNG.
PiperOrigin-RevId: 803453518
2025-09-05 07:24:02 -07:00
ad7d7a2713
Further adjust dot_test threshold (numerics)
...
PiperOrigin-RevId: 803428406
2025-09-05 05:50:16 -07:00
2b4c16e243
Remove Griffin support
...
Also add IsObsolete helper
PiperOrigin-RevId: 803376921
2025-09-05 02:35:40 -07:00
56186193c1
Replace mt19937 with new generator to enable parallel sampling
...
Split it into immutable AesCtrEngine and RngStream
Also add RowSpan and Logits span
PiperOrigin-RevId: 803336423
2025-09-04 23:49:10 -07:00
5d1693e806
Internal change
...
PiperOrigin-RevId: 803083229
2025-09-04 10:31:20 -07:00
afd82376a5
Add AES-CTR RNG for parallel sampling (not yet used)
...
PiperOrigin-RevId: 802991142
2025-09-04 05:58:42 -07:00
4be4799727
Remove kMaxPackages and per-package-related code
...
matmul: remove kMaxClusters, dynamic allocation
PiperOrigin-RevId: 802950348
2025-09-04 03:33:12 -07:00
7263ab8445
MatMul simplification, threading strategy improvements
...
remove MatMul f32 special case (smaller code),
types: Add u32/u64 for use by Activations
move renamed ParallelismStrategy to threading_context so can pass ctx
ensure worker index is unique across clusters
matmul.h: const member functions for renamed policy classes (easier to call)
PiperOrigin-RevId: 802848086
2025-09-03 21:45:07 -07:00
74ffe079c4
Create separate MMStorage objects per cluster.
...
PiperOrigin-RevId: 802588625
2025-09-03 09:35:48 -07:00
b7b3d353db
Simplify MatMul: remove F32 special case (build time)
...
Also move kMaxM into separate kMaxBatchSize
PiperOrigin-RevId: 802086590
2025-09-02 04:29:21 -07:00
1e3c853e80
Add ParallelFor wrapper function and one new mode
...
Move ParallelismType from matmul.h to threading.h
Replace SmallParallelFor with ParallelFor and the new mode
PiperOrigin-RevId: 802038452
2025-09-02 01:40:09 -07:00
3737224132
Add in-cluster parallel policy. Update policy to include cluster_idx.
...
PiperOrigin-RevId: 802016308
2025-09-02 00:16:00 -07:00
27cb8e12d9
Handle non-threading parallel policy.
...
PiperOrigin-RevId: 802012517
2025-09-02 00:02:57 -07:00
0d2e74d74a
Add MMOptions as an argument to Matmul.
...
PiperOrigin-RevId: 802008198
2025-09-01 23:46:39 -07:00
229bd078a1
1.29x speedup: bf16 C1/C2. Extend most ops to any type, expand test coverage.
...
Also increase dot_test.cc range for Zen4, and matmul_test tolerance (failing in some configs)
PiperOrigin-RevId: 801789922
2025-09-01 06:34:04 -07:00
bc0c0bac8b
Add non-threading parallel policy.
...
PiperOrigin-RevId: 800913294
2025-08-29 08:39:06 -07:00
00b70f69c5
Include parallelism type in DoMatMul. Also remove package handling.
...
PiperOrigin-RevId: 800902568
2025-08-29 08:04:52 -07:00
0ae8646731
Fix remainder handling for Paligemma
...
No longer attempt to skip the remainder handling because B might also be a non-padded view.
PiperOrigin-RevId: 800890805
2025-08-29 07:25:52 -07:00
973e284ed6
Refactor Matmul to use a policy class for parallelization.
...
PiperOrigin-RevId: 800864489
2025-08-29 05:40:39 -07:00
Ray Smith
Jan Wassenberg
Nitin Gangahar
Charles Zhao
Marie White