4.7 kB 37 lines
eb15e062
Round 31 blitz: mm_gu3 IQ2_S/IQ2_XXS/IQ2_XS/IQ3_S tab+hh & shift (+6.8..+8.2% / -57.9 ms in kprof), moe_down/gemv Q2_0R 0x6400 half-mantissa decode (-7.6..-17.4%), dq_f16 IQ4 tab init & perm=4 inverse contiguous write
dawn
⋯
audits/argon_r31_dq.md
sourcerendered
download raw history
at main Markdown

Round 31 Track C Audit: dq_f16 (dense_mm.metal) & DW<F> (mma.metal) #

1. Summary & Scope #

  • Owned files:
    • metal/kernels/dense_mm.metal (dq_f16 lines 28–66)
    • metal/kernels/mma.metal (DW<T_IQ4_NL>, DW<T_IQ4_XS> lines 26–69)
    • metal/tests/dense_mm_parity.cpp (--dq-ab <base_dir> [type] [shape] CLI support & uncontended pair filter)
  • Patch: /tmp/argon_r31_dq/argon_r31_dq.patch (verified via git -C vendor/strata apply --check)
  • Parity: 100% bitwise identical (memcmp == 0) across all 14 formats and all 13 shapes in dense_mm_parity --dq-ab, default dense_mm_parity, and mma_parity.

2. Key Findings & Shipped Optimizations #

  1. DRAM Bandwidth Ceiling Diagnosis on Large Shapes (2560x12288, 2560x10240, 2560x6144, 6144x2560):

    • On M5 (153.6 GB/s theoretical peak unified memory bandwidth), dq_f16 on the large matrices already achieves 125–142 GB/s (81%–92% of theoretical DRAM peak), dominated by writing the 62.9 MB (2560x12288) / 52.4 MB (2560x10240) FP16 output buffer Wh.

  2. Eliminating Divergent constant-Memory Table Init in dq_f16<T_IQ4_XS / T_IQ4_NL> (dense_mm.metal):

    • load_tab<T_IQ4_NL> previously executed a dynamic for (int i = tid; i < 256; i += nth) loop in every threadgroup (10,240+ threadgroups per matrix) with divergent constant-memory reads kvalues_iq4nl[i & 15] and kvalues_iq4nl[i >> 4], which serialize across SIMD lanes on Apple GPUs.
    • Because nth.x >= 64 for every dq_f16 dispatch, threads tid < 64 now populate all 256 half2 entries of tab in 4 unrolled bank-conflict-free stores (tab[tid + {0, 64, 128, 192}]) directly from immediate register constants (0xbfad9881u, 0xf6eaddcfu, 0x26190d01u, 0x71594535u), sharing the invariant low nibble h0 across all 4 writes with zero constant-memory loads and zero loop overhead.

  3. Guarded a.perm & Inverse Row Mapping (ro = gid.y) for Sub-8-Bit Formats (dense_mm.metal):

    • Guarded the a.perm division/modulo inside if (a.perm) so a.perm == 0 (~430 of 530 dispatches) executes zero integer divisions/modulos, and folded k0 = 0 into DW<F>::run(..., u, 0, w, p0, p1, tab).
    • For a.perm != 0 (320x10240p4, 96 dispatches/chunk) on sub-8-bit formats (fmt<F>::BB * 16 < fmt<F>::QK * 17), inverted the bijection so output rows ro = gid.y are written contiguously in Wh and input row r is computed via shift/mask (ro & 3u) * (a.n_out >> 2) + (ro >> 2) when a.perm == 4 (1.03x–1.05x speedup on 320x10240p4 for IQ4_NL, Q4_0, Q5_0, Q5_1).

  4. Decoupled Scale Load in DW<T_IQ4_XS> and DW<T_IQ4_NL> (mma.metal):

    • Unpacking the 32 tab values into w[0..31] before loading/applying the block scale d hides the device memory latency of the block scale header (b[0..7]) behind the threadgroup table lookups, matching the two-phase pattern of DW<F>.

3. Ablations Tested & Rejected (Not Shipped) #

  • Dedicated DW<T_Q8_0> & DW<T_IQ3_S> fusing * d into the unpack loop: Slightly regressed 6144x2560 (0.98x) compared to generic DW<F> + deq<F>, because deq<F> already unrolls fixed trip counts (16 and 2) while deferring rdh(b) until after the integer unpack loop.
  • DW<T_Q2_0> 0x6400 FP16 mantissa-OR trick & rd32: Bitwise identical (memcmp == 0) once byte-masked (q[g]), but 0.99x (57.84 -> 58.15 ms) vs ushort p[0], p[1] in DW<T_Q2_0> because DW<T_Q2_0> outputs float w[16] (requiring an extra half-to-float conversion per weight after the half4 FMA).
  • packed_uint4 / uint* p in DW<T_IQ4_XS>: 0.99x vs ushort* p because wider 32/128-bit loads delay the first tab lookup and increase live register pressure.
  • 32x simd_shuffle nibble LUT in DW<T_IQ4_XS / T_IQ4_NL>: 0.66x on 2560x640 / 2560x512 (32 shuffles per thread vs 16 half2 threadgroup loads).

4. Quiet-GPU Paired Benchmark Results (dense_mm_parity --dq-ab) #

  • IQ4_XS: 58.83 -> 58.56 ms (memcmp=0 on all shapes; 6144x2560 1.01x, 2560x640 1.01x, 2560x512 1.01x, 2560x48 1.03x, 10240x4 1.07x, 2560x1 1.03x)
  • IQ4_NL: 64.70 -> 64.50 ms (memcmp=0 on all shapes; 320x10240p4 44.9 -> 42.6 us 1.05x, 6144x2560 1.01x, 2560x640 1.01x, 640x2560 1.01x, 2560x512 1.01x, 2560x48 1.05x, 10240x4 1.04x, 2560x1 1.06x)
  • IQ3_S: 57.02 -> 56.92 ms (memcmp=0 on all shapes; 2560x12288 610.0 -> 606.7 us 1.01x, 2560x128 1.07x, 2560x48 1.02x)
  • 320x10240p4 (perm=4): Q4_0 43.8 -> 42.4 us (1.03x), Q5_0 46.2 -> 45.2 us (1.02x), Q5_1 47.1 -> 46.5 us (1.01x), IQ4_NL 44.9 -> 42.6 us (1.05x), all memcmp=0.