From 8493d0bcfb0fdfd9c3eceae84b40410819756279 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 6 Aug 2026 08:20:27 +0000 Subject: [PATCH] =?UTF-8?q?feat(kernel):=20dense-template=20marlin=20port?= =?UTF-8?q?=20=E2=80=94=20byte-preserving=20E=3D1=20W4A16=20GEMM=20(gated?= =?UTF-8?q?=20OFF)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Row KERNEL-MARLIN-DENSE-PORT. Vendors vLLM's OWN dense marlin GEMM (the direct-A, tile-per-CTA W4A16 kernel it ships for a16 weight-only linears) as a new vt::MarlinDenseGemm op, and routes the E=1 dense NVFP4/MXFP4 projections (dense_nvfp4_gemm.h MatmulNvfp4MarlinD / MatmulMxfp4W4A16D / GateUpFusedMarlinD) through it behind VT_MARLIN_DENSE (default OFF). WHY: #54 proved clamping the single-expert MoE-marlin route to 48 CTAs (VT_MARLIN_E1_PAR1) recovers 81% of the marlin residual (per-call 114.7 vs vLLM 113.1) but the par regrouping of the fp32 C_tmp reduce costs one bf16 ULP that flips a strict token on the 64-layer 32B-NVFP4A16 (test_qwen3_32b_nvfp4a16:344). The dense template is vLLM's OWN dense reduce structure — the byte-preserving fix — so at M<=8 it runs the sms-wide (48-CTA) grid WITHOUT that ULP. Vendored (src/vt/cuda/marlin/libtorch_stable/quantization/marlin/), all cited from vLLM @ 555967922 csrc/libtorch_stable/quantization/marlin/: * kernel.h <- kernel.h (verbatim; namespace marlin, +lda param) * marlin_template.h <- marlin_template.h:1-2081 (verbatim dense kernel; the SHARED marlin.cuh/marlin_dtypes.cuh/dequant.h/ marlin_mma.h it includes are diff-verified byte-identical to our existing vendored copies) * marlin_mm_dense.{h,cu} <- marlin.cu:326-541 marlin::marlin_mm (+ config helpers); torch::stable marlin_gemm wrapper stripped; only the redundant inner is_a_8bit shadow dropped (identical value, avoids -Wshadow); STD_TORCH_CHECK via vt_marlin_check.h * kernel_selector.h, sm80_kernel_...fe2m1f...cu <- generate_kernels.py output. KEY: the dense kernel is a DISTINCT kernel body but the SAME 12-param Marlin<> template as the MoE TUs, so the instantiation set is shared; the dense body + namespace marlin come from the local kernel.h. New op + launcher: * vt::OpId::kMarlinDenseGemm + MarlinDenseArgs + MarlinDenseGemmFn (ops.h, appended before kCount — no id shift); dispatch shim (ops.cpp) * src/vt/cuda/cuda_marlin_dense.cu — vt::Tensor launcher mirroring cuda_moe_marlin.cu (graph-safe c_tmp pool, dense c_tmp sizing marlin.cu:713). Routing reuses the EXISTING resident weights + workspace (same marlin_permute repack for dense and MoE — CONFIRMED, no shim) with rank-2 operand views and NO moe_align gather. dense_gemms execution counter added (the "path RAN" signal). Gates: CPU -fsyntax-only CLEAN on ops.cpp and the VT_MARLIN_NVFP4 routing header. GPU compile: all 3 new dense .cu compile CLEAN on dgx GB10 sm_121a under the EXACT production flags (-Werror=all-warnings, -static-global-template-stub=false, --generate-code=...sm_121a). RED-first unit battery WRITTEN (test_ops_moe_grouped.cpp: NVFP4+MXFP4, M=1..8 x 3 shapes, dense-vs-CPU-ref AND dense-vs-grouped-route, row-shifted stride RED-injection). GPU EXEC gates (unit run + strict token battery dense-ON vs oracle incl. 32B-NVFP4A16:344 + launch-counter + nsys 48-CTA + binding c1..c8 x3) are the scoped dgx follow-up; default stays OFF until the strict battery proves oracle byte-match and the binding beats the MoE route. Records: state.md KERNEL-MARLIN-DENSE-PORT, parity-ledger, porting-inventory §10, kernel-matrix, STATUS/BENCHMARKS/FEATURES. FOLLOWING_AGENTS_PROTOCOL Assisted-by: Claude Code:claude-opus-4-8 [ClaudeCode] --- .agents/NOW.md | 6 +- .agents/kernel-matrix.md | 2 +- .agents/parity-ledger.md | 1 + .agents/porting-inventory.md | 32 + .agents/state.md | 67 + CMakeLists.txt | 9 +- docs/BENCHMARKS.md | 2 +- docs/ENVIRONMENT.md | 1 + docs/FEATURES.md | 2 +- docs/STATUS.md | 2 +- .../model_executor/models/dense_nvfp4_gemm.h | 69 +- include/vt/ops.h | 44 + src/vt/cuda/cuda_marlin_dense.cu | 159 ++ .../quantization/marlin/kernel.h | 43 + .../quantization/marlin/kernel_selector.h | 62 + .../quantization/marlin/marlin_mm_dense.cu | 528 +++++ .../quantization/marlin/marlin_mm_dense.h | 36 + .../quantization/marlin/marlin_template.h | 2081 +++++++++++++++++ .../sm80_kernel_bfloat16_fe2m1f_bfloat16.cu | 70 + src/vt/ops.cpp | 19 + tests/vt/test_ops_moe_grouped.cpp | 248 ++ 21 files changed, 3473 insertions(+), 10 deletions(-) create mode 100644 src/vt/cuda/cuda_marlin_dense.cu create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel.h create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel_selector.h create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.h create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_template.h create mode 100644 src/vt/cuda/marlin/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu diff --git a/.agents/NOW.md b/.agents/NOW.md index 4c7197532..eefa55580 100644 --- a/.agents/NOW.md +++ b/.agents/NOW.md @@ -16,15 +16,15 @@ checkpoint on `upstream/main` at `59674cf1d`. |---|---|---| | Laguna NVFP4 decode speed | **Closed: PARITY+ 1.03x** (byte-exact, default; `VT_LAGUNA_RESIDENT_BF16W` bf16-residency). Benchmark record | vLLM K-run set when convenient | | DeepSeek-V4-Flash decode | **Closed: BEATS ds4 1.144x** (`VT_V4_RESIDENT_W`, byte-exact). Phase-2 routed-expert residency NEGATIVE (−3.4%), default-OFF | — | -| f32-out GEMV audit | Only laguna + deepseek_v4 bf16 tower affected; gate models & on-framework dense unaffected (bf16-out, e2e-verified) | Re-verify deepseek_v4 bf16 tower same-tool | +| f32-out GEMV audit | Only laguna + deepseek_v4 bf16 tower affected; gate/on-framework dense unaffected | Re-verify deepseek_v4 tower same-tool | | Invocation-parity prevention | CI guard (`check-gemv-invocation-consistency.py`) + AGENTS.md checklist landing | Review + merge; CUDA build-verify `kGemvHeuristicAlgos` on dgx | | MiniMax-H3 lane | Portable path complete; e2e prompt-conditioned video on real weights (Thor). Speed = NVFP4 FP4 device path, sm_121-gated | PR #26 rebase + supports-audit synthesis | | Kimi-Linear-48B (KDA+NoPE-MLA+MoE) | **Full-model GB10 e2e RUNS** (bf16-resident §13): CPU+CUDA 13/13·656, no OOM. **Token gate NEAR-TIE 106/128** (6/8 token-exact) | device GDN/MLA islands + bf16 stream; 1.59 tok/s; default OFF | | 35B fresh grid | **BOUND** @`1ea26427`: tput 0.93-1.03x, c16 0.93x. INTAKE + Option A both **RESOLVED NEGATIVE** (H2D-out-of-capture tput WASH) | Real lever left: prefill glue (task #61) | | Qwen3.5-4B revalidation | 0.9971x @`59674cf1` (#35); TTFT/PSS pass, TPOT/ITL open | `docs/bench-evidence/` | -| MXFP4 parity (Qwen3-8B) | **`MARLIN-STRUCT`: decode-graph + gate_up FUSION default-ON (marlin 180→144 GEMM/step = vLLM-structural); #44 3/3, 0.6B/4B 184/184, 32B-NVFP4A16 142/142** | residual = marlin CTA + flash | +| MXFP4 parity (Qwen3-8B) | **`MARLIN-STRUCT`: decode-graph + gate_up FUSION default-ON (180→144 GEMM/step); #44 3/3, 0.6B/4B 184/184, 32B 142/142** | residual = marlin CTA + flash | | ROW-SERVE-ASYNC-DENSE-MIRROR | **LANDED+dgx-VERIFIED** (`f9c969ae`): #31 async mirror on classic dense Qwen3; gate RED→GREEN, SACRED 184/184 | Residual: sibling scope one-liner | -| MXFP4 parity goal | graph+fuse default-ON. c8 residual: marlin CTA 144 vs 48 = DOMINANT +1,177us (`VT_MARLIN_E1_PAR1` opt-in → near-parity, but flips strict 32B token → default-OFF), flash +784, glue +195 | NEXT: dense-template marlin port + full binding (oracle) | +| MXFP4 parity goal | graph+fuse default-ON; c8 dominant residual = marlin CTA 144 vs 48. **`KERNEL-MARLIN-DENSE-PORT` gated-OFF** (`VT_MARLIN_DENSE`): vLLM's own dense marlin = byte-preserving E=1; 3 dense `.cu` compile-clean dgx; unit WRITTEN | NEXT (dgx): strict dense-ON vs oracle + nsys + binding (state) | In-flight branches (default-OFF, not pushed): `laguna-fp4proj-prod` (fp4), laguna bf16/legacy/pipeline-gemv, `ds4-hc-expand-fuse`. Records: diff --git a/.agents/kernel-matrix.md b/.agents/kernel-matrix.md index 7b8cfac91..c273e05a5 100644 --- a/.agents/kernel-matrix.md +++ b/.agents/kernel-matrix.md @@ -118,7 +118,7 @@ host/sched. Detail: state `KERNEL-FA2-GQA-SWAP-FLIP`. | `KERNEL-GEMM-BF16` | BF16 dense GEMM, including torch-Linear/TN-equivalent layout and Qwen GDN merged input projections | Qwen mapper/packing `vllm/model_executor/models/qwen3_5.py:200-210,278-288`; merged linear `vllm/model_executor/layers/linear.py:580-808`; GDN calls `vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py:908-943` (qkvz projection + mixed/z split `:923-936`, construction `:481-496`); unquantized dispatch `vllm/model_executor/layers/utils.py:92-99,332-338`; runtime cuBLASLt heuristic | W1 one-owner BA + **W2A one-owner QKVZ** [loader:165](../src/vllm/model_executor/models/qwen3_5_dense_weights.cpp#L165), merged/fallback qkvz dispatch [forward:2372](../src/vllm/model_executor/models/qwen3_5.cpp#L2372), eligibility seam [internal:52](../src/vllm/model_executor/models/qwen3_5_internal.h#L52), stride-aware [consumers:930](../src/vt/ops.cpp#L930), unchanged [cuBLASLt path:188](../src/vt/cuda/cuda_matmul.cu#L188), exact [oracle generator:1](../tools/bench/gdn_ba_projection_oracle.py#L1), and explicit [trace contracts:113](../tools/bench/online_gate.py#L113) | W1 BA closed through the packed-decode equivalence closure (`f344dec` 235/235 both arms; 35B/GGUF inert). **W2A qkvz implemented 2026-07-15 (test-first)**: merged owner in exact [q,k,v,z] rows, ONE BF16 GEMM + strided mixed/z views on CUDA default, split rollback from the same owner (`VT_GDN_MERGED_QKVZ=0` / `VT_GDN_MERGED_PROJ=0`); CPU tier green — [merged-view battery:2367](../tests/vt/test_ops_gdn.cpp#L2367), [loader/CPU-exactness:456](../tests/vllm/models/test_qwen27_dense_forward.cpp#L456), [eligibility:451](../tests/vllm/models/test_qwen27_paged_forward.cpp#L451), full CTest 107/107, tools 162/162, clean -Werror rebuild. DGX gates at `baea3ec`: default/2a-qkvz-rollback/35B-inertness arms PASS; the `VT_GDN_MERGED_PROJ=0` arm exposed a non-mode-aware gate-test expectation (engine correct — master-off deselects packed decode by the designed BA coupling) → fixed test-first via `detail::PackedGdnDecodeEnvSelected` ([truth table:491](../tests/vllm/models/test_qwen27_paged_forward.cpp#L491), 16/16); 2b re-run, memcheck (first run PATH-only) and the 145→97 BF16 trace pending; `benchmark_binding=false`, no speed credit | [merged GDN projections](specs/gdn-merged-input-projections.md); [packed decode](specs/gdn-packed-decode.md) | `ANCHOR-BACKFILL` | CLAIM-GDN-BA-ROUNDING-1 | | `KERNEL-GEMM-FP8` | FP8/INT8 scaled-mm plus static activation quant | C2x/C3x dispatch `CMakeLists.txt:737-863`; stable quant sources `:383-388`; vLLM cuBLASLt fp8 reuses an in-graph plan (nvjet_sm121_qqtst_* kernels) — no per-call heuristic | [cuda_matmul_fp8_cutlass.cu:312](../src/vt/cuda/cuda_matmul_fp8_cutlass.cu#L312), [cuda_matmul.cu:345](../src/vt/cuda/cuda_matmul.cu#L345); **Hopper `sm_90a` arch coverage (2026-07-28, `CLAIM-CUDA-SM90-C3X`):** the Hopper C3x FP8 scaled-mm is a SEPARATE build-verify TU [cuda_scaled_mm_c3x_sm90.cu](../src/vt/cuda/cuda_scaled_mm_c3x_sm90.cu) (faithful 1:1 port of vLLM `cutlass_3x_gemm_sm90_fp8`, `ArchTag=Sm90`+`KernelTmaWarpSpecialized*FP8FastAccum` → wgmma/TMA), gated by its own `scaledmm-c3x-sm90` FEATURE-TABLE cell (90a-only) — DERIVED+BUILD-VERIFIED, no H100/H200 board ran it; **datacenter-Blackwell `sm_100a` arch coverage (2026-07-28, `CLAIM-CUDA-SM100-C3X`, DC3):** the sm100 C3x FP8 scaled-mm is a SEPARATE build-verify TU [cuda_scaled_mm_c3x_sm100.cu](../src/vt/cuda/cuda_scaled_mm_c3x_sm100.cu) (faithful 1:1 port of vLLM `cutlass_3x_gemm_sm100_fp8`, `ArchTag=Sm100`+`KernelScheduleAuto` → 5th-gen tcgen05 collective; `sm100_fp8_config_{default,M256,M64}`, 2SM `ClusterShape<_2,_2,_1>` default), gated by its own `scaledmm-c3x-sm100` FEATURE-TABLE cell (100a-only) — DERIVED+BUILD-VERIFIED (cuobjdump `sm_100a` cubin + `Sm100TmaUmma`/`SM100_MMA_F8F6F4`/`TMEM` symbols), no B200 board ran it (see backend-matrix `BACKEND-CUDA-SM100`); the sm_12x production body here (`ArchTag=Sm120`) is UNCHANGED (see backend-matrix `BACKEND-CUDA-SM090`); **opt-in per-device plan cache** [fp8_plan_cache.h](../src/vt/cuda/fp8_plan_cache.h) + [GetOrBuildCachedFp8Plan/BuildFp8Plan](../src/vt/cuda/cuda_matmul.cu#L358) (`VT_FP8_PLAN_CACHE=1` opt-in; DEFAULT OFF — bit-exact but measured production-NEUTRAL, premise not reproduced, see ledger 2026-07-18) | [FP8 tests](../tests/vt/test_ops_fp8_cutlass.cpp#L188); **byte-exact cached==fresh** [test_ops_fp8_cutlass.cpp#L387](../tests/vt/test_ops_fp8_cutlass.cpp#L387) + `VT_FP8_PLAN_CACHE=1` on-arm ctest; **CPU key/flag** [test_fp8_plan_cache.cpp](../tests/vt/test_fp8_plan_cache.cpp); 27B 235/235 + 35B 315/315 both flags. **Merged-QKV FP8 sub-lever (`CLAIM-FP8-MERGED-QKV-1`, `VT_FP8_MERGED_QKV` opt-in):** extend the fp4-only merged-QKV fusion to 35B FP8 — ONE fp8 GEMM over the N-concatenated Q/K/V operand + per-column dequant, replacing 3 separate per-shard GEMMs (10 attn layers, 30→10 GEMMs/step). fp8 is PER-TENSOR scaled so a single-alpha concat is incorrect; realized as concat RAW bytes + GEMM alpha=1 + resident per-column alpha vector via NEW `vt::MulColVecF32` ([ops.h](../include/vt/ops.h), [cpu_ops.cpp](../src/vt/cpu/cpu_ops.cpp), [cuda_glue.cu](../src/vt/cuda/cuda_glue.cu)); model glue `ResidentFp8Qkv`/`MergedFp8QkvD`/`MergedFp8QkvEligible` + `ProjectFullAttnQkv` branch ([qwen3_5.cpp](../src/vllm/model_executor/models/qwen3_5.cpp)), resident fields ([qwen3_5_weights.h](../include/vllm/model_executor/models/qwen3_5_weights.h)); byte-exact CPU tests [test_ops_glue.cpp](../tests/vt/test_ops_glue.cpp). CPU gates GREEN (glue 10/10, fp8_cutlass 6/6, matmul 7/7, clean -Werror); **DGX GREEN @ `e9ce593`** (clean CUDA -Werror 0 warn): 35B **315/315 token-exact both arms** + 27B **235/235 both arms** (inert), merge proven to fire; in-situ TPOT A/B **NEUTRAL** (c1/c8 ~0%, c2/c4 −0.5%, all ≤0.9% within rep noise) ⇒ **landed OPT-IN** (`VT_FP8_MERGED_QKV` default OFF, token-exact but not measurably faster; the merged-QKV sub-lever is complete — the broad FP8 row stays `ANCHOR-BACKFILL` for its remaining scope). Spec [fp8-merged-qkv-projection.md](specs/fp8-merged-qkv-projection.md) | [inventory](specs/kernel-family-inventory.md); [fp8-merged-qkv](specs/fp8-merged-qkv-projection.md) | `ANCHOR-BACKFILL` | - | | `KERNEL-GEMM-NVFP4-W4A4` | NVFP4 W4A4 dense quant, merged/fused projections, runtime bucketing, SM12 tactics, v0.25 persistent plan selection and model-owned alpha | SM10/11/12 FP4 families `CMakeLists.txt:940-1002`; CT alpha parameter `compressed_tensors_w4a4_nvfp4.py:95-141`; executed FlashInfer pass-through `kernels/linear/nvfp4/flashinfer.py:97-176`; fused selection `act_quant_fusion.py:36-40,128-181,283-300`; stable fused body `activation_nvfp4_quant_fusion_kernels.cu:30-163`, packed helpers `nvfp4_utils.cuh:25-36,118-329`, vector loads `cuda_vec_utils.cuh:123-175,264-288`; FlashInfer 0.6.13 device pointer `gemm_base.py:1307-1350`, `fp4_gemm_cutlass_sm120.cu:52-77,82-105,135-175`; v0.25 cache lifecycle sources retained | Existing W3-C/W3-F anchors remain. W3-I1 adds the default-off packed [fused producer and dispatch](../src/vt/cuda/cuda_matmul_nvfp4.cu#L1207), public zero-lifecycle [contract](../include/vt/ops.h#L518), and candidate/fallback/graph/alignment [tests](../tests/vt/test_ops_nvfp4_fp4.cpp#L959). **Datacenter-Blackwell arch coverage (2026-07-28, `CLAIM-CUDA-SM100-NVFP4`):** the sm_100a tcgen05 block-scaled NVFP4 GEMM is a SEPARATE build-verify TU [cuda_matmul_nvfp4_sm100.cu](../src/vt/cuda/cuda_matmul_nvfp4_sm100.cu) (faithful 1:1 port of vLLM `Fp4GemmSm100`, `ArchTag=Sm100`+`KernelScheduleAuto`), gated by its own `cutlass-nvfp4-sm100` FEATURE-TABLE cell (100a-only) — DERIVED+BUILD-VERIFIED, no B200 board ran it; the sm_12x production body here is UNCHANGED (see backend-matrix `BACKEND-CUDA-SM100`). The trace-only [controller](../include/vt/cuda/cuda_profiler_control.h#L13), [driver](../scripts/dgx-online-serving.sh#L14), batch-keyed [validator](../tools/bench/online_gate.py#L99), fail-closed [c2 finalizer](../tools/bench/finalize_low_batch_trace.py#L1), and [finalizer tests](../tests/tools/test_low_batch_trace_summary.py#L1) support exact c2 without changing production builds | Clean W3-I1 remains default-off after **27/40 timing + 3/8 memory**. Finalized `179a0fc` proves all 12 local ranges and 1,522 steady oracle windows resolve the same **128 Stream-K 128x64x256 + 80 static-persistent 128x32x256** split. Diagnostic local/oracle FP4 medians are **52.508720 / 52.734326 ms**, so FP4 GEMM is not the positive c2 residual. Status `9e0143fa…7b57` is `complete-diagnostic`; `3f256ab` stays 55/124 and no new speed credit exists. **2026-07-17 (`CLAIM-FP4-QUANT-FAST-1`): two NUMERICS-NEUTRAL decode-glue vectorization sub-levers landed OPT-IN** (default OFF; the underlying quant kernels stay implemented, no row-state change). Bit-identical vectorized-load+store fast kernels behind `VT_FP4_QUANT_FAST` ([ScaledFp4QuantFastKernel](../src/vt/cuda/cuda_matmul_nvfp4.cu)) + `VT_SILU_FP4_FAST` ([SiluAndMulFp4QuantFastKernel](../src/vt/cuda/cuda_matmul_nvfp4.cu)): each thread does ONE 16-byte `uint4` load (vs 16 scalar) + ONE 64-bit packed store (vs eight 1-byte), memory-access-pattern change ONLY (exact `CastToFp4NibbleDev`/`F32ToFp8Dev`/`fmaxf`-amax/bf16-SiLU math unchanged), grounded 1:1 in vLLM `nvfp4_quant_kernels.cu:56-80,98` @ `e24d1b24`; the numerics-changing hw-cvt + `__hmax2` reduction (vLLM's other ~0.4× of the edge) stays out of scope (that is the non-bit-identical `VT_FP4_FUSED_VEC` native kernel). BIT-IDENTITY PROVEN: byte-exact new-vs-old nibbles+scales, [60/60 adversarial parity asserts](../tests/vt/test_ops_nvfp4_fp4.cpp) + full suite 24/24 (26,976), [flag header](../src/vt/cuda/fp4_quant_fast.h) + [CPU test](../tests/vt/test_fp4_quant_fast.cpp) 20/20. Isolated nsys per-launch (swizzled bf16): ScaledFp4Quant K=5120 1.12-1.18× / K=17408 1.44-1.62×, SiluAndMul I=17408 1.14× (c2) / 1.38× (c16-c32) — PARTIAL vs the ≥1.3× flip bar (clears the larger shapes, misses the dominant K=5120 / c2; swizzled small-M is padding-thread-dominated), so BOTH stay default OFF; the orchestrator owns the combined in-situ A/B. Engine token gate both-flags-ON PASSED: 27B 235/235 (16/16 token-exact vs vLLM) + 35B 315/315. `benchmark_binding=false`, binding 52/124. **2026-07-18 (`CLAIM-CONV-UPDATE-FAST-1`): both flags flipped DEFAULT OFF→ON** per the parity-enabler policy (bit-identical ⇒ never-slower + token-safe; under the strict ≥1.0 gate every fraction counts). Predicate parse in [fp4_quant_fast.h](../src/vt/cuda/fp4_quant_fast.h) inverted to default-ON `=0`-rollback; [CPU flag test](../tests/vt/test_fp4_quant_fast.cpp) RED→GREEN 20/20; [CUDA byte-exact test](../tests/vt/test_ops_nvfp4_fp4.cpp) scalar baseline arm → `=0`, re-verified byte-exact 25/25 (26,976). No kernel-body change. Full default set (both flags default ON) 27B 235/235 + 35B 315/315; combined `=0` rollback 235/235 + 315/315. Binding grid re-measures the combined in-situ effect. **2026-07-19 (`CLAIM-SIGMOID-GATE-FOLD-1`): full-attention sigmoid-gate → o_proj activation-quant fusion landed OPT-IN** (`VT_FUSE_SIGMOID_QUANT=1`, default OFF). NEW `vt::SigmoidGateFp4Quant` op ([cuda `SigmoidGateFp4QuantKernel`](../src/vt/cuda/cuda_matmul_nvfp4.cu), [cpu composite](../src/vt/cpu/cpu_ops.cpp), [ops decl](../include/vt/ops.h)) folds `attn*sigmoid(gate)` into the o_proj NVFP4 activation quant — one kernel, no bf16 `gated` intermediate — mirroring vLLM Inductor `triton_poi_fused_mul_scaled_fp4_quant_sigmoid` and the `SiluMulFp4Quant` precedent; model dispatch `SigmoidGateOProjD` ([qwen3_5.cpp](../src/vllm/model_executor/models/qwen3_5.cpp)) fires only on the 27B true-W4A4 o_proj (35B W4A16-Marlin/fp8 reads bf16 acts ⇒ inert, keeps `SigmoidGateBf16`). BIT-IDENTICAL to `SigmoidGateBf16`+`ScaledFp4Quant`: byte-exact [op test](../tests/vt/test_ops_nvfp4_fp4.cpp) **14/14** (CPU f32/bf16 + CUDA, 3 shapes). DGX-GREEN (production flags, clean CUDA `-Werror` 0 warn): 27B **235/235 both arms** + 35B **315/315** (inert). In-situ 27B TTFT A/B (input-1024, 3 reps): c1 −0.15% / c2 −0.03% — NEUTRAL within rep noise (o_proj is a small slice of 27B prefill) ⇒ OPT-IN. Spec [glue-fusion](specs/glue-fusion-2026-07-19.md) — anchor `tests/vt/test_ops_nvfp4_fp4.cpp:113` | [small-M spike](specs/nvfp4-small-m-dispatch.md); [W3-E spike](specs/nvfp4-direct-swizzled-scales.md); [W3-C spike](specs/nvfp4-persistent-plan-cache.md); [W3-F device-alpha spike](specs/nvfp4-device-alpha.md); [W3-G spike](specs/fa2-gqa-split-kv-decode.md); [W3-H normal-producer spike](specs/nvfp4-bf16-producer-vectorization.md); [W3-I fused-producer spike](specs/nvfp4-fused-silu-producer.md) | `ANCHOR-BACKFILL` | CLAIM-SERVE-GATE-1 | -| `KERNEL-GEMM-MARLIN-W4A16` | FP4 W4A16 Marlin dense/grouped GEMM and repack | Marlin generation/targets `CMakeLists.txt:548-679,1168-1274`; capability floor `marlin_utils_fp4.py:29-35` | [cuda_marlin_repack.cu:129](../src/vt/cuda/cuda_marlin_repack.cu#L129), [cuda_moe_marlin.cu:156](../src/vt/cuda/cuda_moe_marlin.cu#L156) | [dense tests](../tests/vt/test_ops_nvfp4_matmul.cpp#L201), [MoE tests](../tests/vt/test_ops_moe_grouped.cpp#L453); 35B gate | [inventory](specs/kernel-family-inventory.md) | `ANCHOR-BACKFILL` | - | +| `KERNEL-GEMM-MARLIN-W4A16` | FP4 W4A16 Marlin dense/grouped GEMM and repack | Marlin generation/targets `CMakeLists.txt:548-679,1168-1274`; capability floor `marlin_utils_fp4.py:29-35`; **DENSE marlin_gemm** `csrc/libtorch_stable/quantization/marlin/marlin.cu:326-541,545` | [cuda_marlin_repack.cu:129](../src/vt/cuda/cuda_marlin_repack.cu#L129), [cuda_moe_marlin.cu:156](../src/vt/cuda/cuda_moe_marlin.cu#L156); **`KERNEL-MARLIN-DENSE-PORT` (gated OFF `VT_MARLIN_DENSE`): vLLM's OWN dense W4A16 GEMM** [cuda_marlin_dense.cu](../src/vt/cuda/cuda_marlin_dense.cu), lifted dispatcher [marlin_mm_dense.cu](../src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu), op [ops.h `kMarlinDenseGemm`](../include/vt/ops.h), routing [dense_nvfp4_gemm.h](../include/vllm/model_executor/models/dense_nvfp4_gemm.h) — byte-preserving E=1 replacement for the single-expert MoE route (no par-regroup ULP) | [dense tests](../tests/vt/test_ops_nvfp4_matmul.cpp#L201), [MoE + DENSE tests](../tests/vt/test_ops_moe_grouped.cpp#L453); 35B gate; **dense-port unit RED-first battery (dense-vs-ref-vs-grouped, M=1..8, NVFP4+MXFP4, wrong-stride RED)**; 3 new dense `.cu` compile-clean dgx sm_121a; GPU exec PENDING | [inventory](specs/kernel-family-inventory.md) | `ANCHOR-BACKFILL` | `KERNEL-MARLIN-DENSE-PORT` | | `KERNEL-GEMM-INT-LOWBIT` | AWQ/GPTQ/integer Marlin, Machete, and AllSpark | `CMakeLists.txt:468-731`; Marlin types/capability `marlin_utils.py:43-149` | - | - | [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | | `KERNEL-GEMM-W4A8` | CUTLASS W4A8 dense/MoE | `CMakeLists.txt:1004-1035`; upstream test `tests/kernels/quantization/test_cutlass_w4a8_moe.py` | - | - | [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | | `KERNEL-GEMM-QUTLASS-MX` | QuTLASS NVFP4/MXFP4 block-scaled GEMM | `cmake/external_projects/qutlass.cmake:56-147`; upstream tests `test_mxfp4_qutlass.py`, `test_nvfp4_qutlass.py` | - | - | [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | diff --git a/.agents/parity-ledger.md b/.agents/parity-ledger.md index 19a7f116c..047e3612f 100644 --- a/.agents/parity-ledger.md +++ b/.agents/parity-ledger.md @@ -888,3 +888,4 @@ Columns: | 2026-08-01 (`SERVE-CLI-CHAT` W0 contract spike; `CLAIM-SERVE-CLI-CHAT-SPIKE`; CPU-only records/spec) | Accepts `.agents/specs/cli-chat-complete.md`, corrects the inventory from “no direct commands” to the actual pinned `chat`/`complete` surface, and decomposes a dual-mode port: exact remote OpenAI HTTP/SSE commands plus preservation of the existing in-process invocation as a compatibility alias. No production, test, CMake, model, kernel, fixture, or generated file changes. | Pinned vLLM `5559679229`: command registration `vllm/entrypoints/cli/main.py:17-37,73-98`; model/auth resolution and stream shaping `vllm/entrypoints/cli/openai.py:30-100`; chat `:155-234`; complete `:237-312`. The local compatibility baseline is `examples/cli/main.cpp:1-207`. | CPU record/doc gates only; benchmark `NOT APPLICABLE`, `benchmark_binding=false`. Implementation remains absent and the row moves `INVENTORIED` -> `SPIKE`. W1-W5 name parse, transport, complete, chat, and packaging gates, including fake-server request/SSE transcript parity, Release `-Werror`, ASan+UBSan, and TSan. | | 2026-08-04 (`HARDEN-DETECTOR-LANES` PR #28 CI repair; `CLAIM-HARDEN-SANITIZER-DISK`; Ordino task `t-e19dc73f`; CPU-only, lifecycle unchanged; closing commit: this checkpoint) | Repairs the hosted ASan+UBSan build's filesystem exhaustion without weakening detector coverage. Sanitizer tests share one internal fully instrumented engine image instead of force-linking another static copy into every executable, and `-g1` retains file/line traces without full type/local-variable DWARF. CI enables the existing `VT_POOL_BYPASS=1` exact-allocation/real-free detector mode. The newly reachable leak survey removes a real minja `MacroNode` ownership cycle by weakly capturing the context that owns the callable. It also closes the two Nix-only suite portability gaps: resolve `true` and the active Python executable instead of assuming `/usr/bin`, and remove inherited `PYTHONHASHSEED` only from the unconfigured control suite. Default build linkage and packaged C ABI exports remain unchanged. | No vLLM behavioral counterpart: this is local build/test infrastructure plus a vendored minja lifetime repair. Anchors: `CMakeLists.txt`, `tests/CMakeLists.txt`, `.github/workflows/ci.yml`, `third_party/minja/minja.hpp`, `tests/tools/test_gdn_packed_component.py`, and `tests/tools/test_online_gate_client.py`. Remote root-cause evidence: GitHub run `30819266647`, job `91704728276`, 99 MiB free then `ld: No space left on device`. | **PASS.** GCC 15.2.0 ASan+UBSan full suite **331/331** under leak detection and pool bypass; TSan full suite **331/331** under pool bypass; affected plain GCC 15 `-Werror` suites **3/3**. ASan+UBSan tree **93 GiB -> 5.6 GiB** (about 94% smaller); TSan tree **1.9 GiB**. `benchmark_binding=false`, performance **NOT APPLICABLE**. Hosted PR confirmation is the next external gate; `continue-on-error` stays until that confirmation. | | 2026-08-08 (`row/KERNEL-FA2-GQA-SWAP`; `CLAIM-KERNEL-FA2-GQA-SWAP`; kernel `KERNEL-ATTN-FA2`; gated default-OFF, lifecycle unchanged) | Ports vLLM's FA2 `seqlenq_ngroups_swapped` decode optimization into the d128 varlen decode launcher (`LaunchDecodeVarlenFA2Bf16`, gate `VT_FA2_DECODE_GQA_SWAP`): the Qwen3-dense decode grid becomes `(batch, kv_heads)` not `(batch, hq)` — the ngroups query heads pack into seqlen_q, KV read once/group, presented WITHOUT a materialized transpose via kv-major-group-minor strides (a 1:1 mirror of the already-shipped d256 `LaunchDecodeFA2Bf16` swap). OFF path byte-identical to the prior plain-varlen reduction; ON is non-byte-exact only when num_splits>1 (split reduction order → near-tie, toward vLLM's own numerics). | Mirrors `flash-attention @ 2c839c33` `mha_fwd_kvcache` seqlenq_ngroups_swapped + `set_params_splitkv` and vLLM v0.25.0 `flash_attn.py flash_attn_varlen_func` decode (#47 measured vLLM's swapped grid `(1,6,16)` = batch×kv_heads vs ours `(1,3,64)` = batch×query_heads). The vendored `flash_fwd_kernel.h` `get_lse_tile`/combine already honor the flag in both the num_splits==1 direct-write and >1 combine paths (the d256 arm is the proof). | GB10 sm_121a CUDA 13.0: op RED-first test 280/280 (both GQA ratios × batch{1,2,4,8} × short+long ctx; `swap_launches==1` proves the grid engaged; swap-vs-plain near-tie; MHA-inert) — RED proven (wrong swapped stride → 26,528 violations); full binary 28/28·454,679 no regression; compute-sanitizer 0-err/0-leak; #44 MXFP4 e2e smoke swap-ON 3/3 deterministic TOKEN-EXACT + coherent, byte-identical to swap-OFF. `benchmark_binding=false` (c1-c8 x3 re-bench + default flip = recorded next step; #47 projects flash ~28%@c2 / ~55%@c8 of the gap). | +| 2026-08-09 (`row/KERNEL-MARLIN-DENSE-PORT`; `CLAIM-KERNEL-MARLIN-DENSE-PORT`; kernel `KERNEL-GEMM-MARLIN-W4A16`; gated default-OFF, lifecycle unchanged) | Vendors vLLM's OWN dense marlin W4A16 GEMM as a new `vt::MarlinDenseGemm` op (`VT_MARLIN_DENSE`, default OFF) and routes the E=1 dense NVFP4/MXFP4 projections (`dense_nvfp4_gemm.h` `MatmulNvfp4MarlinD`/`MatmulMxfp4W4A16D`/`GateUpFusedMarlinD`) through it. The dense kernel is direct-A + tile-per-CTA with vLLM's OWN dense fp32-C_tmp reduce, so at M<=8 it runs the sms-wide (48-CTA) grid WITHOUT the one-bf16-ULP shift the `VT_MARLIN_E1_PAR1` MoE-route par-regroup costs (#54: that ULP flips a strict 32B-NVFP4A16 token). Reuses the EXISTING marlin resident + workspace (same `marlin_permute` repack for dense and MoE — confirmed, no shim); rank-2 operand views, no moe_align gather. | 1:1 lift of vLLM @ `555967922` `csrc/libtorch_stable/quantization/marlin/`: `marlin.cu:326-541` (`marlin::marlin_mm` + config helpers) → `marlin_mm_dense.cu`; the torch::stable `marlin_gemm` wrapper (`:545-894`) → torch-free `cuda_marlin_dense.cu` launcher (mirrors `cuda_moe_marlin.cu`, dense c_tmp sizing `:713-716`); `kernel.h`/`marlin_template.h:1-2081` verbatim (the DENSE kernel — DISTINCT from the moe one, but SAME 12-param `Marlin<>` template so the generated `kernel_selector.h`+`sm80_*.cu` instantiation set is shared, namespace `marlin` from the local kernel.h). Shared `marlin.cuh`/`marlin_dtypes.cuh`/`dequant.h`/`marlin_mma.h` diff-verified byte-identical. Forced-Marlin a16 selection `kernels/linear/__init__.py:879-881`. | CPU `-fsyntax-only` CLEAN (`ops.cpp` + the `VT_MARLIN_NVFP4` routing header). GPU compile: all 3 new dense `.cu` compile CLEAN on dgx GB10 sm_121a under exact production flags (`-Werror=all-warnings`, `-static-global-template-stub=false`, `--generate-code=…sm_121a`). RED-first unit battery WRITTEN (`test_ops_moe_grouped.cpp`: NVFP4+MXFP4, M=1..8 × 3 shapes, dense-vs-CPU-ref AND dense-vs-grouped-route, row-shifted stride RED-injection). `benchmark_binding=false`; GPU EXEC gates (unit run + strict token battery dense-ON vs oracle incl. 32B-NVFP4A16:344 + launch-counter + nsys 48-CTA + binding c1..c8 x3) are the scoped dgx follow-up; default stays OFF until the strict battery proves oracle byte-match and the binding beats the MoE route (state `KERNEL-MARLIN-DENSE-PORT`). | diff --git a/.agents/porting-inventory.md b/.agents/porting-inventory.md index c41c496a8..6c71cedb2 100644 --- a/.agents/porting-inventory.md +++ b/.agents/porting-inventory.md @@ -568,6 +568,38 @@ Examples: `examples/cli` ✅ (C-API client), `examples/server` ✅ (OpenAI serve `nvfp4_marlin_process_scales`/`_global_scale`), the `moe_align_block_size` port, the 35B forward wiring, 16/16 parity, and the A/B TFLOPS measurement. + **DENSE Marlin (row `KERNEL-MARLIN-DENSE-PORT`, gated OFF `VT_MARLIN_DENSE`, + 2026-08-06)**: the byte-preserving E=1 route. `src/vt/cuda/marlin/libtorch_stable/ + quantization/marlin/` now also vendors vLLM's OWN dense marlin (a DISTINCT + kernel from the moe one — direct-A, `lda`, no sorted_token_ids/expert_ids/top_k + gather, its own par-split fp32 C_tmp reduce), all from vLLM @ `555967922` + `csrc/libtorch_stable/quantization/marlin/`: + * `kernel.h` ← `kernel.h` (verbatim; `namespace marlin`, dense `MARLIN_KERNEL_PARAMS` with `lda`) + * `marlin_template.h` ← `marlin_template.h:1-2081` (verbatim dense kernel; the + SHARED `marlin.cuh`/`marlin_dtypes.cuh`/`dequant.h`/`marlin_mma.h` it includes + are byte-identical to our existing vendored copies — diff-verified) + * `marlin_mm_dense.{h,cu}` ← `marlin.cu:326-541` `marlin::marlin_mm` + config + helpers (`get_marlin_kernel`/`determine_exec_config`/`is_valid_config`/…); + the torch::stable `marlin_gemm` host wrapper (`marlin.cu:545-894`) is stripped, + replaced by the torch-free launcher `vt::MarlinDenseGemm` + (`src/vt/cuda/cuda_marlin_dense.cu`, mirrors `cuda_moe_marlin.cu`). Only the + original's redundant inner `is_a_8bit` shadow is dropped (identical value, + avoids `-Wshadow`). `STD_TORCH_CHECK` → `vt_marlin_check.h` as for the moe TU. + * `kernel_selector.h`, `sm80_kernel_bfloat16_fe2m1f_bfloat16.cu` ← + `generate_kernels.py` output. KEY: the dense kernel is the SAME 12-param + `Marlin<>` template as the moe one, so these are the SAME instantiation set — + the dense kernel BODY + `namespace marlin` come from the local dense + `kernel.h`/`marlin_template.h` this TU includes. + New op `vt::OpId::kMarlinDenseGemm` + `MarlinDenseArgs` (ops.h, appended before + `kCount` — no id shift); routing in `dense_nvfp4_gemm.h` reuses the EXISTING + resident weights + workspace (same `marlin_permute` repack for dense and moe — + CONFIRMED via the shared repack ops; no shim needed) with rank-2 operand views + and NO moe_align. **Compile-VERIFIED GB10 sm_121a (2026-08-06): all 3 new dense + `.cu` compile clean under the exact production flags** (`-Werror=all-warnings`, + `-static-global-template-stub=false`, `--generate-code=…sm_121a`). GPU exec gates + (unit RED-first battery, strict token battery dense-ON, nsys 48-CTA, binding + c1..c8) are the dgx follow-up; extracted tree kept at dgx `~/dense_check/vllm.cpp`. + 11. **Vendored FlashAttention-2 (head-dim-256 GQA prefill implemented; ratio-6 split-KV decode `ACTIVE`)**: `src/vt/cuda/flash_attn/` is a byte-identical, torch-free vendor of vllm-project/flash-attention @ `2c839c33`, still the diff --git a/.agents/state.md b/.agents/state.md index 599c9c376..03ef90e26 100644 --- a/.agents/state.md +++ b/.agents/state.md @@ -37379,3 +37379,70 @@ paged_engine` 142/142 at the pure shipping default. Evidence dgx:~/mxfp4-nsys/{o step1_gate,step1_smoke}.log. Env: `VLLM_CPP_QWEN3_DENSE_DECODE_GRAPH` (now default ON), `VT_MOE_FUSED_W13` (default ON, now covers mxfp4), `VT_MARLIN_E1_PAR1` (NEW, default OFF opt-in). +## KERNEL-MARLIN-DENSE-PORT: vLLM's OWN dense marlin GEMM vendored + wired to the E=1 NVFP4/MXFP4 dense projections (byte-preserving, no par-regroup ULP), gated OFF; 3 new dense .cu compile-CLEAN on GB10 sm_121a; RED-first unit battery WRITTEN; GPU exec/strict/nsys/binding gates SCOPED for the operator + + +WHAT + WHY. #54 (`QUANT-CT-MXFP4-MARLIN-STRUCT`) settled the c8 marlin residual as CTA count: +the single-expert MoE-marlin route the dense E=1 projections use pads M<=8 into a 144-CTA grid, +while vLLM's dense marlin covers the same tile set with 48 CTAs. `VT_MARLIN_E1_PAR1` (clamp par→1) +recovers 81% of the residual (per-call 114.7 vs vLLM 113.1; +226us vs vLLM at c8) but regroups the +fp32 C_tmp reduce, costing one bf16 ULP that FLIPS a strict token on the 64-layer 32B-NVFP4A16 +(`test_qwen3_32b_nvfp4a16_paged_engine` REQUIRE :344) → default-OFF. This row is the byte-preserving +fix the #54 comment named: port vLLM's OWN dense marlin template (its own reduce = its own numerics). + +PORTED (all cited from vLLM @ `555967922` `csrc/libtorch_stable/quantization/marlin/`; provenance in +porting-inventory §10). KEY finding: the dense marlin_template.h (2081L) is a DISTINCT kernel from the +moe one (direct-A `A0`, `lda`, no sorted_token_ids/expert_ids/top_k gather, no atomic-add-only path) but +uses the SAME 12-param `Marlin` template — so the +generated `kernel_selector.h` + `sm80_kernel_*.cu` are the SAME instantiation set as the moe TUs, only +the namespace (`marlin` vs `marlin_moe_wna16`) + kernel body differ (they come from the local +`kernel.h`/`marlin_template.h`). The SHARED `marlin.cuh`/`marlin_dtypes.cuh`/`dequant.h`/`marlin_mma.h` +were diff-verified byte-identical to our existing vendored copies. Vendored files (in +`src/vt/cuda/marlin/libtorch_stable/quantization/marlin/`): `kernel.h` (verbatim), `marlin_template.h` +(verbatim `:1-2081`), `marlin_mm_dense.{h,cu}` (= `marlin.cu:326-541` `marlin::marlin_mm` + helpers, +torch wrapper stripped, only the redundant inner `is_a_8bit` shadow dropped → `-Wshadow`), +`kernel_selector.h` + `sm80_kernel_bfloat16_fe2m1f_bfloat16.cu` (generator output). New op: +`vt::OpId::kMarlinDenseGemm` + `MarlinDenseArgs` + `MarlinDenseGemmFn` (ops.h, appended before `kCount` +— no id shift), dispatch shim (ops.cpp), launcher `src/vt/cuda/cuda_marlin_dense.cu` (mirrors +`cuda_moe_marlin.cu`: graph-safe per-stream c_tmp pool, dense c_tmp sizing `marlin.cu:713-716` += `sms * min(ceil(M/16)*16,64) * max_thread_n`). CMake: 3 TUs added to `_MARLIN_SRCS` with the same +nvcc options. ROUTING (`dense_nvfp4_gemm.h`, gated `VT_MARLIN_DENSE` default OFF): `MatmulNvfp4MarlinD` ++ `GateUpFusedMarlinD` branch to `vt::MarlinDenseGemm` when the gate is on AND the op is registered — +REUSING the existing resident (mr.w/mr.s/mr.g) + workspace with rank-2 views and NO moe_align gather. +Repack layout CONFIRMED shared: dense and moe both consume the SAME `MarlinRepackExpertWeight` + +`MarlinProcessExpertScales`(`Mxfp4`) + `MarlinNvfp4ProcessGlobalScale` residents (vLLM's shared +`marlin_permute`), so NO shim was needed. `dense_gemms` execution counter added (the "path RAN" signal). + +GATES DONE. CPU: `-fsyntax-only` CLEAN on `src/vt/ops.cpp` (op+shim) and on the `VT_MARLIN_NVFP4` +routing header (`dense_nvfp4_gemm.h`). GPU compile: all 3 new dense `.cu` compile CLEAN on dgx +GB10 sm_121a under the EXACT production flags (`-Werror=all-warnings`, `-static-global-template-stub= +false`, `--expt-relaxed-constexpr`, `-diag-suppress=20280`, `--generate-code=…sm_121a`) — validating +the hand-lift, the dense kernel body + instantiations in `namespace marlin`, and the launcher API +(`RetireGraphScratch`/`RegisterOp`/`marlin::marlin_mm`). Extracted tree kept at dgx +`~/dense_check/vllm.cpp` for the follow-up build. RED-first unit battery WRITTEN +(`tests/vt/test_ops_moe_grouped.cpp`, under `VT_MARLIN_NVFP4`): two TEST_CASEs (NVFP4 + MXFP4) loop +M=1..8 × model shapes {256×64, 512×128, 128×256}, assert `vt::MarlinDenseGemm` matches BOTH the +independent CPU-dequant reference AND the single-expert grouped route, plus a row-shifted (stride-class) +RED-injection that the comparison must discriminate. + +NOT DONE (operator dgx follow-up, precisely scoped — full from-scratch GB10 build is the long pole, +not attempted here): + (a) UNIT EXEC — build `~/dense_check/vllm.cpp` (Release, `-DVLLM_CPP_CUTLASS_DIR=$HOME/cutlass-4.5.0`, + arch 121a, MARLIN+FLASH+TRITON ON) and run the two dense TEST_CASEs + memcheck. RED-first: they + must be RED on a broken stride and GREEN on the port. + (b) STRICT BATTERY dense-ON (`VT_MARLIN_DENSE=1`), eager AND graphed, vs the ORACLE goldens: + `test_qwen3_32b_nvfp4a16_paged_engine` (the 142/142 suite that blocked par1 — THE decider), + `test_qwen3_paged_engine` 0.6B/4B, #44 MXFP4-8B smoke, the async gate. Add a launch-counter + assertion there (`GetW4A16Stats().dense_gemms>0 && marlin_gemms==0` under the gate). If tokens + match the oracle everywhere strict (PLAUSIBLE — same reduce as vLLM), FLIP DEFAULT ON per + parity-enablers with the flipped-default proof; if a golden legitimately shifts, apply the + near-tie teacher-force razor BEFORE any regen. + (c) nsys c8 same-tool: confirm 48-CTA grids on the dense route + per-call ~113-115us; target the + per-step marlin delta ~+226us or better (the #54 par1 measurement). + (d) BINDING c1..c8 x3 dense-ON vs #51 (1.005/0.925/0.939/0.953) on an on-REAL-disk RelWithDebInfo + build (tmpfs fails mincore); `free -g >= 90` gate + memory monitor + sequential arms; the oracle + host-RAM reservation is the OOM-reboot risk. Updated parity verdict + residual map (flash +784us + expected next-dominant; glue +195 the tail). +Default stays OFF until (b) proves byte-match vs the oracle and (d) beats the MoE route. SHA of this +landing: `row/KERNEL-MARLIN-DENSE-PORT` (code commit + this records commit). + diff --git a/CMakeLists.txt b/CMakeLists.txt index 73a426e63..650fab2a1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1108,7 +1108,14 @@ if(VLLM_CPP_CUDA) src/vt/cuda/marlin/libtorch_stable/moe/marlin_moe_wna16/marlin_mm_moe.cu src/vt/cuda/marlin/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu src/vt/cuda/cuda_moe_marlin.cu - src/vt/cuda/cuda_marlin_repack.cu) + src/vt/cuda/cuda_marlin_repack.cu + # DENSE marlin (row KERNEL-MARLIN-DENSE-PORT): the dense marlin_mm dispatcher + # (`namespace marlin`) + its bf16 NVFP4/MXFP4 kernel instantiations (SAME 12-param + # Marlin<> template as the MoE TUs, distinct namespace + direct-A kernel body) + + # the vt::Tensor launcher. Same nvcc options as the MoE TUs below. + src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu + src/vt/cuda/marlin/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu + src/vt/cuda/cuda_marlin_dense.cu) target_sources(vllm PRIVATE ${_MARLIN_SRCS}) target_compile_definitions(vllm PUBLIC VT_MARLIN_NVFP4=1) set_source_files_properties(${_MARLIN_SRCS} PROPERTIES diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 5b5d0fdbc..0eddcf295 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -309,7 +309,7 @@ built on it rather than keeping the flattering one. | Qwen3-dense decode CUDA-graph | Token-exact pass, ~4.3% e2e directional | Steady-state per-step tok/s | | Kimi-Linear-48B-A3B (KDA+MLA+MoE) | Full-model GB10 e2e RUNS (bf16-resident §13), NEAR-TIE 106/128, pool math CLOSES; default OFF | Full model RUNS on GB10 (bf16-resident, RSS peak 1.7 GiB, min-avail 21 GiB, no OOM). Token NEAR-TIE 106/128 (6/8 prompts exact, numerics vs deterministic oracle). 1.59 tok/s. Detail: spec §13 | | vLLM 0.26 re-benchmark | Pending | Re-run the binding grids on the advanced pin | -| MXFP4 Qwen3-8B (W4A16 Marlin) | #51 x3: c1 1.005, c2/c4/c8 0.925/0.939/0.953, mem 2.18x. `MARLIN-STRUCT`: decode-graph + gate_up FUSION default-ON, marlin 180->144 GEMM/step (vLLM-structural); #44 3/3, 32B-NVFP4A16 142/142 | nsys c8 residual: marlin +1,177us (CTA 144 vs 48, dominant), flash +784, glue +195. `VT_MARLIN_E1_PAR1` opt-in (E=1 grid to 48 CTAs) near-parity but flips a strict 32B token (default-OFF). Detail in benchmark-record | +| MXFP4 Qwen3-8B (W4A16 Marlin) | #51 x3: c1 1.005, c2/c4/c8 0.925/0.939/0.953, mem 2.18x. `MARLIN-STRUCT`: decode-graph + gate_up FUSION default-ON, marlin 180->144 GEMM/step (vLLM-structural); #44 3/3, 32B-NVFP4A16 142/142 | nsys c8: marlin +1,177us (CTA 144 vs 48, dominant); `VT_MARLIN_E1_PAR1` E=1->48 CTAs near-parity but flips a strict 32B token (OFF). Byte-preserving `KERNEL-MARLIN-DENSE-PORT` landed gated-OFF; GPU binding pending | | SGLang floor arms | Never ran | Both arms of the SGLang comparison | | cuBLAS invocation-parity guard | CI guard landed (CPU); `kGemvHeuristicAlgos` refactor build-verify owed | `nvcc` rebuild + SACRED gate on dgx | | Pre-Ampere breadth (Turing `sm_75` / Volta `sm_70` / Pascal) | **NO NUMBER OWED, nothing executes on these arches.** 2026-08-06 sm_75 compile audit (nvcc 13.0.88): 20 unconditional sm_80+ constructs enumerated; detail in .agents/benchmark-record.md | Port the llama.cpp `fattn-tile`/`fattn-vec` fp16 body. Perf floor when a card exists is **llama.cpp on the same card** (vLLM does not run there) | diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 480fde71d..03ee433e3 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -59,6 +59,7 @@ portable/reference path. In normal operation leave them unset. | `VLLM_CPP_QWEN3_DENSE_DECODE_GRAPH` | **on** | Routes pure-decode steps for the SHARED pure-dense forward (`Qwen3DenseModel`, i.e. Qwen3 / Llama / InternLM3 / Mistral / InternLM2 `ForCausalLM`) through the captured decode CUDA graph. **DEFAULT-ON since `QUANT-CT-MXFP4-MARLIN-STRUCT`** (parity-enabler; banks ~+1.3% TPOT@c8 by collapsing the eager inter-kernel launch gap). Token-exact with eager — `=0` opts out to the byte-identical eager decode; gated dgx SACRED `test_qwen3_paged_engine` 184/184 (graph ON==OFF, Qwen3-0.6B near-tie + 4B), async 82/82, Qwen3-8B-MXFP4 #44 smoke 3/3, Qwen3-32B-NVFP4A16 142/142. Honors `VLLM_CPP_CUDAGRAPH=0` | | `VT_MOE_FUSED_W13` | on | Runs a dense/shared MLP gate+up PAIR as ONE Marlin GEMM over the N-concatenated `[2I,H]` operand + `SiluAndMul` (vLLM's merged `gate_up_proj` structure), instead of two split GEMMs + `MoeSiluMul`. **Generalized NVFP4→MXFP4 in `QUANT-CT-MXFP4-MARLIN-STRUCT`** (drops the classic-dense Qwen3-8B-MXFP4 decode from 180→144 marlin GEMM/step = vLLM-structural parity). `=0` is the split A/B fallback. Numerically equivalent to split (the fused fp32 split-K reduce regroups by 1 bf16 ULP; token-exact vs the oracle — #44 fused==split 3/3, unit `test_linear_method` 99.9% bit-exact) | | `VT_MARLIN_E1_PAR1` | off (opt-in) | `=1` clamps the single-expert (`num_experts==1`) DECODE marlin grid to `sms×1` (48 CTAs on GB10) instead of the persistent `sms×par` grid `determine_exec_config` auto-picks (par=3 → 144 CTAs), matching vLLM's dense marlin tile-per-CTA count. Measured same-tool in-model (Qwen3-8B-MXFP4 c8): marlin 17,463→16,512 us/step (−5.4%, per-call 121.3→114.7 vs vLLM 113.1 = near-parity), TPOT 37.22→36.23 ms, token-exact on 8B-MXFP4. **DEFAULT OFF**: `par` regroups the fp32 C_tmp reduce, so the E=1 output differs by 1 bf16 ULP — on the 64-layer Qwen3-32B-NVFP4A16 that accumulates into a strict-token flip vs its committed SACRED anchor (`test_qwen3_32b_nvfp4a16_paged_engine` REQUIRE :344). Real MoE (`num_experts>1`) and prefill (`thread_m_blocks>1`) are untouched (byte-identical) | +| `VT_MARLIN_DENSE` | off (opt-in) | `=1` routes the E=1 dense NVFP4/MXFP4 projections (`dense_nvfp4_gemm.h` `MatmulNvfp4MarlinD`/`GateUpFusedMarlinD`) through vLLM's OWN dense marlin GEMM (`vt::MarlinDenseGemm`) instead of the single-expert MoE-marlin route. The dense kernel is direct-A + tile-per-CTA with vLLM's own dense fp32-C_tmp reduce, so at M≤8 it runs the sms-wide (48-CTA) grid `VT_MARLIN_E1_PAR1` targets but WITHOUT that flag's par-regroup ULP — the byte-preserving fix for the strict-32B token flip (row `KERNEL-MARLIN-DENSE-PORT`). Reuses the same marlin resident + workspace (shared `marlin_permute` repack). **DEFAULT OFF** until the strict token battery proves oracle byte-match and the binding beats the MoE route; then flipped ON per the parity-enabler policy. CUDA-only (needs `VT_MARLIN_NVFP4`) | | `VT_MM_DECODE_EAGER` | off (graph on) | Set to `1` to force the eager per-step multimodal (Qwen3.6-27B image/video) decode instead of routing it through the captured dense decode graph. Rollback / A-B knob; the graphed path is token-exact with the eager path | | `VT_KIMI_DEVICE_COMPUTE` | off (opt-in) | `=1` routes the Kimi-Linear-48B-A3B runner path (`KimiLinearModel::ForwardDevice`) through the W7 DBuf-resident device COMPUTE (`ForwardDeviceCompute`, the whole KDA/NoPE-MLA + MoE hybrid over pooled DBufs via the shared vt:: ops) instead of the default W6 host-reference compose. Default OFF keeps the CPU-verified host-ref-compose seam as production until the device compute is GPU-verified against the SACRED oracle; the device compute is CPU-gated (`test_kimi_linear_forward`, device==W2 reference within f32-accumulation tolerance, greedy-token-identical) but its GPU numerics are a NAMED pending. The flag exists so the device path CAN be exercised as the runner path for that verification | | `VT_WHISPER_ENC_EAGER` | off (flash-tiled attention on) | Set to `1` to force the naive per-key block-reduction attention in the Voxtral/Whisper audio encoder instead of the default flash-tiled kernel. Rollback / A-B knob; token-identical to the default path | diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 36106aa44..d44ac3418 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -65,7 +65,7 @@ are our reading of their documented behavior, not measurements. | GGUF k-quants and i-quants | ✅ | ☐ | ☐ | ✅ | | AWQ | ◐ CPU dequant | ✅ | ✅ | ☐ | | GPTQ | ◐ CPU dequant | ✅ | ✅ | ☐ | -| MXFP4 compressed-tensors | ◐ W4A16 Marlin compute proven; mem 2.6x less. gate_up FUSION + decode-graph default-ON (marlin 180->144 GEMM/step, vLLM-structural); #44 3/3, 32B 142/142. Residual = marlin CTA (`VT_MARLIN_E1_PAR1` opt-in), <1.0x | ✅ | ✅ | ☐ | +| MXFP4 compressed-tensors | ◐ W4A16 Marlin, mem 2.6x less. gate_up FUSION + decode-graph default-ON (180->144 GEMM/step); #44 3/3, 32B 142/142. Residual = marlin CTA; byte-preserving `KERNEL-MARLIN-DENSE-PORT` landed gated-OFF, GPU gates pending | ✅ | ✅ | ☐ | | fp8 weights | ✅ | ✅ | ✅ | ☐ | | bf16 / fp16 | ✅ | ✅ | ✅ | ✅ | | Safetensors direct load, no conversion | ✅ | ✅ | ✅ | ☐ | diff --git a/docs/STATUS.md b/docs/STATUS.md index cd182608c..9ff46e260 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -84,7 +84,7 @@ token-for-token correctness against the pinned oracle. | Safetensors loading | Supported | Both gate models plus every registered dense/MoE family | | GGUF loading (F32/F16/BF16/Q4_0/Q8_0/Q2_K/Q3_K/Q4_K/Q5_K/Q6_K/IQ2_XXS/IQ3_XXS/IQ2_S/MXFP4/NVFP4) | Supported; compute-in-quant (keep-quant) on CPU AND now CUDA for the six K-block encodings PLUS Q2_K/IQ2_XXS/IQ3_XXS (DeepSeek-V4 W8, 2026-07-29 - the FIRST CUDA keep-quant GGUF k-quant GEMM `KERNEL-QUANT-CIQ-GEMM-CUDA`, MMVQ-style dequant-in-kernel, GB10-gated 92401/92401 vs the CPU oracle, so a CUDA runner keeps blocks compressed and dots them on the GPU instead of the ARM cores); **NVFP4 now COMPUTES IN FP4 on CUDA for the dense-MLP and full-attention projections (2026-07-29, `CLAIM-GGUF-NVFP4-COMPUTE`), no longer materialize-only** | Weights in six block encodings stay compressed from file to matmul on CPU (no BF16 expansion). NVFP4 (ggml type 40) DEQUANTIZES, including the per-tensor (per-expert) `.scale` sidecar the container keeps outside the blocks; gated BIT-EXACT against the compressed-tensors NVFP4 path on real Qwen3.6-27B bytes from both containers. **It no longer expands to bf16 on CUDA:** an NVFP4 matmul/expert weight is REPACKED at load into the same (`weight_packed [N,K/2]`, `weight_scale [N,K/16]`) operand pair the compressed-tensors path produces - a pure byte permutation, gated BYTE-IDENTICAL against that container - and the existing `vt::MatmulNvfp4*` kernels run on it, so no new kernel exists and no numerics are re-derived. Covers the dense MLP + full-attention q/k/v/o and the MoE shared/routed experts; the GDN `in_proj_*` family and `ssm_out` still expand (the V-head reorder rewrites their layout) and a CPU build still expands everything - the documented `part` subset. **MEASURED on GB10 (2026-07-29), same-binary A/B, one `flock`, idle box, 2 reps per arm:** peak RSS **50.8 -> 25.7 GiB**, load-and-generate **1:58 -> 0:41**, and the 256 projections that move cost 35 840 MiB expanded against 10 080 MiB fp4-resident (3.56x). **The divergence against the safetensors sibling CLOSES:** the fp4 arm is token-IDENTICAL over 24 greedy tokens where the bf16 arm of the same binary diverges at index 4, which retires the reading that that divergence was permanent. It is REPORTED, not gated: the two containers are not the same model - the GGUF NVFP4-quantizes 192 GDN `in_proj` tensors the safetensors keeps BF16 (mean relative weight error ~0.18) and their activation global scales differ - so identity is not guaranteed and a cross-container throughput arm is not valid. SACRED gates unmoved: `test_qwen27_paged_engine` 235/235, `test_qwen36_paged_engine` 315/315. **The MoE (35B) stacked-expert arm is now HARDWARE-GATED too (2026-07-29)**, superseding the gap recorded here: the real 35B A3B NVFP4 GGUF loads and generates through the fp4 path, its 120 routed-expert stacks x 256 experts repack to the modelopt safetensors' own operands with ZERO differing bytes over 840 sampled (tensor, expert) slabs, and all 840 per-expert `.scale[e]` are bit-identical to that expert's `weight_scale_2` - the per-expert scale INDEXING, mutation-proved against both a `scales[0]`-for-every-expert and an expert-0-slab-for-every-expert mutant. Same-binary A/B: peak RSS 68.5 -> 22.7 GiB (3.01x), load-and-generate 1:51.9 -> 0:28.8, tokens IDENTICAL (correct here, since the 35B routed experts run the W4A16 grouped GEMM in both arms). Recorded as OPEN, not smoothed over: this case's 24-token greedy stream is NOT run-to-run stable (one of three `use_a16` runs and one of four safetensors-reference runs differed), so the binding results are the weight-level byte identity and the residency audit, not a token-exactness claim; `test_qwen36_paged_engine` is token-exact at ITS engine params, so the instability belongs to this case's configuration and attributing it is owed work. That run also found and FIXED a latent defect the MoE arm made reachable: the two fp4 fused MoE blocks issued the router GEMM assuming the safetensors `[K,N]` gate layout and threw `matmul: inner dims mismatch` on the GGUF's `[N,K]` one; `MoeRouterLogits` now branches on `nk` (inert for the safetensors path, SACRED gates unmoved). **Q2_K (id 10) + IQ2_XXS (id 16) DEQUANTIZE (2026-07-29, `CLAIM-DSV4-GGUF-LOADER`):** the ~2-bit types the single-Spark `DeepSeek-V4-Flash-GGUF UD-IQ2_XXS`/`UD-Q2_K_XL` vehicles use, ported 1:1 from llama.cpp `ggml-quants.c` (`iq2xxs_grid` codebook + signs; Q2_K nibble sub-scale/min), unit-gated on hand-derived bytes (`test_gguf_dequant` 15/15). Dequant-only (no vec_dot -> expand-bf16). A V4-GGUF model still cannot RUN: the V4-GGUF name map (tensor-manifest-blocked) + the V4 forward (W3-W8) remain. **Multi-shard split GGUF READING landed (2026-08-03, `CLAIM-GGUF-SPLIT-SHARDS`):** `GgufFile::Open` now transparently stitches llama.cpp `gguf-split` shards (`...-00001-of-00003.gguf`) — every shard mmap'd, tensor tables merged, KV metadata taken from shard `00001`, and the sibling shard mappings kept alive by the primary so keep-quant mmap-borrows stay valid across shards (`OwnsSpan` is shard-aware); `VT_GGUF_NO_SPLIT=1` opts out; unit-gated (`test_gguf` split-merge / no-split / count-mismatch cases, 33/33 local). This unblocks the real 3-shard `unsloth/DeepSeek-V4-Flash-0731 UD-IQ2_M` (~91 GiB), whose layout is the NATIVE `deepseek4` arch — per-block `ffn_gate_tid2eid` hash tables (hash layers 0/1/2) + `hc_*` MHC + DSA compressor/indexer are all PRESENT (name-map 1328/1328), `vocab_size` derives from `token_embd` — NOT a standard llama.cpp conversion, so no loader-layout change is owed. It now loads THROUGH 1324/1328 tensors; the sole remaining gap is 4 routed-expert slabs quantized with IQ2_S (id 22, ×2) + MXFP4 (id 39, ×2) — encodings we have GGUF block traits for but no keep-quant vec_dot, so they hit the expand→dequant path which lacks them. Dequant-expanding those 4 big expert tensors to bf16 would add ~17 GiB (~106 GiB total → GB10 OOM-reboot risk), so the memory-safe fix is an IQ2_S+MXFP4 keep-quant kernel (CPU dequant dispatch + the `iq2s_grid` codebook + a CUDA `DotSuperblock`), spec'd as the next brick **IQ2_S (id 22) + MXFP4 (id 39) DEQUANTIZE + KEEP-QUANT on CPU (2026-08-03, `CLAIM-DSV4-UDIQ2M-QUANT`, off-GPU):** the extra per-tensor "dynamic" encodings the `unsloth/DeepSeek-V4-Flash-GGUF UD-IQ2_M` checkpoint mixes into its last routed-expert slabs (IQ2_S `ffn_gate/up` dotting Q8_K, MXFP4 `ffn_down` dotting Q8_0) — ported 1:1 from llama.cpp `ggml-quants.c` @ 237ad9b96 (`iq2s_grid` 1024-entry codebook + DIRECT sign bytes; MXFP4 `kvalues_mxfp4` + `e8m0_to_fp32_half` micro-scaling, distinct from the compressed-tensors `E8M0ToF32` NVFP4 path). CPU dequant + keep-quant `vec_dot`, unit-gated on hand-derived golden bytes (`test_gguf_dequant` 17/17), an INDEPENDENT f64 dequant-then-dot + GEMM NMSE (`test_ops_quant_dot` 19/19), and keep-quant routing (`test_gguf_keep_quant` 37/37) — all CPU-green, so UD-IQ2_M's four previously-`unsupported ggml type 22/39` slabs now load COMPRESSED (no ~17 GiB bf16 expansion that OOM-reboots the box). CUDA: the IQ2_S device `DotSuperblock` is wired into the Q8_K grouped-MoE GEMM and now **CUDA-BUILT + LINKED on GB10 (sm_121a, CUDA 13.0, `-Werror`, 2026-08-03 integration)** — it compiles clean and the merged binary links; MXFP4's device dot (`DotMXFP4`) is written but NOT wired (Q8_0-activation needs a separate 32-block GEMM) so it is marked `[[maybe_unused]]` to keep the ready math without tripping nvcc #177-D, and on GPU MXFP4 CPU-fallbacks like Q4_0/Q8_0. The V4-GGUF forward + a real UD-IQ2_M GPU load/coherence run are owed | | AWQ / GPTQ quantization | W0 spike + W1 CPU INT4 dequant primitive; not yet loadable end to end | INT4 unpack+dequant-to-bf16 for BOTH community formats, mirroring vLLM 1:1 (AWQ reverse-order `awq_triton.py`; GPTQ `qdq_4.cuh` with zero_offset v1/v2 + act-order g_idx). Unit-gated RED-first (hand-computed known bytes + double-precision roundtrip). NOT wired to a loader, no GPU Marlin compute, no model run yet: config recognizer (W2), Marlin GPU GEMM riding the vendored NVFP4 Marlin (W4), CPU e2e (W3), GPTQ 8/2/3-bit (W5) and MoE (W6) are named next bricks. See [.agents/specs/awq-gptq-quant.md](../.agents/specs/awq-gptq-quant.md) | -| MXFP4 (compressed-tensors `mxfp4-pack-quantized`) | Compute PROVEN (#38); GQA-swap ON (#49). **`MARLIN-STRUCT`: decode-graph (`VLLM_CPP_QWEN3_DENSE_DECODE_GRAPH`) + gate_up FUSION (NVFP4 to MXFP4) DEFAULT-ON.** Decode marlin 180 to 144 GEMM/step (vLLM-structural), byte-token-exact: 0.6B/4B 184/184 (graph ON==OFF), async 82/82, #44 3/3 (fused==split), 32B-NVFP4A16 142/142. nsys c8: marlin CTA count 144 vs vLLM 48 is the dominant residual (+1,177us; W4A16 memory-bound). `VT_MARLIN_E1_PAR1` opt-in clamps E=1 grid to 48 CTAs (near-parity, 8B token-exact) but flips a strict 32B-NVFP4 token so DEFAULT-OFF; byte-preserving route is the dense-template port (#50 NO-GO). Detail in state.md | Shared with DeepSeek-V4-Flash + Kimi-K3 MXFP4-expert paths. CPU E8M0 dequant unit-gated 5/5·1142. GPU W4A4 fp4 GEMM + MoE-expert e2e remain later bricks | +| MXFP4 (compressed-tensors `mxfp4-pack-quantized`) | Compute PROVEN (#38); GQA-swap ON (#49). **`MARLIN-STRUCT`: decode-graph (`VLLM_CPP_QWEN3_DENSE_DECODE_GRAPH`) + gate_up FUSION (NVFP4 to MXFP4) DEFAULT-ON.** Decode marlin 180->144 GEMM/step (vLLM-structural), byte-token-exact (0.6B/4B 184/184 graph ON==OFF, async 82/82, #44 3/3, 32B-NVFP4A16 142/142). nsys c8: marlin CTA 144 vs vLLM 48 dominates (+1,177us). `VT_MARLIN_E1_PAR1` clamps E=1->48 CTAs (near-parity, 8B token-exact) but flips a strict 32B token so OFF. **`KERNEL-MARLIN-DENSE-PORT` (gated OFF `VT_MARLIN_DENSE`): vLLM's OWN dense marlin = byte-preserving E=1; 3 dense `.cu` compile-clean dgx; unit battery WRITTEN; GPU exec/strict/nsys/binding PENDING.** Detail in state.md | Shared with DeepSeek-V4-Flash + Kimi-K3 MXFP4 paths. CPU E8M0 dequant 5/5·1142. GPU W4A4 + MoE-expert e2e later | | CPU backend vs llama.cpp | At or ahead on every axis (GGUF) | Prefill 1.18x ahead, decode at parity, peak memory 1.01x, byte-identical greedy tokens. Single-stream only; no concurrent-serving comparison has been measured | | Paged KV cache + prefix caching | Supported | Block-paged full attention, hybrid full-attention + GDN state groups, automatic prefix caching (APC) on by default for dense models (cache-ON gated end to end: token-identical output, cache hits, faster TTFT) | | fp8 KV cache (`cache_dtype=fp8`) | In progress (W1 CPU brick), not yet usable end-to-end | HIGH-priority memory/throughput lever (halves the KV footprint). W0 spike + W1 CPU brick landed (`KV-FP8` ACTIVE): fp8-e4m3 K/V STORE (`Quantize(hp/scale)`) + the paged-attention READ dequant (`Dequant(fp8)*scale`) + the `cache_dtype` config parse, all CPU-gated RED-first (`test_ops_fp8_kv_cache` 8/8·511; a wrong store direction fails 3/480). Storage is 1-byte fp8 (`DType::kI8`) + a `Fp8KVCacheDataType` interpretation enum, per-tensor k/v scales (mirroring vLLM `BaseKVCacheMethod`). The CUDA store + fp8 paged-attention read (the GPU memory-halving path, DGX-blocked), the runner/spec integration (half-sized KV blocks + checkpoint-scale threading + `--kv-cache-dtype`/`--calculate-kv-scales`), fp8_e5m2 and per-head scales are named W2-W5 in [.agents/specs/fp8-kv-cache.md](../.agents/specs/fp8-kv-cache.md). No model can run with an fp8 KV cache yet | diff --git a/include/vllm/model_executor/models/dense_nvfp4_gemm.h b/include/vllm/model_executor/models/dense_nvfp4_gemm.h index 2ec9ff395..d7bd43d41 100644 --- a/include/vllm/model_executor/models/dense_nvfp4_gemm.h +++ b/include/vllm/model_executor/models/dense_nvfp4_gemm.h @@ -103,6 +103,24 @@ inline bool FusedGateUpEnabled() { return on; } +// VT_MARLIN_DENSE (default OFF): route the E=1 dense NVFP4/MXFP4 projections through +// vLLM's OWN dense marlin GEMM (vt::MarlinDenseGemm) instead of the single-expert +// MoE-marlin route. The dense kernel is direct-A + tile-per-CTA with vLLM's dense +// fp32-C_tmp reduce, so at M<=8 it naturally runs the 48-CTA (sms-wide) grid the MoE +// path only reaches with the VT_MARLIN_E1_PAR1 clamp — WITHOUT that clamp's par +// regrouping, which costs one bf16 ULP vs the oracle and flips a strict 32B token +// (row QUANT-CT-MXFP4-MARLIN-STRUCT / #50 / #54). Same resident weights + workspace; +// the repack permute is vLLM's shared marlin_permute for both dense and MoE. Default +// OFF until the strict-gate battery + binding prove it byte-matches the oracle +// everywhere and beats the MoE route; then flipped ON per parity-enablers. +inline bool MarlinDenseEnabled() { + static const bool on = [] { + const char* e = std::getenv("VT_MARLIN_DENSE"); + return e != nullptr && e[0] == '1'; + }(); + return on; +} + // --- Execution counters (the "this path actually RAN" positive signal) ------ // A passing correctness gate does NOT prove a new code path was exercised — a // mis-wired dispatch that silently fell back to the BF16 arm would also pass if @@ -113,6 +131,7 @@ struct Nvfp4W4A16Stats { uint64_t marlin_gemms = 0; // MatmulNvfp4MarlinD launches uint64_t fused_gate_up = 0; // GateUpFusedMarlinD launches (one per MLP) uint64_t fallback_gemms = 0; // naive vt::MatmulNvfp4 / CPU dequant launches + uint64_t dense_gemms = 0; // vt::MarlinDenseGemm launches (VT_MARLIN_DENSE route) }; inline Nvfp4W4A16Stats& MutableW4A16Stats() { @@ -336,9 +355,33 @@ inline DBuf MatmulNvfp4MarlinD(Dev d, const Tensor& x, const Nvfp4Weight& w, const int64_t M = x.shape[0], K = x.shape[1], N = w.n; MarlinDenseResident& mr = MarlinDenseResidentFor(&w); if (!mr.ready) BuildMarlinDenseResident(d, w, mr); - DenseAlignCache& ac = DenseAlignFor(d, static_cast(M)); int sms = 0; void* ws = DenseMarlinWorkspace(d, &sms); // zeroed once; kernel self-resets + + // VT_MARLIN_DENSE (default OFF): route through vLLM's OWN dense marlin GEMM. + // Same resident (mr.w/mr.s/mr.g) + workspace; rank-2 operand views (the dense + // launcher wants [K/16, N*2] / [K/gs, N], not the MoE rank-3 [1, ...]); NO + // moe_align cache (direct-A). Byte-preserving vs the oracle (its own dense + // fp32-C_tmp reduce). Only when the op is realized for this device. + if (MarlinDenseEnabled() && + vt::OpRegistered(vt::OpId::kMarlinDenseGemm, d.q.device.type)) { + ++MutableW4A16Stats().dense_gemms; + DBuf outbf(d, DType::kBF16, {M, N}); + Tensor wqd = MakeTensor(mr.w, DType::kI32, d.q.device, {K / 16, N * 2}); + Tensor scd = MakeTensor(mr.s, DType::kI8, d.q.device, {K / w.group_size, N}); + Tensor ggd = MakeTensor(mr.g, DType::kF32, d.q.device, {1}); + Tensor wstd = MakeTensor(ws, DType::kI32, d.q.device, {sms * 4}); + vt::MarlinDenseArgs dargs{static_cast(M), static_cast(N), static_cast(K)}; + dargs.group_size = static_cast(w.group_size); + dargs.mxfp4 = w.is_mxfp4; + vt::MarlinDenseGemm(d.q, outbf.t(), x, wqd, scd, ggd, wstd, dargs); + if (out_dtype == DType::kBF16) return outbf; + DBuf out(d, DType::kF32, {M, N}); + vt::CastF32(d.q, out.t(), outbf.t()); + return out; + } + + DenseAlignCache& ac = DenseAlignFor(d, static_cast(M)); ++MutableW4A16Stats().marlin_gemms; // Marlin's output is bf16 (c_type=kBFloat16); an f32 result is the bf16 output @@ -473,9 +516,31 @@ inline DBuf GateUpFusedMarlinD(Dev d, const Tensor& x, const Nvfp4Weight& gw, const int64_t M = x.shape[0], K = x.shape[1], N = gw.n; MarlinDensePairResident& mr = MarlinDensePairResidentFor(&gw); if (!mr.ready) BuildMarlinDensePairResident(d, gw, uw, mr); - DenseAlignCache& ac = DenseAlignFor(d, static_cast(M)); int sms = 0; void* ws = DenseMarlinWorkspace(d, &sms); // zeroed once; kernel self-resets + + // VT_MARLIN_DENSE (default OFF): fused gate_up over the 2N-concatenated operand + // via vLLM's OWN dense marlin GEMM. Same merged resident (mr.w/mr.s/mr.g), rank-2 + // views, no moe_align; byte-preserving reduce. Only when the op is realized here. + if (MarlinDenseEnabled() && + vt::OpRegistered(vt::OpId::kMarlinDenseGemm, d.q.device.type)) { + ++MutableW4A16Stats().dense_gemms; + DBuf gud(d, DType::kBF16, {M, 2 * N}); + Tensor wqd = MakeTensor(mr.w, DType::kI32, d.q.device, {K / 16, 2 * N * 2}); + Tensor scd = MakeTensor(mr.s, DType::kI8, d.q.device, {K / gw.group_size, 2 * N}); + Tensor ggd = MakeTensor(mr.g, DType::kF32, d.q.device, {1}); + Tensor wstd = MakeTensor(ws, DType::kI32, d.q.device, {sms * 4}); + vt::MarlinDenseArgs dargs{static_cast(M), static_cast(2 * N), + static_cast(K)}; + dargs.group_size = static_cast(gw.group_size); + dargs.mxfp4 = gw.is_mxfp4; + vt::MarlinDenseGemm(d.q, gud.t(), x, wqd, scd, ggd, wstd, dargs); + DBuf actd(d, DType::kBF16, {M, N}); + vt::SiluAndMul(d.q, actd.t(), gud.t()); + return actd; + } + + DenseAlignCache& ac = DenseAlignFor(d, static_cast(M)); ++MutableW4A16Stats().fused_gate_up; DBuf gu(d, DType::kBF16, {M, 2 * N}); diff --git a/include/vt/ops.h b/include/vt/ops.h index 6ac0815f3..3d21daf9c 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -331,6 +331,13 @@ enum class OpId : uint8_t { // BYTE-EXACT (sequential reductions) to the host Laguna forward. Additive: only // LagunaForwardResidentDecode dispatches it. Appended before kCount (no id shift). kLaguna, + // DENSE Marlin W4A16 GEMM (lift of vLLM's own dense marlin.cu marlin_gemm; see + // MarlinDenseGemm below). Byte-preserving replacement for the single-expert + // MoE-marlin route the dense E=1 NVFP4/MXFP4 projections use today — direct-A, + // tile-per-CTA, vLLM's own dense fp32-C_tmp reduce (no par regrouping ULP). + // CUDA-only (Blackwell sm_12xa; vendored dense marlin TUs, VT_MARLIN_NVFP4). + // Appended before kCount (no existing op's id shifts). + kMarlinDenseGemm, kCount }; @@ -738,6 +745,25 @@ struct MoeMarlinArgs { using MoeGroupedGemmNvfp4MarlinFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, const MoeMarlinArgs&); +// DENSE Marlin W4A16 GEMM (lift of vLLM's own dense marlin_gemm; see +// MarlinDenseGemm below). Scalar params travel in MarlinDenseArgs. Unlike the +// MoE path there is NO moe_align gather (sorted_token_ids/expert_ids/top_k): +// `a` is a plain [size_m, size_k] contiguous activation (lda = size_k). +struct MarlinDenseArgs { + int size_m = 0; // number of tokens (rows of `a`) + int size_n = 0; // output features + int size_k = 0; // input features (contraction; multiple of 16) + // Block-scale format selector, identical semantics to MoeMarlinArgs: default = + // NVFP4 (fp8-e4m3 scales, group 16, per-tensor global scale). group_size 32 + + // mxfp4=true selects the MXFP4 path (E8M0 scales, group_blocks 2, NO global + // scale; the `global_scale` tensor is ignored). Mirrors vLLM's is_nvfp4 branch. + int group_size = 16; + bool mxfp4 = false; +}; +using MarlinDenseGemmFn = + void (*)(Queue&, Tensor& /*c*/, const Tensor& /*a*/, const Tensor& /*b_q_weight*/, + const Tensor& /*b_scales*/, const Tensor& /*global_scale*/, Tensor& /*workspace*/, + const MarlinDenseArgs&); using MoeSiluMulFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&); // --- Qwen3.6 elementwise "glue" ops (M0.9 forward). These replace host-side // loops so the decode step can run entirely on-device (CUDA-graph capture). @@ -1356,6 +1382,24 @@ void MoeGroupedGemmNvfp4Marlin(Queue& q, Tensor& c, const Tensor& a, const Tenso const Tensor& expert_ids, const Tensor& num_tokens_past_padded, const Tensor& topk_weights, const MoeMarlinArgs& args); +// MarlinDenseGemm (lift of vLLM's DENSE marlin_gemm, marlin.cu:545 -> marlin_mm +// at :326 — the byte-preserving dense W4A16 kernel vLLM itself ships for a16 +// weight-only linears). One launch computes y = a @ dequant(b).T with vLLM's own +// direct-A, tile-per-CTA layout and dense fp32-C_tmp reduce — NOT the MoE +// single-expert route, whose par regrouping of the reduce costs one bf16 ULP. +// c [size_m, size_n] bf16 (out) +// a [size_m, size_k] bf16 (token hidden; contiguous, lda=size_k) +// b_q_weight [size_k/16, size_n*8/pack] i32 — Marlin-interleaved fp4 (SAME +// repack as the MoE path: marlin_permute; a shim is added only if +// a layout divergence is proven — see dense_nvfp4_gemm.h) +// b_scales [size_k/group_size, size_n] fp8 (processed marlin scales) +// global_scale [1] f32 (nvfp4 only; ignored for mxfp4) +// workspace [>= sms] i32 (zeroed reduction locks) +// CUDA-only (Blackwell sm_12xa; needs the vendored dense Marlin TUs, VT_MARLIN_NVFP4). +void MarlinDenseGemm(Queue& q, Tensor& c, const Tensor& a, const Tensor& b_q_weight, + const Tensor& b_scales, const Tensor& global_scale, Tensor& workspace, + const MarlinDenseArgs& args); + // out[R,I] = silu(gate[R,I]) * up[R,I] (moe-semantics.md §4; the fused-MoE // element-wise activation between the grouped gate/up and down GEMMs). gate/up // f32 or bf16, out f32/bf16; silu/mul computed in f32, rounded on store. Unlike diff --git a/src/vt/cuda/cuda_marlin_dense.cu b/src/vt/cuda/cuda_marlin_dense.cu new file mode 100644 index 000000000..05fd0b60b --- /dev/null +++ b/src/vt/cuda/cuda_marlin_dense.cu @@ -0,0 +1,159 @@ +// vllm.cpp — DENSE Marlin NVFP4/MXFP4 W4A16 GEMM drop-in (vt::Tensor launcher). +// +// Torch-free host launcher for the vendored DENSE Marlin kernel (src/vt/cuda/ +// marlin/libtorch_stable/quantization/marlin/, a 1:1 lift of vLLM's dense +// marlin.cu @ 555967922). It mirrors the a16 (weight-only) branch of vLLM's +// `marlin_gemm` (marlin.cu:545): b_type=kFE2M1f + s_type=kFE4M3fn (NVFP4) or +// s_type=kFE8M0fnu (MXFP4), bf16 activation/output, no act-order/zero-point/bias. +// All those irrelevant branches are dropped; the compute call into +// marlin::marlin_mm is the verbatim vendored dispatcher (marlin_mm_dense.cu). +// +// This is the BYTE-PRESERVING replacement for the single-expert MoE-marlin route +// the dense E=1 projections use today (dense_nvfp4_gemm.h): the DENSE kernel's +// direct-A, tile-per-CTA grid + its own dense fp32 C_tmp reduce ARE vLLM's own +// numerics, so it does not incur the one-bf16-ULP shift the MoE par regrouping +// does (row QUANT-CT-MXFP4-MARLIN-STRUCT / #50 / #54). +// +// Weights MUST be pre-repacked into Marlin's interleaved layout with processed +// fp8 block scales + per-tensor global scale — the SAME resident the MoE route +// builds (dense_nvfp4_gemm.h BuildMarlinDenseResident); the repack permute is +// vLLM's shared marlin_permute for both dense and MoE. +// +// Isolated TU (heavy templated kernel). Gated by VT_MARLIN_NVFP4. + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "core/scalar_type.hpp" +#include "libtorch_stable/quantization/marlin/marlin_mm_dense.h" + +#include "vt/cuda/graph_safe_scratch.h" +#include "vt/ops.h" + +namespace vt::cuda { +namespace { + +// max_thread_n from the vendored marlin.cuh:28 (dense C_tmp reduce scratch upper +// bound; vLLM marlin.cu:716 sizes c_tmp as sms * max_m_block_size * max_thread_n). +constexpr int kMarlinMaxThreadN = 256; + +void Check(cudaError_t err, const char* what) { + if (err != cudaSuccess) { + throw std::runtime_error(std::string("vt cuda: marlin_dense: ") + what + ": " + + cudaGetErrorString(err)); + } +} + +cudaStream_t AsStream(const Queue& q) { return static_cast(q.handle); } + +// Persistent per-stream C_tmp workspace pool (VT_MARLIN_WS_POOL, default ON). +// Same rationale as cuda_moe_marlin.cu: vLLM allocates c_tmp per call through +// PyTorch's CACHING allocator (a cheap pool hit, not a raw cudaMalloc); a raw +// per-GEMM cudaMallocAsync/cudaFreeAsync serializes on the forward host thread +// and is a steady-state decode idle. This mirrors the caching allocator with a +// grown-on-demand per-stream buffer. c_tmp is scratch the kernel fully writes +// before it reads (vLLM uses new_empty; no zero-on-entry invariant), so reuse is +// race-free under the forward's single-stream ordering. The RETIRE-not-free on +// regrow keeps a captured decode graph's baked c_tmp pointer valid across a later +// larger forward (graph_safe_scratch.h). VT_MARLIN_WS_POOL=0 restores per-call. +bool MarlinWsPoolEnabled() { + static const bool on = [] { + const char* e = std::getenv("VT_MARLIN_WS_POOL"); + return !(e != nullptr && e[0] == '0'); + }(); + return on; +} + +float* EnsureCtmp(cudaStream_t s, size_t bytes) { + struct Scratch { + void* p = nullptr; + size_t cap = 0; + }; + static std::mutex mu; + static std::unordered_map pool; + std::lock_guard lk(mu); + Scratch& sc = pool[s]; + if (bytes > sc.cap) { + RetireGraphScratch(sc.p); + Check(cudaMallocAsync(&sc.p, bytes, s), "cudaMallocAsync c_tmp (pool)"); + sc.cap = bytes; + } + return static_cast(sc.p); +} + +// vt::MarlinDenseGemm registered kernel. +void MarlinDenseGemmKernelCuda(Queue& q, Tensor& c, const Tensor& a, const Tensor& b_q_weight, + const Tensor& b_scales, const Tensor& global_scale, + Tensor& workspace, const MarlinDenseArgs& args) { + cudaStream_t s = AsStream(q); + const int dev = q.device.index; + + // NVFP4 W4A16, bf16 activation/output; OR MXFP4 W4A16 when args.mxfp4 (E8M0 + // scales => s_type kFE8M0fnu, group_size 32 => group_blocks 2, NO global scale). + const vllm::ScalarType a_type = vllm::kBFloat16; + const vllm::ScalarType b_type = vllm::kFE2M1f; + const vllm::ScalarType c_type = vllm::kBFloat16; + const vllm::ScalarType s_type = args.mxfp4 ? vllm::kFE8M0fnu : vllm::kFE4M3fn; + + const int size_m = args.size_m; + const int size_n = args.size_n; + const int size_k = args.size_k; + const int group_size = args.group_size; // 16 (nvfp4) or 32 (mxfp4) + const int num_groups = size_k / group_size; + // MXFP4 has NO global scale — the dense kernel reads global_scale_ptr only under + // (b_type==kFE2M1f && s_type==kFE4M3fn), so pass nullptr on the mxfp4 path. + void* global_scale_ptr = args.mxfp4 ? nullptr : global_scale.data; + + int sms = -1; + Check(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev), + "cudaDeviceGetAttribute(sms)"); + + // C_tmp for the fp32 global reduce (use_fp32_reduce && !use_atomic_add). Size + // per vLLM marlin.cu:713-716: sms * min(ceil(size_m/16)*16, 64) * max_thread_n. + const bool use_atomic_add = false; + const bool use_fp32_reduce = true; + int max_m_block_size = (size_m + 16 - 1) / 16 * 16; + if (max_m_block_size > 64) max_m_block_size = 64; + const int64_t c_tmp_elems = + static_cast(sms) * max_m_block_size * kMarlinMaxThreadN; + const size_t c_tmp_bytes = static_cast(c_tmp_elems) * sizeof(float); + float* c_tmp = nullptr; + bool c_tmp_pooled = false; + if (MarlinWsPoolEnabled()) { + c_tmp = EnsureCtmp(s, c_tmp_bytes); // persistent per-stream, reused + c_tmp_pooled = true; + } else { + Check(cudaMallocAsync(&c_tmp, c_tmp_bytes, s), "cudaMallocAsync c_tmp"); + } + + // lda = A.stride(0) = size_k (a is [size_m, size_k] contiguous, row-major). + marlin::marlin_mm( + a.data, b_q_weight.data, c.data, c_tmp, /*b_bias=*/nullptr, /*a_s=*/nullptr, + b_scales.data, global_scale_ptr, /*zp=*/nullptr, /*g_idx=*/nullptr, /*perm=*/nullptr, + /*a_tmp=*/nullptr, size_m, size_n, size_k, /*lda=*/size_k, workspace.data, a_type, b_type, + c_type, s_type, /*has_bias=*/false, /*has_act_order=*/false, /*is_k_full=*/true, + /*has_zp=*/false, num_groups, group_size, dev, s, /*thread_k=*/-1, /*thread_n=*/-1, sms, + use_atomic_add, use_fp32_reduce, /*is_zp_float=*/false); + + if (c_tmp && !c_tmp_pooled) Check(cudaFreeAsync(c_tmp, s), "cudaFreeAsync c_tmp"); + Check(cudaGetLastError(), "marlin_dense marlin_mm launch"); +} + +struct Registrar { + Registrar() { + RegisterOp(OpId::kMarlinDenseGemm, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&MarlinDenseGemmKernelCuda))); + } +}; +Registrar g_registrar; + +} // namespace +} // namespace vt::cuda diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel.h b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel.h new file mode 100644 index 000000000..8c9cec88b --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel.h @@ -0,0 +1,43 @@ + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin +#endif + +#include "marlin.cuh" +#include "marlin_dtypes.cuh" +#include "core/scalar_type.hpp" + +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, \ + int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, \ + const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, \ + const float *__restrict__ global_scale_ptr, \ + const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ + int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ + bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \ + int max_shared_mem + +namespace MARLIN_NAMESPACE_NAME { +template shared + // fetch pipeline + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin(MARLIN_KERNEL_PARAMS); + +} diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel_selector.h b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel_selector.h new file mode 100644 index 000000000..0f8c19390 --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/kernel_selector.h @@ -0,0 +1,62 @@ +// auto generated by generate_kernels.py +// clang-format off +if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE8M0fnu && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu new file mode 100644 index 000000000..f375eecf9 --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.cu @@ -0,0 +1,528 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +// [vt lift] 1:1 lift of vLLM csrc/libtorch_stable/quantization/marlin/marlin.cu +// (@ 555967922) — the DENSE marlin_mm dispatcher + its config helpers. The +// torch::stable `marlin_gemm` host wrapper (marlin.cu:545-894) is removed; the +// vt::Tensor launcher (src/vt/cuda/cuda_marlin_dense.cu) does that job. The +// `marlin::marlin_mm` compute path below (marlin.cu:326-541) is kept verbatim. +// STD_TORCH_CHECK resolves to the torch-free throwing shim (vt_marlin_check.h, +// pulled in transitively via core/scalar_type.hpp) — the ONLY torch coupling. + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin +#endif + +#include "kernel.h" + +#include +#include +#include +#include +// [vt lift] torch host launcher removed; marlin_mm dispatcher kept verbatim. +// See src/vt/cuda/cuda_marlin_dense.cu for the vt::Tensor launcher. + +// [vt lift] removed: #include +// [vt lift] removed: #include +// [vt lift] removed: #include +// [vt lift] removed: #include +// [vt lift] removed: #include +// [vt lift] removed: #include +// [vt lift] removed: #include "libtorch_stable/torch_utils.h" + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +namespace MARLIN_NAMESPACE_NAME { + +__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; + +using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); + +// For a given "a" of size [M,K] performs a permutation of the K columns based +// on the given "perm" indices. (act_order only; unused on the W4A16 dense path +// but kept for verbatim fidelity with vLLM's marlin_mm.) +__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int lda, int block_rows) { + auto start_row = block_rows * blockIdx.x; + int finish_row = start_row + block_rows; + if (finish_row > size_m) { + finish_row = size_m; + } + int cur_block_rows = finish_row - start_row; + + int input_row_stride = lda * sizeof(half) / 16; + int output_row_stride = size_k * sizeof(half) / 16; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int input_offset = row * input_row_stride; + int output_offset = row * output_row_stride; + + half const* a_row_half = + reinterpret_cast(a_int4_ptr + input_offset); + half* out_half = reinterpret_cast(out_int4_ptr + output_offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + auto cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + auto cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int i = 0; i < cur_block_rows; i++) { + int cur_row = start_row + i; + if (cur_row < size_m) { + permute_row(cur_row); + } + } +} + +typedef struct { + int thread_k; + int thread_n; + int num_threads; +} thread_config_t; + +thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}}; + +thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + {64, 128, 128}, + {128, 64, 128}}; + +typedef struct { + int blocks_per_sm; + thread_config_t tb_cfg; +} exec_config_t; + +int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full, int stages) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + } else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * stages; + } +} + +int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int has_zp, bool is_zp_float, bool is_a_8bit, + int stages) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + int tb_m = thread_m_blocks * 16; + int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2); + int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; + int sh_red_size = tb_m * (tb_n + 8) * 2; + int sh_bias_size = tb_n * 2; + int tmp_size = + (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; + tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); + + int sh_s_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full, stages); + int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; + int sh_zp_size = 0; + if (has_zp) { + if (is_zp_float) + sh_zp_size = sh_s_size; + else if (num_bits == 4) + sh_zp_size = sh_s_size / 4; + else if (num_bits == 8) + sh_zp_size = sh_s_size / 2; + } + + int total_size = + tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; + + return total_size; +} + +bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int has_zp, bool is_zp_float, bool is_a_8bit, int stages, + int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Check that pipeline fits into cache + int cache_size = get_kernel_cache_size( + th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + return cache_size <= max_shared_mem; +} + +MarlinFuncPtr get_marlin_kernel( + const vllm::ScalarType a_type, const vllm::ScalarType b_type, + const vllm::ScalarType c_type, const vllm::ScalarType s_type, + int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, + bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, + int threads, bool is_zp_float, int stages) { + int num_bits = b_type.size_bits(); + auto kernel = MarlinDefault; + +#include "kernel_selector.h" + + return kernel; +} + +exec_config_t determine_exec_config( + const vllm::ScalarType& a_type, const vllm::ScalarType& b_type, + const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m, + int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, + int num_bits, int group_size, bool has_act_order, bool is_k_full, + bool has_zp, bool is_zp_float, int is_a_8bit, int stages, + int max_shared_mem, int sms) { + exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; + thread_config_t* thread_configs = thread_m_blocks > 1 + ? large_batch_thread_configs + : small_batch_thread_configs; + int thread_configs_size = + thread_m_blocks > 1 + ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) + : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); + + for (int i = 0; i < thread_configs_size; i++) { + thread_config_t th_config = thread_configs[i]; + + if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, has_zp, + is_zp_float, is_a_8bit, stages, + max_shared_mem - 512)) { + continue; + } + + int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, has_zp, + is_zp_float, is_a_8bit, stages); + + int group_blocks = 0; + if (!has_act_order) { + group_blocks = group_size == -1 ? -1 : group_size / 16; + } + + auto kernel = + get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, + th_config.thread_n / 16, th_config.thread_k / 16, + m_block_size_8, has_act_order, has_zp, group_blocks, + th_config.num_threads, is_zp_float, stages); + + if (kernel == MarlinDefault) continue; + + return {1, th_config}; + } + + return exec_cfg; +} + +void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, + void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k, + int lda, void* workspace, vllm::ScalarType const& a_type, + vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, + vllm::ScalarType const& s_type, bool has_bias, + bool has_act_order, bool is_k_full, bool has_zp, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k_init, + int thread_n_init, int sms, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { + bool is_a_8bit = a_type.size_bits() == 8; + STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + STD_TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } else { + STD_TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + } else { + if (group_size == -1) { + group_blocks = -1; + } else { + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + int num_bits = b_type.size_bits(); + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + int4* C_tmp_ptr = (int4*)C_tmp; + + const int4* bias_ptr = (const int4*)b_bias; + const float* a_s_ptr = (const float*)a_s; + const int4* b_s_ptr = (const int4*)b_s; + const float* g_s_ptr = (const float*)g_s; + + const int4* zp_ptr = (const int4*)zp; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + int block_rows = div_ceil(prob_m, sms); + // avoid ">>>" being formatted to "> > >" + // clang-format off + permute_cols_kernel<<>>( + A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); + // clang-format on + A_ptr = a_tmp_ptr; + lda = prob_k; + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by + // having a full K, we have full original groups) + if (is_k_full) has_act_order = false; + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + STD_TORCH_CHECK(max_shared_mem > 0); + + int major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + dev); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + dev); + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); + int stages = 4; + if (major_capability == 7 && minor_capability == 5) { + stages = 2; + STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); + } + if (a_type == vllm::kFE4M3fn) { + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); + STD_TORCH_CHECK( + major_capability * 10 + minor_capability == 89 || + major_capability == 12, + "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " + "Marlin W4A16 on other devices)."); + } + + int max_par = 16; + if (prob_n <= 4096) max_par = 16 * 8; + int max_shared_mem_new = max_shared_mem; + int rest_m = prob_m; + int max_thread_m_blocks = 4; + while (rest_m) { + int par_count = rest_m / (max_thread_m_blocks * 16); + if (par_count > max_par) par_count = max_par; + int prob_m_split = + par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; + + int thread_k = thread_k_init; + int thread_n = thread_n_init; + + int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); + int m_block_size_8 = prob_m_split <= 8 && a_type.size_bits() == 16; + + // Set thread config + exec_config_t exec_cfg; + thread_config_t thread_tfg; + if (thread_k != -1 && thread_n != -1) { + thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; + exec_cfg = exec_config_t{1, thread_tfg}; + STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + } else { + // Auto config + exec_cfg = determine_exec_config( + a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k, + thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem, + sms); + thread_tfg = exec_cfg.tb_cfg; + if (thread_tfg.thread_n != -1) { + if (prob_n / thread_tfg.thread_n * + div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= + sms) { + if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split, + prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, has_zp, is_zp_float, + is_a_8bit, stages, max_shared_mem_new)) { + thread_tfg = {128, 64, 128}; + exec_cfg = {1, thread_tfg}; + } + } + } + + if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { + max_thread_m_blocks--; + continue; + } + } + + int num_threads = thread_tfg.num_threads; + thread_k = thread_tfg.thread_k; + thread_n = thread_tfg.thread_n; + int blocks = sms * exec_cfg.blocks_per_sm; + if (exec_cfg.blocks_per_sm > 1) + max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + STD_TORCH_CHECK( + is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, + prob_k, num_bits, group_size, has_act_order, is_k_full, + has_zp, is_zp_float, is_a_8bit, stages, + max_shared_mem_new), + "Invalid thread config: thread_m_blocks = ", thread_m_blocks, + ", thread_k = ", thread_tfg.thread_k, + ", thread_n = ", thread_tfg.thread_n, + ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, + ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", prob_m_split = ", prob_m_split, ", group_size = ", group_size, + ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, + ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, + ", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new); + + auto kernel = get_marlin_kernel( + a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, + thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, + num_threads, is_zp_float, stages); + + if (kernel == MarlinDefault) { + STD_TORCH_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, + ", prob_m_split = ", prob_m_split, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, + ", num_threads = ", num_threads, ", num_bits = ", num_bits); + } + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_shared_mem_new); + + bool part_use_atomic_add = + use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; + + // avoid ">>>" being formatted to "> > >" + // clang-format off + kernel<<>>( + A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, + g_idx_ptr, num_groups, + prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add, + use_fp32_reduce, max_shared_mem_new); + // clang-format on + + A_ptr += prob_m_split * (lda / (is_a_8bit ? 16 : 8)); + a_s_ptr += prob_m_split; + C_ptr += prob_m_split * (prob_n / 8); + rest_m -= prob_m_split; + } +} + +} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.h b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.h new file mode 100644 index 000000000..d7c118110 --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_mm_dense.h @@ -0,0 +1,36 @@ +// vt lift — declaration of the vendored DENSE Marlin dispatcher (marlin::marlin_mm), +// defined in marlin_mm_dense.cu (= vLLM csrc/libtorch_stable/quantization/marlin/ +// marlin.cu:326-541 `marlin::marlin_mm`, with the torch::stable `marlin_gemm` +// host wrapper stripped). The vt::Tensor launcher (src/vt/cuda/cuda_marlin_dense.cu) +// calls this directly. +// +// This is the DIRECT-A, tile-per-CTA dense GEMM — vLLM's OWN dense W4A16 numerics, +// distinct from the moe/marlin_moe_wna16 dispatcher (marlin_mm.h): no +// sorted_token_ids / expert_ids / top_k gather, an `lda` (A.stride(0)) parameter, +// and its own par-split fp32 C_tmp reduce structure. It is the byte-preserving +// replacement for the single-expert MoE-marlin route the dense E=1 projections use +// today (dense_nvfp4_gemm.h), whose par regrouping of the fp32 C_tmp reduce costs +// one bf16 ULP vs the oracle (row QUANT-CT-MXFP4-MARLIN-STRUCT / #50 / #54). +#pragma once + +#include + +#include "core/scalar_type.hpp" + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin +#endif + +namespace MARLIN_NAMESPACE_NAME { + +void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, + void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k, + int lda, void* workspace, vllm::ScalarType const& a_type, + vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, + vllm::ScalarType const& s_type, bool has_bias, bool has_act_order, + bool is_k_full, bool has_zp, int num_groups, int group_size, + int dev, cudaStream_t stream, int thread_k, int thread_n, int sms, + bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); + +} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_template.h b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_template.h new file mode 100644 index 000000000..32b8f8bdd --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/marlin_template.h @@ -0,0 +1,2081 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin +#endif + +#include "marlin.cuh" +#include "marlin_dtypes.cuh" +#include "dequant.h" +#include "marlin_mma.h" +#include "core/scalar_type.hpp" + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +namespace MARLIN_NAMESPACE_NAME { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks, // extra global storage for barrier synchronization + bool use_fp32_reduce // whether to use fp32 global reduce +) {} + +} // namespace marlin + +#else + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm(typename MarlinScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (count == 4) { + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } else if constexpr (count == 2) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" + : "=r"(a[0]), "=r"(a[1]) + : "r"(smem)); + } else if constexpr (count == 1) { + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(a[0]) + : "r"(smem)); + } else { + static_assert(count == 1 || count == 2 || count == 4, "invalid count"); + } +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale(typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::FragS& frag_s, + int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 s = MarlinScalarType::num2num2( + reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +template +__device__ inline void scale_and_sub( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::scalar_t s, + typename MarlinScalarType::scalar_t zp) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 s2 = MarlinScalarType::num2num2(s); + scalar_t2 zp2 = MarlinScalarType::num2num2(zp); + frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); + frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); +} + +template +__device__ inline void sub_zp( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::scalar_t2& frag_zp, int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 zp = MarlinScalarType::num2num2( + reinterpret_cast(&frag_zp)[i]); + frag_b[0] = __hsub2(frag_b[0], zp); + frag_b[1] = __hsub2(frag_b[1], zp); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::FragS& frag_s_1, + typename MarlinScalarType::FragS& frag_s_2, + typename MarlinScalarType::FragS& frag_s_3, + typename MarlinScalarType::FragS& frag_s_4, int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float( + float* c, typename MarlinScalarType::FragS& s) { + using scalar_t = typename MarlinScalarType::scalar_t; + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], MarlinScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], MarlinScalarType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } +} + +// Wait until value of lock to be negative, and then add 1 +__device__ inline void wait_negative_and_add(int* lock) { + if (threadIdx.x == 0) { + int state = 0; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state >= 0); + atomicAdd(lock, 1); + } + __syncthreads(); +} + +template shared + // fetch pipeline + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin( + const int4* __restrict__ A0, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C0, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ b_bias_ptr, + // float scales of input matrix, only used when is_a_8bit == true. + // shape (m,) + const float* __restrict__ a_scales_ptr, + // fp16 quantization scales. shape (k/groupsize, n) + const int4* __restrict__ scales_ptr, + // float global scale (for nvfp4// only) + const float* __restrict__ global_scale_ptr, + // 4bit packed zero-points of shape + // (k/groupsize, n/pack_factor) + const int4* __restrict__ zp_ptr, + // int32 group indices of shape k + const int* __restrict__ g_idx, + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int lda, // A.stride(0), equal to prob_k is A is contiguous + int* locks, // extra global storage for barrier synchronization + bool has_bias, + bool use_atomic_add, // whether to use atomic add to reduce + bool use_fp32_reduce, // whether to use fp32 global reduce + int max_shared_mem) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890 + // FP8 computation is only supported for Ada Lovelace or newer architectures. + if constexpr (a_type_id == vllm::kFE4M3fn.id()) return; + #endif + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + // Turing TensorCore only supports fp16 and int8 + if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id()) + return; + #endif + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits(); + // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 && + // num_bits == 4 + constexpr bool use_fp16_accum = + a_type_id == vllm::kFloat16.id() && + (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) && + !(group_blocks == -1 && num_bits == 4)); + #else + constexpr bool use_fp16_accum = false; + #endif + using Adtype = MarlinScalarType; + using Cdtype = MarlinScalarType; + const int4* A = A0; + int4* C = C0; + + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + using scalar_32bit_t = typename MarlinScalarType::scalar_32bit_t; + + using c_scalar_t = typename MarlinScalarType::scalar_t; + using c_scalar_t2 = typename MarlinScalarType::scalar_t2; + + using FragA = typename MarlinScalarType::FragA; + using FragB = typename MarlinScalarType::FragB; + using FragC = typename MarlinScalarType::FragC; + using FragS = typename MarlinScalarType::FragS; + using FragZP = typename MarlinScalarType::FragZP; + + static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id); + static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id); + static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id); + static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id); + if constexpr (b_type == vllm::kFE2M1f) { + static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || + s_type == vllm::kFE8M0fnu && group_blocks == 2); + } else if constexpr (s_type == vllm::kFE8M0fnu) { + // MXFP8: FP8 weights with e8m0 microscaling block scales + static_assert(b_type == vllm::kFE4M3fn && group_blocks == 2); + } else if constexpr (std::is_same::value) { + static_assert(s_type == vllm::kBFloat16); + } else if constexpr (std::is_same::value) { + static_assert(s_type == vllm::kFloat16); + } + + constexpr bool is_a_8bit = a_type.size_bits() == 8; + constexpr bool is_8bit_scale = s_type.size_bits() == 8; + if constexpr (!is_a_8bit) { + static_assert(std::is_same::value); + } + constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8; + constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 || + b_type == vllm::kS4 || b_type == vllm::kS8 || + b_type == vllm::kU4B8 || b_type == vllm::kU8B128; + // see comments of dequant.h for more details + constexpr bool dequant_skip_flop = + is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) || + b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || + has_zp && !is_zp_float && !std::is_same::value || + has_zp && !is_zp_float && !(b_type == vllm::kU8); + + float global_scale_f32 = 1.0f; + + if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { + global_scale_f32 = global_scale_ptr[0]; + } + + constexpr bool has_act_order = group_blocks == 0; + constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); + + extern __shared__ int4 sh[]; + float* sh_a_s = reinterpret_cast(sh); + int4* sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); + constexpr int pack_factor = 32 / b_type.size_bits(); + static_assert(thread_m_blocks == 1 || !m_block_size_8); + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > m_block_size) { + parallel = prob_m / m_block_size; + prob_m = m_block_size; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + + int global_mn_tiles = parallel * n_tiles; + int part2_mn_tiles = global_mn_tiles; + int part1_mn_iters = 0; + bool in_part2 = false; + + if (global_mn_tiles > gridDim.x) { + part2_mn_tiles = global_mn_tiles % gridDim.x; + if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x; + part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; + } + + int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = 0; + int slice_col_par = blockIdx.x; + int slice_col; + int slice_iters = + k_tiles; // number of threadblock tiles in the current slice + // total number of active threadblocks in the current slice + int slice_count = 1; + // index of threadblock in current slice; numbered bottom to top + int slice_idx = 0; + + int par_id = 0; + int locks_off = 0; + + if (part2_mn_tiles >= gridDim.x) { + // when part2_mn_tiles >= sms + // then there are at most $sms$ conflict tile blocks + locks_off = blockIdx.x; + } else { + locks_off = (iters * blockIdx.x) / k_tiles - 1; + } + + // Compute all information about the current slice which is required for + // synchronization. + bool first_init = true; + auto init_part2_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (part2_mn_tiles >= gridDim.x) { + if (slice_count > 1 && slice_idx == slice_count - 1) { + locks_off++; + } + } else { + locks_off++; + } + + if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { + constexpr int threads_per_m = 16 * thread_n_blocks / 8; + int m_per_thread = + div_ceil(thread_m_blocks * 16, threads / threads_per_m); + if (m_block_size_8) m_per_thread = div_ceil(8, threads / threads_per_m); + for (int i = 0; i < m_per_thread; i++) { + int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; + if (row < prob_m) { + int col = slice_col * 16 * thread_n_blocks / 8 + + threadIdx.x % threads_per_m; + C[row * prob_n / 8 + col] = {0, 0, 0, 0}; + } + } + // After write zero to output, write a negative value to lock. + // Every SM that processes the same slice would wait for + // the negative value, and then atomicAdd 1 to it. + // After all SMs are processed, the lock value would back to 0 again. + __syncthreads(); + if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count; + } + + if (slice_col == n_tiles) { + A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); + C += 16 * thread_m_blocks * prob_n / 8; + slice_col = 0; + par_id++; + } + if (is_a_8bit && (first_init || slice_col == 0)) { + __syncthreads(); + int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], + threadIdx.x < prob_m); + } + }; + + auto init_part1_slice = [&]() { + if (part1_mn_iters) { + part1_mn_iters--; + par_id = slice_col_par / n_tiles; + slice_col = slice_col_par % n_tiles; + slice_iters = k_tiles; + A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; + C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; + if (is_a_8bit) { + __syncthreads(); + int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], + threadIdx.x < prob_m); + } + } + }; + + auto init_slice = [&]() { + if (!in_part2 && !part1_mn_iters) { + in_part2 = true; + slice_col_par = (iters * blockIdx.x) / k_tiles; + slice_row = (iters * blockIdx.x) % k_tiles; + slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; + par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; + A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; + C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; + } + if (!in_part2) { + init_part1_slice(); + } else { + init_part2_slice(); + first_init = false; + } + }; + + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = lda / (is_a_8bit ? 16 : 8); + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * m_block_size; + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); + constexpr int b_sh_stride = + ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); + constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = + b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); + constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + constexpr int act_s_max_num_groups = 32; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + + constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Zero-points sizes/strides + int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; + constexpr int zp_sh_stride = is_zp_float + ? 16 * thread_n_blocks / 8 + : ((16 * thread_n_blocks) / pack_factor) / 4; + constexpr int zp_tb_groups = s_tb_groups; + constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; + int zp_gl_rd_delta = zp_gl_stride; + + // Global A read index of current thread. + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); + a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; + + int b_gl_rd; + if (threads <= b_sh_stride) { + b_gl_rd = threadIdx.x; + } else { + b_gl_rd = + b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + } + + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + auto b_sh_rd = threadIdx.x * b_thread_vecs; + b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); + + // For act_order + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + } + } + auto s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stage; + + // Zero-points + int zp_gl_rd; + if constexpr (has_zp) { + if constexpr (group_blocks == -1) { + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } else { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + auto zp_sh_wr = threadIdx.x; + bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (is_a_8bit) { + s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); + } else if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; + else if constexpr (group_blocks == -1 && + (m_block_size_8 || (has_zp && !dequant_skip_flop))) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; + + int bias_sh_rd; + if constexpr (m_block_size_8) { + bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + } else { + bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + + (threadIdx.x % 32) % 4; + } + + int bias_sh_wr = threadIdx.x; + int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; + + // Zero-points have the same read layout as the scales + // (without column-wise case) + constexpr int num_col_threads = 8; + constexpr int num_row_threads = 4; + constexpr int num_ints_per_thread = 8 / pack_factor; + int zp_sh_rd; + if constexpr (has_zp) { + if constexpr (is_zp_float) { + if constexpr (group_blocks != -1) { + zp_sh_rd = + 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; + } + } else if (is_a_8bit) { + zp_sh_rd = num_ints_per_thread * num_col_threads * + ((threadIdx.x / 32) % tb_n_warps / 2) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } else { + zp_sh_rd = num_ints_per_thread * num_col_threads * + ((threadIdx.x / 32) % tb_n_warps) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + } + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + + // Shared memory storage for global fetch pipelines. + constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; + constexpr int sh_b_size = stages * b_sh_stage; + int4* sh_b = sh_new; + int4* sh_red = sh_new; + constexpr int sh_size_b_red_min = + (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_size_b_red_max = + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); + constexpr int sh_b_red_bias_size = + sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) + ? sh_size_b_red_max + : (sh_size_b_red_min + sh_bias_size); + + int4* sh_bias = sh_new + sh_size_b_red_min; + int4* sh_g_idx = sh_new + sh_b_red_bias_size; + int4* sh_zp = sh_g_idx + (stages * g_idx_stage); + constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) + : (stages * s_sh_stage); + int4* sh_s = sh_zp + (stages * zp_sh_stage); + int4* sh_a = sh_s + sh_s_size; + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + FragS frag_s[2][4]; // No act-order + FragS frag_bias[2][4]; + FragS act_frag_s[2][4][4]; // For act-order + int frag_qzp[2][num_ints_per_thread]; // Zero-points + FragZP frag_zp; // Zero-points in fp16 + FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ + + if constexpr (is_a_8bit) { + #pragma unroll + for (int j = 0; j < 2; j++) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + + // Zero accumulators. + auto zero_accums = [&]() { + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + + auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups > act_s_max_num_groups) { + sh_num_groups = act_s_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + cp_async4_pred( + &sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], + a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) { + constexpr int count = div_ceil(b_sh_stride, threads); + int b_gl_idx = + b_gl_rd + (i % count) * threads + + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); + + cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); + } + + b_gl_rd += b_gl_rd_delta_o; + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta * s_tb_groups; + } + } + + if constexpr (has_zp && group_blocks != -1) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + auto fetch_col_zp_to_shared = [&]() { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + }; + + auto fetch_col_scale_to_shared = [&]() { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + ldsm( + frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + + #pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + using IT1 = typename std::conditional_t; + using IT0 = typename std::conditional_t; + constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks == -1) { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 && dequant_skip_flop) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } else if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) { + if (k % b_sh_wr_iters == 0) { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } + } + } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } else { + reinterpret_cast(&frag_s[k % 2])[0] = + reinterpret_cast( + sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } + } else if (group_blocks >= b_sh_wr_iters) { + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k % b_sh_wr_iters; + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + int warp_col = warp_id % tb_n_warps; + + cur_k += warp_row * 16 * b_sh_wr_iters; + + auto th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, + 9}; // Tensor core offsets per thread + + #pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + auto fetch_zp_to_registers = [&](int k, int full_pipe) { + // This code does not handle group_blocks == 0, + // which signifies act_order. + // has_zp implies AWQ, which doesn't have act_order, + static_assert(!has_zp || group_blocks != 0); + + if constexpr (has_zp && !is_zp_float) { + int pipe = full_pipe % stages; + + if constexpr (group_blocks == -1) { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 || is_a_8bit) { + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; + } + } + } else if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = + (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + } + } + } else { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + sh_zp_stage += cur_group_id * zp_sh_stride; + + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = + (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + } + } + } + + else if constexpr (has_zp && is_zp_float) { + int pipe = full_pipe % stages; + + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_zpf[k % 2])[0] = + sh_zp_stage[zp_sh_rd]; + } + } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + reinterpret_cast(&frag_zpf[k % 2])[0] = + sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; + } + } + } + }; + + auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) { + if constexpr (a_type.size_bits() != b_type.size_bits()) { + if constexpr (is_a_8bit && has_zp) { + sub_zp_and_dequant( + q, frag_b_ptr, zp); + } else { + dequant(q, frag_b_ptr); + } + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + bool is_first_matmul_in_slice = true; + auto matmul = [&](int k, int pipe) { + if (is_a_8bit) return; + int k2 = k % 2; + constexpr int g = + group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; + const bool is_new_zp = + (group_blocks == 0) || + ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && + (pipe % g == 0) || + (group_blocks == -1 && is_first_matmul_in_slice); + if constexpr (has_zp && !is_zp_float) { + if (is_new_zp) { + if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; + int zp_quant_0, zp_quant_1; + + if constexpr (b_type.size_bits() == 4) { + zp_quant_0 = frag_qzp[k2][0]; + zp_quant_1 = zp_quant_0 >> 8; + } else { + static_assert(b_type.size_bits() == 8); + zp_quant_0 = frag_qzp[k2][0]; + zp_quant_1 = frag_qzp[k2][1]; + } + + dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); + dequant_data(zp_quant_1, + reinterpret_cast(&frag_zp) + 2); + } + } + if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { + if (is_new_zp) { + reinterpret_cast(&frag_zp)[0] = + reinterpret_cast(&frag_zpf[k2])[0]; + } + } + + if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) { + int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; + int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; + + dequant_fp8_scales( + s_quant_0, reinterpret_cast(&frag_s[k2])); + dequant_fp8_scales( + s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); + } + + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. + #pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + int b_quant_0, b_quant_1; + + if constexpr (b_type_id == vllm::kFE2M1f.id()) { + b_quant_1 = frag_b_quant[k2][0][j]; + b_quant_0 = b_quant_1 << 8; + } else if constexpr (b_type.size_bits() == 4) { + b_quant_0 = frag_b_quant[k2][0][j]; + b_quant_1 = b_quant_0 >> 8; + } else { + static_assert(b_type.size_bits() == 8); + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); + b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + } + + dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); + dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); + + if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) { + sub_zp(frag_b0, frag_zp[j], 0); + sub_zp(frag_b1, frag_zp[j], 1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order && !is_a_8bit) { + static_assert(group_blocks != -1); + scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], + act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); + scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], + act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); + } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && + group_blocks == -1 && !is_a_8bit) { + int idx = (threadIdx.x / 4) % 2; + scalar_t2 s2 = Adtype::nums2num2( + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); + if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); + scale_and_sub(frag_b0, s2.x, frag_zp[j].x); + scale_and_sub(frag_b1, s2.y, frag_zp[j].y); + } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && + !is_a_8bit) { + if (is_new_zp) + frag_zp[j] = __hmul2(frag_zp[j], + *reinterpret_cast(&frag_s[k2][j])); + scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); + scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); + } else if constexpr (group_blocks != -1 && !is_a_8bit) { + scale(frag_b0, frag_s[k2][j], 0); + scale(frag_b1, frag_s[k2][j], 1); + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + if constexpr (m_block_size_8) { + mma_trans(frag_a[k2][i], frag_b0, frag_b1, + frag_c[i][j][0]); + } else { + mma(frag_a[k2][i], frag_b0, + frag_c[i][j][0]); + mma(frag_a[k2][i], frag_b1, + frag_c[i][j][1]); + } + } + } + }; + + auto matmul_a8 = [&](int k) { + int k2 = k % 2; + #pragma unroll + for (int j = 0; j < 2; j++) { + FragB frag_b[2]; + + if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) { + dequant_data(frag_b_quant[k2][0][j * 2], + reinterpret_cast(&frag_b)); + dequant_data(frag_b_quant[k2][0][j * 2 + 1], + reinterpret_cast(&frag_b) + 2); + } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) { + int off = (threadIdx.x / 32) % 2 * 2 + j; + int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2], + reinterpret_cast(&frag_b), zp); + zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2 + 1], + reinterpret_cast(&frag_b) + 2, zp); + } else { + reinterpret_cast(&frag_b)[0] = + reinterpret_cast(&frag_b_quant[k2][j])[0]; + reinterpret_cast(&frag_b)[1] = + reinterpret_cast(&frag_b_quant[k2][j])[1]; + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma( + frag_a[k2][i], frag_b[0], + (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); + mma( + frag_a[k2][i], frag_b[1], + (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); + } + + if constexpr (group_blocks != -1) { + if (group_blocks == 2 || k == 1) { + if constexpr (a_type == vllm::kS8) { + int2 s_vals[2]; + s_vals[0] = { + (int)reinterpret_cast(&frag_s[k2][j * 2][0])[0], + (int)reinterpret_cast(&frag_s[k2][j * 2][0])[1]}; + s_vals[1] = { + (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[0], + (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[1]}; + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + int scale = reinterpret_cast(&s_vals[0])[g % 2]; + *reinterpret_cast(&frag_c[i][j][0][g]) += + *reinterpret_cast(&frag_c_tmp[i][j][0][g]) * + scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + int scale = reinterpret_cast(&s_vals[1])[g % 2]; + *reinterpret_cast(&frag_c[i][j][1][g]) += + *reinterpret_cast(&frag_c_tmp[i][j][1][g]) * + scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } else { + float2 s_vals[2]; + if constexpr (s_type_id != vllm::kFE8M0fnu.id()) { + static_assert(a_type.size_bits() == 16 || + s_type.size_bits() == 16); + s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]); + s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]); + } else { + int32_t* s_vals_int = reinterpret_cast(&s_vals[0]); + int32_t s_vals_e8m0 = + *reinterpret_cast(&frag_s[k2][j][0]); + + s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23; + s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15; + s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7; + s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1; + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&s_vals[0])[g % 2]; + frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&s_vals[1])[g % 2]; + frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + auto red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = + b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + + #pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { + #pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { + #pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; + j += (m_block_size_8 ? 2 : 1)) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = reinterpret_cast( + &sh_red[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); + #pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast( + frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh_red[red_sh_wr] = reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { + #pragma unroll + for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; + i += (m_block_size_8 ? 2 : 1)) { + float* c_rd = + reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); + #pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast( + frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce_fp16 = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * tb_n_warps; + if (threadIdx.x < active_threads) { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr; + if constexpr (m_block_size_8) { + c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + } else { + c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); + } + constexpr int c_sh_wr_delta = active_threads; + auto c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) { + // Interestingly, doing direct global accesses here really seems to mess up + // the compiler and lead to slowdowns, hence we also use async-copies even + // though these fetches are not actually asynchronous. + #pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { + if constexpr (m_block_size_8) { + cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + i * c_gl_stride + + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], + (threadIdx.x % 4) * 2 + i < prob_m); + } else if constexpr (is_a_8bit) { + int2* sh_red_int2 = reinterpret_cast(sh_red); + int2* c_int2 = reinterpret_cast(C); + cp_async2_ca_pred( + &sh_red_int2[c_sh_wr + c_sh_wr_delta * i], + &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } else { + cp_async4_pred( + &sh_red[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + } + cp_async_fence(); + cp_async_wait<0>(); + } + + #pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { + bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || + 8 * (i / 2) + row < prob_m) || + (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); + if (mask) { + if (!first) { + c_scalar_t* c_red_f16; + if constexpr (is_a_8bit) { + int2 tmp = + reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } else { + int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } + #pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { + int delta = 0; + if constexpr (m_block_size_8) { + delta = j % 2 == 1 ? -2 : 0; + } + reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + + (i % 4) + delta] += Cdtype::num2float(c_red_f16[j]); + } + } + if (!last) { + c_scalar_t c_f16[is_a_8bit ? 4 : 8]; + #pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { + int delta = 0; + if constexpr (m_block_size_8) { + delta = j % 2 == 1 ? -2 : 0; + } + c_f16[j] = Cdtype::float2num(reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + + (i % 4) + delta]); + } + if constexpr (m_block_size_8) { + C[c_gl_wr + i * c_gl_stride + + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = + *reinterpret_cast(c_f16); + } else if constexpr (is_a_8bit) { + int2* c_int2 = reinterpret_cast(C); + c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)] = + *reinterpret_cast(c_f16); + } else { + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast(c_f16); + } + } + } + } + } + }; + + // Globally reduce over threadblocks that compute the same column block. + // We use a tmp C buffer to reduce in full fp32 precision. + auto global_reduce_fp32 = [&](bool first = false, bool last = false) { + constexpr int tb_m = thread_m_blocks * 16; + constexpr int tb_n = thread_n_blocks * 16; + + constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; + + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + + constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; + constexpr int th_size = num_floats * sizeof(float) / 16; + + int c_cur_offset = locks_off * c_size; + + if (!is_th_active) { + return; + } + + if (!first) { + float* frag_c_ptr = reinterpret_cast(&frag_c); + #pragma unroll + for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { + sh_red[threadIdx.x] = + C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; + + float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); + #pragma unroll + for (int f = 0; f < 4; f++) { + frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; + } + } + } + + if (!last) { + int4* frag_c_ptr = reinterpret_cast(&frag_c); + #pragma unroll + for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { + C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&](bool last) { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr; + if constexpr (m_block_size_8) { + c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + + (threadIdx.x % 32) / 4; + c_sh_wr += 64 * (threadIdx.x / 32); + } else { + c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); + } + + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) { + if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { + c0 *= global_scale_f32; + c1 *= global_scale_f32; + } + c_scalar_t2 res = + Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && + b_type.size_bits() == 4 && + (has_zp && dequant_skip_flop || !has_zp)) { + c_scalar_t2 tmp_scale = s[0]; + if constexpr (m_block_size_8) { + tmp_scale = Cdtype::num2num2( + reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); + } + res = __hmul2(res, tmp_scale); + } + if (has_bias && last) { + c_scalar_t2 tmp_bias = b_bias[0]; + if constexpr (m_block_size_8) { + tmp_bias = Cdtype::num2num2( + reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + } + res = __hadd2(res, tmp_bias); + } + + if constexpr (m_block_size_8) { + ((c_scalar_t*)sh_red)[idx] = res.x; + ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y; + } else { + ((c_scalar_t2*)sh_red)[idx] = res; + } + }; + + if (threadIdx.x / 32 < tb_n_warps) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) { + if constexpr (m_block_size_8) { + int wr = c_sh_wr + 16 * j; + write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], + frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], + frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + } else { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + } + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + + #pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + if (c_gl_wr < c_gl_wr_end) { + if (use_atomic_add && slice_count > 1) { + c_scalar_t2* C_half2 = reinterpret_cast(&C[c_gl_wr]); + c_scalar_t2* sh_red_half2 = + reinterpret_cast(&sh_red[c_sh_rd]); + #pragma unroll + for (int a = 0; a < 4; a++) { + atomicAdd(&C_half2[a], sh_red_half2[a]); + } + } else { + C[c_gl_wr] = sh_red[c_sh_rd]; + } + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + __syncthreads(); + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + + #pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], + g_idx[last_g_idx]); + } + + if constexpr (has_zp && !is_zp_float && group_blocks == -1) { + if (i == 0) { + fetch_col_zp_to_shared(); + if constexpr (!dequant_skip_flop) { + fetch_col_scale_to_shared(); + } + } + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + fetch_zp_to_registers(0, 0); + a_gl_rd += a_gl_rd_delta_o * (stages - 1); + if constexpr (has_act_order) { + slice_k_start_shared_fetch += tb_k * (stages - 1); + } + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + + #pragma unroll + for (int pipe = 0; pipe < stages;) { + #pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + fetch_zp_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + + if constexpr (!is_a_8bit) { + matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); + } else { + static_assert(group_blocks != 0 && group_blocks != 1); + matmul_a8(k); + } + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd += a_gl_rd_delta_o * stages; + + if constexpr (has_act_order) { + slice_k_start += tb_k * stages; + + if (slice_k_start < prob_k) { + slice_k_start_shared_fetch += tb_k * stages; + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_act_order_scales_to_shared(false, first_group_id, + last_group_id); + __syncthreads(); + } + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + // convert fp16 accum to fp32 for reduction + if constexpr (use_fp16_accum) { + #pragma unroll + for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) { + float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; + scalar_t* frag_c_part_half = + reinterpret_cast(frag_c_part_float); + + #pragma unroll + for (int i = 3; i >= 0; i--) { + frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]); + } + } + } + + if constexpr (is_a_8bit) { + float frag_a_s[2 * thread_m_blocks]; + + for (int i = 0; i < 2 * thread_m_blocks; i++) + frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; + + #pragma unroll + for (int j = 0; j < 2; j++) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float c_val = frag_c[i][j][0][g]; + + if constexpr (a_type == vllm::kS8) { + c_val = __int2float_rn(*reinterpret_cast(&c_val)); + } + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][0][g] = c_val * s_val; + } + #pragma unroll + for (int g = 0; g < 4; g++) { + float c_val = frag_c[i][j][1][g]; + + if constexpr (a_type == vllm::kS8) { + c_val = __int2float_rn(*reinterpret_cast(&c_val)); + } + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][1][g] = c_val * s_val; + } + } + } + } + + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1 && + (has_zp && dequant_skip_flop || !has_zp)) { + if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + + thread_block_reduce(); + + if (has_bias && last) { + __syncthreads(); + cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], + threadIdx.x < 16 * thread_n_blocks / 8); + cp_async_fence(); + } + + if constexpr (!has_act_order && group_blocks == -1 && + (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { + if constexpr (is_a_8bit) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + } + } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + if constexpr (m_block_size_8) { + int idx = (threadIdx.x / 4) % 2; + c_scalar_t2* frag_s_half2 = + reinterpret_cast(frag_s); + #pragma unroll + for (int i = 0; i < 8; i++) { + frag_s_half2[i] = Cdtype::num2num2( + reinterpret_cast(&frag_s_half2[i])[idx]); + } + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 aa[2]; + aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]); + aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]); + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&aa[0])[g % 2]; + frag_c[i][j][0][g] *= scale; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&aa[1])[g % 2]; + frag_c[i][j][1][g] *= scale; + } + } + } + } else if (!has_act_order && group_blocks == -1 && + b_type.size_bits() == 8 && + (has_zp && dequant_skip_flop || !has_zp)) { + if (threadIdx.x / 32 < tb_n_warps) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); + + if constexpr (!m_block_size_8) { + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + } + + if (slice_count > 1 && !use_atomic_add) { + // only globally reduce if there is more than one block in a slice + barrier_acquire(&locks[locks_off], slice_idx); + if (use_fp32_reduce) { + global_reduce_fp32(slice_idx == 0, last); + } else { + global_reduce_fp16(slice_idx == 0, last); + } + barrier_release(&locks[locks_off], last); + } + + if (has_bias && last) { + cp_async_wait<0>(); + __syncthreads(); + reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; + if constexpr (!is_a_8bit) + reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; + __syncthreads(); + } + + if (use_atomic_add && slice_count > 1 && slice_idx != 0) + wait_negative_and_add(&locks[locks_off]); + if (last || use_atomic_add) + // only the last block in a slice actually writes the result + write_result(last); + slice_row = 0; + if (!in_part2) { + slice_col_par += gridDim.x; + } else { + slice_col_par++; + slice_col++; + } + is_first_matmul_in_slice = true; + init_slice(); + + if (slice_iters) { + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + + (threadIdx.x % b_sh_stride); + b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; + + bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + } else { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + start_pipes(); + } + } + } +} + +} // namespace MARLIN_NAMESPACE_NAME + +#endif diff --git a/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu new file mode 100644 index 000000000..8688df862 --- /dev/null +++ b/src/vt/cuda/marlin/libtorch_stable/quantization/marlin/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu @@ -0,0 +1,70 @@ +// auto generated by generate_kernels.py +// clang-format off + +#include "kernel.h" +#include "marlin_template.h" + +namespace MARLIN_NAMESPACE_NAME { + + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +} diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index 36afd5924..89a569976 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -843,6 +843,25 @@ void MoeGroupedGemmNvfp4Marlin(Queue& q, Tensor& c, const Tensor& a, const Tenso num_tokens_past_padded, topk_weights, args); } +void MarlinDenseGemm(Queue& q, Tensor& c, const Tensor& a, const Tensor& b_q_weight, + const Tensor& b_scales, const Tensor& global_scale, Tensor& workspace, + const MarlinDenseArgs& args) { + VT_CHECK(a.rank == 2 && c.rank == 2, "marlin_dense: a/c must be rank-2"); + VT_CHECK(a.dtype == DType::kBF16 && c.dtype == DType::kBF16, "marlin_dense: a/c must be bf16"); + VT_CHECK(args.size_k % 16 == 0, "marlin_dense: size_k must be a multiple of 16 (group size)"); + VT_CHECK(args.group_size == 16 || args.group_size == 32, + "marlin_dense: group_size must be 16 (nvfp4) or 32 (mxfp4)"); + VT_CHECK(a.shape[0] == args.size_m && a.shape[1] == args.size_k, + "marlin_dense: a shape must be [size_m, size_k]"); + VT_CHECK(c.shape[0] == args.size_m && c.shape[1] == args.size_n, + "marlin_dense: c shape must be [size_m, size_n]"); + VT_CHECK(b_q_weight.rank == 2, "marlin_dense: b_q_weight must be rank-2 [K/16, N*8/pack]"); + VT_CHECK(global_scale.dtype == DType::kF32, "marlin_dense: global_scale must be f32"); + VT_CHECK(workspace.dtype == DType::kI32, "marlin_dense: workspace must be i32 (reduction locks)"); + reinterpret_cast(GetOp(OpId::kMarlinDenseGemm, q.device.type))( + q, c, a, b_q_weight, b_scales, global_scale, workspace, args); +} + void MoeSiluMul(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up) { VT_CHECK(gate.Numel() == out.Numel() && up.Numel() == out.Numel(), "moe_silu_mul: out/gate/up must have the same element count"); diff --git a/tests/vt/test_ops_moe_grouped.cpp b/tests/vt/test_ops_moe_grouped.cpp index 74fb64540..a92f30824 100644 --- a/tests/vt/test_ops_moe_grouped.cpp +++ b/tests/vt/test_ops_moe_grouped.cpp @@ -1239,4 +1239,252 @@ TEST_CASE("CUDA moe_align parallel == serial (expert_ids/num_pad exact, per-expe } } } + +// ─── DENSE Marlin (row KERNEL-MARLIN-DENSE-PORT) ──────────────────────────── +// vt::MarlinDenseGemm is vLLM's OWN dense W4A16 GEMM (direct-A, tile-per-CTA, +// dense fp32-C_tmp reduce) — the byte-preserving replacement for the +// single-expert MoE-marlin route the dense E=1 NVFP4/MXFP4 projections use today +// (dense_nvfp4_gemm.h). The gate (mission gate a): the dense op must match, per +// output element, BOTH an INDEPENDENT CPU-dequant reference AND the grouped +// (single-expert MoE) route it replaces, across M=1..8, NVFP4 + MXFP4, and a set +// of model-representative shapes; plus a wrong-stride RED-injection proof that the +// comparison actually discriminates (not a vacuous pass). RED-first: on a wrong +// launcher stride (lda / operand layout) the reference match must FAIL. + +// Count elements outside tolerance (RED-injection helper: 0 == match, >0 == diff). +size_t MismatchCount(const std::vector& got, const std::vector& want, float atol, + float rtol) { + size_t bad = 0; + for (size_t i = 0; i < got.size(); ++i) { + const float tol = atol + rtol * std::fabs(want[i]); + if (!(std::fabs(got[i] - want[i]) <= tol)) ++bad; + } + return bad; +} + +// CPU reference y[M,N] = act(bf16)[M,K] @ dequant(w)[N,K]^T, accumulated in f32. +std::vector DenseRefNvfp4(const Nvfp4Weight& w, const std::vector& act_bf16, + int64_t M, int64_t N, int64_t K) { + std::vector deq(static_cast(N * K)); + vllm::DequantNvfp4ToBf16(w.packed.data(), w.scale.data(), w.scale2, N, K, deq.data()); + std::vector ref(static_cast(M * N), 0.0f); + for (int64_t m = 0; m < M; ++m) + for (int64_t n = 0; n < N; ++n) { + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) + acc += vt::BF16ToF32(act_bf16[static_cast(m * K + k)]) * + vt::BF16ToF32(deq[static_cast(n * K + k)]); + ref[static_cast(m * N + n)] = acc; + } + return ref; +} + +std::vector DenseRefMxfp4(const Mxfp4Weight& w, const std::vector& act_bf16, + int64_t M, int64_t N, int64_t K) { + std::vector deq(static_cast(N * K)); + vllm::DequantMxfp4ToBf16(w.packed.data(), w.scale.data(), N, K, deq.data()); + std::vector ref(static_cast(M * N), 0.0f); + for (int64_t m = 0; m < M; ++m) + for (int64_t n = 0; n < N; ++n) { + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) + acc += vt::BF16ToF32(act_bf16[static_cast(m * K + k)]) * + vt::BF16ToF32(deq[static_cast(n * K + k)]); + ref[static_cast(m * N + n)] = acc; + } + return ref; +} + +TEST_CASE("CUDA marlin DENSE gemm matches CPU-dequant ref AND the grouped route (NVFP4)") { + if (!HasCuda()) { + MESSAGE("no CUDA backend registered; skipping"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kCUDA); + REQUIRE(vt::OpRegistered(vt::OpId::kMarlinDenseGemm, DeviceType::kCUDA)); + QueueGuard gq(gpu); + void* stream = gq.q.handle; + const int dev = gq.q.device.index; + const int sms = vt::cuda::MarlinDeviceSms(dev); + + // Model-representative shapes (K % 128 == 0, N % 64 == 0: the marlin tile + // constraints the dense projections satisfy). The 48-CTA E=1 regime is M<=8. + const std::vector> shapes = {{256, 64}, {512, 128}, {128, 256}}; + for (auto [K, N] : shapes) { + const Nvfp4Weight w = MakeNvfp4Weight(N, K, 9100 + static_cast(K + N)); + // Resident (repack ONCE): wq [K/16, N*2] i32, sc [K/16, N] fp8, gg [1] f32. + std::vector sc_bufs{w.scale.data()}; + std::vector sc_lens{w.scale.size()}; + const float sf = vt::cuda::MarlinNvfp4CombinedScaleFactor(sc_bufs, sc_lens); + DeviceTensor dpacked(gpu, gq.q, DType::kI8, {N, K / 2}, w.packed.data()); + DeviceTensor dscale(gpu, gq.q, DType::kI8, {N, K / 16}, w.scale.data()); + DeviceTensor wq(gpu, gq.q, DType::kI32, {K / 16, N * 2}); + DeviceTensor sc(gpu, gq.q, DType::kI8, {K / 16, N}); + vt::cuda::MarlinRepackExpertWeight(stream, dev, static_cast(wq.ptr()), + static_cast(dpacked.ptr()), + static_cast(K), static_cast(N)); + vt::cuda::MarlinProcessExpertScales(stream, static_cast(dscale.ptr()), + static_cast(sc.ptr()), static_cast(K), + static_cast(N), sf); + const float g = vt::cuda::MarlinNvfp4ProcessGlobalScale(w.scale2, sf); + DeviceTensor gg(gpu, gq.q, DType::kF32, {1}, &g); + gpu.Synchronize(gq.q); + + // Rank-3 resident views for the single-expert MoE route (SAME memory). + Tensor wq3 = MakeTensor(wq.ptr(), DType::kI32, gq.q.device, {1, K / 16, N * 2}); + Tensor sc3 = MakeTensor(sc.ptr(), DType::kI8, gq.q.device, {1, K / 16, N}); + + DeviceTensor ws(gpu, gq.q, DType::kI32, {sms * 4}); + + for (int64_t M = 1; M <= 8; ++M) { + const auto act_f = RandomF32(static_cast(M * K), 700 + static_cast(M)); + const auto act_bf16 = ToBf16(act_f); + const auto ref = DenseRefNvfp4(w, act_bf16, M, N, K); + DeviceTensor dact(gpu, gq.q, DType::kBF16, {M, K}, act_bf16.data()); + + // (1) DENSE route. + DeviceTensor dout(gpu, gq.q, DType::kBF16, {M, N}); + gpu.Memset(gq.q, ws.ptr(), 0, static_cast(sms) * 4 * sizeof(int32_t)); + vt::MarlinDenseArgs dargs{static_cast(M), static_cast(N), static_cast(K)}; + vt::MarlinDenseGemm(gq.q, dout.tensor(), dact.tensor(), wq.tensor(), sc.tensor(), + gg.tensor(), ws.tensor(), dargs); + std::vector h_dense(static_cast(M * N)); + dout.Download(gq.q, h_dense.data()); + std::vector got_dense(static_cast(M * N)); + for (size_t i = 0; i < got_dense.size(); ++i) got_dense[i] = vt::BF16ToF32(h_dense[i]); + CheckClose(got_dense, ref, 3e-2f, 3e-2f); + + // (2) GROUPED single-expert MoE route (all M tokens -> expert 0). + const int block = (M <= 8) ? 8 : vt::cuda::MarlinMoeAlignBlockSizeSelect(static_cast(M), 1, 1); + int max_tok = 0, max_blk = 0; + vt::cuda::MarlinMoeAlignSizes(static_cast(M), 1, 1, block, &max_tok, &max_blk); + std::vector tids(static_cast(M), 0); + DeviceTensor dtid(gpu, gq.q, DType::kI32, {M}, tids.data()); + DeviceTensor sorted_ids(gpu, gq.q, DType::kI32, {max_tok}); + DeviceTensor expert_ids(gpu, gq.q, DType::kI32, {max_blk}); + DeviceTensor num_pad(gpu, gq.q, DType::kI32, {1}); + vt::cuda::MarlinMoeAlignBlockSize(stream, static_cast(dtid.ptr()), + static_cast(M), 1, 1, block, + static_cast(sorted_ids.ptr()), + static_cast(expert_ids.ptr()), + static_cast(num_pad.ptr())); + std::vector ones(static_cast(M), 1.0f); + DeviceTensor topkw(gpu, gq.q, DType::kF32, {M}, ones.data()); + DeviceTensor mout(gpu, gq.q, DType::kBF16, {M, N}); + gpu.Memset(gq.q, ws.ptr(), 0, static_cast(sms) * 4 * sizeof(int32_t)); + vt::MoeMarlinArgs margs{block, 1, static_cast(M), static_cast(N), + static_cast(K), false}; + vt::MoeGroupedGemmNvfp4Marlin(gq.q, mout.tensor(), dact.tensor(), wq3, sc3, + gg.tensor(), ws.tensor(), sorted_ids.tensor(), + expert_ids.tensor(), num_pad.tensor(), topkw.tensor(), margs); + std::vector h_moe(static_cast(M * N)); + mout.Download(gq.q, h_moe.data()); + std::vector got_moe(static_cast(M * N)); + for (size_t i = 0; i < got_moe.size(); ++i) got_moe[i] = vt::BF16ToF32(h_moe[i]); + CheckClose(got_moe, ref, 3e-2f, 3e-2f); + // Dense and grouped agree to within a couple bf16 ULPs (both marlin; the + // dense reduce differs from the par-regrouped grouped reduce by ~1 ULP — + // exactly the point of the port). + CheckClose(got_dense, got_moe, 5e-2f, 5e-2f); + + // (3) WRONG-STRIDE RED injection: a row-shifted reference is a stride-class + // perturbation; the correct-vs-shifted comparison MUST discriminate (else + // the whole gate is vacuous). Only meaningful when M*N gives real spread. + if (M >= 2) { + std::vector ref_shift(ref.size()); + for (size_t i = 0; i < ref.size(); ++i) + ref_shift[i] = ref[(i + static_cast(N)) % ref.size()]; + CHECK(MismatchCount(got_dense, ref_shift, 3e-2f, 3e-2f) > 0); + } + } + } +} + +TEST_CASE("CUDA marlin DENSE gemm matches CPU-dequant ref AND the grouped route (MXFP4)") { + if (!HasCuda()) { + MESSAGE("no CUDA backend registered; skipping"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kCUDA); + REQUIRE(vt::OpRegistered(vt::OpId::kMarlinDenseGemm, DeviceType::kCUDA)); + QueueGuard gq(gpu); + void* stream = gq.q.handle; + const int dev = gq.q.device.index; + const int sms = vt::cuda::MarlinDeviceSms(dev); + + const std::vector> shapes = {{256, 64}, {512, 128}}; + for (auto [K, N] : shapes) { + const Mxfp4Weight w = MakeMxfp4Weight(N, K, 9500 + static_cast(K + N)); + // MXFP4 resident: group_size 32 => sc [K/32, N]; E8M0 passthrough (no global). + DeviceTensor dpacked(gpu, gq.q, DType::kI8, {N, K / 2}, w.packed.data()); + DeviceTensor dscale(gpu, gq.q, DType::kI8, {N, K / 32}, w.scale.data()); + DeviceTensor wq(gpu, gq.q, DType::kI32, {K / 16, N * 2}); + DeviceTensor sc(gpu, gq.q, DType::kI8, {K / 32, N}); + vt::cuda::MarlinRepackExpertWeight(stream, dev, static_cast(wq.ptr()), + static_cast(dpacked.ptr()), + static_cast(K), static_cast(N)); + vt::cuda::MarlinProcessExpertScalesMxfp4(stream, static_cast(dscale.ptr()), + static_cast(sc.ptr()), static_cast(K), + static_cast(N)); + const float g = 1.0f; // unused (kernel skips global for E8M0) + DeviceTensor gg(gpu, gq.q, DType::kF32, {1}, &g); + gpu.Synchronize(gq.q); + + Tensor wq3 = MakeTensor(wq.ptr(), DType::kI32, gq.q.device, {1, K / 16, N * 2}); + Tensor sc3 = MakeTensor(sc.ptr(), DType::kI8, gq.q.device, {1, K / 32, N}); + + DeviceTensor ws(gpu, gq.q, DType::kI32, {sms * 4}); + + for (int64_t M = 1; M <= 8; ++M) { + const auto act_f = RandomF32(static_cast(M * K), 800 + static_cast(M)); + const auto act_bf16 = ToBf16(act_f); + const auto ref = DenseRefMxfp4(w, act_bf16, M, N, K); + DeviceTensor dact(gpu, gq.q, DType::kBF16, {M, K}, act_bf16.data()); + + DeviceTensor dout(gpu, gq.q, DType::kBF16, {M, N}); + gpu.Memset(gq.q, ws.ptr(), 0, static_cast(sms) * 4 * sizeof(int32_t)); + vt::MarlinDenseArgs dargs{static_cast(M), static_cast(N), static_cast(K)}; + dargs.group_size = 32; + dargs.mxfp4 = true; + vt::MarlinDenseGemm(gq.q, dout.tensor(), dact.tensor(), wq.tensor(), sc.tensor(), + gg.tensor(), ws.tensor(), dargs); + std::vector h_dense(static_cast(M * N)); + dout.Download(gq.q, h_dense.data()); + std::vector got_dense(static_cast(M * N)); + for (size_t i = 0; i < got_dense.size(); ++i) got_dense[i] = vt::BF16ToF32(h_dense[i]); + CheckClose(got_dense, ref, 4e-2f, 4e-2f); + + const int block = (M <= 8) ? 8 : vt::cuda::MarlinMoeAlignBlockSizeSelect(static_cast(M), 1, 1); + int max_tok = 0, max_blk = 0; + vt::cuda::MarlinMoeAlignSizes(static_cast(M), 1, 1, block, &max_tok, &max_blk); + std::vector tids(static_cast(M), 0); + DeviceTensor dtid(gpu, gq.q, DType::kI32, {M}, tids.data()); + DeviceTensor sorted_ids(gpu, gq.q, DType::kI32, {max_tok}); + DeviceTensor expert_ids(gpu, gq.q, DType::kI32, {max_blk}); + DeviceTensor num_pad(gpu, gq.q, DType::kI32, {1}); + vt::cuda::MarlinMoeAlignBlockSize(stream, static_cast(dtid.ptr()), + static_cast(M), 1, 1, block, + static_cast(sorted_ids.ptr()), + static_cast(expert_ids.ptr()), + static_cast(num_pad.ptr())); + std::vector ones(static_cast(M), 1.0f); + DeviceTensor topkw(gpu, gq.q, DType::kF32, {M}, ones.data()); + DeviceTensor mout(gpu, gq.q, DType::kBF16, {M, N}); + gpu.Memset(gq.q, ws.ptr(), 0, static_cast(sms) * 4 * sizeof(int32_t)); + vt::MoeMarlinArgs margs{block, 1, static_cast(M), static_cast(N), + static_cast(K), false}; + margs.group_size = 32; + margs.mxfp4 = true; + vt::MoeGroupedGemmNvfp4Marlin(gq.q, mout.tensor(), dact.tensor(), wq3, sc3, + gg.tensor(), ws.tensor(), sorted_ids.tensor(), + expert_ids.tensor(), num_pad.tensor(), topkw.tensor(), margs); + std::vector h_moe(static_cast(M * N)); + mout.Download(gq.q, h_moe.data()); + std::vector got_moe(static_cast(M * N)); + for (size_t i = 0; i < got_moe.size(); ++i) got_moe[i] = vt::BF16ToF32(h_moe[i]); + CheckClose(got_moe, ref, 4e-2f, 4e-2f); + CheckClose(got_dense, got_moe, 6e-2f, 6e-2f); + } + } +} #endif // VT_MARLIN_NVFP4