Skip to content

cuda: warp-per-row E8 kernels, lattice table in shared memory - #935

Open
Nanetnounou wants to merge 1 commit into
JustVugg:devfrom
Nanetnounou:pr/e8-warp-kernels
Open

cuda: warp-per-row E8 kernels, lattice table in shared memory#935
Nanetnounou wants to merge 1 commit into
JustVugg:devfrom
Nanetnounou:pr/e8-warp-kernels

Conversation

@Nanetnounou

Copy link
Copy Markdown
Contributor

What

grouped_hidden_e8_dual and grouped_down_e8 used one thread block per output element: width/32 useful threads out of 256 launched, an 8-stage __syncthreads tree to reduce a single float, and per-element lattice lookups through __constant__ memory. __constant__ is built for broadcast; here every lane reads a different index, so the constant cache serializes up to 32 accesses per warp. Measured effect on GB10: ~21 GB/s effective weight bandwidth out of 273 available.

This PR replaces them with warp-per-output-row kernels (__shfl_down_sync reduction, E8 lattice table in shared memory pre-multiplied by 0.5, activation tile with stride 33 against bank conflicts, sign parity via __popc, 4 independent accumulators per lane).

Measured (GB10, GLM-5.2 E8-IQ3, decode, nsys cuda_gpu_kern_sum over 23 forwards)

step routed GPU critical path
original block kernels 10.01 s
+ warp/shuffle reduction 8.40 s (-16%)
+ lattice table in shared memory 1.58 s (-81%)

The dominant cost was the serialized __constant__ cache, not the reduction. End to end: expert matmul 435 β†’ 69 ms/token (~120 GB/s, 44% of peak); a 60-token prompt went 25.72 s β†’ 4.00 s (6.4x).

Precision: a single accumulator regresses RMS vs a float64 oracle by 45%; 4 accumulators restore it (2.796e-07 vs 2.719e-07 for the old block kernel) at identical throughput.

Tests

test_backend_cuda.cu gains parametrable fmt=6 widths (T6_I/T6_O/T6_S) and T6_VERBOSE=1 RMS output, so both kernel families can be checked against the float64 oracle at real dimensions.

make check: 415 tests OK. Note: make -C c cuda-test CUDA_ARCH=native currently exits 1 on GB10 on a pristine dev checkout as well β€” a v1.5.0-based build of the same test passes on the same machine and driver, so this looks like a dev-branch regression on sm_121 unrelated to this PR; we can file a separate issue with details.

Disclosure

Developed with AI assistance (Anthropic Claude), human-reviewed; all figures were measured on the hardware described.

grouped_hidden_e8_dual and grouped_down_e8 used one thread block per
output element: width/32 useful threads out of 256 launched, an
8-stage __syncthreads tree to reduce a single float, and per-element
lattice lookups through __constant__ memory. __constant__ is built
for broadcast; here every lane reads a different index, so the
constant cache serializes up to 32 accesses per warp, 8 times per
sub-block. Net effect measured on GB10: ~21 GB/s effective weight
bandwidth out of 273 available, in the kernel family the decode path
lives in. grouped_down_e8 was the worst case: 3.7M blocks per token,
784 bytes read per block, 64 useful threads out of 256, 3 useful
bytes per launched thread.

Replacement kernels (grouped_hidden_e8_dual_warp, grouped_down_e8_warp):

  - one warp per output row, __shfl_down_sync reduction: no
    __syncthreads, no shared-memory round trip for the reduction;
  - E8 lattice table staged in shared memory, pre-multiplied by 0.5;
  - activation tile in shared memory with a stride of 33, which
    removes the 32-way bank conflicts of a power-of-two stride;
  - sign parity via __popc;
  - 4 independent accumulators per lane: a single accumulator
    regresses precision by 45% vs a float64 oracle; 4 restore it
    (2.796e-07 vs 2.719e-07 for the old block kernel) at identical
    throughput.

Isolating the contributions (routed GPU critical path over 23
forwards, GB10, GLM-5.2 E8-IQ3, decode): original block kernels
10.01 s; +warp/shuffle 8.40 s (-16%); +lattice table in shared
memory 1.58 s (-81%). The dominant cost was the serialized
__constant__ cache, not the reduction.

