SM90 native BF16/FP8 GEMM kernels, fused-SiLU quantization, and linear test harness - #4795
Open
lzhangzz wants to merge 14 commits into
Open
SM90 native BF16/FP8 GEMM kernels, fused-SiLU quantization, and linear test harness#4795lzhangzz wants to merge 14 commits into
lzhangzz wants to merge 14 commits into
Conversation
…nd quant-aware Forward scale companions
Full-suite linear validation (890/922 cases) shows dense/MoE BF16 picks the native GMMA kernels or cuBLAS, FP16 dense picks cuBLAS, and FP16 indexed MoE was already unsupported with these kernels present. The remaining failures are pre-existing and involve other kernel TUs (U4/FP4 paths).
65536-token problems hit two pre-existing legacy-kernel bugs (grid.y over the 65535 hardware limit in SchedulerSm70, int32 gmem offset overflow). fp16 indexed/blocked MoE has no SM90 kernel (GMMA grouped is bf16-only).
Operands over 2^31 elements (e.g. 64k tokens x 36k K in the MLA o-proj weight-as-A path) overflowed the int32 tile base offset, yielding illegal addresses or silent numeric corruption. Before/after bench on e2m1/u4 paths shows no performance change (within +/-2% noise).
Split gemm2_kernels into per-family static libs (sm70/75/80/90) gated on CMAKE_CUDA_ARCHITECTURES buckets matching arch.h device ranges, following the linear_attn pattern. cublas.cu stays compiled for all archs; the SM100 -> SM90 H100-compat auto-enable and ENABLE_CUBLAS_GROUPED are preserved. sm_100 and sm_120 now compile and dispatch to the SM80 s16816 kernels (kernel arch 800 accepts Sm100/Sm120 devices); sm_120 no longer claims SM90 kernel compatibility since 90a cubins cannot load on it.
The gate/up pairing of the fused-SiLU epilogue is defined along the n axis and does not survive the m/n swap performed by TransposedKernel; a transposed launch would silently compute wrong results. This also covers the fp8 per-token activation-scale output (param_W), which reaches the FP8 weight-as-A kernels exclusively through the silu+e4m3 feasibility path. Non-uniform MoE problems (indexed/ragged A) need indexed/ragged B after transposition, which no registered kernel provides, so they remain naturally infeasible. Blocked (uniform grouped) and flat problems transpose exactly and keep their twins.
_fused_silu_block returned the SM90-native [g128|u128] block-pack for fp8 weights on every arch, but only SM90 kernels consume that layout; on A100/SM100 the legacy interleaved epilogue would have read block-packed weights as element-interleaved. Off-SM90 fp8 now falls back to element interleave as before.
Neither header is included anywhere; the FP8 pool is served by v3 and the weight-as-A family.
double-quote-string-fixer rewrites f-string quotes on Python < 3.12 (CI lint runs 3.10); converting to match.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR modernizes TurboMind’s GEMM/linear stack for Hopper (SM90) by introducing native GMMA (+TMA) BF16/FP8 kernel families (including fused SiLU + in-kernel FP8 quantization), updating dispatch/registration and cache formats accordingly, and adding a case-driven Python benchmark/validation harness for linear correctness and tuning records.
Changes:
- Added new SM90 native BF16/FP8 kernel families, fused-SiLU FP8 quantization output support, and updated Llama FFN/MoE call paths to thread scale companions through
LlamaLinear::Forward. - Reworked GEMM kernel registration to self-register via
Registrar(whole-archive linked), and versioned the dispatch-cache binary format (header + expanded key fields). - Added
tests/turbomind/linearharness utilities (bench_linear.py, reports, scan drivers) plus FFN builder tests for fused layouts and SM90/non-SM90 behavior.
Reviewed changes
Copilot reviewed 84 out of 85 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/turbomind/test_ffn_builder.py | New unit tests covering w1/w3 fusion layouts and SM90-specific fused-SiLU packing behavior. |
| tests/turbomind/linear/validate_records.py | New multi-GPU validator that replays exported dispatch records against the PyTorch reference. |
| tests/turbomind/linear/test_linear.py | New pytest smoke test running the linear fixture reference-vs-TurboMind checks. |
| tests/turbomind/linear/sched_cmp_report.py | New log parser to compare row/col scheduler policies and summarize tuned results. |
| tests/turbomind/linear/run_sm90_scan.py | New chunked driver to scan full-suite kernel usage without OOM. |
| tests/turbomind/linear/run_sched_cmp.py | New runner producing tune-verbose logs for scheduler comparisons. |
| tests/turbomind/linear/regtune_cmp_report.py | New report script comparing sweep dirs for regressions/improvements. |
| tests/turbomind/linear/reference.py | New reference helpers for block-packed fused-SiLU and FP8 quantization modeling. |
| tests/turbomind/linear/model_specs.py | New YAML-driven model/weight spec schema for test case generation. |
| tests/turbomind/linear/kernel_usage_report.py | New aggregation tool to summarize used/unused kernels from bench logs. |
| tests/turbomind/linear/benchmark.py | New bench_linear implementation (tuning, validation, export/import records). |
| tests/turbomind/linear/bench_linear.py | New thin entrypoint wrapper for bench_linear. |
| tests/turbomind/linear/init.py | Package marker for the new linear test harness modules. |
| src/turbomind/python/CMakeLists.txt | Adds linear_bind.cpp to the Python extension build. |
| src/turbomind/python/bind.cpp | Extends bindings (fp8 dlpack export as uint8, ContextGuard variants, Tensor transpose helper, Param.set, linear binding hook). |
| src/turbomind/models/moe_weight.cc | Simplifies MoE expert table linking to use strided pointer tables and standard descriptors. |
| src/turbomind/models/llama/moe_ffn_layer.cc | Updates MoE FFN forward path to support FP8 fused-SiLU outputs with scale companions. |
| src/turbomind/models/llama/LlamaLinear.h | Extends LlamaLinear::Forward API with input/output scale companions. |
| src/turbomind/models/llama/LlamaLinear.cu | Implements scale-companion plumbing, FP8 output scale allocation, and refined indexed-vs-flat operand handling. |
| src/turbomind/models/llama/LlamaFfnLayer.cc | Threads fused-SiLU FP8 output scales into the downstream w2 linear call. |
| src/turbomind/models/linear_weight.h | Adds helper to mark FP8 fused-SiLU outputs as FP8 + group-128 scales. |
| src/turbomind/models/linear_weight.cc | Adds pack disable env-gate and SM90 BF16/FP16 “physical (N,K)” transpose for native GMMA+TMA path. |
| src/turbomind/models/ffn_weight.cc | Enables SM90 FP8 fused-SiLU output format (FP8 + dynamic group-128 scales). |
| src/turbomind/kernels/gemm/test/test_gemm_v2.cc | Removes legacy standalone GEMM v2 test executable source. |
| src/turbomind/kernels/gemm/test/models.h | Removes legacy benchmark model shape list header. |
| src/turbomind/kernels/gemm/test/gemm_bench.cu | Removes legacy NVBench-based GEMM benchmark source. |
| src/turbomind/kernels/gemm/sm90_v3_traits.h | Adds tile/pipeline trait definitions for SM90 v3 kernel family. |
| src/turbomind/kernels/gemm/sm90_utils.h | Adds SM90 TMA descriptor stride replacement helper and includes for SM90 desc ops. |
| src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h | Adds trait definitions for SM90 FP8 weight-as-A family. |
| src/turbomind/kernels/gemm/sm90_bf16_traits.h | Adds trait definitions for SM90 BF16 native kernel family. |
| src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h | Generalizes GMMA atom selection with MAX_OP_N and fixes scale application iteration to avoid OOB/scale-boundary issues. |
| src/turbomind/kernels/gemm/registry.h | Updates registry interface toward kernel self-registration and modern kernel base class usage. |
| src/turbomind/kernels/gemm/registry.cu | Switches to registrar-based kernel discovery and adds optional occupancy debug printing. |
| src/turbomind/kernels/gemm/registrar.h | New self-registration mechanism (Registrar, Collector, global factory list). |
| src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h | New grouped-MoE FP8 TMA descriptor prepare kernel for SM90. |
| src/turbomind/kernels/gemm/kernel/sm90_16816_8.cu | Migrates kernel registration to Registrar self-registration. |
| src/turbomind/kernels/gemm/kernel/sm90_16816_16.cu | Removes SM90 F16/BF16 s16816 kernel registrations (legacy path). |
| src/turbomind/kernels/gemm/kernel/sm80_16816_8.cu | Migrates SM80 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm80_16816_16.cu | Migrates SM80 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm75_16816_8.cu | Migrates SM75 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm75_16816_4.cu | Migrates SM75 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm75_16816_16.cu | Migrates SM75 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm70_884_8.cu | Migrates SM70 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel/sm70_884_16.cu | Migrates SM70 kernel registration to Registrar. |
| src/turbomind/kernels/gemm/kernel.h | Extends kernel launch signature to accept an extra output (W, Wdesc) for scale companions. |
| src/turbomind/kernels/gemm/kernel.cu | Extends kernel naming and feasibility behavior for fused-SiLU and scheduler/raster descriptors. |
| src/turbomind/kernels/gemm/kernel_impl.h | Updates base kernel impl signature (accepts W/Wdesc; unused for non-SM90 paths). |
| src/turbomind/kernels/gemm/kernel_impl_sm90.h | Extends SM90 kernel impl to support fused-SiLU gating, output scales (W), grouped TMA prepare, and improved feasibility checks. |
| src/turbomind/kernels/gemm/kernel_impl_sm90_bf16.h | New SM90 BF16 native kernel implementation layer (TMA-based). |
| src/turbomind/kernels/gemm/iterator_sm80.h | Fixes potential int32 overflow in gmem base offset calculation (uses int64). |
| src/turbomind/kernels/gemm/iterator_sm70.h | Fixes potential int32 overflow in gmem base offset calculation (uses int64). |
| src/turbomind/kernels/gemm/gmma_bf16_sm90.h | New helper to build GMMA BF16 traits without pulling Cutlass internal .inl headers. |
| src/turbomind/kernels/gemm/gemm.h | Extends GEMM run signature to accept output scale companions (W, Wdesc). |
| src/turbomind/kernels/gemm/gemm.cu | Adds optional tune-verbose logging and threads W/Wdesc through kernel launches. |
| src/turbomind/kernels/gemm/gemm_universal_sm90.h | Fixes WARPGROUPS naming and related barrier sizing. |
| src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h | Refactors grouped TMA handling to a host prepare-kernel and cleans up in-kernel descriptor rewriting. |
| src/turbomind/kernels/gemm/dispatch_cache.cu | Adds dispatch cache header/versioning and expands keying to include epilogue and new kernel desc fields. |
| src/turbomind/kernels/gemm/desc.h | Extends kernel descriptor (algo/raster/atom_layout/fused-silu support) and updates op-class naming. |
| src/turbomind/kernels/gemm/cublas.cu | Migrates cublas kernels to registrar-based registration and updates signature for W/Wdesc. |
| src/turbomind/kernels/gemm/convert.h | Adds WeightPackEnv() and removes blocked pointer table helper. |
| src/turbomind/kernels/gemm/convert_v3.cu | Implements TM_GEMM_WEIGHT_PACK env logic and updates converter selection rules. |
| src/turbomind/kernels/gemm/context.cu | Improves grouped tiling cost estimation by using average group size along the group axis. |
| src/turbomind/kernels/gemm/CMakeLists.txt | Buckets kernel compilation by arch family, adds whole-archive linking for registrar self-registration, and preserves SM100 grouped support. |
| src/turbomind/kernels/gemm/arch/mma_sm80.h | Renames MMA op-class enum constants to the new naming. |
| src/turbomind/kernels/gemm/arch/mma_sm70.h | Renames MMA op-class enum constants to the new naming. |
| src/turbomind/kernels/gemm/arch.h | Adjusts arch compatibility (SM120 falls back to SM80 family; SM90 no longer claims SM120 compatibility). |
| src/turbomind/core/module.h | Adds Param::set() to replace parameter slots from C++ (used by new linking/alloc flows). |
| lmdeploy/turbomind/builders/ffn.py | Adds SM90 fused-SiLU block-pack packing logic and updated fusion feasibility checks. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
1
to
6
| // Copyright (c) OpenMMLab. All rights reserved. | ||
|
|
||
| #include "src/turbomind/kernels/gemm/arch.h" | ||
| #include "src/turbomind/kernels/gemm/registrar.h" | ||
| #include "src/turbomind/kernels/gemm/registry.h" | ||
|
|
Comment on lines
+58
to
+62
| elif "{'case':" in line: | ||
| mrow = ROW_RE.search(line) | ||
| if mrow: | ||
| d = eval(mrow.group(0)) # noqa: S307 - benchmark's own row dict | ||
| if d.get('tflops'): |
Comment on lines
+72
to
+75
| """Match QuantizeSymm: per-row group-absmax → e4m3 + float scales [ceil(N/gs), M]. | ||
|
|
||
| Returns (fp8_tensor as uint8 view for storage, scales float32, dequant bf16). | ||
| """ |
Matches the linear_attn pattern; pytest collects test_linear.py as linear.test_linear and absolute 'tests.turbomind.linear' imports break when a site-packages 'tests' package shadows the repo tree (CI UT import error). Scripts now run via 'python -m tests.turbomind.linear.bench_linear'.
- registry.cu: include <cstdlib>/<iostream> for getenv/cout - kernel_impl_sm90_bf16: align dynamic smem to 1024 like the FP8 path (GMMA/TMA SW128 smem descriptors) - sched_cmp_report: parse row dicts with ast.literal_eval, not eval - reference: fix quantize_symm_row_fp8 return docstring (e4m3 tensor, not a uint8 view)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
SM90 native BF16/FP8 GEMM kernels, fused-SiLU quantization, and linear test harness
Summary
Reworks the SM90 GEMM stack around new native GMMA (wgmma + TMA) kernel families — BF16 (
64n16) and FP8 (64n32, v3 and weight-as-A) — with fused-SiLU epilogues that quantize to e4m3 + dynamic group-128 scales in-kernel, replacing the legacy s16816 F16 path on Hopper. Ships a case-drivenbench_linear.pytest harness and threads activation/output scale companions through theLlamaLinearForward API to support them.Kernel & dispatch
sm90_64n16_16, FP8sm90_64n32_8(v3 + weight-as-A), with fused-SiLU variants; FP8 fused SiLU emits per-token activation scales (param_W) for chained FP8 FFN.Registrar+ whole-archive link, replacing the hardcodedRegistryctor list; dispatch cache format versioned (magic/version/record-size header) withepilogueadded to the cache key.sm90_16816_16.cu, full-suite validated GMMA/cuBLAS-only picks) and FP8 v2/v4 headers (unreferenced); U4/FP8 s16816 kernels retained.Arch support
ENABLE_CUBLAS_GROUPEDpreserved.Weights & model glue
LlamaLinear::Forwardgains input/output scale-companion tensors; FP8 inputs carry dynamic group scales, FP8 fused-SiLU outputs are e4m3 + group-128 scales consumed by the next linear (dense FFN and MoE FFN paths).TM_GEMM_WEIGHT_PACK=1restores old behavior). FP8 MoE linking unified on the strided-ptr expert table (MakeBlockedPtrsremoved).[g128|u128]block-pack for FP8,[g64|u64]for SM90 BF16, element interleave elsewhere — with TP alignment via lcm(block_out, pack_block). FP8 block-pack is SM90-gated (off-SM90 keeps interleave).Testing
tests/turbomind/linear: case-driven harness (smoke/quick/fullsuites) validating quant-vs-dequant numerics across shapes × formats × batch sizes × TP/EP, with record validation and ffn-builder tests. Full-batch axis capped at 32768 (65536 hits two pre-existing legacy-kernel limits); fp16 MoE excluded on SM90 (pre-existing gap, fails identically on base).