Ray Smith
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
Ray Smith
2f6cbde8ff
Added a smaller tile size to flash attention for smaller batch sizes
...
PiperOrigin-RevId: 813226193
2025-09-30 05:49:20 -07:00
Ray Smith
d15731d201
Used hn::BroadcastLane instead of Set(..., x.raw)
...
PiperOrigin-RevId: 811386295
2025-09-25 09:42:03 -07:00
Jan Wassenberg
501fdf000e
Remove no longer used MatVec
...
PiperOrigin-RevId: 809059409
2025-09-19 09:03:22 -07:00
Jan Wassenberg
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
Jan Wassenberg
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
Ray Smith
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
Jan Wassenberg
24b1760f03
Refactor: move Worker to ThreadingContext, factor out MMDecompress
...
PiperOrigin-RevId: 804909921
2025-09-09 07:56:12 -07:00
Jan Wassenberg
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
Jan Wassenberg
34ceee6c30
Update MatMul comments, removing mention of partial.
...
PiperOrigin-RevId: 804872289
2025-09-09 05:57:33 -07:00
Jan Wassenberg
a5ab99e4ba
Memory use reduction: smaller/single MMStorage
...
PiperOrigin-RevId: 804865029
2025-09-09 05:32:46 -07:00
Jan Wassenberg
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
Jan Wassenberg
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
Jan Wassenberg
ad7d7a2713
Further adjust dot_test threshold (numerics)
...
PiperOrigin-RevId: 803428406
2025-09-05 05:50:16 -07:00
Jan Wassenberg
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
Jan Wassenberg
4be4799727
Remove kMaxPackages and per-package-related code
...
matmul: remove kMaxClusters, dynamic allocation
PiperOrigin-RevId: 802950348
2025-09-04 03:33:12 -07:00
Jan Wassenberg
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
Marie White
74ffe079c4
Create separate MMStorage objects per cluster.
...
PiperOrigin-RevId: 802588625
2025-09-03 09:35:48 -07:00
Jan Wassenberg
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
Jan Wassenberg
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
Marie White
3737224132
Add in-cluster parallel policy. Update policy to include cluster_idx.
...
PiperOrigin-RevId: 802016308
2025-09-02 00:16:00 -07:00
Marie White
27cb8e12d9
Handle non-threading parallel policy.
...
PiperOrigin-RevId: 802012517
2025-09-02 00:02:57 -07:00
Marie White
0d2e74d74a
Add MMOptions as an argument to Matmul.
...
PiperOrigin-RevId: 802008198
2025-09-01 23:46:39 -07:00
Jan Wassenberg
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
Marie White
bc0c0bac8b
Add non-threading parallel policy.
...
PiperOrigin-RevId: 800913294
2025-08-29 08:39:06 -07:00
Marie White
00b70f69c5
Include parallelism type in DoMatMul. Also remove package handling.
...
PiperOrigin-RevId: 800902568
2025-08-29 08:04:52 -07:00
Jan Wassenberg
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
Marie White
973e284ed6
Refactor Matmul to use a policy class for parallelization.
...
PiperOrigin-RevId: 800864489
2025-08-29 05:40:39 -07:00
Jan Wassenberg
7288891439
Remove F64 partial storage in matmul.
...
Also remove no longer used kMaxN; row_ptrs only used for C
PiperOrigin-RevId: 800774757
2025-08-29 00:12:08 -07:00
Jan Wassenberg
31c09cca4c
f32 LoopKC: 1.37x(M=512), 1.19(M=128) single-K F32,BF16 matmul speedup on SKX
...
Add a special case for A=F32,B=BF16, used when there is no native bf16 dot product.
dot-inl: ensure bf16,f32 and f32,bf16 both get promoted to float before f64 summation
matmul.cc: update autotuning to reflect actual A size
matmul_test: add all combinations of bf16/f32, report all results, not just first difference, check non-vector-aligned K
PiperOrigin-RevId: 800487817
2025-08-28 08:55:50 -07:00
Jan Wassenberg
5411fd846d
Minor: batched NotifyGenerate, fix comment/dep
...
PiperOrigin-RevId: 799889802
2025-08-26 23:33:17 -07:00
Jan Wassenberg
86afd53076
1.04x speedup: Parallelize SoftCap
...
Also require opt-in constexpr flag for observer callbacks, update zones
PiperOrigin-RevId: 799655163
2025-08-26 11:55:20 -07:00
Jan Wassenberg
71406cf6d0
More profiler interface fixes: hwy:: plus avoid ADD_ZONE
...
PiperOrigin-RevId: 794493165
2025-08-13 03:15:48 -07:00
Jan Wassenberg
faa4102992
(Resubmit) Prepare profiler annotations for new API
...
Pass hwy::Profiler& to low-level functions.
Used ThreadingContext arg instead of NestedPools.
Use new PROFILER_ZONE3.
PiperOrigin-RevId: 794461159
2025-08-13 01:38:24 -07:00
The gemma.cpp Authors
a2d9133f7d
Prepare profiler annotations for new API
...
Pass hwy::Profiler& to low-level functions.
Used ThreadingContext arg instead of NestedPools.
Use new PROFILER_ZONE3.
PiperOrigin-RevId: 793865287
2025-08-11 17:51:38 -07:00
Jan Wassenberg
4cbf63e6f0
Prepare profiler annotations for new API
...
Pass hwy::Profiler& to low-level functions.
Used ThreadingContext arg instead of NestedPools.
Use new PROFILER_ZONE3.
PiperOrigin-RevId: 793821255
2025-08-11 15:34:52 -07:00
Jan Wassenberg
eef564e8f0
Prepare profiler annotations for new API
...
PiperOrigin-RevId: 792808391
2025-08-08 16:51:29 -07:00
Jan Wassenberg
701841897b
Default to disabling per-socket parallelization
...
weights: default to Read for small-batch (only look at qbatch, not the larger prefill tbatch)
PiperOrigin-RevId: 790787643
2025-08-04 09:49:14 -07:00
Jan Wassenberg
d1638587f0
1.14x batch decode speedup: parallelize RMSNorm ops
...
Activations was over-parallelized, use single pool instead.
Also improve profiler zone annotations,
pass through worker args (for tracking concurrency), now non-optional.
PiperOrigin-RevId: 788790976
2025-07-30 00:55:45 -07:00
Jan Wassenberg
e76e29ce11
De-singleton ThreadingContext so callers can pass in their own
...
weights.cc: fix BindB argument for bf16 tensors
threading_test: enable autotune
PiperOrigin-RevId: 785763618
2025-07-22 02:08:46 -07:00
Jan Wassenberg
349c86f2d9
Fix bench_matmul perf regression: A input should be padded
...
PiperOrigin-RevId: 781976414
2025-07-11 07:36:36 -07:00
Jan Wassenberg
0f70f285e0
1.1x prefill and decode speedup (attention/activations)
...
Optimizations
- Better load-balancing in attention threading
(Previously, clusters were limited by #heads)
- Add MulByConstTo to avoid zero-init
- Parallel activations
Cleanup
- Prepare for RowPtr in A or B
- Pass through thread_id to ops
- Avoid warning in bench_matmul
PiperOrigin-RevId: 773723423
2025-06-20 08:59:53 -07:00
Jan Wassenberg
4f5785b0fd
Update instrumentation for new Highway wall-time profiler
...
Pass the thread index through and use new zone_id.
PiperOrigin-RevId: 773344242
2025-06-19 07:46:04 -07:00
Jan Wassenberg
7f62c2606e
Fix bf16 KV recompression and Rope(), fixes #608
...
Also add more helpful error message for prompt > seq_len
Also update ops_test, adding coverage for Rope().
PiperOrigin-RevId: 772945644
2025-06-18 09:14:20 -07:00
Jan Wassenberg
343482c7ef
1.02x batch decode speedup: BF16 KV cache
...
ops-inl.h: Vectorize Rope(), template
Remove unused MulBy, and extra-arg overloads of MulByConst and Softmax
Fix for DecompressAndZeroPad: ensure second vector filled
PiperOrigin-RevId: 772779163
2025-06-17 23:21:59 -07:00
Jan Wassenberg
cd80d8b24d
Speed up builds by skipping rarely used targets
...
Centralize previous code into GEMMA_DISABLED_TARGETS
PiperOrigin-RevId: 772433723
2025-06-17 05:44:20 -07:00
Jan Wassenberg
6773e4517c
Split Activations into Griffin/Attention to reduce memory usage for attention-only tests.
...
PiperOrigin-RevId: 772025282
2025-06-16 07:52:59 -07:00
Jan Wassenberg
2c72ff2aa5
Fix MatMul issue caused by autotuning bucketing, refs #608 , thanks @ufownl
...
PiperOrigin-RevId: 771077158
2025-06-13 06:58:42 -07:00
Jan Wassenberg
01cdefeda7
1.64x batch=1 prefill speedup: nested parallelization for Attention
...
(DotSoftmaxWeightedSum)
Also fix tsan error in matmul (atomic_flag instead of static)
PiperOrigin-RevId: 770241705
2025-06-11 11:28:46 -07:00
Jan Wassenberg
bd98b43cea
Rename RowPtr->StridedView, CRows->RowPtrs
...
PiperOrigin-RevId: 770046362
2025-06-11 02:30:53 -07:00