Measured end to end on the same setup (nsys cuda_gpu_kern_sum):

  expert matmul         435 -> 69 ms/token (~120 GB/s, 44% of peak)
  60-token prompt, e2e  25.72 s -> 4.00 s  (6.4x)

test_backend_cuda.cu: T6_I/T6_O/T6_S make the fmt=6 test widths
parametrable and T6_VERBOSE=1 prints the RMS, so both kernel families
can be checked against the float64 oracle at real dimensions.

Developed with AI assistance (Anthropic Claude), human-reviewed; all
figures above were measured on the hardware described.
@monotophic

Copy link
Copy Markdown
Contributor

Authored by Fable 5 in Claude Code, analysis in partnership with @monotophic.

Hi! We're running a review pass on this branch now (great work btw),
since we just took #1037 through review on directly adjacent ground
(warp-per-row fmt=8 kernels), and GB10 measurements alongside it β€” my pair is
sm_121 / CUDA 13 / unified-memory NVIDIA, so the branch gets decode
correctness plus throughput at S ∈ {1, 4, 8, 32} on hardware this queue hasn't
had access to. Two methods from the #1037 cycle come with it: a 256-value
exhaustive decode sweep (catches per-element drift that RMS checks average
away) with mutation-bites proving the tests fail when the kernel is wrong, and
%-of-peak normalized against a measured peak (on GB10,
cudaDevAttrMemoryClockRate already returns the effective LPDDR5X rate, so a
2Γ—-DDR-derived roofline overstates peak 2Γ—).

Results land on this thread when the pass completes β€” no action needed from
anyone in the meantime.


Also, in case it saves a round of CI frustration:
the HIP syntax check is failing on __shfl_down_sync at backend_cuda.cu:638 β€”
hipcc has no _sync variants. The form that passed the same check in #1037:

#if defined(__HIP_PLATFORM_AMD__) || defined(__HIP__)
#define f8_shfl_down(v,o) __shfl_down((v),(o),32)
#else
#define f8_shfl_down(v,o) __shfl_down_sync(0xffffffffu,(v),(o),32)
#endif

Worth pinning the width to 32 explicitly (as above) rather than taking the
compiler's __shfl_down suggestion: the default width compiles fine but folds
two 32-lane rows into one reduction on wave64 HIP devices β€” silent wrong
answers, no error. Two smaller walls from the same check, in case they're
next: new test TUs in the gpu-compile target can't include
<cuda_runtime.h> directly (route through backend_cuda.cu /
backend_gpu_compat.h), and host-side isnan needs std:: qualification
under ROCm.

@monotophic

Copy link
Copy Markdown
Contributor

Authored by Fable 5 in Claude Code, analysis in partnership with @monotophic

We ran this branch on my GB10 (spark, sm_121) and reviewed the kernels against the CPU reference. Short version: the kernels are correct, the speedup is real and slightly better than advertised on our box, and we found one test-harness bug worth fixing plus one piece of good news about your failing cuda-test.

Reproduction (merged with current dev, GLM-5.2 decode shape D=6144/I=2048, 8 experts x S rows, kernel-window timing via COLI_CUDA_PROFILE, 5 reps):

S block kernels warp kernels speedup eff. weight BW % of measured peak
1 5.752 Β± 0.147 ms 0.822 Β± 0.001 ms 7.00x 140.6 GB/s 52.8%
4 23.644 Β± 0.803 ms 3.096 Β± 0.152 ms 7.64x 149.4 GB/s 56.1%
8 44.508 Β± 0.509 ms 6.060 Β± 0.099 ms 7.34x 152.6 GB/s 57.3%
32 176.770 Β± 0.722 ms 23.956 Β± 0.108 ms 7.38x 154.4 GB/s 58.0%

The %-of-peak column is against a measured 266.2 GB/s read ceiling (float4 grid-stride sum, 512 MB) rather than the 273 GB/s spec β€” on GB10 cudaDevAttrMemoryClockRate reports the effective rate, so derived peaks mislead. Your ~21 GB/s diagnosis of the block kernels reproduces exactly (we measure 19.6–20.9). We also confirmed the decode is term-by-term identical to quant.h's e8_expand_sub (the 0.5 prescale is a power of two, so the shared-memory codebook is exact; __popc parity == the XOR chain), verified the zero-filled tail at odd widths (2050, 2080, 300 β€” both kernel families pass the double oracle; compute-sanitizer memcheck clean), and confirmed determinism (fixed-order fold, warp-uniform exits, no atomics).

