Skip to content

perf(quant): multi-row tile + AVX-512 + AMX arms for the K1b grouped IDOT - #1606

Open
heyyyyjude wants to merge 1 commit into
JustVugg:devfrom
heyyyyjude:xeon-k1b-tile
Open

heyyyyjude wants to merge 1 commit into
JustVugg:devfrom
heyyyyjude:xeon-k1b-tile

Conversation

@heyyyyjude

Copy link
Copy Markdown

The fmt=4 (gs64/gs128) family was the worst-served format on the fastest CPUs: the planar gate hard-disabled the family on AVX-512 builds, the K1b kernel was 256-bit only, and every x86 int4 path re-unpacked the weight block once per activation row — the multi-row bandwidth loss the decode failure ledger's reopen gate names (33–35 vs 66.88 GB/s at 8 sessions).

What changes (all opt-in; default behavior is unchanged on every platform):

  • Planar gate: AVX-512 builds now enable the planar family for the integer K1b path only; fmt=2 keeps the pair layout there (the f32 planar twin mirrors the AVX2 accumulation order, not dot_i4f_avx512's), so no existing bit-exactness claim changes. IDOT_GS=1 stays the opt-in it was.
  • K1b vector kernel: 1×4 row tile (the fmt=2 K2 idea brought to the grouped family) pays each block's load+mask once per 4 rows, plus AVX-512-VNNI arms — one vpdpbusd zmm per whole 64-element group.
  • K1c: AMX int8 tile kernel (Sapphire Rapids+). B tile = 16 output rows' group block unpacked once to signed v−8, so tdpbssd returns the same int32 the vector path computes, amortized over up to 16 activation rows. Linux arms tile state via ARCH_REQ_XCOMP_PERM, Windows 11 via EnableProcessOptionalXStateFeatures, others fail closed. AMX=0 kills it, AMX_S_MIN (default 8) sets the row threshold.

Correctness: all three shapes keep one contract — exact int32 group dots and the same per-row ascending-g fmaf chain — so they are bit-identical to each other and to the pure-C reference on every ISA. test_int_kernel_exact now covers the tile remainders (S=3/6/9/18) and an O%16 output tail:

  • native AVX-VNNI: 11,312 checks, 0 failures
  • Intel SDE -spr (avx512-vnni): 10,544 checks, 0 failures, with the -mix histogram confirming TDPBSSD/TILELOADD/TILESTORED actually executed
  • engine builds warning-free for native, -march=sapphirerapids (Windows), and x86_64-linux-gnu (zig cc cross-compile)

Performance: tests/bench_i4p_gidot.c (not a gate) measures per-row vs multi-row: 1.20–1.50× on a 256-bit AVX-VNNI laptop (Core Ultra 9 185H). The 512-bit and AMX arms need a real AMX Xeon A/B per the benchmark protocol — I don't have that hardware; treat this PR as correctness-gated enablement, and server datapoints (including negative ones) are invited. Happy to file the experiment manifest once numbers exist.

…IDOT

The fmt=4 (gs64/gs128) family was the worst-served format on the fastest
CPUs: the planar gate hard-disabled the whole family on AVX-512 builds, the
K1b kernel was 256-bit only, and every x86 int4 path re-unpacked the weight
block once per activation row -- the multi-row bandwidth loss the decode
failure ledger's reopen gate names (33-35 vs 66.88 GB/s at 8 sessions).

- planar gate: AVX-512 builds now enable the planar family for the INTEGER
  K1b path only; fmt=2 keeps the pair layout there (the f32 planar twin
  mirrors the AVX2 accumulation order, not dot_i4f_avx512's), so no existing
  bit-exactness claim changes. IDOT_GS=1 stays the opt-in it was.
- K1b vector kernel: 1x4 row tile (the fmt=2 K2 idea brought to the grouped
  family) pays each block's load+mask once per 4 rows, plus AVX-512-VNNI
  arms -- one vpdpbusd zmm per whole 64-element group.
- K1c: AMX int8 tile kernel (Sapphire Rapids+). B tile = 16 output rows'
  group block unpacked once to signed v-8, so tdpbssd returns the same
  int32 the vector path computes; amortized over up to 16 activation rows.
  Linux arms tile state via ARCH_REQ_XCOMP_PERM, Windows 11 via
  EnableProcessOptionalXStateFeatures, others fail closed. AMX=0 kills it,
  AMX_S_MIN (default 8) sets the row threshold.

All three shapes keep one contract: exact int32 group dots and the same
per-row ascending-g fmaf chain, so they are bit-identical to each other and
to the pure-C reference on every ISA. test_int_kernel_exact now covers the
tile remainders (S=3/6/9/18) and an O%16 output tail, and passes native
(avx-vnni, 11312 checks), and under Intel SDE -spr (avx512-vnni, 10544
checks) where the -mix histogram confirms TDPBSSD/TILELOADD executed.
tests/bench_i4p_gidot.c (not a gate) measures per-row vs multi-row: 1.20-
1.50x on a 256-bit AVX-VNNI laptop; the 512-bit and AMX arms await a real
Xeon A/B per the benchmark protocol.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@JustVugg
JustVugg changed the base branch from main to dev September 18, 2026 22:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants