Fix H20 packed-K GEMM correctness

[guzekai01]

Summary

This PR fixes three correctness issues in the H20 packed-K GEMM path used by dense W4AFP8/FP8 workloads:

• keep packed-K warp-N large enough for the double-buffered S2R pipeline;
• preserve trailing accumulator vectors when the K-warp reduction partition has a partial tail;
• establish the required shared-memory proxy ordering and TMA store completion before epilogue buffer reuse.

The changes are limited to Humming kernel selection, reduction, and epilogue synchronization. They do not add or modify checkpoint quantization schemas.

Root cause

Packed-K warp-N selection

Packed-K maps each warp iteration to a 16-column N fragment. The S2R pipeline alternates between two buffers, so it requires at least two N iterations. The H20 heuristic could select warp_shape_n=16, leaving only one iteration and violating that pipeline invariant.

Packed-K configurations now clamp warp-N to 32. Non-packed layouts keep their previous heuristic result.

Partial accumulator tail in K-warp reduction

The shared-memory K-warp reducer partitioned the flattened accumulator register array using floor division for both the reduction iteration count and the vectors handled per iteration. Whenever the register count was not evenly divisible by floor(warp_shape_m / 16), the trailing vectors were never moved through shared memory and therefore never accumulated across K warps.

Warp-M 56 is the concrete H20 reproducer, not a special-cased shape. With warp-N 32 and FP32 accumulation, the register array contains 14 int4 vectors per thread. The old calculation used three iterations of four vectors and processed only 12. The two omitted vectors collectively represent the final 8 rows of the 56x128 output tile. Other shapes can exhibit the same truncation whenever their flattened register count leaves a remainder under the old partitioning.

The reducer now uses ceiling division for the iteration count and per-iteration width, with bounds checks for the partial final chunk. This handles the general remainder case without a block-M-specific branch.

TMA-C shared-memory proxy ordering

The epilogue writes an output tile to shared memory through the generic proxy, then TMA reads the same tile through the async proxy. A thread barrier coordinates the writers but does not make generic-proxy writes visible to the TMA async proxy.

Every math writer now issues fence.proxy.async.shared::cta before the existing synchronization. TMA output operations are also committed as store groups, and outstanding stores are waited on where required before the shared-memory source buffer is reused.

Without this ordering, the affected server configuration produced corrupted logits and repetitive output even though the individual kernel configuration could appear correct in isolated runs.

Validation

• H20 heuristic tests: 20 passed.
• Targeted H20 GPU tests passed for:
  • the representative warp-M 56 partial reduction tail;
  • TMA-C shared-memory buffer reuse;
  • the exact M=160, warp-M 56, TMA, stream-K dense configuration selected on H20.
• Server regression before the proxy fence: 0/9 GSM8K examples correct.
• Short server validation after the fix: 8/9 examples correct.
• 200-example validation after the fix: 0.96 accuracy.
• Full gsm8k_train8_test1319.jsonl validation after the fix: 1275/1319 correct, or 0.966641 accuracy. The matching baseline was 1263/1319, or 0.957544.
• Full-run server logs contained no CUDA, assertion, traceback, or runtime errors.

The proxy-ordering failure is timing-sensitive: the focused GPU tests exercise the affected path, but removing only the proxy fence does not currently make them fail deterministically. The server-level accuracy regression is therefore the effective fail-before/pass-after validation for that part of the fix.

[gemini-code-assist]

Code Review

This pull request introduces several fixes and enhancements to the humming library's epilogue and tuning logic. Key changes include adding a TMA store commit and stream-K wait group logic in gmem_writer.cuh, introducing tma_fence_async_shared to prevent shared memory reuse issues in pipeline.cuh, and adding bounds checks in smem_reducer.cuh to avoid out-of-bounds accesses during reduction. Additionally, the tuning heuristics in sm90_h20.py were updated to enforce a minimum warp shape N of 32 for packed K layouts, and comprehensive tests were added to validate these behaviors. There are no review comments, and I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.