Good news about your failing cuda-test: it is not an sm_121 regression. At this branch's base the failure is the stale "fmt=7 must be rejected" negative case at test_backend_cuda.cu:295 β€” fmt=7 became real and #981 fixed that expectation on dev after you branched. We verified backend_cuda_test passes on current dev on GB10, and on this PR merged with dev it passes with both COLI_E8_WARP=1 and =0. A rebase makes your test suite green on this hardware. (Pristine dev's make cuda-test does still fail today, but in test_mxfp4_cuda β€” a separate issue, not this PR's.)

Must-fix β€” T6_S=1 corrupts the heap in the harness: the expert-group case hardwires erows={1,1} (total=2) while x/group_e are sized by S, and coli_cuda_expert_group transfers total rows. At the invocation the new comment recommends (T6_I=6144 T6_O=2048 T6_S=1) that's a 24 KB read past x and a 24 KB write past group_e β€” on our box it aborts with free(): invalid next size right after printing your exact RMS figures (2.796e-07/2.719e-07). T6_S=3 fails the other way (false FAIL 0.45 for both families β€” 3 oracle rows vs 2 computed). Suggest sizing those buffers by total (or deriving erows from S), and recommending T6_S=2 in the comment. Your precision conclusion survives: at the valid T6_S=2 we measure warp 2.663e-07 vs block 2.648e-07.

Should-fix:

  • e8_group_launch has no shared-memory guard: e8_warp_smem(width) exceeds the 48 KB static limit for width > 10912, and every fmt=6 group launch then fails (caught, but the engine silently drops the whole grouped path to CPU, where the block kernels used to work). One width check falling back to the block branch closes it. GLM-5.2's shapes are safe.
  • The HIP CI failure (__shfl_down_sync, no HIP equivalent): the fix needs the shuffle width pinned to 32, not just the intrinsic swapped β€” HIP's default width is the wavefront size, so on wave64 a bare __shfl_down folds two output rows into one reduction. We use a guarded macro for the same situation in CUDA: warp-per-row fmt=8 (fp8-e4m3) kernels, reference-mirroring accumulationΒ #1037 (fmt=8 warp kernels); happy to align on one shared macro if both land.

Nits: COLI_E8_WARP breaks the file's COLI_CUDA_* convention for runtime switches (and shadows the compile-time COLI_E8_* macros) β€” COLI_CUDA_E8_WARP? And since the default flips numerics ON (accumulation order changes, as your comment honestly warns), a release-note line may save someone a bisect.

Nice work β€” the constant-cache diagnosis, the stride-33 pad derivation, and the "never test the two kernels against each other" discipline are all solid, and T6_VERBOSE made our independent verification cheap.

@monotophic

Copy link
Copy Markdown
Contributor

Authored by Fable 5 in Claude Code, analysis in partnership with @monotophic.

A correction to our review above, from a belated verification pass: in the
shared-memory should-fix, we wrote that an oversized grouped launch "silently
drops the whole grouped path to CPU." The destination is wrong. The engine's
group-failure fallback retries each expert via coli_cuda_expert_mlp β€” the
per-expert GPU block kernels β€” and reaches the CPU only if those also fail.
We confirmed empirically that expert_mlp succeeds at width 10944 where the
grouped warp launch fails. Everything else in that bullet stands as verified:
e8_warp_smem crosses the 48 KB static limit exactly at width > 10912
(10912 passes, 10944 fails under COLI_E8_WARP=1 only), the failure is
caught via cudaGetLastError with no guard ahead of it, and the one-width-check
fallback to the block branch remains the suggested fix. The practical impact is
milder than we stated: an unguarded wide launch costs the warp-kernel speedup,
not GPU execution.

Small heads-up while we're here: on the branch as it sits (pre-rebase), the
recommended T6_S=1 invocation exits at the stale fmt=7 negative case before
ever reaching T6 β€” so anyone reproducing the T6 behavior should do it on the
branch merged with current dev (where #981's fix is in), which is the
environment our Reproduction section used.

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