5.1 kB 43 lines
8bc5f038
Round 8 blitz: prompt_chunk sigs/cuts shared-expert+MoE-plan overlap, W8 IQ2_XS/IQ2_XXS opt, gr_norm_h_stream_out, dq_f16+act_f16 half4 (385.1 tok/s 2K prefill, 122.9 tok/s short prefill)
dawn
⋯
audits/argon_r8_dq.md
sourcerendered
download raw history
at main Markdown

Round 8 Track C: Barrier-Free Constant-Table dq_f16<F> + Vectorized act_f16 #

1. Summary of Changes (metal/kernels/dense_mm.metal & metal/tests/dense_mm_parity.cpp) #

  1. Zero-Threadgroup-Memory, Barrier-Free dq_f16<F> (DWC<F>):

    • Previously, dq_f16<F> allocated threadgroup uint tab[fmt<F>::TAB], invoked load_tab<F>(tab, tid, 64), and executed threadgroup_barrier(mem_flags::mem_threadgroup) in every single threadgroup ((n_in / UNIT / 64) * n_out threadgroups per matrix, e.g., 30,720 threadgroups for a 2560 x 10240 IQ3_S matrix), even though each thread decodes at most one unit (16 or 32 weights) before exiting:
      • For T_IQ3_S (TAB = 544, 2,176 bytes static threadgroup memory), each thread loaded 8 uints from iq3s_grid + 2 uchars from ksigns_iq2xs (which deq<T_IQ3_S> does not even use!) into threadgroup memory and waited on a threadgroup barrier just to perform 4 grid lookups in deq<T_IQ3_S>.
      • For T_IQ4_NL and T_IQ4_XS (TAB = 256, 1,024 bytes static threadgroup memory), each thread constructed 4 half2 codebook entries in threadgroup memory and waited on a barrier.
      • For all other 10 formats (TAB = 1), tab was unused (static_tg_mem = 0), yet every threadgroup still executed threadgroup_barrier(mem_flags::mem_threadgroup).
    • Added DWC<F> in metal/kernels/dense_mm.metal with direct constant-memory specializations for T_IQ3_S (iq3s_grid), T_IQ4_NL (kvalues_iq4nl), and T_IQ4_XS (kvalues_iq4nl), forwarding the remaining 10 formats to DW<F>::run(..., nullptr).
    • Removed threadgroup uint tab[fmt<F>::TAB], load_tab<F>, and threadgroup_barrier from dq_f16<F>. Static threadgroup memory is now 0 bytes across all 13 dq_f16<F> specializations.

  2. Vectorized half4 Output Stores in dq_f16<F>:

    • Since H = UNIT / 2 is always a multiple of 4 (H = 8 for UNIT = 16, H = 16 for UNIT = 32) and chunk offsets p0, p1 are always 4-element (8-byte) aligned, replaced the scalar half store loop with half4 vector stores (*(device half4*)(o + p0 + j) and *(device half4*)(o + p1 + j)), cutting store instructions by 4x (4 half4 stores instead of 16 scalar stores for UNIT = 16; 8 half4 stores instead of 32 scalar stores for UNIT = 32).

  3. Vectorized float4 -> half4 act_f16:

    • Updated act_f16 to convert 4 elements per thread via ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]) with a tail loop at gid == (n >> 2) for any remainder elements. This remains backward-compatible with existing host dispatch grids (n + 255) / 256 while allowing (n + 1023) / 1024 dispatches.

2. Bitwise Parity & Benchmark Results on Apple M5 GPU #

All 13 formats (F32, BF16, Q8_0, Q4_0, Q5_0, IQ4_NL, IQ4_XS, Q2_0, IQ3_S, Q3_K, Q4_K, Q5_K, Q6_K) were verified against ref::dequant + ref::f32_to_f16 with std::memcmp in both dense_mm_parity (3.65 s) and dense_mm_parity --bench (1.48 s):

  • 100% bitwise identity (memcmp == 0) across all 13 formats and all tested matrix shapes (2560 x 10240, 2560 x 12288, 2560 x 6144, 10240 x 324, plus 2560/640/6144 -> 100/48).
Format UNIT static_tg_mem (Old -> New) 2560 x 10240 (us) 2560 x 12288 (us) 2560 x 6144 (us) 10240 x 324 (us) memcmp
F32 16 0 B -> 0 B 1191.7 1436.9 677.9 124.2 0 (PASS)
BF16 16 0 B -> 0 B 777.6 936.1 447.0 85.7 0 (PASS)
Q8_0 32 0 B -> 0 B 582.1 699.9 324.8 62.9 0 (PASS)
Q4_0 32 0 B -> 0 B 521.3 625.4 269.5 48.4 0 (PASS)
Q5_0 32 0 B -> 0 B 533.3 640.8 290.4 51.6 0 (PASS)
IQ4_NL 32 1024 B -> 0 B 520.4 624.3 268.0 49.3 0 (PASS)
IQ4_XS 32 1024 B -> 0 B 516.1 626.3 262.6 47.4 0 (PASS)
Q2_0 16 0 B -> 0 B 460.0 571.6 208.3 41.1 0 (PASS)
IQ3_S 16 2176 B -> 0 B 495.6 602.1 237.0 43.2 0 (PASS)
Q3_K 32 0 B -> 0 B 506.7 615.3 243.4 44.8 0 (PASS)
Q4_K 32 0 B -> 0 B 529.5 637.2 274.6 46.9 0 (PASS)
Q5_K 32 0 B -> 0 B 542.1 649.9 291.0 49.0 0 (PASS)
Q6_K 32 0 B -> 0 B 547.7 659.1 297.8 56.0 0 (PASS)
  • Note on DRAM bandwidth saturation on M5 (~100–130 GB/s): For large matrices (2560 x 10240 writes 52.4 MB of fp16 output and reads 14.1 MB of IQ3_S weights = 66.5 MB per matrix), 495 us corresponds to 134 GB/s effective DRAM bandwidth, which is at the M5 hardware memory bandwidth ceiling. On smaller shapes (2560 x 48 ssm_alpha/ssm_beta), removing the 2,176-byte threadgroup table load + barrier in dq_f16_21 (IQ3_S) speeds up each dispatch by 1.56x (3.3 us -> 2.1 us).

3. Artifacts #

  • Clean unified diff: /tmp/argon_r8_dq/argon_r8_dq.patch (verified with git -C vendor/strata apply --check).