diff --git a/CMakeLists.txt b/CMakeLists.txt index 9b0f7ddf..1d82ee28 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -82,18 +82,38 @@ message(STATUS "Using gencode flag: ${GPU_GENCODE}") # SM86/87/89 share the same source instantiations; we emit arch-specific # codegen per target so Jetson Orin (SM87) is not forced through an # incompatible cubin. SM120 (5090) uses PTX JIT from -# compute_80. Thor SM110 has its own attention path (fvk.attention_qkv_fp16 -# cuBLAS decomposed), hand-tuned for unified LPDDR memory — building -# FA2 there would add ~10 MB to the .so and unused compile time. +# compute_80. +# +# Thor SM110 was excluded on the grounds that it has its own attention path +# (fvk.attention_qkv_fp16, cuBLAS decomposed) and that FA2 would cost ~10 MB +# of .so for nothing. That trade does not survive a long prefill: the +# decomposed path materialises an (S * heads, S_kv) score buffer, which is +# 3.4 GB per layer at ten thousand tokens, and the model that needs it here +# is bf16 at head_dim 256 -- one instantiation, not the twelve the size +# estimate assumed. FA2_HDIMS/FA2_DTYPES are narrowed for this target below. +# +# It stays opt-in all the same: every other Thor model still uses its own +# attention path and would only be paying the compile time and the binary, so +# the target that wants it asks for it. Default OFF keeps a Thor build's +# sources and symbols exactly what they were. +option(FLASHRT_ENABLE_THOR_FA2 + "Build the vendored FA2 attention kernels on Jetson AGX Thor (sm_110)" OFF) +if(FLASHRT_ENABLE_THOR_FA2 AND NOT GPU_ARCH STREQUAL "110") + message(FATAL_ERROR + "FLASHRT_ENABLE_THOR_FA2 is the Thor (sm_110) FA2 gate; current " + "GPU_ARCH=${GPU_ARCH} decides FA2 by architecture and needs no flag.") +endif() if(GPU_ARCH STREQUAL "80" OR GPU_ARCH STREQUAL "86" OR GPU_ARCH STREQUAL "87" OR - GPU_ARCH STREQUAL "89" OR GPU_ARCH STREQUAL "120" OR + GPU_ARCH STREQUAL "89" OR + (GPU_ARCH STREQUAL "110" AND FLASHRT_ENABLE_THOR_FA2) OR + GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121") set(ENABLE_FA2 ON) message(STATUS "FA2 in-SO attention: ENABLED (sm_${GPU_ARCH})") else() set(ENABLE_FA2 OFF) - message(STATUS "FA2 in-SO attention: DISABLED (Thor SM110 uses fvk.attention_qkv_fp16 cuBLAS path)") + message(STATUS "FA2 in-SO attention: DISABLED (sm_${GPU_ARCH})") endif() # SM80-family CUTLASS INT8 kernels used by the Jetson Orin SM87 Pi0.5 fast @@ -137,6 +157,21 @@ set(FA2_HDIMS "64;96;128;256" CACHE STRING "Semicolon-separated FA2 head_dim instantiations to build. Qwen3-VL 2B vision uses 64; other shipped models use 96 and 256.") set(FA2_DTYPES "fp16;bf16" CACHE STRING "Semicolon-separated FA2 dtype instantiations to build. Pi0 uses fp16; pi0.5/groot use bf16.") + +# Thor runs one model family through FA2 -- Qwen3.6 full attention, bf16 at +# head_dim 256 -- so the default there is that one instantiation rather than +# the twelve-file matrix the RTX targets distribute. This is why enabling FA2 +# on this arch does not cost what the original exclusion assumed. Both remain +# cache variables: an explicit -DFA2_HDIMS on the command line still wins. +if(GPU_ARCH STREQUAL "110" AND ENABLE_FA2) + if(NOT DEFINED CACHE{FA2_HDIMS} OR FA2_HDIMS STREQUAL "64;96;128;256") + set(FA2_HDIMS "256" CACHE STRING "" FORCE) + endif() + if(NOT DEFINED CACHE{FA2_DTYPES} OR FA2_DTYPES STREQUAL "fp16;bf16") + set(FA2_DTYPES "bf16" CACHE STRING "" FORCE) + endif() +endif() + option(FLASHRT_ENABLE_NATIVE_CPP "Build Python-free operation libraries for native C++ consumers" OFF) option(FLASHRT_BUILD_FA2_PYTHON_ADAPTER @@ -163,13 +198,47 @@ option(FLASHRT_ENABLE_LINGBOT "Build LingBot-VLA Thor (sm_110) model kernels and bindings" ON) option(FLASHRT_ENABLE_MOTUS "Build Motus-specific RTX SM120 kernels and bindings" ON) -# qwen3_5_moe family (Nex-N2-mini / Qwen3.5-3.6-35B-A3B) RTX SM120 kernels. -# OFF by default: these are NVFP4/sm_120a-only block-scaled kernels and must -# not enter the default flash_rt_kernels build on SM89/87/110 or for models -# that don't use them. Enable explicitly (-DFLASHRT_ENABLE_QWEN35MOE=ON) on a -# Blackwell build to get the Nex-N2 / Qwen3.6-35B-A3B path. +# qwen3_5_moe family (Nex-N2-mini / Qwen3.5-3.6-35B-A3B) kernels, split into +# three tiers by the hardware each tier actually requires. All are OFF by +# default so they never enter the flash_rt_kernels build for models that do +# not use them. +# +# _CORE Architecture-neutral: layout/split, bf16 matvec, router top-k, +# activation fusion, GDN recurrence, weighted-sum reducer, and the +# bf16 GEMM. Uses only bf16 intrinsics plus cp.async and +# mma.m16n8k16.bf16, i.e. SM80 and newer. +# _W4A16 Weight-only 4-bit: converts packed operands with the +# __nv_cvt_fp4x2_to_halfraw2 intrinsic and accumulates in bf16, so +# it needs no block-scaled MMA. SM89 and newer emit the hardware +# conversion; older targets take the intrinsic's software path. +# _W4A4 Block-scaled 4-bit MMA through CUTLASS. Requires the +# CUTE_ARCH_F8F6F4_MMA_ENABLED path, which only sm_120a/sm_121a +# provide. On other targets CUTLASS still compiles these +# translation units but replaces the MMA with +# CUTE_INVALID_CONTROL_PATH, so the gate must be explicit: a +# successful build would otherwise produce kernels that fail at +# run time. +# +# FLASHRT_ENABLE_QWEN35MOE remains the single switch for a full Blackwell +# build and turns on all three tiers, so existing configure lines are +# unchanged. option(FLASHRT_ENABLE_QWEN35MOE - "Build qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) RTX SM120 kernels" OFF) + "Build all qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) kernel tiers" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_CORE + "Build architecture-neutral qwen3_5_moe kernels" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_W4A16 + "Build weight-only 4-bit qwen3_5_moe kernels" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_W4A4 + "Build block-scaled 4-bit MMA qwen3_5_moe kernels (sm_120a/sm_121a)" OFF) +if(FLASHRT_ENABLE_QWEN35MOE) + set(FLASHRT_ENABLE_QWEN35MOE_CORE ON) + set(FLASHRT_ENABLE_QWEN35MOE_W4A16 ON) + set(FLASHRT_ENABLE_QWEN35MOE_W4A4 ON) +endif() +# The upper tiers reuse the core layout, reducer, and activation kernels. +if(FLASHRT_ENABLE_QWEN35MOE_W4A16 OR FLASHRT_ENABLE_QWEN35MOE_W4A4) + set(FLASHRT_ENABLE_QWEN35MOE_CORE ON) +endif() # Qwen3-VL adds a handful of kernels (rotate_half RoPE, ...) that the shared # binary does not yet need. Built into a SEPARATE flash_rt_qwen3_vl_kernels # module so flash_rt_kernels.so stays stable; OFF by default. @@ -628,6 +697,7 @@ if(GPU_ARCH STREQUAL "110") POSITION_INDEPENDENT_CODE ON ) target_include_directories(cutlass_nvfp4_w4a16_sm100_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/csrc ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm/fp4 ${CUTLASS_INCLUDE} @@ -642,6 +712,37 @@ if(GPU_ARCH STREQUAL "110") message(STATUS "SM100 CUTLASS NVFP4 W4A16 GEMM (Thor): ENABLED") endif() +# ── Grouped NVFP4 MoE GEMM (qwen3_5_moe weight-only tier, Thor SM110) ── +# One launch per layer over every routed expert, with the per-group shapes read +# from device memory. Only the qwen3_5_moe MoE prefill calls it, so it is its +# own object gated on that model's tier rather than a second source in the +# W4A16 object above: a Thor build that does not ask for this model must not +# pay its CUTLASS grouped-kernel compile time or carry its symbols. The +# matching bindings are guarded on FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100. +if(GPU_ARCH STREQUAL "110" AND FLASHRT_ENABLE_QWEN35MOE_W4A16) + add_library(qwen35moe_nvfp4_grouped_sm100_obj OBJECT + csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu) + set_target_properties(qwen35moe_nvfp4_grouped_sm100_obj PROPERTIES + CUDA_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + ) + target_include_directories(qwen35moe_nvfp4_grouped_sm100_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/csrc + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm/fp4 + ${CUTLASS_INCLUDE} + ${CUTLASS_DIR}/tools/util/include + ) + target_compile_options(qwen35moe_nvfp4_grouped_sm100_obj PRIVATE + $<$: + --expt-relaxed-constexpr -O3 --use_fast_math + ${GPU_GENCODE} + > + ) + set(ENABLE_QWEN35MOE_GROUPED_SM100 ON) + message(STATUS "qwen3_5_moe grouped NVFP4 MoE GEMM (Thor): ENABLED") +endif() + # ── Dedicated M=1 decode GEMV (FP8 + BF16) ── # Origin sm_120a, but neither kernel uses an SM120-specific instruction, so both # compile for Thor (sm110) and Ada (sm89) too. Independent of @@ -954,6 +1055,15 @@ if(ENABLE_FA2 AND > ) message(STATUS "FA2 vendor arch: sm_${GPU_ARCH} AOT only (FA2_ARCH_NATIVE_ONLY=ON)") + elseif(GPU_ARCH STREQUAL "110") + # Native SASS via the same gencode the rest of this build uses, not + # compute_80 PTX. The 5090 learned that lesson: routing it through + # compute_80 PTX produced SASS that drifted from the native build by + # about an fp16 ULP a layer. + target_compile_options(fa2_vendor_obj PRIVATE + $<$:${GPU_GENCODE}> + ) + message(STATUS "FA2 vendor arch: sm_110a AOT only (FA2_ARCH_NATIVE_ONLY=ON)") elseif(GPU_ARCH STREQUAL "120") target_compile_options(fa2_vendor_obj PRIVATE $<$: @@ -975,16 +1085,26 @@ if(ENABLE_FA2 AND # build). See commit history for the reason each gencode matters — # routing 5090 through compute_80 PTX produced a subtle sm_120 SASS # drift that accumulated to cos 0.98 under Pi0 FP8. - target_compile_options(fa2_vendor_obj PRIVATE - $<$: - "SHELL:-gencode arch=compute_80,code=sm_80" - "SHELL:-gencode arch=compute_120,code=sm_120" - "SHELL:-gencode arch=compute_120,code=compute_120" - "SHELL:-gencode arch=compute_121,code=sm_121" - "SHELL:-gencode arch=compute_121,code=compute_121" - > - ) - message(STATUS "FA2 vendor arch: sm_80 + sm_120/sm_121 AOT + Blackwell PTX fallback (default)") + # sm_110 is not in the consumer family this multi-arch list distributes + # for, and compute_80 SASS does not run on it, so it takes its own + # gencode here as well rather than silently getting no cubin. + if(GPU_ARCH STREQUAL "110") + target_compile_options(fa2_vendor_obj PRIVATE + $<$:${GPU_GENCODE}> + ) + message(STATUS "FA2 vendor arch: sm_110a AOT (Thor)") + else() + target_compile_options(fa2_vendor_obj PRIVATE + $<$: + "SHELL:-gencode arch=compute_80,code=sm_80" + "SHELL:-gencode arch=compute_120,code=sm_120" + "SHELL:-gencode arch=compute_120,code=compute_120" + "SHELL:-gencode arch=compute_121,code=sm_121" + "SHELL:-gencode arch=compute_121,code=compute_121" + > + ) + message(STATUS "FA2 vendor arch: sm_80 + sm_120/sm_121 AOT + Blackwell PTX fallback (default)") + endif() endif() # Macros the wrapper reads to #ifdef-guard optional hdim/dtype @@ -1520,30 +1640,72 @@ else() message(STATUS "NVFP4 swizzle/quantize kernels: SKIPPED (slim build)") endif() -# ── qwen3_5_moe family (Nex-N2 / Qwen3.6-35B-A3B) SM120 kernels ── -# Gated: only compiled when explicitly enabled AND on a Blackwell (NVFP4) -# build. The bindings are #ifdef FLASHRT_HAVE_QWEN35MOE in bindings.cpp, so -# the symbols are absent (and the .cu never compiled) on every other target. -if(FLASHRT_ENABLE_QWEN35MOE AND ENABLE_NVFP4) +# ── qwen3_5_moe family (Nex-N2 / Qwen3.6-35B-A3B) kernels ── +# Three tiers, each with its own compile definition. The matching bindings in +# csrc/bindings.cpp are guarded on the same macros, so a tier that is off +# contributes neither symbols nor translation units. +if(FLASHRT_ENABLE_QWEN35MOE_CORE) target_sources(flash_rt_kernels PRIVATE csrc/kernels/qwen35moe_layout.cu - csrc/kernels/moe_grouped_gemv_sm120.cu csrc/kernels/bf16_matvec_sm120.cu - csrc/kernels/w4a16_matvec_sm120.cu - csrc/kernels/moe_grouped_w4a16_sm120.cu csrc/kernels/gdn_recurrent_seq_sm120.cu + csrc/kernels/gdn_wy_prefill_edge.cu + csrc/kernels/causal_conv1d_rows_edge.cu csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu - csrc/kernels/moe_m16_mma_sm120.cu - csrc/kernels/moe_m64_mma_sm120.cu - csrc/kernels/moe_blocktile_mma_sm120.cu + csrc/kernels/moe_route_prefill_edge.cu + csrc/kernels/moe_shared_combine_edge.cu csrc/kernels/moe_weighted_sum_sm120.cu + csrc/kernels/w16a16_gemm_sm120.cu + csrc/kernels/qwen35moe_e0m3_dequant.cu) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_CORE=1) + message(STATUS "qwen3_5_moe core kernels: ENABLED (sm_${GPU_ARCH})") +endif() + +if(FLASHRT_ENABLE_QWEN35MOE_W4A16) + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/w4a16_matvec_sm120.cu + csrc/kernels/moe_grouped_w4a16_sm120.cu + csrc/kernels/w4a16_edge_sm120.cu + csrc/kernels/w4a16_mrows_edge_sm120.cu csrc/kernels/w4a16_gemm_sm120.cu - csrc/kernels/w16a16_gemm_sm120.cu) - target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE=1) - message(STATUS "qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) SM120 kernels: ENABLED") -elseif(FLASHRT_ENABLE_QWEN35MOE) - message(STATUS "qwen3_5_moe SM120 kernels: SKIPPED (requires Blackwell NVFP4)") + csrc/kernels/qwen35moe_grouped_quant.cu) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_W4A16=1) + # Packed-weight loads in flight per row in the W4A16 GEMVs. The kernels + # default to 4, which is the value the SM120 path was validated with. Thor + # measures faster at 2 because the kernel is register-limited there (8 blocks + # per SM where every other limit allows 24); that is a property of a 20-SM + # part and is not assumed to transfer, so it is set only for that target. + # Either value produces bit-identical output; see w4a16_edge_sm120.cu. + if(GPU_ARCH STREQUAL "110") + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_W4A16_EDGE_UNROLL=2) + endif() + message(STATUS "qwen3_5_moe weight-only 4-bit kernels: ENABLED (sm_${GPU_ARCH})") +endif() + +# Block-scaled 4-bit MMA needs the sm_120a/sm_121a CUTLASS path. Refuse rather +# than build silently broken kernels on other targets. +if(FLASHRT_ENABLE_QWEN35MOE_W4A4) + if(NOT ENABLE_NVFP4) + message(FATAL_ERROR + "FLASHRT_ENABLE_QWEN35MOE_W4A4 requires a block-scaled-MMA target " + "(current GPU_ARCH=${GPU_ARCH}). CUTLASS compiles these kernels on " + "other architectures but replaces the MMA with an invalid control " + "path, so they would fail at run time. Use " + "FLASHRT_ENABLE_QWEN35MOE_CORE and FLASHRT_ENABLE_QWEN35MOE_W4A16 " + "instead.") + endif() + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/moe_grouped_gemv_sm120.cu + csrc/kernels/moe_m16_mma_sm120.cu + csrc/kernels/moe_m64_mma_sm120.cu + csrc/kernels/moe_blocktile_mma_sm120.cu) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_W4A4=1) + message(STATUS "qwen3_5_moe block-scaled 4-bit kernels: ENABLED (sm_${GPU_ARCH})") endif() # ── MelBandRoformer custom fused kernels (BF16/FP8, gated) ── @@ -1642,6 +1804,15 @@ if(GPU_ARCH STREQUAL "110") ENABLE_CUTLASS_SM100_NVFP4_W4A16=1) endif() +# The grouped MoE GEMM built above, linked and declared only when the model +# tier that calls it is on. +if(ENABLE_QWEN35MOE_GROUPED_SM100) + target_sources(flash_rt_kernels PRIVATE + $) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100=1) +endif() + if(GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121" OR GPU_ARCH STREQUAL "110") target_sources(flash_rt_kernels PRIVATE diff --git a/README.md b/README.md index 937bf5b5..5ddaf5d1 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,32 @@ DGX Spark / GB10: | NVFP4, 128 | **170.1 ms** | **40.42 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) | | NVFP4, 16 K | **8.545 s** | **54.94 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) | +#### Qwen3.6-35B-A3B + +RTX 5090: + +| Mode | Prefill | Decode | Source | +|---|---:|---:|---| +| NVFP4, 64 | **40.42 ms** | **257.95 tok/s** | [Qwen3.6-MoE usage](docs/qwen36_moe_usage.md#validation) | + +Jetson AGX Thor, against vLLM 0.26.0 on the same part and protocol: + +| Mode | Prefill | Decode | Source | +|---|---:|---:|---| +| NVFP4, 20 | **89.5 ms** (vLLM 102.3) | | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | +| NVFP4, 1 K | **216.0 ms** (vLLM 319.4) | **87.1 tok/s** (vLLM 31.6) | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | +| NVFP4, 2 K | **379.6 ms** (vLLM 495.0) | **86.3 tok/s** (vLLM 31.5) | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | +| NVFP4, 32 K | **7207.5 ms** (vLLM 7231.8) | | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | + +TTFT leads at every length from 20 to 32768 tokens; 128 K context reaches the +board at 2470 tok/s of prefill. The decode column was taken before a later +round that moved the steady step from 89.0 to 102.6 tok/s, and vLLM's side is +unaffected by it, so the ratios shown are lower bounds. + +Speculative decode with the MTP head reaches **106.74 tok/s** against 100.35 +plain in the same process, emitting the same tokens as greedy decoding. See +[speculative decode](docs/qwen36_moe_usage.md#speculative-decode). + #### Qwen3-8B | Hardware | Mode | Prefill | Decode | Source | @@ -732,7 +758,7 @@ extension modules: | Artifact | Size | What it contains | |---|---|---| | `flash_rt/flash_rt_kernels.so` | ~3 MB | Hand-written memory-bound kernels (norm, activation, fusion, FP8 quant, cuBLASLt wrappers, Thor FMHA). **Always built.** | -| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built only on RTX targets** — Thor skips it and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) for attention instead. | +| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built automatically on RTX targets.** Thor skips it by default and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) instead; `-DFLASHRT_ENABLE_THOR_FA2=ON` builds it there for the one model whose long prefill needs it (Qwen3.6, bf16 head_dim 256 — a single instantiation). | **Crucially — no `pip install flash-attn` required.** The FA2 kernel is vendored at source level and built into `flash_rt_fa2.so` during @@ -862,7 +888,7 @@ CMake reads `nvidia-smi --query-gpu=compute_cap` to pick the target arch. Override for cross-compilation or when auto-detect fails: ```bash -cmake -B build -S . -DGPU_ARCH=110 # Jetson AGX Thor (FA2 skipped, CUTLASS SM100 path ON) +cmake -B build -S . -DGPU_ARCH=110 # Jetson AGX Thor (FA2 opt-in, CUTLASS SM100 path ON) cmake -B build -S . -DGPU_ARCH=121 # DGX Spark / GB10 (FA2 sm_121 AOT, NVFP4 ON) cmake -B build -S . -DGPU_ARCH=120 # RTX 5090 (FA2 sm_120 AOT, NVFP4 ON) cmake -B build -S . -DGPU_ARCH=89 # RTX 4090 (FA2 sm_80 AOT natively runs on Ada) @@ -870,10 +896,12 @@ cmake -B build -S . -DGPU_ARCH=86 # RTX 3090 / A10 (FA2 sm_80 AOT) cmake -B build -S . -DGPU_ARCH=80 # A100 (FA2 sm_80 AOT) ``` -FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`. Other -arches (notably Thor SM110 and SM90 Hopper) route attention through -the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of FA2 — -`flash_rt_fa2.so` simply isn't built, and no runtime error results. +FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`, and on +Thor SM110 when `-DFLASHRT_ENABLE_THOR_FA2=ON` is passed. Other arches +(notably SM90 Hopper, and Thor without that flag) route attention +through the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of +FA2 — `flash_rt_fa2.so` simply isn't built, and no runtime error +results. ### Build timing (one-time) diff --git a/benchmarks/qwen36_moe_edge_decode.py b/benchmarks/qwen36_moe_edge_decode.py new file mode 100644 index 00000000..64c6eba3 --- /dev/null +++ b/benchmarks/qwen36_moe_edge_decode.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +"""First-light latency probe for the Qwen3.6-MoE (qwen3_5_moe) edge path. + +Reports the numbers quoted in ``docs/qwen36_moe_usage.md``: weight load time, +resident and peak allocation, prefill latency, and decode throughput on the +eager and the captured-graph paths. The two decode paths are compared token for +token, because a throughput number for a path that emits different text is not +a throughput number for the same work. + +Usage: + + PYTHONPATH=. python benchmarks/qwen36_moe_edge_decode.py \\ + --checkpoint /path/to/Qwen3.6-35B-A3B \\ + --prompt-tokens 64 --max-new-tokens 32 +""" + +from __future__ import annotations + +import argparse +import statistics +import time + +import torch + +GIB = 2 ** 30 + + +def _sync(device: str) -> None: + torch.cuda.synchronize(device) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True, + help="path to the BF16 checkpoint directory") + parser.add_argument("--prompt-tokens", type=int, default=64) + parser.add_argument("--max-new-tokens", type=int, default=64) + parser.add_argument("--prefill-reps", type=int, default=5) + parser.add_argument("--decode-reps", type=int, default=8) + parser.add_argument("--max-seq", type=int, default=512) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend + + # Select and initialise the device before touching the memory stats: they + # are per-device counters and are not addressable until then. + torch.cuda.set_device(args.device) + torch.cuda.init() + torch.cuda.reset_peak_memory_stats(args.device) + t0 = time.perf_counter() + frontend = Qwen36MoeTextFrontend( + args.checkpoint, device=args.device, max_seq=args.max_seq) + _sync(args.device) + load_s = time.perf_counter() - t0 + + print(f"runtime weight load {load_s:8.2f} s") + print(f"resident allocated after load " + f"{torch.cuda.memory_allocated(args.device) / GIB:8.2f} GiB") + print(f"peak allocated during load " + f"{torch.cuda.max_memory_allocated(args.device) / GIB:8.2f} GiB") + + base = frontend.tokenizer.encode( + "The quick brown fox jumps over the lazy dog. ") + ids = (base * (args.prompt_tokens // len(base) + 2))[:args.prompt_tokens] + + # Prefill: the first call carries warmup and lazy weight packing, so it is + # reported separately rather than averaged into the steady-state figure. + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + frontend.generate(max_new_tokens=1) + _sync(args.device) + first_ms = (time.perf_counter() - t0) * 1e3 + + warm = [] + for _ in range(args.prefill_reps): + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + frontend.generate(max_new_tokens=1) + _sync(args.device) + warm.append((time.perf_counter() - t0) * 1e3) + + print(f"first prefill, including warmup {first_ms:8.2f} ms") + print(f"subsequent prefill " + f"{min(warm):8.2f}-{max(warm):.2f} ms") + + def run(fn) -> tuple[list[float], list[int]]: + # Median and range over every repetition, not a best-of: a single best + # sample hides both contention and variance, and the baseline this is + # compared against reports the same shape. + rates, toks = [], None + for _ in range(args.decode_reps): + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + out = fn() + _sync(args.device) + rates.append(args.max_new_tokens / (time.perf_counter() - t0)) + toks = list(out) + return sorted(rates), toks + + def report(label: str, rates: list[float]) -> None: + med = statistics.median(rates) + print(f"{label:<44}{med:8.2f} tok/s " + f"(range {rates[0]:.2f}-{rates[-1]:.2f} over {len(rates)} runs)") + + state = frontend._decode_state + from flash_rt.frontends.torch import _nexn2_rtx_decode as dec + + def eager(): + t = torch.tensor(ids, dtype=torch.long, device=args.device) + with torch.no_grad(): + return dec.generate_greedy( + state, t, args.max_new_tokens, frontend._fvk, args.device) + + eager_rate, eager_toks = run(eager) + graph_rate, graph_toks = run( + lambda: frontend.generate(max_new_tokens=args.max_new_tokens)) + + report(f"{args.prompt_tokens}/{args.max_new_tokens} eager decode", + eager_rate) + report(f"{args.prompt_tokens}/{args.max_new_tokens} warm graph decode", + graph_rate) + free, total = torch.cuda.mem_get_info(args.device) + print(f"{'device free memory at exit':<44}{free / GIB:8.2f} GiB " + f"of {total / GIB:.2f} -- a shared device invalidates the timings") + same = eager_toks == graph_toks + print(f"eager and graph emit the same tokens {str(same):>8}" + f" ({sum(a == b for a, b in zip(eager_toks, graph_toks))}" + f"/{len(graph_toks)})") + if not same: + raise SystemExit("eager and captured decode disagree; " + "the throughput numbers are not comparable") + + +if __name__ == "__main__": + main() diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 94781d44..60679c64 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -39,6 +39,9 @@ #ifdef ENABLE_CUTLASS_SM100_NVFP4_W4A16 #include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh" #endif +#ifdef FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 +#include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" +#endif #ifdef ENABLE_ACTION_FFN_MEGAKERNEL_V6T #include "kernels/megakernel/action_ffn_megakernel_v6t_sm120.cuh" #endif @@ -179,22 +182,34 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #ifdef FLASHRT_HAVE_QWEN36_KERNELS #include "kernels/qwen36_misc.cuh" #endif -#ifdef FLASHRT_HAVE_QWEN35MOE +#ifdef FLASHRT_HAVE_QWEN35MOE_CORE #include "kernels/qwen35moe_layout.cuh" -#include "kernels/moe_grouped_gemv_sm120.cuh" #include "kernels/bf16_matvec_sm120.cuh" -#include "kernels/w4a16_matvec_sm120.cuh" -#include "kernels/moe_grouped_w4a16_sm120.cuh" #include "kernels/gdn_recurrent_seq_sm120.cuh" +#include "kernels/gdn_wy_prefill_edge.cuh" +#include "kernels/causal_conv1d_rows_edge.cuh" #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" +#include "kernels/moe_route_prefill_edge.cuh" +#include "kernels/moe_shared_combine_edge.cuh" +#include "kernels/moe_weighted_sum_sm120.cuh" +#include "kernels/w16a16_gemm_sm120.cuh" +#include "kernels/qwen35moe_e0m3_dequant.cuh" +#endif // FLASHRT_HAVE_QWEN35MOE_CORE +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 +#include "kernels/w4a16_matvec_sm120.cuh" +#include "kernels/moe_grouped_w4a16_sm120.cuh" +#include "kernels/w4a16_edge_sm120.cuh" +#include "kernels/w4a16_mrows_edge_sm120.cuh" +#include "kernels/w4a16_gemm_sm120.cuh" +#include "kernels/qwen35moe_grouped_quant.cuh" +#endif // FLASHRT_HAVE_QWEN35MOE_W4A16 +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 +#include "kernels/moe_grouped_gemv_sm120.cuh" #include "kernels/moe_m16_mma_sm120.cuh" #include "kernels/moe_m64_mma_sm120.cuh" #include "kernels/moe_blocktile_mma_sm120.cuh" -#include "kernels/moe_weighted_sum_sm120.cuh" -#include "kernels/w4a16_gemm_sm120.cuh" -#include "kernels/w16a16_gemm_sm120.cuh" -#endif // FLASHRT_HAVE_QWEN35MOE +#endif // FLASHRT_HAVE_QWEN35MOE_W4A4 #include "kernels/bf16_matvec_qwen36.cuh" #include "kernels/bf16_matmul_bf16.cuh" #ifdef FLASHRT_HAVE_QWEN36_KERNELS @@ -4679,17 +4694,21 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("W"), py::arg("out"), py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); + // max_algos=0 (the default) keeps the environment-driven autotune this + // entry has always had; a caller passes 1 to take the heuristic's own pick + // and get a run-to-run reproducible reduction order. See the header. m.def("bf16_matmul_cublaslt_bf16", [](uintptr_t x, uintptr_t W, uintptr_t out, - int M, int N, int K, uintptr_t stream) { + int M, int N, int K, uintptr_t stream, int max_algos) { flash_rt::kernels::bf16_matmul_cublaslt_bf16( reinterpret_cast(x), reinterpret_cast(W), reinterpret_cast<__nv_bfloat16*>(out), - M, N, K, to_stream(stream)); + M, N, K, to_stream(stream), max_algos); }, py::arg("x"), py::arg("W"), py::arg("out"), - py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); + py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0, + py::arg("max_algos") = 0); #ifdef FLASHRT_HAVE_QWEN36_KERNELS m.def("bf16_matmul_qwen36_bf16", @@ -5469,6 +5488,36 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("head_k_dim"), py::arg("head_v_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + m.def("gated_deltanet_recurrent_edge_qwen36_bf16", + [](uintptr_t q, uintptr_t k, uintptr_t v, + uintptr_t g, uintptr_t beta, + uintptr_t state, uintptr_t out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, uintptr_t stream) { + const int rc = + flash_rt::kernels::gated_deltanet_recurrent_edge_qwen36_bf16( + to_ptr(q), to_ptr(k), to_ptr(v), + to_ptr(g), to_ptr(beta), + to_ptr(state), to_ptr(out), + B, num_v_heads, head_k_dim, head_v_dim, + use_qk_l2norm, to_stream(stream)); + if (rc != 0) { + throw std::runtime_error( + "gated_deltanet_recurrent_edge_qwen36_bf16 failed with " + + std::to_string(rc) + " for B=" + std::to_string(B) + + " num_v_heads=" + std::to_string(num_v_heads) + + " head_k_dim=" + std::to_string(head_k_dim) + + " head_v_dim=" + std::to_string(head_v_dim) + + " (this entry supports head dims of 128 only)"); + } + }, + py::arg("q"), py::arg("k"), py::arg("v"), + py::arg("g"), py::arg("beta"), + py::arg("state"), py::arg("out"), + py::arg("B"), py::arg("num_v_heads"), + py::arg("head_k_dim"), py::arg("head_v_dim"), + py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + // In/out-state variant for K-iter chained per-step save (A2c-3). m.def("gated_deltanet_recurrent_inout_qwen36_bf16", [](uintptr_t q, uintptr_t k, uintptr_t v, @@ -5588,7 +5637,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("S"), py::arg("stream") = 0); #endif // FLASHRT_HAVE_QWEN36_KERNELS (gated_deltanet_qwen36 part 1) -#ifdef FLASHRT_HAVE_QWEN35MOE +#ifdef FLASHRT_HAVE_QWEN35MOE_CORE m.def("qwen35moe_lin_split_qkv_broadcast_bf16", [](uintptr_t conv_out, uintptr_t q32, uintptr_t k32, uintptr_t v32, @@ -5621,61 +5670,6 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("W"), py::arg("out"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); - m.def("w4a16_matvec_sm120_bf16", - [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, - int N, int K, float alpha, uintptr_t stream) -> int { - return flash_rt::kernels::w4a16_matvec_sm120_bf16( - to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), - N, K, alpha, to_stream(stream)); - }, - py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), - py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); - - m.def("moe_m16_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_m16_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - - m.def("moe_m64_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_m64_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - - m.def("moe_blocktile_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_blocktile_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - m.def("moe_weighted_sum_sm120_bf16", [](uintptr_t d_dn, uintptr_t rows, uintptr_t tw, uintptr_t out, int S, int TOPK, int HID, int dn_stride, uintptr_t stream) -> int { @@ -5687,17 +5681,6 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("S"), py::arg("TOPK"), py::arg("HID"), py::arg("dn_stride"), py::arg("stream") = 0); - m.def("w4a16_gemm_sm120_bf16", - [](uintptr_t X, uintptr_t W, uintptr_t SFB, uintptr_t Y, - int M, int N, int K, float alpha, uintptr_t stream) -> int { - return flash_rt::gemm::w4a16_gemm_sm120_bf16( - to_ptr(X), to_ptr(W), to_ptr(SFB), to_ptr(Y), - M, N, K, alpha, to_stream(stream)); - }, - py::arg("X"), py::arg("W"), py::arg("SFB"), py::arg("Y"), - py::arg("M"), py::arg("N"), py::arg("K"), - py::arg("alpha") = 1.0f, py::arg("stream") = 0); - m.def("w16a16_gemm_sm120_bf16", [](uintptr_t X, uintptr_t W, uintptr_t Y, int M, int N, int K, float alpha, uintptr_t stream) -> int { @@ -5709,6 +5692,16 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("M"), py::arg("N"), py::arg("K"), py::arg("alpha") = 1.0f, py::arg("stream") = 0); + m.def("moe_router_topk_warp_sm120_bf16", + [](uintptr_t logits, uintptr_t out_idx, uintptr_t out_val, + int n_experts, int k, uintptr_t stream) { + return flash_rt::kernels::moe_router_topk_warp_sm120_bf16( + to_ptr(logits), to_ptr(out_idx), to_ptr(out_val), + n_experts, k, to_stream(stream)); + }, + py::arg("logits"), py::arg("out_idx"), py::arg("out_val"), + py::arg("n_experts"), py::arg("k"), py::arg("stream") = 0); + m.def("moe_router_topk_sm120_bf16", [](uintptr_t logits, uintptr_t out_idx, uintptr_t out_val, int n_experts, int k, uintptr_t stream) -> int { @@ -5737,6 +5730,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("gate"), py::arg("out"), py::arg("n"), py::arg("stream") = 0); + m.def("qwen35moe_e0m3_dequant_bf16", + [](uintptr_t packed, uintptr_t scale, uintptr_t out, + int rows, int cols, int group_size, float global_scale, + uintptr_t stream) -> int { + return flash_rt::kernels::qwen35moe_e0m3_dequant_bf16( + to_ptr(packed), to_ptr(scale), to_ptr(out), + rows, cols, group_size, global_scale, to_stream(stream)); + }, + py::arg("packed"), py::arg("scale"), py::arg("out"), + py::arg("rows"), py::arg("cols"), py::arg("group_size"), + py::arg("global_scale"), py::arg("stream") = 0); + m.def("gdn_recurrent_seq_sm120_bf16", [](uintptr_t q, uintptr_t k, uintptr_t v, uintptr_t g, uintptr_t beta, uintptr_t state, uintptr_t out, int S, int num_v_heads, @@ -5751,6 +5756,132 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + m.def("moe_shared_gate_combine_edge_bf16", + [](uintptr_t routed, uintptr_t shared, uintptr_t gate, uintptr_t out, + int S, int dim, uintptr_t stream) { + flash_rt::kernels::moe_shared_gate_combine_edge_bf16( + to_ptr(routed), to_ptr(shared), to_ptr(gate), to_ptr(out), + S, dim, to_stream(stream)); + }, + py::arg("routed"), py::arg("shared"), py::arg("gate"), + py::arg("out"), py::arg("S"), py::arg("dim"), + py::arg("stream") = 0); + + m.def("moe_route_prefill_bf16", + [](uintptr_t logits, uintptr_t ti, uintptr_t tw, uintptr_t se, + uintptr_t stok, uintptr_t inv, uintptr_t group_off, uintptr_t ws, + int ws_bytes, int S, int n_experts, int topk, + uintptr_t stream) -> int { + return flash_rt::kernels::moe_route_prefill_bf16( + to_ptr(logits), to_ptr(ti), to_ptr(tw), to_ptr(se), + to_ptr(stok), to_ptr(inv), to_ptr(group_off), to_ptr(ws), + ws_bytes, S, n_experts, topk, to_stream(stream)); + }, + py::arg("logits"), py::arg("ti"), py::arg("tw"), py::arg("se"), + py::arg("stok"), py::arg("inv"), py::arg("group_off"), py::arg("ws"), + py::arg("ws_bytes"), py::arg("S"), py::arg("n_experts"), + py::arg("topk"), py::arg("stream") = 0); + + m.def("moe_route_prefill_workspace_bytes", + [](int S, int topk, int n_experts) -> int { + return flash_rt::kernels::moe_route_prefill_workspace_bytes( + S, topk, n_experts); + }, + py::arg("S"), py::arg("topk"), py::arg("n_experts")); + + m.def("moe_route_sfa_offsets", + [](uintptr_t group_off, uintptr_t sfa_off, int n_experts, int n_col, + uintptr_t stream) { + flash_rt::kernels::moe_route_sfa_offsets( + to_ptr(group_off), to_ptr(sfa_off), n_experts, n_col, + to_stream(stream)); + }, + py::arg("group_off"), py::arg("sfa_off"), py::arg("n_experts"), + py::arg("n_col"), py::arg("stream") = 0); + + m.def("causal_conv1d_qwen36_rows_hist_bf16", + [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t hist, + uintptr_t out, int B, int S, int conv_dim, int k, bool apply_silu, + uintptr_t stream) { + flash_rt::kernels::causal_conv1d_qwen36_rows_hist_bf16( + to_ptr(x), to_ptr(w), to_ptr(bias), to_ptr(hist), to_ptr(out), + B, S, conv_dim, k, apply_silu, to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("hist"), + py::arg("out"), py::arg("B"), py::arg("S"), py::arg("conv_dim"), + py::arg("k"), py::arg("apply_silu") = true, py::arg("stream") = 0); + + m.def("causal_conv1d_qwen36_rows_bf16", + [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t out, + int B, int S, int conv_dim, int k, bool apply_silu, + uintptr_t stream) { + flash_rt::kernels::causal_conv1d_qwen36_rows_bf16( + to_ptr(x), to_ptr(w), to_ptr(bias), to_ptr(out), + B, S, conv_dim, k, apply_silu, to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("out"), + py::arg("B"), py::arg("S"), py::arg("conv_dim"), py::arg("k"), + py::arg("apply_silu") = true, py::arg("stream") = 0); + + m.def("w4a16_mrows_edge_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t SFB, uintptr_t out, + int M, int N, int K, double alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_mrows_edge_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(SFB), to_ptr(out), + M, N, K, static_cast(alpha), to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("SFB"), py::arg("out"), + py::arg("M"), py::arg("N"), py::arg("K"), py::arg("alpha"), + py::arg("stream") = 0); + + m.def("gdn_wy_norm_pack_q_cumsum_edge_bf16", + [](uintptr_t q, uintptr_t k, uintptr_t g, uintptr_t k_l2, + uintptr_t q_pack, uintptr_t g_cumsum, int S, int num_k_heads, + int num_v_heads, int head_dim, int qk_group, uintptr_t stream) { + flash_rt::kernels::gdn_wy_norm_pack_q_cumsum_edge_bf16( + to_ptr(q), to_ptr(k), to_ptr(g), to_ptr(k_l2), + to_ptr(q_pack), to_ptr(g_cumsum), S, num_k_heads, + num_v_heads, head_dim, qk_group, to_stream(stream)); + }, + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("k_l2"), + py::arg("q_pack"), py::arg("g_cumsum"), py::arg("S"), + py::arg("num_k_heads"), py::arg("num_v_heads"), py::arg("head_dim"), + py::arg("qk_group"), py::arg("stream") = 0); + + m.def("gdn_wy_pack_v_edge_bf16", + [](uintptr_t v, uintptr_t v_pack, int S, int num_v_heads, + int head_dim, uintptr_t stream) { + flash_rt::kernels::gdn_wy_pack_v_edge_bf16( + to_ptr(v), to_ptr(v_pack), S, num_v_heads, head_dim, + to_stream(stream)); + }, + py::arg("v"), py::arg("v_pack"), py::arg("S"), + py::arg("num_v_heads"), py::arg("head_dim"), + py::arg("stream") = 0); +#endif // FLASHRT_HAVE_QWEN35MOE_CORE + +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 + m.def("w4a16_matvec_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, + int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_matvec_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), + N, K, alpha, to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), + py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); + + m.def("w4a16_gemm_sm120_bf16", + [](uintptr_t X, uintptr_t W, uintptr_t SFB, uintptr_t Y, + int M, int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::gemm::w4a16_gemm_sm120_bf16( + to_ptr(X), to_ptr(W), to_ptr(SFB), to_ptr(Y), + M, N, K, alpha, to_stream(stream)); + }, + py::arg("X"), py::arg("W"), py::arg("SFB"), py::arg("Y"), + py::arg("M"), py::arg("N"), py::arg("K"), + py::arg("alpha") = 1.0f, py::arg("stream") = 0); + m.def("moe_grouped_w4a16_sm120_bf16", [](uintptr_t A, uintptr_t W, uintptr_t sfb, uintptr_t alpha, uintptr_t eidx, uintptr_t D, int slots, int N, int K, @@ -5766,6 +5897,126 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); + // Bitwise-identical variants tuned for a part where these two are compute + // bound rather than bandwidth bound. See w4a16_edge_sm120.cuh. + m.def("w4a16_matvec_edge_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, + int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_matvec_edge_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), + N, K, alpha, to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), + py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); + + m.def("moe_grouped_w4a16_edge_sm120_bf16", + [](uintptr_t A, uintptr_t W, uintptr_t sfb, uintptr_t alpha, + uintptr_t eidx, uintptr_t D, int slots, int N, int K, + long a_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::kernels::moe_grouped_w4a16_edge_sm120_bf16( + to_ptr(A), to_ptr(W), to_ptr(sfb), to_ptr(alpha), to_ptr(eidx), + to_ptr(D), slots, N, K, a_stride, w_stride, sfb_stride, + to_stream(stream)); + }, + py::arg("A"), py::arg("W"), py::arg("sfb"), py::arg("alpha"), + py::arg("eidx"), py::arg("D"), py::arg("slots"), py::arg("N"), + py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + + // Grouped NVFP4 activation quantisers (csrc/kernels/qwen35moe_grouped_quant.cu). + // Only the MoE prefill of this model calls them, and they write the + // grouped GEMM's per-group scale-factor layout rather than the general + // quantiser's, so they are built and declared with this tier. + m.def("moe_grouped_silu_quant_nvfp4_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + + m.def("moe_grouped_silu_quant_nvfp4_warp_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_warp_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + + m.def("moe_grouped_quant_nvfp4_bf16", + [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t src_row, + uintptr_t out_packed, uintptr_t out_sf, + int slots, int K, uintptr_t stream) -> int { + return moe_grouped_quant_nvfp4_bf16( + to_ptr(A), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(src_row), + to_ptr(out_packed), to_ptr(out_sf), + slots, K, to_stream(stream)); + }, + py::arg("A"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("src_row"), + py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("K"), py::arg("stream") = 0); +#endif // FLASHRT_HAVE_QWEN35MOE_W4A16 + +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 + m.def("moe_m16_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_m16_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + + m.def("moe_m64_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_m64_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + + m.def("moe_blocktile_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_blocktile_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + m.def("moe_grouped_gemv_sm120_bf16", [](uintptr_t A_stack, uintptr_t B_stack, uintptr_t D, uintptr_t SFA_stack, uintptr_t SFB_stack, @@ -5787,7 +6038,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("a_stride"), py::arg("sfa_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); -#endif // FLASHRT_HAVE_QWEN35MOE +#endif // FLASHRT_HAVE_QWEN35MOE_W4A4 #ifdef FLASHRT_HAVE_QWEN36_KERNELS m.def("qwen36_gdn_gating_bf16", @@ -7497,7 +7748,42 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; m.def("nvfp4_sf_swizzled_bytes", &flash_rt::fp4::nvfp4_sf_swizzled_bytes, py::arg("rows"), py::arg("D")); -#endif +#endif // ENABLE_CUTLASS_SM100_NVFP4_W4A16 + +// Grouped NVFP4 MoE GEMM (qwen3_5_moe weight-only tier on Thor). Its own +// object library and its own gate: a Thor build that does not ask for this +// model neither compiles it nor exports these names. +#ifdef FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 + // Every routed expert of a layer in one launch, with the per-group shapes + // taken from device memory so the routing never reaches the host. + m.def("moe_grouped_gemm_nvfp4_sm100_bf16out", + [](uintptr_t A_packed, uintptr_t SFA, uintptr_t W_stack, + uintptr_t SFB_stack, uintptr_t alpha_dev, uintptr_t D, + uintptr_t group_off, uintptr_t sfa_off, + int groups, int N, int K, long w_stride, long sfb_stride, + uintptr_t scratch, size_t scratch_bytes, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_bf16out( + to_ptr(A_packed), to_ptr(SFA), to_ptr(W_stack), + to_ptr(SFB_stack), to_ptr(alpha_dev), to_ptr(D), + to_ptr(group_off), to_ptr(sfa_off), + groups, N, K, w_stride, sfb_stride, + to_ptr(scratch), scratch_bytes, to_stream(stream)); + }, + py::arg("A_packed"), py::arg("SFA"), py::arg("W_stack"), + py::arg("SFB_stack"), py::arg("alpha_dev"), py::arg("D"), + py::arg("group_off"), py::arg("sfa_off"), py::arg("groups"), + py::arg("N"), py::arg("K"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("scratch"), + py::arg("scratch_bytes"), py::arg("stream") = 0); + + m.def("moe_grouped_gemm_nvfp4_sm100_scratch_bytes", + [](int groups) -> size_t { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_scratch_bytes( + groups); + }, + py::arg("groups")); +#endif // FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 #ifdef ENABLE_ACTION_FFN_MEGAKERNEL_V6T // Action FFN megakernel V6tuned (ku256_sd4_su3 tile). Fused FP8 diff --git a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu new file mode 100644 index 00000000..49d78997 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu @@ -0,0 +1,312 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Grouped NVFP4 block-scaled GEMM for sm_100-class Blackwell. See header. + +#include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" +#include +#include + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "cute/tensor.hpp" + +namespace flash_rt { +namespace gemm { + +namespace { + +using namespace cute; + +using ElementA = cutlass::float_e2m1_t; +using ElementB = cutlass::float_e2m1_t; +using ElementC = cutlass::bfloat16_t; +using ElementD = cutlass::bfloat16_t; +using ElementAccumulator = float; +using ElementSF = cutlass::float_ue4m3_t; + +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::RowMajor; + +using ElementPairA = cutlass::nv_float4_t; +using ElementPairB = cutlass::nv_float4_t; + +constexpr int AlignmentA = 32; +constexpr int AlignmentB = 32; +constexpr int AlignmentC = 8; +constexpr int AlignmentD = 8; + +using ProblemShape = cutlass::gemm::GroupProblemShape>; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using ClusterShape = Shape; +// N of 128 rather than 256, measured. A prefill routes about sixty-four rows +// to the average expert across 256 groups of unequal size, and the narrower +// tile gives the scheduler twice as many blocks to balance them across twenty +// SMs. Paired against N=256 on the same machine state: 377.9/378.2/377.9 ms +// against 429.6/386.3/387.6 -- the worst run of this tile beats the best run +// of the other, and the spread goes from 43 ms to 0.3. +using MmaTileShape = Shape<_128, _128, _256>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + Shape<_128, _64>, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutC*, AlignmentC, + ElementD, LayoutC*, AlignmentD, + cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementPairA, LayoutA*, AlignmentA, + ElementPairB, LayoutB*, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +using StrideA = typename Gemm::GemmKernel::InternalStrideA; +using StrideB = typename Gemm::GemmKernel::InternalStrideB; +using StrideC = typename Gemm::GemmKernel::InternalStrideC; +using StrideD = typename Gemm::GemmKernel::InternalStrideD; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; +using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; +using Sm1xxBlkScaledConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + +using ProblemShapeMNK = typename ProblemShape::UnderlyingProblemShape; + +// Everything the launch needs, laid out back to back in one scratch buffer so +// the caller allocates once. Filled by a device kernel from the routing, which +// is what keeps the whole thing free of a host round trip. +struct GroupArgs { + ProblemShapeMNK* shapes; + const ElementA** ptr_A; + const ElementB** ptr_B; + const ElementSF** ptr_SFA; + const ElementSF** ptr_SFB; + ElementD** ptr_D; + const float** ptr_alpha; + StrideA* stride_A; + StrideB* stride_B; + StrideC* stride_C; + StrideD* stride_D; + LayoutSFA* layout_SFA; + LayoutSFB* layout_SFB; +}; + +constexpr size_t align_up(size_t v, size_t a) { return (v + a - 1) / a * a; } + +size_t args_bytes(int g) { + size_t n = 0; + n = align_up(n + sizeof(ProblemShapeMNK) * g, 256); + n = align_up(n + sizeof(void*) * g * 6, 256); // A B SFA SFB D alpha + n = align_up(n + sizeof(StrideA) * g, 256); + n = align_up(n + sizeof(StrideB) * g, 256); + n = align_up(n + sizeof(StrideC) * g, 256); + n = align_up(n + sizeof(StrideD) * g, 256); + n = align_up(n + sizeof(LayoutSFA) * g, 256); + n = align_up(n + sizeof(LayoutSFB) * g, 256); + return n; +} + +GroupArgs carve(void* base, int g) { + auto* p = static_cast(base); + size_t o = 0; + auto take = [&](size_t bytes) { + void* r = p + o; + o = align_up(o + bytes, 256); + return r; + }; + GroupArgs a{}; + a.shapes = static_cast(take(sizeof(ProblemShapeMNK) * g)); + auto* ptrs = static_cast(take(sizeof(void*) * g * 6)); + auto at = [&](int i) { return static_cast(ptrs + i * g); }; + a.ptr_A = static_cast(at(0)); + a.ptr_B = static_cast(at(1)); + a.ptr_SFA = static_cast(at(2)); + a.ptr_SFB = static_cast(at(3)); + a.ptr_D = static_cast(at(4)); + a.ptr_alpha = static_cast(at(5)); + a.stride_A = static_cast(take(sizeof(StrideA) * g)); + a.stride_B = static_cast(take(sizeof(StrideB) * g)); + a.stride_C = static_cast(take(sizeof(StrideC) * g)); + a.stride_D = static_cast(take(sizeof(StrideD) * g)); + a.layout_SFA = static_cast(take(sizeof(LayoutSFA) * g)); + a.layout_SFB = static_cast(take(sizeof(LayoutSFB) * g)); + return a; +} + +// One thread per group. Reads the routing (prefix sums of the per-expert token +// counts) and writes the descriptor arrays CUTLASS reads. No host involvement, +// which is the point: the launch shape below depends only on the group count. +__global__ void fill_group_args( + GroupArgs a, + const uint8_t* __restrict__ A_packed, + const uint8_t* __restrict__ SFA, + const uint8_t* __restrict__ W_stack, + const uint8_t* __restrict__ SFB_stack, + const float* __restrict__ alpha, + uint8_t* __restrict__ D, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + int groups, int N, int K, long w_stride, long sfb_stride) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= groups) return; + + const int off = group_off[e]; + const int m = group_off[e + 1] - off; + + a.shapes[e] = cute::make_shape(m, N, K); + a.ptr_A[e] = reinterpret_cast( + A_packed + static_cast(off) * (K / 2)); + a.ptr_B[e] = reinterpret_cast(W_stack + e * w_stride); + a.ptr_SFA[e] = reinterpret_cast(SFA + sfa_off[e]); + a.ptr_SFB[e] = reinterpret_cast(SFB_stack + e * sfb_stride); + a.ptr_D[e] = reinterpret_cast( + D + static_cast(off) * N * sizeof(ElementD)); + a.ptr_alpha[e] = alpha + e; + + a.stride_A[e] = cutlass::make_cute_packed_stride( + StrideA{}, cute::make_shape(m, K, 1)); + a.stride_B[e] = cutlass::make_cute_packed_stride( + StrideB{}, cute::make_shape(N, K, 1)); + a.stride_C[e] = cutlass::make_cute_packed_stride( + StrideC{}, cute::make_shape(m, N, 1)); + a.stride_D[e] = cutlass::make_cute_packed_stride( + StrideD{}, cute::make_shape(m, N, 1)); + a.layout_SFA[e] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA( + cute::make_shape(m, N, K, 1)); + a.layout_SFB[e] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB( + cute::make_shape(m, N, K, 1)); +} + +} // namespace + +size_t moe_grouped_gemm_nvfp4_sm100_scratch_bytes(int groups) { + if (groups <= 0) return 0; + // The CUTLASS workspace for a grouped launch scales with the group count and + // the scheduler, not with the token counts, so a bound taken at construction + // stays valid however the routing falls. + return args_bytes(groups) + static_cast(groups) * 1024 + (1u << 20); +} + +int moe_grouped_gemm_nvfp4_sm100_bf16out( + const void* A_packed, + const void* SFA, + const void* W_stack, + const void* SFB_stack, + const void* alpha_dev, + void* D, + const void* group_off, + const void* sfa_off, + int groups, + int N, + int K, + long w_stride, + long sfb_stride, + void* scratch, + size_t scratch_bytes, + cudaStream_t stream) { + if (!A_packed || !SFA || !W_stack || !SFB_stack || !alpha_dev || !D + || !group_off || !sfa_off || !scratch) return 1; + if (groups <= 0 || N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const size_t need = args_bytes(groups); + if (scratch_bytes < need) return 3; + + GroupArgs ga = carve(scratch, groups); + const int threads = 128; + fill_group_args<<<(groups + threads - 1) / threads, threads, 0, stream>>>( + ga, + static_cast(A_packed), + static_cast(SFA), + static_cast(W_stack), + static_cast(SFB_stack), + static_cast(alpha_dev), + static_cast(D), + static_cast(group_off), + static_cast(sfa_off), + groups, N, K, w_stride, sfb_stride); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + // The cluster is a runtime shape, so it was swept without recompiling: + // (1,1) (2,1) (1,2) (2,2) (4,1) (1,4) at the prefill shape, then the best + // two alternated three times each. The within-pair differences (+1.4%, + // -0.4%, +0.8%) came out smaller than the drift between runs (427 to 385 ms + // for the same setting), so the cluster shape does not move this. Left as a + // knob with the result written down rather than as a knob to try again. + static const dim3 kCluster = [] { + const char* v = std::getenv("FLASHRT_MOE_GROUPED_CLUSTER"); + int x = 1, y = 1; + if (v && std::sscanf(v, "%d,%d", &x, &y) == 2 && x >= 1 && y >= 1) { + return dim3(x, y, 1); + } + return dim3(1, 1, 1); + }(); + hw_info.cluster_shape = kCluster; + hw_info.cluster_shape_fallback = dim3(1, 1, 1); + + typename Gemm::Arguments args_proto{}; + // The fusion argument type is reachable only through an Arguments instance, + // which is how the CUTLASS example spells it too. + decltype(args_proto.epilogue.thread) fusion_args; + fusion_args.alpha = 0.0f; + fusion_args.alpha_ptr_array = ga.ptr_alpha; + fusion_args.dAlpha = {_0{}, _0{}, 1}; + fusion_args.beta = 0.0f; + fusion_args.beta_ptr_array = nullptr; + fusion_args.dBeta = {_0{}, _0{}, 0}; + + typename Gemm::GemmKernel::TileSchedulerArguments scheduler{}; + + // Host-side problem shapes are passed as nullptr deliberately: the shapes + // live only on device, so nothing here depends on the routing and the call + // is safe to capture. + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {groups, ga.shapes, nullptr}, + {ga.ptr_A, ga.stride_A, ga.ptr_B, ga.stride_B, + ga.ptr_SFA, ga.layout_SFA, ga.ptr_SFB, ga.layout_SFB}, + {fusion_args, nullptr, ga.stride_C, ga.ptr_D, ga.stride_D}, + hw_info, scheduler}; + + Gemm gemm; + const size_t ws = Gemm::get_workspace_size(args); + if (need + ws > scratch_bytes) return 4; + void* ws_ptr = static_cast(scratch) + need; + + auto status = gemm.can_implement(args); + if (status != cutlass::Status::kSuccess) { + std::fprintf(stderr, + "[moe_grouped_gemm_nvfp4_sm100] can_implement FAIL groups=%d N=%d " + "K=%d status=%d\n", groups, N, K, static_cast(status)); + return 5; + } + status = gemm.initialize(args, ws_ptr, stream); + if (status != cutlass::Status::kSuccess) return 6; + status = gemm.run(stream); + return status == cutlass::Status::kSuccess ? 0 : 7; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh new file mode 100644 index 00000000..84616804 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Grouped NVFP4 block-scaled GEMM for sm_100-class Blackwell (datacenter +// SM100 / Jetson AGX Thor SM110): every routed expert of a MoE layer in one +// launch. +// +// Why this exists. Prefill routes S*8 tokens across 256 experts. Serving that +// as one GEMV per (token, expert) slot re-reads an expert's weight once per +// token that chose it -- 9.7 GB a layer at S=1024 -- and serving it as one +// GEMM per expert costs 256 launches and 256 Python iterations a layer, whose +// host time exceeded the device time. Neither scales: the first is bounded by +// L2 bandwidth, the second by the host. +// +// A grouped GEMM is bounded by neither. One launch covers every expert, each +// weight is read once, and -- because CUTLASS accepts the per-group problem +// shapes from device memory (the host-side array is optional) -- the launch +// geometry is host-known and the routing is not, which is what a CUDA-graph +// capture requires. That is the property this is really for: capture a prefill +// chunk once and replay it for any context length, rather than tuning a +// threshold per prompt length. +// +// Wire format matches cutlass_nvfp4_w4a16_gemm_sm100: e2m1 nibbles, UE4M3 +// block scales of 16, Sm1xx block-scaled atom layout, BF16 out. The per-expert +// global scale enters as the epilogue's per-group alpha. + +#pragma once + +#include +#include + +namespace flash_rt { +namespace gemm { + +// Scratch the entry point needs, in bytes, for `groups` groups. Holds the +// per-group pointer/stride/layout arrays it fills on device, plus the CUTLASS +// workspace. Allocate once and reuse; it does not depend on the token counts. +size_t moe_grouped_gemm_nvfp4_sm100_scratch_bytes(int groups); + +// D[off_e : off_e + cnt_e, :] = A[off_e : off_e + cnt_e, :] @ W[e].T * alpha[e] +// +// A_packed (slots, K/2) u8 rows sorted by expert +// SFA per-group block-scaled atom layouts, group e at byte offset +// sfa_offsets[e] +// W_stack (E, N, K/2) u8 +// SFB_stack (E, sfb_bytes) u8 +// alpha_dev (E,) f32 device +// D (slots, N) bf16 +// group_off (E + 1,) i32 device, prefix sums of the per-expert counts +// sfa_off (E,) i32 device, byte offset of group e's SFA block +// +// Nothing is read to the host: the group shapes are derived on device from +// group_off. Returns 0 on success, nonzero on argument or CUTLASS error. +int moe_grouped_gemm_nvfp4_sm100_bf16out( + const void* A_packed, + const void* SFA, + const void* W_stack, + const void* SFB_stack, + const void* alpha_dev, + void* D, + const void* group_off, + const void* sfa_off, + int groups, + int N, + int K, + long w_stride, + long sfb_stride, + void* scratch, + size_t scratch_bytes, + cudaStream_t stream); + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/bf16_matmul_bf16.cu b/csrc/kernels/bf16_matmul_bf16.cu index 32d32cc7..6017dbb5 100644 --- a/csrc/kernels/bf16_matmul_bf16.cu +++ b/csrc/kernels/bf16_matmul_bf16.cu @@ -50,9 +50,14 @@ struct Bf16LtKey { int M; int N; int K; + // Part of the key, not just of the search: a caller that asked for a + // deterministic pick and one that let the timing loop run must not share + // whichever plan happened to be built first. + int max_algos; bool operator==(const Bf16LtKey& other) const { - return M == other.M && N == other.N && K == other.K; + return M == other.M && N == other.N && K == other.K + && max_algos == other.max_algos; } }; @@ -61,6 +66,7 @@ struct Bf16LtKeyHash { size_t h = static_cast(key.M); h = h * 1315423911u + static_cast(key.N); h = h * 1315423911u + static_cast(key.K); + h = h * 1315423911u + static_cast(key.max_algos); return h; } }; @@ -72,7 +78,10 @@ static std::mutex g_bf16_mu; static std::unordered_map g_bf16_plans; -static int get_bf16_autotune_algos() { +// requested > 0 is the caller's own bound; 0 falls back to the environment, +// then to the historical default of 8. +static int get_bf16_autotune_algos(int requested) { + if (requested > 0) return std::clamp(requested, 1, 32); const char* env = std::getenv("FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS"); if (!env || !*env) return 8; return std::clamp(std::atoi(env), 1, 32); @@ -94,10 +103,11 @@ static void ensure_bf16_lt() { } } -static Bf16LtPlan& get_bf16_lt_plan(int M, int N, int K) { +static Bf16LtPlan& get_bf16_lt_plan(int M, int N, int K, + int max_algos) { std::lock_guard lock(g_bf16_mu); ensure_bf16_lt(); - Bf16LtKey key{M, N, K}; + Bf16LtKey key{M, N, K, max_algos}; auto it = g_bf16_plans.find(key); if (it != g_bf16_plans.end()) return it->second; @@ -159,9 +169,10 @@ static void autotune_bf16_lt_plan( int M, int N, int K, - cudaStream_t stream) { + cudaStream_t stream, + int max_algos) { if (plan.autotuned) return; - const int num_algos = get_bf16_autotune_algos(); + const int num_algos = get_bf16_autotune_algos(max_algos); if (num_algos <= 1) { plan.autotuned = true; return; @@ -434,12 +445,13 @@ void bf16_matmul_cublaslt_bf16( const __nv_bfloat16* W, __nv_bfloat16* out, int M, int N, int K, - cudaStream_t stream) { + cudaStream_t stream, + int max_algos) { if (M <= 0 || N <= 0 || K <= 0) return; - Bf16LtPlan& plan = get_bf16_lt_plan(M, N, K); + Bf16LtPlan& plan = get_bf16_lt_plan(M, N, K, max_algos); if (!plan.autotuned) { std::lock_guard lock(g_bf16_mu); - autotune_bf16_lt_plan(plan, x, W, out, M, N, K, stream); + autotune_bf16_lt_plan(plan, x, W, out, M, N, K, stream, max_algos); } const float alpha = 1.0f; const float beta = 0.0f; diff --git a/csrc/kernels/bf16_matmul_bf16.cuh b/csrc/kernels/bf16_matmul_bf16.cuh index 9fd36aa4..0fb626bd 100644 --- a/csrc/kernels/bf16_matmul_bf16.cuh +++ b/csrc/kernels/bf16_matmul_bf16.cuh @@ -36,6 +36,18 @@ void bf16_matmul_bf16( int K, cudaStream_t stream); +// ``max_algos`` bounds how many cuBLASLt candidates the first call for a shape +// times before it commits to one. 0 (the default, and what every existing call +// site passes) keeps the current behaviour: the count comes from +// FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS, defaulting to 8. +// +// A caller passes 1 to take the heuristic's own first choice and skip the +// timing loop. That matters beyond speed: timing is noisy, so different +// processes pick different algorithms, different algorithms reduce in +// different orders, and a model whose output is compared token for token then +// disagrees with itself across runs. Plans are cached per (M, N, K, max_algos), +// so one caller asking for a deterministic pick does not decide the algorithm +// for another that did not. void bf16_matmul_cublaslt_bf16( const __nv_bfloat16* x, const __nv_bfloat16* W, @@ -43,6 +55,7 @@ void bf16_matmul_cublaslt_bf16( int M, int N, int K, - cudaStream_t stream); + cudaStream_t stream, + int max_algos = 0); } // namespace flash_rt::kernels diff --git a/csrc/kernels/causal_conv1d_rows_edge.cu b/csrc/kernels/causal_conv1d_rows_edge.cu new file mode 100644 index 00000000..3e92e13d --- /dev/null +++ b/csrc/kernels/causal_conv1d_rows_edge.cu @@ -0,0 +1,138 @@ +#include "causal_conv1d_rows_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kMaxK = 4; +constexpr int kThreadsX = 256; // matches the existing prefill entry + +// Written the same way as the existing entry's, not merely equivalent to it: +// the two are meant to agree to the bit. +__device__ __forceinline__ float rows_silu(float v) { + return v / (1.0f + __expf(-v)); +} + +// One thread, one channel, `kRows` consecutive tokens. The k-1 inputs a token +// shares with the next are kept in registers and shifted along, so the reads +// are one element per output rather than k. +template +__global__ void causal_conv1d_rows_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + const __nv_bfloat16* __restrict__ hist, + __nv_bfloat16* __restrict__ out, + int B, int S, int conv_dim, int k, + bool apply_silu) +{ + const int c = blockIdx.x * kThreadsX + threadIdx.x; + if (c >= conv_dim) return; + const int s0 = blockIdx.y * kRows; + if (s0 >= S) return; + const int b = blockIdx.z; + + float wv[kMaxK]; + #pragma unroll + for (int i = 0; i < kMaxK; ++i) { + wv[i] = (i < k) ? static_cast(w[c * k + i]) : 0.0f; + } + const float b0 = (bias != nullptr) ? static_cast(bias[c]) : 0.0f; + + const size_t base = static_cast(b) * S * conv_dim + c; + + // win[j] holds x[s0 - (k-1) + j], the window the first output needs. Before + // the start of this block that is the previous block's trailing inputs when + // there are any, and zero when the sequence itself starts here. + float win[kMaxK]; + #pragma unroll + for (int j = 0; j < kMaxK; ++j) { + const int t = s0 - (k - 1) + j; + if (j >= k) { win[j] = 0.0f; continue; } + if (t >= 0 && t < S) { + win[j] = static_cast(x[base + static_cast(t) * conv_dim]); + } else if (t < 0 && hist != nullptr) { + // hist is (B, conv_dim, k-1), newest last: t == -1 is the final column. + const int hj = t + (k - 1); + win[j] = static_cast( + hist[(static_cast(b) * conv_dim + c) * (k - 1) + hj]); + } else { + win[j] = 0.0f; + } + } + + #pragma unroll + for (int r = 0; r < kRows; ++r) { + const int s = s0 + r; + if (s >= S) break; + if (r > 0) { + // Shift by one and pull in the token that just became current. + #pragma unroll + for (int j = 0; j < kMaxK - 1; ++j) win[j] = win[j + 1]; + win[k - 1] = static_cast( + x[base + static_cast(s) * conv_dim]); + } + float acc = b0; + #pragma unroll + for (int i = 0; i < kMaxK; ++i) { + if (i < k) acc = fmaf(win[i], wv[i], acc); + } + if (apply_silu) acc = rows_silu(acc); + out[base + static_cast(s) * conv_dim] = __float2bfloat16(acc); + } +} + +} // namespace + +void causal_conv1d_qwen36_rows_bf16( + const void* x, + const void* w, + const void* bias, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream) +{ + if (B <= 0 || S <= 0 || conv_dim <= 0 || k <= 0 || k > kMaxK) return; + + causal_conv1d_qwen36_rows_hist_bf16(x, w, bias, nullptr, out, B, S, + conv_dim, k, apply_silu, stream); +} + +void causal_conv1d_qwen36_rows_hist_bf16( + const void* x, + const void* w, + const void* bias, + const void* hist, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream) +{ + if (B <= 0 || S <= 0 || conv_dim <= 0 || k <= 0 || k > kMaxK) return; + + constexpr int kRows = 8; + const dim3 block(kThreadsX); + const dim3 grid((conv_dim + kThreadsX - 1) / kThreadsX, + (S + kRows - 1) / kRows, + B); + causal_conv1d_rows_kernel<<>>( + reinterpret_cast(x), + reinterpret_cast(w), + reinterpret_cast(bias), + reinterpret_cast(hist), + reinterpret_cast<__nv_bfloat16*>(out), + B, S, conv_dim, k, apply_silu); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/causal_conv1d_rows_edge.cuh b/csrc/kernels/causal_conv1d_rows_edge.cuh new file mode 100644 index 00000000..82788f32 --- /dev/null +++ b/csrc/kernels/causal_conv1d_rows_edge.cuh @@ -0,0 +1,65 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// Causal depthwise conv1d over a whole prompt, several output tokens per +// thread. +// +// The existing prefill entry gives one thread one (channel, token), so each +// input element is fetched once for every output that needs it -- k times -- +// and the channel's weight row is fetched once per token. At the Qwen3.6 +// prefill shape that is 134 MB of reads for 67 MB of data, and the kernel +// measures about four times off what its traffic implies. +// +// Here a thread walks `rows` consecutive tokens of one channel, holding the +// last k inputs in registers, so each input is read once and the weight row +// once per thread rather than once per token. +// +// x (B, S, conv_dim) bf16 +// w (conv_dim, k) bf16 +// bias (conv_dim,) bf16 or null +// out (B, S, conv_dim) bf16 +// hist (B, conv_dim, k-1) bf16 or null -- the previous block's last k-1 +// inputs, channel-major with the newest last, which is the layout the +// decode conv state already carries. Null means the sequence starts +// here and the reads before it are zero. +// +// `hist` is what lets a chunked prefill stop concatenating. Prepending the +// history to the activations and slicing the result back off copies the whole +// block twice per layer -- 691 ms of a 32768-token prefill, in `cat` and its +// batched copy -- to supply three tokens of context. +// +// k must be at most 4. Layout, causality and the optional silu match the +// existing entry exactly; this is the same function computed with less +// traffic. +void causal_conv1d_qwen36_rows_bf16( + const void* x, + const void* w, + const void* bias, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream); + +// Same, continuing from a previous block's trailing inputs. +void causal_conv1d_qwen36_rows_hist_bf16( + const void* x, + const void* w, + const void* bias, + const void* hist, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/fp4_e2m1_compat.cuh b/csrc/kernels/fp4_e2m1_compat.cuh new file mode 100644 index 00000000..b4fe15ef --- /dev/null +++ b/csrc/kernels/fp4_e2m1_compat.cuh @@ -0,0 +1,68 @@ +// Packed-E2M1 to half2 conversion that does not require . +// +// The CUDA header only exists from 12.8 onwards, and on architectures without +// the cvt.rn.f16x2.e2m1x2 instruction it decodes each nibble in software +// anyway. E2M1 has sixteen representable values, so a table gives the same +// result on every target and removes the toolkit dependency: a Jetson image +// pinned to CUDA 12.6 can still build the weight-only 4-bit kernels. +// +// Where the header is present it is used, so the emitted code on those targets +// is unchanged. + +#pragma once + +#include +#include + +// Define FLASHRT_FP4_FORCE_TABLE to take the portable path even where the +// header exists. Used to check the two against each other. +#if !defined(FLASHRT_FP4_FORCE_TABLE) && defined(__has_include) +#if __has_include() +#define FLASHRT_HAVE_CUDA_FP4_HEADER 1 +#endif +#endif + +#ifdef FLASHRT_HAVE_CUDA_FP4_HEADER +#include +#endif + +namespace flash_rt { +namespace fp4 { + +// Decode one byte holding two E2M1 values, low nibble first. +__device__ __forceinline__ __half2_raw cvt_e2m1x2_to_halfraw2(uint8_t pair) { +#ifdef FLASHRT_HAVE_CUDA_FP4_HEADER + return __nv_cvt_fp4x2_to_halfraw2( + static_cast<__nv_fp4x2_storage_t>(pair), __NV_E2M1); +#else + // The sixteen E2M1 values as raw half bit patterns, indexed by the 4-bit + // code: one sign bit, two exponent bits, one mantissa bit, giving 0, + // +/-0.5, +/-1, +/-1.5, +/-2, +/-3, +/-4, +/-6. Function-local so no + // translation unit owns a device symbol. + constexpr unsigned short kAsHalfRaw[16] = { + 0x0000, // 0.0 + 0x3800, // 0.5 + 0x3C00, // 1.0 + 0x3E00, // 1.5 + 0x4000, // 2.0 + 0x4200, // 3.0 + 0x4400, // 4.0 + 0x4600, // 6.0 + 0x8000, // -0.0 + 0xB800, // -0.5 + 0xBC00, // -1.0 + 0xBE00, // -1.5 + 0xC000, // -2.0 + 0xC200, // -3.0 + 0xC400, // -4.0 + 0xC600, // -6.0 + }; + __half2_raw out; + out.x = kAsHalfRaw[pair & 0x0F]; + out.y = kAsHalfRaw[(pair >> 4) & 0x0F]; + return out; +#endif +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/kernels/gated_deltanet_qwen36.cu b/csrc/kernels/gated_deltanet_qwen36.cu index c80754c7..e8cd0a57 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cu +++ b/csrc/kernels/gated_deltanet_qwen36.cu @@ -150,6 +150,135 @@ __global__ void gated_deltanet_recurrent_kernel( __float2bfloat16(out_t); } + +// Same recurrence, without the local-memory round trip. +// +// The kernel above keeps the thread's whole state column in `float col[HD]`. +// A 128-iteration loop unrolled by 16 leaves the index non-constant, so the +// array cannot live in registers -- ncu measures 39 registers per thread for a +// 128-float array, which means it is in local memory, read and written across +// five passes. Against 2 MB of real state traffic that is roughly 6 MB of +// spill, and the kernel lands at 108 GB/s on a 244 GB/s part. +// +// The column never needs to be held. The recurrence reads the state twice -- +// once to form the k-weighted sum, once to update and emit -- and the whole +// state is 1 MB, so the second read is an L2 hit. Arithmetic, and the order of +// every accumulation, is identical to the kernel above; only where the +// intermediate lives changes. +template +__global__ void gated_deltanet_recurrent_edge_kernel( + const __nv_bfloat16* __restrict__ q_in, + const __nv_bfloat16* __restrict__ k_in, + const __nv_bfloat16* __restrict__ v_in, + const __nv_bfloat16* __restrict__ g_in, + const __nv_bfloat16* __restrict__ beta_in, + __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ out_, + int num_v_heads, + bool use_qk_l2norm) +{ + static_assert(HD == 128, "HD must be 128 for Qwen3.6 (single instantiation)"); + const int h = blockIdx.x; + const int b = blockIdx.y; + const int t = threadIdx.x; + if (t >= HD) return; + + __shared__ float smem[2 * HD + 32]; + float* qs = smem; + float* ks = smem + HD; + float* scratch = smem + 2 * HD; + + const size_t qkv_off = ((size_t)b * num_v_heads + h) * HD + t; + qs[t] = static_cast(q_in[qkv_off]); + ks[t] = static_cast(k_in[qkv_off]); + __syncthreads(); + + if (use_qk_l2norm) { + float q_sq = qs[t] * qs[t]; + float k_sq = ks[t] * ks[t]; + q_sq = block_reduce_sum(q_sq, scratch); + __syncthreads(); + k_sq = block_reduce_sum(k_sq, scratch); + const float q_inv = rsqrtf(q_sq + kEps); + const float k_inv = rsqrtf(k_sq + kEps); + qs[t] *= q_inv; + ks[t] *= k_inv; + __syncthreads(); + } + + qs[t] *= rsqrtf(static_cast(HD)); + __syncthreads(); + + const float g_t = + __expf(static_cast(g_in[b * num_v_heads + h])); + const float beta_t = + static_cast(beta_in[b * num_v_heads + h]); + + const size_t state_h_off = (((size_t)b * num_v_heads + h)) * HD * HD; + + // kv_mem[t] = sum_i (state[i][t] * g_t) * ks[i], accumulated in i order -- + // the same order, and the same rounded product, as the version that stored + // the column first. + float kv_mem = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + const float c = + static_cast(state[state_h_off + (size_t)i * HD + t]) * g_t; + kv_mem = fmaf(c, ks[i], kv_mem); + } + + const float v_t = + static_cast(v_in[(size_t)b * num_v_heads * HD + h * HD + t]); + const float delta = (v_t - kv_mem) * beta_t; + + // Second pass: re-derive the decayed column, apply the rank-one update, + // store it, and accumulate the output -- one read and one write per element. + float out_t = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + const size_t off = state_h_off + (size_t)i * HD + t; + const float c = fmaf(ks[i], delta, static_cast(state[off]) * g_t); + state[off] = __float2bfloat16(c); + out_t = fmaf(c, qs[i], out_t); + } + out_[(size_t)b * num_v_heads * HD + h * HD + t] = + __float2bfloat16(out_t); +} + +} // namespace + +int gated_deltanet_recurrent_edge_qwen36_bf16( + const void* q, + const void* k, + const void* v, + const void* g, + const void* beta, + void* state, + void* out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, + cudaStream_t stream) +{ + constexpr int kHD = 128; + if (head_k_dim != kHD || head_v_dim != kHD) return 2; + if (!q || !k || !v || !g || !beta || !state || !out) return 1; + if (B <= 0 || num_v_heads <= 0) return 3; + dim3 grid(num_v_heads, B); + dim3 block(kHD); + gated_deltanet_recurrent_edge_kernel<<>>( + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + reinterpret_cast(g), + reinterpret_cast(beta), + reinterpret_cast<__nv_bfloat16*>(state), + reinterpret_cast<__nv_bfloat16*>(out), + num_v_heads, use_qk_l2norm); + return 0; +} + +namespace { + } // namespace void gated_deltanet_recurrent_qwen36_bf16( diff --git a/csrc/kernels/gated_deltanet_qwen36.cuh b/csrc/kernels/gated_deltanet_qwen36.cuh index f497eff9..2e11c577 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cuh +++ b/csrc/kernels/gated_deltanet_qwen36.cuh @@ -60,6 +60,26 @@ void gated_deltanet_recurrent_qwen36_bf16( // to state_out (different buffer). Caller chains state_in[k+1] := // state_out[k] to support per-step state save without an extra // .copy_(state_save, state) launch per step. +// Spill-free variant of the above: identical arithmetic and accumulation +// order, but the thread's state column is re-read rather than held in a +// 128-float local array. Same arguments, same results, bit for bit. +// +// Shape-specialized: head_k_dim and head_v_dim must both be 128. Returns +// non-zero for a null pointer (1), an unsupported head dim (2) or a +// non-positive batch/head count (3), rather than leaving the output buffer +// undefined -- the binding turns that into an exception. +int gated_deltanet_recurrent_edge_qwen36_bf16( + const void* q, + const void* k, + const void* v, + const void* g, + const void* beta, + void* state, + void* out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, + cudaStream_t stream); + void gated_deltanet_recurrent_inout_qwen36_bf16( const void* q, const void* k, diff --git a/csrc/kernels/gdn_wy_prefill_edge.cu b/csrc/kernels/gdn_wy_prefill_edge.cu new file mode 100644 index 00000000..315c6f70 --- /dev/null +++ b/csrc/kernels/gdn_wy_prefill_edge.cu @@ -0,0 +1,195 @@ +#include "gdn_wy_prefill_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kChunk = 64; +constexpr float kEps = 1e-6f; // matches the sequential scan's l2 eps + +// Butterfly order, the same summation order the sibling WY normalisation uses. +// Reduction order decides the low bits here, so this is not interchangeable +// with the shuffle-down helper in common.cuh. +template +__device__ __forceinline__ float wy_block_sum(float val, float* smem) { + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + if (lane == 0) smem[warp] = val; + __syncthreads(); + if (warp == 0) { + val = (lane < (kHD / 32)) ? smem[lane] : 0.0f; + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + if (lane == 0) smem[0] = val; + } + __syncthreads(); + return smem[0]; +} + +// One block per (unique k-head, token). The block reduces both q and k over +// head_dim, writes the unique-head k, and scatters q into the qk_group v-head +// slots of the packed buffer -- so the GQA broadcast never materialises. +// +// The grid covers chunks * 64 tokens rather than S, so the threads past the +// end of the sequence are the ones that zero the packed tail. +template +__global__ void gdn_wy_norm_pack_q_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ k_l2, + __nv_bfloat16* __restrict__ q_pack, + int S, + int num_k_heads, + int num_v_heads, + int qk_group) +{ + const int t = threadIdx.x; + const int h = blockIdx.x; // unique k-head + const int s = blockIdx.y; // token, may run past S into the pad + if (t >= kHD || h >= num_k_heads) return; + + const int chunk = s / kChunk; + const int tt = s - chunk * kChunk; + + if (s >= S) { + const __nv_bfloat16 zero = __float2bfloat16(0.0f); + for (int r = 0; r < qk_group; ++r) { + const int vh = h * qk_group + r; + q_pack[((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * kHD + t] = zero; + } + return; + } + + // q and k arrive GQA-broadcast, so the group leader carries the value. + const size_t src = (static_cast(s) * num_v_heads + h * qk_group) + * kHD + t; + const float qv = static_cast(q[src]); + const float kv = static_cast(k[src]); + + __shared__ float scratch[32]; + const float q_sq = wy_block_sum(qv * qv, scratch); + __syncthreads(); // scratch is reused by the second reduction + const float k_sq = wy_block_sum(kv * kv, scratch); + __syncthreads(); + + const __nv_bfloat16 q_norm = __float2bfloat16(qv * rsqrtf(q_sq + kEps)); + const __nv_bfloat16 k_norm = __float2bfloat16(kv * rsqrtf(k_sq + kEps)); + + k_l2[(static_cast(s) * num_k_heads + h) * kHD + t] = k_norm; + + for (int r = 0; r < qk_group; ++r) { + const int vh = h * qk_group + r; + q_pack[((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * kHD + t] = q_norm; + } +} + +// One thread per (chunk, v-head): 64 dependent adds, chunks * num_v_heads of +// them in flight. The sibling path runs one block of num_v_heads threads +// serially over the whole sequence, which is fine for a decode step and two +// orders of magnitude off for a prefill. +__global__ void gdn_wy_cumsum_g_chunk_kernel( + const __nv_bfloat16* __restrict__ g, + __nv_bfloat16* __restrict__ g_cumsum, + int S, + int num_v_heads, + int chunks) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= chunks * num_v_heads) return; + const int chunk = idx / num_v_heads; + const int vh = idx - chunk * num_v_heads; + + const int s0 = chunk * kChunk; + const int s1 = min(s0 + kChunk, S); + float acc = 0.0f; + for (int s = s0; s < s1; ++s) { + const size_t off = static_cast(s) * num_v_heads + vh; + acc += static_cast(g[off]); + g_cumsum[off] = __float2bfloat16(acc); + } +} + +__global__ void gdn_wy_pack_v_kernel( + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ v_pack, + int S, + int num_v_heads, + int head_dim) +{ + const int t = threadIdx.x; + const int vh = blockIdx.x; + const int s = blockIdx.y; + if (t >= head_dim || vh >= num_v_heads) return; + + const int chunk = s / kChunk; + const int tt = s - chunk * kChunk; + const size_t dst = + ((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * head_dim + t; + v_pack[dst] = (s < S) + ? v[(static_cast(s) * num_v_heads + vh) * head_dim + t] + : __float2bfloat16(0.0f); +} + +} // namespace + +void gdn_wy_norm_pack_q_cumsum_edge_bf16( + const void* q, + const void* k, + const void* g, + void* k_l2, + void* q_pack, + void* g_cumsum, + int S, + int num_k_heads, + int num_v_heads, + int head_dim, + int qk_group, + cudaStream_t stream) +{ + if (S <= 0 || head_dim != 128) return; + const int chunks = (S + kChunk - 1) / kChunk; + + gdn_wy_norm_pack_q_kernel<128> + <<>>( + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast<__nv_bfloat16*>(k_l2), + reinterpret_cast<__nv_bfloat16*>(q_pack), + S, num_k_heads, num_v_heads, qk_group); + + const int total = chunks * num_v_heads; + gdn_wy_cumsum_g_chunk_kernel<<<(total + 127) / 128, 128, 0, stream>>>( + reinterpret_cast(g), + reinterpret_cast<__nv_bfloat16*>(g_cumsum), + S, num_v_heads, chunks); +} + +void gdn_wy_pack_v_edge_bf16( + const void* v, + void* v_pack, + int S, + int num_v_heads, + int head_dim, + cudaStream_t stream) +{ + if (S <= 0) return; + const int chunks = (S + kChunk - 1) / kChunk; + gdn_wy_pack_v_kernel<<>>( + reinterpret_cast(v), + reinterpret_cast<__nv_bfloat16*>(v_pack), + S, num_v_heads, head_dim); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/gdn_wy_prefill_edge.cuh b/csrc/kernels/gdn_wy_prefill_edge.cuh new file mode 100644 index 00000000..febb2106 --- /dev/null +++ b/csrc/kernels/gdn_wy_prefill_edge.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// WY chunked delta-rule front matter for a batched prefill, with the head +// counts as runtime arguments. +// +// The sibling 27B path has equivalent kernels, but its v-head count is a +// compile-time constant and its gate cumulative sum runs one block of +// `num_v_heads` threads serially over S -- a decode shape. These take the +// counts as arguments and parallelise the cumulative sum over chunks, which is +// what a prefill of a few thousand tokens needs. +// +// Layout conventions match the mma WY kernels: +// packed: (chunks, num_v_heads, 64, head_dim), chunks = ceil(S / 64), +// pack[c, h, i, d] = x[c * 64 + i, h, d], zero past S. +// g_cumsum: (S, num_v_heads), cumulative within each 64-token chunk. + +// Fuses the q/k l2 normalisation, the GQA broadcast of q into v-head slots, +// the chunk-major packing of q, and the gate cumulative sum. +// +// `q` and `k` are read as (S, num_v_heads, head_dim) already broadcast across +// the GQA group -- the form the conv split kernel writes -- and only the group +// leaders are touched, so no strided host-side slice is needed. +// +// q, k (S, num_v_heads, head_dim) bf16, GQA-broadcast +// g (S, num_v_heads) bf16 +// k_l2 (S, num_k_heads, head_dim) bf16 out, unique heads only +// q_pack (chunks, num_v_heads, 64, head_dim) bf16 out +// g_cumsum (S, num_v_heads) bf16 out +// +// head_dim must be 128. qk_group = num_v_heads / num_k_heads. +void gdn_wy_norm_pack_q_cumsum_edge_bf16( + const void* q, + const void* k, + const void* g, + void* k_l2, + void* q_pack, + void* g_cumsum, + int S, + int num_k_heads, + int num_v_heads, + int head_dim, + int qk_group, + cudaStream_t stream); + +// Chunk-major packing of the un-decayed v the chunk_h stage produces. +// v (S, num_v_heads, head_dim) bf16 +// v_pack (chunks, num_v_heads, 64, head_dim) bf16 out, zero past S +void gdn_wy_pack_v_edge_bf16( + const void* v, + void* v_pack, + int S, + int num_v_heads, + int head_dim, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_grouped_w4a16_sm120.cu b/csrc/kernels/moe_grouped_w4a16_sm120.cu index 8a82e89d..6599b6aa 100644 --- a/csrc/kernels/moe_grouped_w4a16_sm120.cu +++ b/csrc/kernels/moe_grouped_w4a16_sm120.cu @@ -7,7 +7,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -34,9 +34,8 @@ __device__ __forceinline__ float blockdot_g(uint64_t b_pack, float acc = 0.0f; #pragma unroll for (int j = 0; j < 8; ++j) { - const __nv_fp4x2_storage_t bb = - static_cast<__nv_fp4x2_storage_t>(b_pack >> (j * 8)); - const __half2_raw wr = __nv_cvt_fp4x2_to_halfraw2(bb, __NV_E2M1); + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); const float2 wf = __half22float2(*reinterpret_cast(&wr)); const float2 xf = __bfloat1622float2(xb2[j]); acc = fmaf(wf.x, xf.x, acc); diff --git a/csrc/kernels/moe_route_prefill_edge.cu b/csrc/kernels/moe_route_prefill_edge.cu new file mode 100644 index 00000000..10434e4e --- /dev/null +++ b/csrc/kernels/moe_route_prefill_edge.cu @@ -0,0 +1,348 @@ +#include "moe_route_prefill_edge.cuh" + +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kMaxExperts = 1024; +constexpr int kMaxTopK = 32; +constexpr int kSlotsPerBlock = 256; // slots a scatter/histogram block owns +constexpr int kRouteThreads = 256; + +// One warp per token, the whole row held in registers: PER_LANE experts per +// lane, strided so the row loads coalesced. A block-wide version of this cost +// eight barriers per top-k round -- sixty-four per token -- and ran seventy +// times off the bandwidth the row needs; there is no barrier here at all. +// +// Softmax first, then the top-k renormalised over itself. The full denominator +// cancels between the two, but it is kept because it only cancels exactly in +// exact arithmetic and this seeds a decode that has to reproduce. +// +// Ties go to the lower expert index. bf16 logits make the tail probabilities +// tie outright often enough to matter (6% of slots at 256 experts), and the +// tensor top-k this replaces does not define which of two equal experts it +// ranks first -- so the rank order inside a token's top-k can differ from it +// while the selected set, which is what the grouped GEMM reads, does not. +template +__global__ void route_topk_warp_kernel( + const __nv_bfloat16* __restrict__ logits, + int* __restrict__ ti, + float* __restrict__ tw, + int S, + int n_experts, + int topk) +{ + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int s = blockIdx.x * (blockDim.x >> 5) + warp; + if (s >= S) return; // whole warp, so the shuffles stay put + + const __nv_bfloat16* row = logits + static_cast(s) * n_experts; + float v[PER_LANE]; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) v[i] = static_cast(row[i * 32 + lane]); + + float m = -CUDART_INF_F; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) m = fmaxf(m, v[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + + float sum = 0.0f; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) { v[i] = __expf(v[i] - m); sum += v[i]; } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sum += __shfl_xor_sync(0xffffffff, sum, off); + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) v[i] /= sum; + + // Lane r keeps rank r, so the results end up spread one per lane and the + // write below is a single coalesced store. + float my_val = 0.0f; + int my_idx = 0; + for (int r = 0; r < topk; ++r) { + float best = -CUDART_INF_F; + int best_i = n_experts; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) { + const int e = i * 32 + lane; + if (v[i] > best || (v[i] == best && e < best_i)) { best = v[i]; best_i = e; } + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, best, off); + const int oi = __shfl_xor_sync(0xffffffff, best_i, off); + if (ov > best || (ov == best && oi < best_i)) { best = ov; best_i = oi; } + } + if (lane == r) { my_val = best; my_idx = best_i; } + // Compile-time indices: a computed one would push v[] into local memory. + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) + if (i * 32 + lane == best_i) v[i] = -CUDART_INF_F; + } + + float tsum = (lane < topk) ? my_val : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + tsum += __shfl_xor_sync(0xffffffff, tsum, off); + + if (lane < topk) { + const size_t base = static_cast(s) * topk; + ti[base + lane] = my_idx; + tw[base + lane] = my_val / tsum; + } +} + +// Per-block expert histogram over a fixed slice of slots. The warp-level match +// gives each lane its rank among the lanes of its warp holding the same +// expert, which is what makes the scatter's ordering fall out without atomics. +__global__ void slot_hist_kernel( + const int* __restrict__ ti, + int* __restrict__ blk_hist, + int slots, + int n_experts) +{ + const int blk = blockIdx.x; + const int t = threadIdx.x; + const int slot = blk * kSlotsPerBlock + t; + + extern __shared__ int s_hist[]; // n_experts + for (int e = t; e < n_experts; e += blockDim.x) s_hist[e] = 0; + __syncthreads(); + + if (slot < slots) atomicAdd(&s_hist[ti[slot]], 1); + __syncthreads(); + + int* out = blk_hist + static_cast(blk) * n_experts; + for (int e = t; e < n_experts; e += blockDim.x) out[e] = s_hist[e]; +} + +// One block per expert: exclusive scan of that expert's per-block counts, so a +// scatter block knows where its own slots for that expert begin. +// +// Scanned across the block in tiles rather than by one thread in a loop. The +// number of slot-blocks grows with the sequence, so a serial walk here is +// O(S) on a single thread -- invisible at two thousand tokens, and the reason +// the prefill rate fell away past four thousand. +template +__global__ void expert_block_scan_kernel( + const int* __restrict__ blk_hist, + int* __restrict__ blk_off, + int* __restrict__ counts, + int n_blocks, + int n_experts) +{ + const int e = blockIdx.x; + const int t = threadIdx.x; + __shared__ int s[kThreads]; + __shared__ int s_carry; + if (t == 0) s_carry = 0; + __syncthreads(); + + for (int base = 0; base < n_blocks; base += kThreads) { + const int b = base + t; + const size_t off = static_cast(b) * n_experts + e; + const int own = (b < n_blocks) ? blk_hist[off] : 0; + s[t] = own; + __syncthreads(); + + // Hillis-Steele inclusive scan; subtracting own value gives the exclusive + // one without a second pass. + for (int d = 1; d < kThreads; d <<= 1) { + const int add = (t >= d) ? s[t - d] : 0; + __syncthreads(); + s[t] += add; + __syncthreads(); + } + const int tile_total = s[kThreads - 1]; + if (b < n_blocks) blk_off[off] = s_carry + s[t] - own; + __syncthreads(); + if (t == 0) s_carry += tile_total; + __syncthreads(); + } + if (t == 0) counts[e] = s_carry; +} + +__global__ void group_off_kernel( + const int* __restrict__ counts, + int* __restrict__ group_off, + int n_experts) +{ + if (threadIdx.x != 0) return; + int acc = 0; + for (int e = 0; e < n_experts; ++e) { + group_off[e] = acc; + acc += counts[e]; + } + group_off[n_experts] = acc; +} + +// Places every slot at group_off[e] + blk_off[blk][e] + its rank within the +// block. Rank comes from the warp match plus the counts of the earlier warps, +// so two runs on the same routing place the same slot in the same row. +__global__ void slot_scatter_kernel( + const int* __restrict__ ti, + const int* __restrict__ group_off, + const int* __restrict__ blk_off, + int* __restrict__ se, + long* __restrict__ stok, + int* __restrict__ inv, + int slots, + int n_experts, + int topk) +{ + const int blk = blockIdx.x; + const int t = threadIdx.x; + const int slot = blk * kSlotsPerBlock + t; + const int warp = t >> 5; + const int lane = t & 31; + const int n_warps = blockDim.x >> 5; + + extern __shared__ int s_warp_hist[]; // n_warps * n_experts + for (int i = t; i < n_warps * n_experts; i += blockDim.x) s_warp_hist[i] = 0; + __syncthreads(); + + const int e = (slot < slots) ? ti[slot] : -1; + const unsigned active = __ballot_sync(0xffffffff, e >= 0); + int rank_in_warp = 0; + if (e >= 0) { + const unsigned same = __match_any_sync(active, e); + const unsigned lower = same & ((1u << lane) - 1u); + rank_in_warp = __popc(lower); + if (rank_in_warp == 0) { + s_warp_hist[warp * n_experts + e] = __popc(same); + } + } + __syncthreads(); + + if (e >= 0) { + int before = 0; + for (int w = 0; w < warp; ++w) before += s_warp_hist[w * n_experts + e]; + const int row = group_off[e] + + blk_off[static_cast(blk) * n_experts + e] + + before + rank_in_warp; + se[row] = e; + stok[row] = slot / topk; // 64-bit: see the note in the header + inv[slot] = row; + } +} + +__global__ void sfa_offsets_kernel( + const int* __restrict__ group_off, + int* __restrict__ sfa_off, + int n_experts, + int n_col) +{ + if (threadIdx.x != 0) return; + int acc = 0; + for (int e = 0; e < n_experts; ++e) { + sfa_off[e] = acc; + const int c = group_off[e + 1] - group_off[e]; + acc += ((c + 127) / 128) * (n_col * 512); + } +} + +int route_blocks(int slots) { + return (slots + kSlotsPerBlock - 1) / kSlotsPerBlock; +} + +} // namespace + +int moe_route_prefill_workspace_bytes(int S, int topk, int n_experts) +{ + const int slots = S * topk; + const int nblk = route_blocks(slots); + // blk_hist + blk_off + counts + return static_cast( + (2 * static_cast(nblk) * n_experts + n_experts) * sizeof(int)); +} + +int moe_route_prefill_bf16( + const void* logits, + void* ti, + void* tw, + void* se, + void* stok, + void* inv, + void* group_off, + void* ws, + int ws_bytes, + int S, + int n_experts, + int topk, + cudaStream_t stream) +{ + if (S <= 0) return 0; + if (n_experts <= 0 || n_experts > kMaxExperts || (n_experts % 32) != 0) + return 1; + if (topk <= 0 || topk > kMaxTopK || topk > n_experts) return 2; + if (ws_bytes < moe_route_prefill_workspace_bytes(S, topk, n_experts)) + return 3; + + const int slots = S * topk; + const int nblk = route_blocks(slots); + int* blk_hist = reinterpret_cast(ws); + int* blk_off = blk_hist + static_cast(nblk) * n_experts; + int* counts = blk_off + static_cast(nblk) * n_experts; + int* ti_i = reinterpret_cast(ti); + + // One warp per token, four warps a block. + const int warps = kRouteThreads / 32; + const int topk_grid = (S + warps - 1) / warps; + float* tw_f = reinterpret_cast(tw); + const __nv_bfloat16* lg = reinterpret_cast(logits); + switch (n_experts / 32) { + case 1: route_topk_warp_kernel<1><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 2: route_topk_warp_kernel<2><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 4: route_topk_warp_kernel<4><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 8: route_topk_warp_kernel<8><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 16: route_topk_warp_kernel<16><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 32: route_topk_warp_kernel<32><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + default: return 4; // expert count is not 32 * a power of two + } + + slot_hist_kernel<<>>( + ti_i, blk_hist, slots, n_experts); + + expert_block_scan_kernel<256><<>>( + blk_hist, blk_off, counts, nblk, n_experts); + + group_off_kernel<<<1, 32, 0, stream>>>( + counts, reinterpret_cast(group_off), n_experts); + + const size_t scatter_smem = + static_cast(kSlotsPerBlock / 32) * n_experts * sizeof(int); + slot_scatter_kernel<<>>( + ti_i, reinterpret_cast(group_off), blk_off, + reinterpret_cast(se), reinterpret_cast(stok), + reinterpret_cast(inv), slots, n_experts, topk); + + return 0; +} + +void moe_route_sfa_offsets( + const void* group_off, + void* sfa_off, + int n_experts, + int n_col, + cudaStream_t stream) +{ + sfa_offsets_kernel<<<1, 32, 0, stream>>>( + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), n_experts, n_col); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_route_prefill_edge.cuh b/csrc/kernels/moe_route_prefill_edge.cuh new file mode 100644 index 00000000..c8536b4f --- /dev/null +++ b/csrc/kernels/moe_route_prefill_edge.cuh @@ -0,0 +1,66 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// Everything a grouped MoE prefill needs from its router logits, as kernels. +// +// The chain this replaces was softmax, top-k, a renormalising divide, a stable +// argsort, two gathers, a bincount, a cumulative sum and a scatter -- ten +// tensor ops per layer, of which the top-k alone cost 25 ms of a 2048-token +// prefill. +// +// The permutation is built as a counting sort with per-block offsets, not an +// atomic scatter: prefill seeds a decode that has to reproduce, so slot order +// within an expert is fixed (ascending slot index, matching a stable argsort) +// rather than left to the order blocks happen to arrive in. +// +// logits (S, n_experts) bf16 +// ti (S, topk) int32 out, expert per (token, rank) +// tw (S, topk) fp32 out, weights renormalised over the top-k +// se (S * topk,) int32 out, expert per sorted slot +// stok (S * topk,) int64 out, token per sorted slot -- 64-bit, alone +// among these, because it is handed to the grouped activation +// quantiser as its gather index and that kernel reads a long. +// Emitting int32 here reads as garbage row indices there, which +// surfaces as an illegal access three kernels later. +// inv (S * topk,) int32 out, sorted row holding slot i +// group_off (n_experts + 1,) int32 out, prefix sums over experts +// ws workspace, moe_route_prefill_workspace_bytes(S, topk, n_experts) +// +// n_experts must be 32 times a power of two, at most 1024, since the top-k +// holds a row across one warp; topk at most 32. Returns 0 on success. +int moe_route_prefill_bf16( + const void* logits, + void* ti, + void* tw, + void* se, + void* stok, + void* inv, + void* group_off, + void* ws, + int ws_bytes, + int S, + int n_experts, + int topk, + cudaStream_t stream); + +int moe_route_prefill_workspace_bytes(int S, int topk, int n_experts); + +// Per-expert scale-factor byte offsets for the block-scaled activation layout, +// derived from the group boundaries the routing kernel already produced. The +// layout blocks rows by 128, so a group of c rows takes ceil(c / 128) super +// blocks of n_col * 512 bytes. +// group_off (n_experts + 1,) int32 +// sfa_off (n_experts,) int32 out +void moe_route_sfa_offsets( + const void* group_off, + void* sfa_off, + int n_experts, + int n_col, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_router_topk_sm120.cu b/csrc/kernels/moe_router_topk_sm120.cu index e3b9e164..1ada0ab5 100644 --- a/csrc/kernels/moe_router_topk_sm120.cu +++ b/csrc/kernels/moe_router_topk_sm120.cu @@ -69,6 +69,79 @@ __global__ void router_topk_kernel(const __nv_bfloat16* __restrict__ logits, } } +// One warp, no barriers. +// +// The kernel above spreads 256 logits over 256 threads and runs k rounds, each +// with a warp reduction plus three block-wide barriers to publish the winner +// and mask it. That is 24 barriers to pick eight of 256 values -- 6.5 us to +// read 512 bytes, which is latency, not work. +// +// Here one warp owns all the logits in registers, eight per lane, and the same +// k rounds run entirely in shuffles. Nothing to synchronise, nothing in shared +// memory. +// +// The result is identical rather than merely equivalent, and for a reason +// worth stating: argmax under a total order -- greater value wins, lower index +// breaks ties -- selects one specific element, so the answer does not depend +// on the shape of the reduction tree the way a floating-point sum does. +// Changing which lane holds which logit therefore cannot change the output. +template +__global__ void router_topk_warp1_kernel(const __nv_bfloat16* __restrict__ logits, + int* __restrict__ out_idx, + float* __restrict__ out_val, + int n, int k) { + const int lane = threadIdx.x; + float v[kPerLane]; +#pragma unroll + for (int j = 0; j < kPerLane; ++j) { + const int i = j * 32 + lane; + v[j] = (i < n) ? static_cast(logits[i]) : -FLT_MAX; + } + + for (int r = 0; r < k; ++r) { + float best = -FLT_MAX; + int bidx = -1; +#pragma unroll + for (int j = 0; j < kPerLane; ++j) { + const int i = j * 32 + lane; + if (v[j] > best || (v[j] == best && i < bidx)) { best = v[j]; bidx = i; } + } + warp_argmax(best, bidx); // butterfly: every lane ends with it + if (lane == 0) { + out_idx[r] = bidx; + out_val[r] = best; + } + // The lane that owns the winner clears it. bidx = j * 32 + owner, so the + // low five bits are the owner and the rest is the slot. + if (bidx >= 0 && (bidx & 31) == lane) v[bidx >> 5] = -FLT_MAX; + } +} + +} // namespace + +int moe_router_topk_warp_sm120_bf16(const void* logits, void* out_idx, + void* out_val, int n_experts, int k, + cudaStream_t stream) { + if (!logits || !out_idx || !out_val) return 1; + if (n_experts <= 0 || k <= 0 || k > 32) return 2; + auto* oi = reinterpret_cast(out_idx); + auto* ov = reinterpret_cast(out_val); + const auto* lg = reinterpret_cast(logits); + if (n_experts <= 32) + router_topk_warp1_kernel<1><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 64) + router_topk_warp1_kernel<2><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 128) + router_topk_warp1_kernel<4><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 256) + router_topk_warp1_kernel<8><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else + return 3; // caller falls back to the block kernel + return 0; +} + +namespace { + } // namespace int moe_router_topk_sm120_bf16(const void* logits, void* out_idx, void* out_val, diff --git a/csrc/kernels/moe_router_topk_sm120.cuh b/csrc/kernels/moe_router_topk_sm120.cuh index cdabc16a..69bbe52b 100644 --- a/csrc/kernels/moe_router_topk_sm120.cuh +++ b/csrc/kernels/moe_router_topk_sm120.cuh @@ -18,5 +18,13 @@ namespace kernels { int moe_router_topk_sm120_bf16(const void* logits, void* out_idx, void* out_val, int n_experts, int k, cudaStream_t stream); +// Single-warp variant, for n_experts <= 256. Same selection rule and the same +// descending order, and identical output -- argmax under a total order does not +// depend on the reduction tree. Returns 3 for a width it cannot hold, so the +// caller can fall back to the block kernel above. +int moe_router_topk_warp_sm120_bf16(const void* logits, void* out_idx, + void* out_val, int n_experts, int k, + cudaStream_t stream); + } // namespace kernels } // namespace flash_rt diff --git a/csrc/kernels/moe_shared_combine_edge.cu b/csrc/kernels/moe_shared_combine_edge.cu new file mode 100644 index 00000000..ed7891b0 --- /dev/null +++ b/csrc/kernels/moe_shared_combine_edge.cu @@ -0,0 +1,59 @@ +#include "moe_shared_combine_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +__global__ void moe_shared_gate_combine_kernel( + const float* __restrict__ routed, + const __nv_bfloat16* __restrict__ shared, + const __nv_bfloat16* __restrict__ gate, + __nv_bfloat16* __restrict__ out, + int S, + int dim) +{ + const int row = blockIdx.x; + if (row >= S) return; + // expf, not the fast intrinsic: routing downstream is discrete, and a gate + // that lands a few ulp away flips ties in later layers. + const float g = 1.0f / (1.0f + expf(-static_cast(gate[row]))); + const size_t base = static_cast(row) * dim; + for (int i = threadIdx.x; i < dim; i += blockDim.x) { + // Multiply and add as two rounded operations, not one contracted fma. + // Written as `routed + shared * g` the compiler contracts it, which is one + // rounding instead of two and therefore a different number -- measured, one + // element in 16384 by one ulp, where the routed sum is small against the + // gated shared term. That is a fine trade in isolation and the wrong one + // here: the decode step computes this as separate tensor ops, and this + // kernel is only allowed to stand in for it if it lands on the same bits. + out[base + i] = __float2bfloat16(__fadd_rn( + routed[base + i], + __fmul_rn(static_cast(shared[base + i]), g))); + } +} + +} // namespace + +void moe_shared_gate_combine_edge_bf16( + const void* routed, + const void* shared, + const void* gate, + void* out, + int S, + int dim, + cudaStream_t stream) +{ + if (S <= 0 || dim <= 0) return; + const int threads = dim < 256 ? 128 : 256; + moe_shared_gate_combine_kernel<<>>( + reinterpret_cast(routed), + reinterpret_cast(shared), + reinterpret_cast(gate), + reinterpret_cast<__nv_bfloat16*>(out), S, dim); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_shared_combine_edge.cuh b/csrc/kernels/moe_shared_combine_edge.cuh new file mode 100644 index 00000000..8f8d52d2 --- /dev/null +++ b/csrc/kernels/moe_shared_combine_edge.cuh @@ -0,0 +1,35 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// The tail of a fine-grained MoE layer: gate the shared expert and add it to +// the routed sum. +// +// out = routed + shared * sigmoid(gate_logit[row]) +// +// Replaces a sigmoid, a broadcast multiply, an add and a cast -- four tensor +// ops and two full (S, hidden) fp32 intermediates per layer. +// +// The existing bf16 gate-mul-residual kernel does not fit here: the routed sum +// arrives in fp32 from the weighted reduction, and taking it through bf16 to +// reach that kernel would change the accumulation rather than merely fuse it. +// This keeps the arithmetic in fp32 and rounds once, at the store. +// +// routed (S, dim) fp32 +// shared (S, dim) bf16 +// gate (S,) bf16, the raw gate logit -- the sigmoid is applied here +// out (S, dim) bf16 +void moe_shared_gate_combine_edge_bf16( + const void* routed, + const void* shared, + const void* gate, + void* out, + int S, + int dim, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/nvfp4_convert.cuh b/csrc/kernels/nvfp4_convert.cuh new file mode 100644 index 00000000..d1d27e85 --- /dev/null +++ b/csrc/kernels/nvfp4_convert.cuh @@ -0,0 +1,112 @@ +#pragma once + +// NVFP4 element and scale-factor conversions. +// +// Moved out of quantize.cu unchanged so a translation unit that produces the +// same wire format without pulling in the whole quantiser -- the gated +// qwen3_5_moe grouped quantiser is the first -- encodes it with the same code +// rather than a second copy of these thresholds. quantize.cu includes this +// header where the definitions used to be, so its own kernels are unaffected. +// +// Everything here is a device-side __forceinline__ helper: including this +// header adds no symbol and no code to a TU that does not call it. +// +// FP4 E2M1 values: +/-{0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} +// UE4M3 (unsigned E4M3): one scale factor per 16-element block + +#include + +#include + +// FP4 E2M1 value table (magnitude only, 3 bits): +// 0b000 = 0.0 (E=0, M=0) +// 0b001 = 0.5 (E=0, M=1, subnormal) +// 0b010 = 1.0 (E=1, M=0) +// 0b011 = 1.5 (E=1, M=1) +// 0b100 = 2.0 (E=2, M=0) +// 0b101 = 3.0 (E=2, M=1) +// 0b110 = 4.0 (E=3, M=0) +// 0b111 = 6.0 (E=3, M=1) + +__device__ __forceinline__ uint8_t float_to_fp4_e2m1(float v) { + uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; + float a = fabsf(v); + uint8_t mag; + if (a < 0.25f) mag = 0; // -> 0.0 + else if (a < 0.75f) mag = 1; // -> 0.5 + else if (a < 1.25f) mag = 2; // -> 1.0 + else if (a < 1.75f) mag = 3; // -> 1.5 + else if (a < 2.5f) mag = 4; // -> 2.0 + else if (a < 3.5f) mag = 5; // -> 3.0 + else if (a < 5.0f) mag = 6; // -> 4.0 + else mag = 7; // -> 6.0 + return sign | mag; +} + +// Branchless equivalent of float_to_fp4_e2m1 — bit-identical, but the 8-way +// if-else (which diverges across a warp and serializes) becomes a sum of +// threshold comparisons (predicated, no divergence). Used by the prefetch _v2 +// quant/norm kernels where the encode is the hot per-element op. +__device__ __forceinline__ uint8_t float_to_fp4_e2m1_branchless(float v) { + float a = fabsf(v); + uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; + uint8_t mag = (uint8_t)((a >= 0.25f) + (a >= 0.75f) + (a >= 1.25f) + + (a >= 1.75f) + (a >= 2.5f) + (a >= 3.5f) + + (a >= 5.0f)); + return sign | mag; +} + +__device__ __forceinline__ float fp4_e2m1_to_float(uint8_t v) { + float mag; + switch (v & 0x7u) { + case 0: mag = 0.0f; break; + case 1: mag = 0.5f; break; + case 2: mag = 1.0f; break; + case 3: mag = 1.5f; break; + case 4: mag = 2.0f; break; + case 5: mag = 3.0f; break; + case 6: mag = 4.0f; break; + default: mag = 6.0f; break; + } + return (v & 0x8u) ? -mag : mag; +} + +// Convert float to UE4M3 (unsigned, 4-bit exponent, 3-bit mantissa) +// Rounds UP (ceil) so that scale >= true_amax / 6.0 (avoids FP4 overflow) +// UE4M3: bias=7, normal = 2^(E-7) * (1 + M/8), subnormal = 2^(-6) * M/8 +// Range: [~0.002, 240] +__device__ __forceinline__ uint8_t float_to_ue4m3_ceil(float v) { + if (v <= 0.0f) return 0; + if (v > 240.0f) return 0xFE; // max finite: E=14, M=7 -> 2^7 * 1.875 = 240 + + uint32_t bits = __float_as_uint(v); + int float_exp = ((bits >> 23) & 0xFF) - 127; // unbiased float exponent + uint32_t frac = bits & 0x7FFFFF; // 23-bit float mantissa + + int ue_exp = float_exp + 7; // UE4M3 bias = 7 + + if (ue_exp <= 0) { + // Subnormal in UE4M3: value = 2^(-6) * M/8 + float scaled = v * 512.0f; // v / (2^(-6) / 8) + int m = (int)ceilf(scaled); + if (m > 7) return (1 << 3) | 0; // smallest normal: E=1, M=0 + if (m < 1) m = 1; + return (uint8_t)m; + } + if (ue_exp >= 15) return 0xFE; // clamp to max + + // Extract top 3 mantissa bits, round up + int m = (int)(frac >> 20); // top 3 of 23 bits + if (frac & 0xFFFFF) m++; // ceil: round up if remaining bits nonzero + if (m >= 8) { m = 0; ue_exp++; } + if (ue_exp >= 15) return 0xFE; + + return (uint8_t)((ue_exp << 3) | m); +} + +__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { + int e = (v >> 3) & 0xF; + int m = v & 0x7; + if (e == 0) return ldexpf((float)m / 8.0f, -6); + return ldexpf(1.0f + (float)m / 8.0f, e - 7); +} diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index a25a7be7..33d72140 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -343,98 +343,10 @@ void quantize_fp8_device_fp16(const __half* input, __nv_fp8_e4m3* output, // UE4M3 (unsigned E4M3): scale factor per 16-element block // ================================================================ -// FP4 E2M1 value table (magnitude only, 3 bits): -// 0b000 = 0.0 (E=0, M=0) -// 0b001 = 0.5 (E=0, M=1, subnormal) -// 0b010 = 1.0 (E=1, M=0) -// 0b011 = 1.5 (E=1, M=1) -// 0b100 = 2.0 (E=2, M=0) -// 0b101 = 3.0 (E=2, M=1) -// 0b110 = 4.0 (E=3, M=0) -// 0b111 = 6.0 (E=3, M=1) - -__device__ __forceinline__ uint8_t float_to_fp4_e2m1(float v) { - uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; - float a = fabsf(v); - uint8_t mag; - if (a < 0.25f) mag = 0; // -> 0.0 - else if (a < 0.75f) mag = 1; // -> 0.5 - else if (a < 1.25f) mag = 2; // -> 1.0 - else if (a < 1.75f) mag = 3; // -> 1.5 - else if (a < 2.5f) mag = 4; // -> 2.0 - else if (a < 3.5f) mag = 5; // -> 3.0 - else if (a < 5.0f) mag = 6; // -> 4.0 - else mag = 7; // -> 6.0 - return sign | mag; -} - -// Branchless equivalent of float_to_fp4_e2m1 — bit-identical, but the 8-way -// if-else (which diverges across a warp and serializes) becomes a sum of -// threshold comparisons (predicated, no divergence). Used by the prefetch _v2 -// quant/norm kernels where the encode is the hot per-element op. -__device__ __forceinline__ uint8_t float_to_fp4_e2m1_branchless(float v) { - float a = fabsf(v); - uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; - uint8_t mag = (uint8_t)((a >= 0.25f) + (a >= 0.75f) + (a >= 1.25f) - + (a >= 1.75f) + (a >= 2.5f) + (a >= 3.5f) - + (a >= 5.0f)); - return sign | mag; -} - -__device__ __forceinline__ float fp4_e2m1_to_float(uint8_t v) { - float mag; - switch (v & 0x7u) { - case 0: mag = 0.0f; break; - case 1: mag = 0.5f; break; - case 2: mag = 1.0f; break; - case 3: mag = 1.5f; break; - case 4: mag = 2.0f; break; - case 5: mag = 3.0f; break; - case 6: mag = 4.0f; break; - default: mag = 6.0f; break; - } - return (v & 0x8u) ? -mag : mag; -} - -// Convert float to UE4M3 (unsigned, 4-bit exponent, 3-bit mantissa) -// Rounds UP (ceil) so that scale >= true_amax / 6.0 (avoids FP4 overflow) -// UE4M3: bias=7, normal = 2^(E-7) * (1 + M/8), subnormal = 2^(-6) * M/8 -// Range: [~0.002, 240] -__device__ __forceinline__ uint8_t float_to_ue4m3_ceil(float v) { - if (v <= 0.0f) return 0; - if (v > 240.0f) return 0xFE; // max finite: E=14, M=7 -> 2^7 * 1.875 = 240 - - uint32_t bits = __float_as_uint(v); - int float_exp = ((bits >> 23) & 0xFF) - 127; // unbiased float exponent - uint32_t frac = bits & 0x7FFFFF; // 23-bit float mantissa - - int ue_exp = float_exp + 7; // UE4M3 bias = 7 - - if (ue_exp <= 0) { - // Subnormal in UE4M3: value = 2^(-6) * M/8 - float scaled = v * 512.0f; // v / (2^(-6) / 8) - int m = (int)ceilf(scaled); - if (m > 7) return (1 << 3) | 0; // smallest normal: E=1, M=0 - if (m < 1) m = 1; - return (uint8_t)m; - } - if (ue_exp >= 15) return 0xFE; // clamp to max - - // Extract top 3 mantissa bits, round up - int m = (int)(frac >> 20); // top 3 of 23 bits - if (frac & 0xFFFFF) m++; // ceil: round up if remaining bits nonzero - if (m >= 8) { m = 0; ue_exp++; } - if (ue_exp >= 15) return 0xFE; - - return (uint8_t)((ue_exp << 3) | m); -} - -__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { - int e = (v >> 3) & 0xF; - int m = v & 0x7; - if (e == 0) return ldexpf((float)m / 8.0f, -6); - return ldexpf(1.0f + (float)m / 8.0f, e - 7); -} +// The element and scale-factor converters live in nvfp4_convert.cuh so the +// gated qwen3_5_moe grouped quantiser encodes the same wire format with the +// same code instead of a second copy. Definitions are unchanged. +#include "nvfp4_convert.cuh" // UE8M0 conversion: 8-bit unsigned exponent, 0 mantissa bits, bias=127 // value = 2^(exp - 127), same as IEEE FP32 exponent extraction diff --git a/csrc/kernels/qwen35moe_e0m3_dequant.cu b/csrc/kernels/qwen35moe_e0m3_dequant.cu new file mode 100644 index 00000000..43c63f9b --- /dev/null +++ b/csrc/kernels/qwen35moe_e0m3_dequant.cu @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Streamed routed-expert block to bf16. See header. + +#include "kernels/qwen35moe_e0m3_dequant.cuh" + +#include +#include +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kThreads = 256; + +// Sign-magnitude: low three bits are the magnitude, bit 3 the sign. Kept as a +// signed integer because the values are exactly the integers 0..7, which is the +// property the quantizer's group scale is chosen against. +__device__ __forceinline__ float decode_nibble(uint8_t code) { + const float magnitude = static_cast(code & 0x07u); + return (code & 0x08u) ? -magnitude : magnitude; +} + +// One thread per packed byte: two output values that always share a scale, +// because group_size is even. +__global__ void dequant_kernel(const uint8_t* __restrict__ packed, + const uint8_t* __restrict__ scale, + __nv_bfloat162* __restrict__ out, + int rows, int cols, int group_size, + float global_scale) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + const int pairs_per_row = cols >> 1; + if (index >= rows * pairs_per_row) return; + + const int row = index / pairs_per_row; + const int pair = index - row * pairs_per_row; + const int group = (pair << 1) / group_size; + const int groups_per_row = cols / group_size; + + const __half_raw raw = __nv_cvt_fp8_to_halfraw( + scale[row * groups_per_row + group], __NV_E4M3); + const float step = + __half2float(*reinterpret_cast(&raw)) * global_scale; + + const uint8_t byte = packed[index]; + out[index] = __floats2bfloat162_rn( + decode_nibble(byte & 0x0Fu) * step, + decode_nibble(byte >> 4) * step); +} + +} // namespace + +int qwen35moe_e0m3_dequant_bf16(const void* packed, const void* scale, + void* out, int rows, int cols, + int group_size, float global_scale, + cudaStream_t stream) { + if (!packed || !scale || !out) return 1; + if (rows <= 0 || cols <= 0) return 2; + if (cols & 1) return 3; + if (group_size <= 0 || (group_size & 1)) return 4; + if (cols % group_size) return 5; + + const long long pairs = static_cast(rows) * (cols >> 1); + const long long blocks = (pairs + kThreads - 1) / kThreads; + if (blocks > 2147483647LL) return 6; + + dequant_kernel<<(blocks), kThreads, 0, stream>>>( + reinterpret_cast(packed), + reinterpret_cast(scale), + reinterpret_cast<__nv_bfloat162*>(out), + rows, cols, group_size, global_scale); + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/qwen35moe_e0m3_dequant.cuh b/csrc/kernels/qwen35moe_e0m3_dequant.cuh new file mode 100644 index 00000000..580564eb --- /dev/null +++ b/csrc/kernels/qwen35moe_e0m3_dequant.cuh @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Decode a streamed routed-expert block into bf16. +// +// The edge bundle stores each expert as sign-magnitude 4-bit values with one +// e4m3 scale byte per group of 16 along K, scaled by a per-tensor float the +// bundle keeps beside the blocks. That is not the format the block-scaled 4-bit +// GEMMs read: they decode E2M1, and they want the scale bytes in the SM1xx +// swizzled tile layout. Neither difference is bridgeable by relabelling -- +// E2M1's sixteen values and this format's sixteen are different sets. +// +// So the streaming path decodes here and hands bf16 to the existing bf16 GEMM, +// which costs bandwidth on an already-resident block but needs no swizzle, no +// second codebook inside a GEMM, and no architecture beyond SM80. Reading the +// bundle's own linear scale layout is the point: it removes the swizzle step +// rather than implementing it. + +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// out[r][c] = value(packed) * e4m3(scale[r][c / group_size]) * global_scale +// +// packed (rows, cols / 2) bytes, low nibble first, each nibble +// magnitude | sign << 3 with magnitude in 0..7 +// scale (rows, cols / group_size) bytes, each an e4m3 magnitude +// out (rows, cols) bf16 +// +// cols must be even and a multiple of group_size, and group_size must be even +// so that a byte's two values always share one scale. +int qwen35moe_e0m3_dequant_bf16(const void* packed, const void* scale, + void* out, int rows, int cols, + int group_size, float global_scale, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/qwen35moe_grouped_quant.cu b/csrc/kernels/qwen35moe_grouped_quant.cu new file mode 100644 index 00000000..d9626cd6 --- /dev/null +++ b/csrc/kernels/qwen35moe_grouped_quant.cu @@ -0,0 +1,295 @@ +// Grouped NVFP4 activation quantisers for the qwen3_5_moe MoE path. +// See qwen35moe_grouped_quant.cuh for the tier this is built under. + +#include "qwen35moe_grouped_quant.cuh" + +#include "nvfp4_convert.cuh" + +#include + +// ── Grouped activation quantiser for the MoE grouped GEMM ── +// +// Same math as quantize_bf16_to_nvfp4_swizzled_kernel, block for block; what +// differs is where the scale factors land. The block-scaled GEMM wants each +// group's scales in the Sm1xx atom layout for that group's own row count, and +// that layout blocks rows by 128, so a group beginning at an arbitrary row of a +// jointly-quantised matrix has no contiguous sub-block to point at. Quantising +// per group is correct but costs a launch and a host iteration per expert. +// +// Here a row reads the expert it was sorted by, subtracts its group's first +// row, and indexes its group's own block. Nothing reaches the host, which is +// what lets the surrounding prefill chunk be captured. +__global__ void moe_grouped_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ input, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + const long* __restrict__ src_row, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int cols, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; // row index inside its group + // Gather while quantising when a permutation is given. Materialising the + // sorted activation first is a full read and a full write of an (S, HID) + // matrix per layer -- 14.5 ms of a 2048-token prefill -- for rows this + // kernel is about to read once anyway. + const size_t in_row = (src_row == nullptr) ? (size_t)row + : (size_t)src_row[row]; + const __nv_bfloat16* row_in = input + in_row * cols; + uint8_t* row_fp4 = fp4_data + (size_t)row * cols / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; + const int tid = threadIdx.x; + + // Per-16-block amax without atomics. One thread takes eight bf16 (a half + // block), reduces them in registers, and pairs with its neighbour through a + // shuffle -- j and j^1 land on lanes t and t^1 because the block size is + // even. The first version of this kernel used one atomicMax per element and + // ran at 58.3 ms for traffic worth 0.8; the atomics were all of it. + const int vec8 = cols >> 3; + for (int j = tid; j < vec8; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + float a = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) a = fmaxf(a, fabsf(__bfloat162float(bf[i]))); + a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, 1)); + if ((j & 1) == 0) smem[j >> 1] = a; + } + __syncthreads(); + + const int rb = local / 128; + const int ri = local % 128; + for (int b = tid; b < num_blocks; b += blockDim.x) { + uint8_t ue_scale = float_to_ue4m3_ceil(smem[b] * (1.0f / 6.0f)); + const int cb = b / 4; + const int ci = b % 4; + sf_base[(rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + ci] = ue_scale; + smem[b] = ue4m3_to_float(ue_scale); + } + __syncthreads(); + + // Pack four bytes at a time: eight bf16 in, one uint32 out, and the eight + // share a 16-block so the scale is read once. + const int quads = cols >> 3; + for (int j = tid; j < quads; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + const float scale = smem[j >> 1]; + const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + uint32_t packed = 0; + #pragma unroll + for (int k = 0; k < 4; ++k) { + uint32_t lo = float_to_fp4_e2m1(__bfloat162float(bf[2 * k]) * inv); + uint32_t hi = float_to_fp4_e2m1( + __bfloat162float(bf[2 * k + 1]) * inv); + packed |= ((hi << 4) | (lo & 0xF)) << (k * 8); + } + *reinterpret_cast(row_fp4 + (j << 2)) = packed; + } +} + +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream) +{ + if (!A || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int num_blocks = K / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const int threads = 256; + const size_t smem = (size_t)num_blocks * sizeof(float); + moe_grouped_quant_nvfp4_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(src_row), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + K, num_blocks, n_col_blocks); + return 0; +} + +// ── Gate and quantise in one pass, for the grouped MoE's down projection ── +// +// The grouped GEMM produces gate and up interleaved in one (slots, 2*inter) +// buffer, and the gate op wants them as two matrices. Slicing columns out of it +// is not free: the halves are strided, so `.contiguous()` copies both -- 67 MB +// a layer at 2048 tokens, to feed an op that then writes another 17 and has it +// read straight back by the quantiser. +// +// Reading the merged buffer directly costs none of that. The silu is computed +// and rounded to bf16 exactly as silu_mul_sm120_bf16 does, so the value that +// reaches the quantiser is the same one it saw before. +__global__ void moe_grouped_silu_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ merged, // (slots, 2 * inter) + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int inter, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; // inter gated values, then scales + float* gated = smem; + float* scales = smem + inter; + + const int tid = threadIdx.x; + for (int i = tid; i < inter; i += blockDim.x) { + const float gv = __bfloat162float(g_in[i]); + const float uv = __bfloat162float(u_in[i]); + // Rounded to bf16 here, as the separate gate kernel does, so the + // quantiser downstream sees the identical value. + gated[i] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + __syncthreads(); + + for (int b = tid; b < num_blocks; b += blockDim.x) { + float a = 0.0f; + #pragma unroll 4 + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[b * 16 + j])); + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + const int rb = local / 128, ri = local % 128; + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + scales[b] = ue4m3_to_float(ue); + } + __syncthreads(); + + const int half = inter >> 1; + for (int p = tid; p < half; p += blockDim.x) { + const int i = p * 2; + const float s = scales[i >> 4]; + const float inv = (s > 0.0f) ? (1.0f / s) : 0.0f; + row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(gated[i + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[i] * inv) & 0x0F)); + } +} + +// Warp-per-row form of the same thing. +// +// The block-per-row kernel above gives 256 threads a row of 512 values -- two +// elements each -- behind three barriers and three passes over shared memory, +// so a block reads two kilobytes and then waits. Measured 2.9x off what that +// traffic implies. +// +// Here a warp owns a row and a lane owns one 16-element scale-factor group: +// it reads its own sixteen gate and up values as vectors, gates them, takes +// its own maximum and packs its own eight bytes. Nothing is shared, so there +// are no barriers and no shared memory at all, and each lane has sixteen +// values in flight instead of two. +// +// The arithmetic is the same in the same order, so the output is identical. +__global__ void moe_grouped_silu_quant_nvfp4_warp_kernel( + const __nv_bfloat16* __restrict__ merged, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int slots, int inter, int num_blocks, int n_col_blocks) +{ + const int warp_in_blk = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int row = blockIdx.x * (blockDim.x >> 5) + warp_in_blk; + if (row >= slots) return; + + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + const int rb = local / 128, ri = local % 128; + + for (int b = lane; b < num_blocks; b += 32) { + float gated[16]; + const int base = b * 16; + #pragma unroll + for (int j = 0; j < 16; ++j) { + const float gv = __bfloat162float(g_in[base + j]); + const float uv = __bfloat162float(u_in[base + j]); + gated[j] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + float a = 0.0f; + #pragma unroll + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[j])); + + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + + const float sc = ue4m3_to_float(ue); + const float inv = (sc > 0.0f) ? (1.0f / sc) : 0.0f; + uint8_t* out8 = row_fp4 + (size_t)b * 8; + #pragma unroll + for (int p = 0; p < 8; ++p) { + out8[p] = (uint8_t)((float_to_fp4_e2m1(gated[2 * p + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[2 * p] * inv) & 0x0F)); + } + } +} + +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + constexpr int kThreads = 256; + const int rows_per_block = kThreads / 32; + const int grid = (slots + rows_per_block - 1) / rows_per_block; + moe_grouped_silu_quant_nvfp4_warp_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + slots, inter, num_blocks, n_col_blocks); + return 0; +} + +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const size_t smem = ((size_t)inter + num_blocks) * sizeof(float); + moe_grouped_silu_quant_nvfp4_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + inter, num_blocks, n_col_blocks); + return 0; +} diff --git a/csrc/kernels/qwen35moe_grouped_quant.cuh b/csrc/kernels/qwen35moe_grouped_quant.cuh new file mode 100644 index 00000000..adddd7de --- /dev/null +++ b/csrc/kernels/qwen35moe_grouped_quant.cuh @@ -0,0 +1,49 @@ +#pragma once + +// Grouped NVFP4 activation quantisers for the qwen3_5_moe MoE path. +// +// Built only with the weight-only 4-bit tier +// (-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON); the matching bindings are guarded on +// FLASHRT_HAVE_QWEN35MOE_W4A16, so a build without that tier contains neither +// these translation units nor their symbols. They live here rather than in +// quantize.cu because the layout they write is the grouped GEMM's, not the +// general quantiser's: scale factors go into the Sm1xx atom layout for each +// group's own row count. + +#include +#include + +// Grouped activation quantiser for the MoE grouped GEMM: every expert's block +// in one launch. Same math as quantize_bf16_to_nvfp4_swizzled; what differs is +// that each group's scale factors go into the Sm1xx atom layout for that +// group's own row count, which is what the block-scaled grouped GEMM reads. +// Quantising per group instead is correct but costs a launch and a host +// iteration per expert -- and a host iteration is what a graph capture cannot +// have. +// +// A (slots, K) bf16, rows already sorted by expert +// expert_of_row (slots,) i32 +// group_off (E + 1,) i32 prefix sums of the per-expert row counts +// sfa_off (E,) i32 byte offset of each group's SF block +// K must be a multiple of 16. Returns 0 on success, nonzero on arg error. +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream); + +// Gate and quantise in one pass: reads the grouped GEMM's merged (slots, +// 2*inter) gate/up output directly, so the strided column halves are never +// copied out. The silu is rounded to bf16 exactly as silu_mul_sm120_bf16 does, +// so the quantiser sees the same value it did when the two were separate. +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); + +// Warp-per-row form of the above: a lane owns one 16-element scale-factor +// group and keeps it in registers, so there is no shared memory and no +// barrier. Same arithmetic in the same order, so the output is identical. +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu new file mode 100644 index 00000000..4d42b01e --- /dev/null +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -0,0 +1,352 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// W4A16 GEMV variants for a bandwidth-poor part. See header for what differs +// and why. + +#include "kernels/w4a16_edge_sm120.cuh" + +#include +#include +#include "kernels/fp4_e2m1_compat.cuh" +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kWarps = 2; // output-row groups per block +constexpr int kThreads = kWarps * 32; // 256 +// Packed-weight loads in flight per row. +// +// The default is four, which is the value the SM120 path was validated with. +// Thor (sm_110) measures faster at two: ncu puts this kernel at 121 registers a +// thread there, which caps it at 8 blocks per SM when shared memory, warps and +// the SM limit all allow 24 -- Block Limit Registers is the only binding one, +// and achieved occupancy is 31%. wv[R][kUnroll] alone is R * kUnroll eight-byte +// values, 64 registers at R=8, so halving it buys back the warps. Swept in the +// captured decode step on that part, one build each: +// +// kUnroll 1 2 3 4 +// step 10.025 9.743 10.360 10.384 ms +// tok/s 99.7 102.6 96.5 96.3 +// +// That trade is a property of a 20-SM part with 244 GB/s, so it is set per +// architecture in CMake rather than globally -- a device with far more SMs and +// bandwidth may well prefer the deeper per-thread parallelism, and this branch +// has no measurement for one. +// +// The accumulation order does not move either way: the main loop advances by +// 32*kUnroll and the tail takes the remainder, so a lane visits the same +// k-blocks in the same sequence for any kUnroll and the result is bit-identical. +#ifndef FLASHRT_W4A16_EDGE_UNROLL +#define FLASHRT_W4A16_EDGE_UNROLL 4 +#endif +constexpr int kUnroll = FLASHRT_W4A16_EDGE_UNROLL; + +// A 16-element NVFP4 block is 16 bf16 of activation, 32 bytes. Held at that +// stride, the eight lanes of a 128-bit shared-load phase land on banks +// 0,8,16,24,0,8,16,24 -- four banks, two-way conflicted. At 48 bytes they land +// on 0,12,24,4,16,28,8,20: eight distinct banks. The 16 spare bytes per block +// cost K/2 bytes of shared memory (12 KB at K=4096) and buy back the 2.41x +// wavefront overhead the conflict was costing. +constexpr int kBlockSlots = 24; // bf16 slots per 16-element block +constexpr int kBlockInt4 = kBlockSlots / 8; // 3 int4 per block, 2 used + +// UE4M3 -> fp32 without a table. +// +// The value is (1 + m/8) * 2^(e-7) for e > 0, which is exactly an fp32 with +// exponent field e+120 and mantissa m<<20, and m * 2^-9 for e == 0. Four +// integer ops and a select, against a __constant__ load whose index differs +// per lane -- and constant memory serves one address per cycle, so a divergent +// index serialises the warp. +// +// Bit 7 is not a sign bit: UE4M3 is unsigned, and the quantizer's saturation +// byte 0xFE must decode to +448. +__device__ __forceinline__ float ue4m3_to_float(uint32_t v) { + const uint32_t e = (v >> 3) & 0xFu; + const uint32_t m = v & 0x7u; + const float normal = __uint_as_float(((e + 120u) << 23) | (m << 20)); + const float subnormal = static_cast(m) * (1.0f / 512.0f); + return e == 0u ? subnormal : normal; +} + +// SF swizzle byte offset, identical packing to bf16_weight_to_nvfp4_swizzled. +__device__ __forceinline__ int sf_off(int rb_ncs, int row_inner, int k_block) { + return (rb_ncs + (k_block >> 2)) * 512 + row_inner + (k_block & 3); +} + +// One NVFP4 block (16 elements / 8 packed bytes) dotted with 16 bf16 acts. +__device__ __forceinline__ float blockdot(uint64_t b_pack, + const __nv_bfloat162* xb2) { + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); + const float2 wf = __half22float2(*reinterpret_cast(&wr)); + const float2 xf = __bfloat1622float2(xb2[j]); + acc = fmaf(wf.x, xf.x, acc); + acc = fmaf(wf.y, xf.y, acc); + } + return acc; +} + +// Stage x into the padded shared layout: block b occupies int4 slots +// 3b and 3b+1, leaving 3b+2 as the padding that separates the banks. +__device__ __forceinline__ void stage_padded( + const __nv_bfloat16* __restrict__ x, __nv_bfloat16* x_sh, int K) { + const int4* x_i4 = reinterpret_cast(x); + int4* sh_i4 = reinterpret_cast(x_sh); + const int n_i4 = K >> 3; // 8 bf16 per int4, 2 per block + for (int j = threadIdx.x; j < n_i4; j += kThreads) + sh_i4[(j >> 1) * kBlockInt4 + (j & 1)] = x_i4[j]; +} + +// The K loop, shared by both entry points: R output rows per warp, kUnroll +// packed-weight loads per row in flight. +// +// R exists because a warp with one row does not have enough memory-level +// parallelism on this part. In situ the dominant stall is the global-load +// dependency (long scoreboard, 5.8-13 cycles per issued instruction) while the +// ALU pipe sits at 27-50%: the loop is waiting on memory it has not asked for +// yet. Each lane keeps R*kUnroll eight-byte loads outstanding instead of +// kUnroll, and the K=512 shapes -- where K_BLOCKS is exactly 32, so the +// unrolled body never runs and the tail leaves ONE load in flight -- get the +// whole factor from R. +// +// The rows a warp takes are consecutive and 32-aligned by construction, so +// their scale offsets differ by a constant and cost no extra registers. The +// per-row arithmetic is untouched: same lane-to-block mapping, same order, same +// reduction, so the result is bit-identical to R = 1. +template +__device__ __forceinline__ void row_dot( + const uint64_t* __restrict__ w_row0, size_t row_stride_u64, + const uint8_t* __restrict__ SFB, const __nv_bfloat16* x_sh, + int K_BLOCKS, int rb_ncs, int row_inner, int lane, float (&acc)[R]) { +#pragma unroll + for (int r = 0; r < R; ++r) acc[r] = 0.0f; + + int kb = lane; + const int step = 32 * kUnroll; + for (; kb + 32 * (kUnroll - 1) < K_BLOCKS; kb += step) { + uint64_t wv[R][kUnroll]; + float sf[R][kUnroll]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + wv[r][u] = w_row0[r * row_stride_u64 + kb + 32 * u]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + sf[r][u] = ue4m3_to_float(__ldg( + SFB + sf_off(rb_ncs, row_inner + 16 * r, kb + 32 * u))); +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + acc[r] += blockdot( + wv[r][u], reinterpret_cast( + x_sh + (size_t)(kb + 32 * u) * kBlockSlots)) + * sf[r][u]; + } + for (; kb < K_BLOCKS; kb += 32) { + uint64_t wv[R]; + float sf[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wv[r] = w_row0[r * row_stride_u64 + kb]; +#pragma unroll + for (int r = 0; r < R; ++r) + sf[r] = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner + 16 * r, kb))); +#pragma unroll + for (int r = 0; r < R; ++r) + acc[r] += blockdot( + wv[r], reinterpret_cast( + x_sh + (size_t)kb * kBlockSlots)) * sf[r]; + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[r] += __shfl_xor_sync(0xffffffff, acc[r], off); +} + +// Rows per warp. More rows means more outstanding loads and more registers, so +// the useful value is where the added parallelism stops paying for the +// occupancy it costs -- and that turned out to differ between the two entry +// points, which is why they do not share a constant. Measured at the shapes +// the decode issues, cold: the dense GEMV peaks at 2 (q_proj 47.9 us against +// 53.9 at 4, lm_head 1205 against 1366) while the grouped one peaks at 4 +// (gate_up 42.9 against 47.8 at 2). The grouped launch carries a slot per grid +// row, so it has fewer blocks per row tile and leans harder on what each +// thread keeps in flight. +constexpr int kRowsDense = 2; // K >= 2048: kUnroll fires, 2 * 4 = 8 +constexpr int kRowsGrouped = 4; // K >= 2048: 4 * 4 = 16 +constexpr int kRowsSmall = 8; // K < 2048: tail only, 8 * 1 = 8 + +// The row block a warp owns must not straddle a 32-row scale group, or +// row_inner + 16 * r stops describing the swizzle. Warps take R consecutive +// rows starting at a multiple of R, so this holds for any R dividing 32. +static_assert(32 % kRowsDense == 0 && 32 % kRowsGrouped == 0 + && 32 % kRowsSmall == 0, + "rows per warp must divide the 32-row scale group"); + +template +__global__ void w4a16_matvec_edge_kernel( + const __nv_bfloat16* __restrict__ x, + const uint8_t* __restrict__ W, + const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ out, + float alpha, int N, int K, int n_col_super) { + extern __shared__ __nv_bfloat16 x_sh[]; + stage_padded(x, x_sh, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; + + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R]; + row_dot( + reinterpret_cast(W + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB, x_sh, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { +#pragma unroll + for (int r = 0; r < R; ++r) + if (row0 + r < N) out[row0 + r] = __float2bfloat16(acc[r] * alpha); + } +} + +// grid = (ceil(N/(8*R)), slots). Block computes 8*R output rows of one slot. +template +__global__ void moe_grouped_w4a16_edge_kernel( + const __nv_bfloat16* __restrict__ A_stack, + const uint8_t* __restrict__ W_stack, + const uint8_t* __restrict__ SFB_stack, + const float* __restrict__ alpha_stack, + const int* __restrict__ expert_idx, + __nv_bfloat16* __restrict__ D, + int N, int K, int n_col_super, + long a_stride, long w_stride, long sfb_stride) { + const int slot = blockIdx.y; + const int e = expert_idx[slot]; + + extern __shared__ __nv_bfloat16 x_sh[]; + stage_padded(A_stack + (long)slot * a_stride, x_sh, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; + + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R]; + row_dot( + reinterpret_cast( + W_stack + (long)e * w_stride + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB_stack + (long)e * sfb_stride, x_sh, K >> 4, + rb * n_col_super, (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { + const float a = alpha_stack[e]; +#pragma unroll + for (int r = 0; r < R; ++r) + if (row0 + r < N) + D[(long)slot * N + row0 + r] = __float2bfloat16(acc[r] * a); + } +} + +// Shared memory for the padded stage: kBlockSlots bf16 per 16 elements. +inline size_t smem_bytes(int K) { + return (size_t)(K >> 4) * kBlockSlots * sizeof(__nv_bfloat16); +} + +// Rows per warp for this K, dropping to 1 when N cannot fill even one warp's +// worth. Above 32 rows a warp would straddle a scale group; the constants +// enforce that, this only picks between them. +inline int rows_per_warp(int N, int K, int rows_big) { + const int r = (K >= 2048) ? rows_big : kRowsSmall; + return (N >= kWarps * r) ? r : 1; +} + +} // namespace + +int w4a16_matvec_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int N, + int K, + float alpha, + cudaStream_t stream) { + if (!x_bf16 || !W_packed || !SFB || !out) return 1; + if (N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int n_col_super = ((K >> 4) + 3) / 4; + const auto* xp = reinterpret_cast(x_bf16); + const auto* wp = reinterpret_cast(W_packed); + const auto* sp = reinterpret_cast(SFB); + auto* op = reinterpret_cast<__nv_bfloat16*>(out); + const size_t smem = smem_bytes(K); +#define FLASHRT_LAUNCH_MATVEC(R) \ + w4a16_matvec_edge_kernel<<>>( \ + xp, wp, sp, op, alpha, N, K, n_col_super) + switch (rows_per_warp(N, K, kRowsDense)) { + case kRowsDense: FLASHRT_LAUNCH_MATVEC(kRowsDense); break; + case kRowsSmall: FLASHRT_LAUNCH_MATVEC(kRowsSmall); break; + default: FLASHRT_LAUNCH_MATVEC(1); break; + } +#undef FLASHRT_LAUNCH_MATVEC + return 0; +} + +int moe_grouped_w4a16_edge_sm120_bf16( + const void* A_stack, + const void* W_stack, + const void* SFB_stack, + const void* alpha_stack, + const void* eidx, + void* D, + int slots, + int N, + int K, + long a_stride, + long w_stride, + long sfb_stride, + cudaStream_t stream) { + if (!A_stack || !W_stack || !SFB_stack || !alpha_stack || !eidx || !D) + return 1; + if (slots <= 0 || N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int n_col_super = ((K >> 4) + 3) / 4; + const auto* ap = reinterpret_cast(A_stack); + const auto* wp = reinterpret_cast(W_stack); + const auto* sp = reinterpret_cast(SFB_stack); + const auto* alp = reinterpret_cast(alpha_stack); + const auto* ep = reinterpret_cast(eidx); + auto* dp = reinterpret_cast<__nv_bfloat16*>(D); + const size_t smem = smem_bytes(K); +#define FLASHRT_LAUNCH_GROUPED(R) \ + moe_grouped_w4a16_edge_kernel \ + <<>>( \ + ap, wp, sp, alp, ep, dp, N, K, n_col_super, \ + a_stride, w_stride, sfb_stride) + switch (rows_per_warp(N, K, kRowsGrouped)) { + case kRowsGrouped: FLASHRT_LAUNCH_GROUPED(kRowsGrouped); break; + case kRowsSmall: FLASHRT_LAUNCH_GROUPED(kRowsSmall); break; + default: FLASHRT_LAUNCH_GROUPED(1); break; + } +#undef FLASHRT_LAUNCH_GROUPED + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/w4a16_edge_sm120.cuh b/csrc/kernels/w4a16_edge_sm120.cuh new file mode 100644 index 00000000..1ede1670 --- /dev/null +++ b/csrc/kernels/w4a16_edge_sm120.cuh @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// W4A16 GEMV variants for a bandwidth-poor part (Jetson-class Blackwell). +// +// Same math, same weight layout, same results as w4a16_matvec_sm120 and +// moe_grouped_w4a16_sm120 -- bit for bit. What differs is two things the +// profiler found on a 20-SM part with ~244 GB/s of memory, where the original +// pair sits at 51% of that while the BF16 GEMV of the same shape reaches 100%: +// +// 1. The staged activation is read from shared memory at a 32-byte stride +// across lanes, which puts eight lanes of a 128-bit load phase on four +// banks. Measured: 432,685 bank conflicts over 98,816 shared loads, 2.41x +// the wavefronts the traffic needs. Padding each block's footprint to 48 +// bytes lands the phase on eight distinct banks. +// +// 2. The UE4M3 block scale is decoded through a 256-entry __constant__ LUT +// indexed by a per-lane byte. Constant memory serves one address per +// cycle, so a divergent index serialises. The decode is four integer ops, +// so it does not need a table at all. +// +// Neither changes an arithmetic result, which is the point: the variant is +// accepted only if it is bitwise identical to the kernel it replaces. + +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// y(1,N) = (x(1,K) bf16) . (W(N,K) NVFP4)^T, fp32 accumulate, bf16 out. +// Arguments and semantics are those of w4a16_matvec_sm120_bf16. +int w4a16_matvec_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int N, + int K, + float alpha, + cudaStream_t stream); + +// Grouped per-slot GEMV: D[s,:] = A[s,:] . W[eidx[s]]^T * alpha[eidx[s]]. +// Arguments and semantics are those of moe_grouped_w4a16_sm120_bf16. +int moe_grouped_w4a16_edge_sm120_bf16( + const void* A, + const void* W, + const void* SFB, + const void* alpha, + const void* eidx, + void* D, + int slots, + int N, + int K, + long a_stride, + long w_stride, + long sfb_stride, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/w4a16_gemm_sm120.cu b/csrc/kernels/w4a16_gemm_sm120.cu index a069bcce..66b8a79a 100644 --- a/csrc/kernels/w4a16_gemm_sm120.cu +++ b/csrc/kernels/w4a16_gemm_sm120.cu @@ -18,7 +18,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -166,10 +166,10 @@ __global__ __launch_bounds__(GM_THREADS) void w4a16_gemm_kernel( const int ncol = warp_n * 32 + jb * 8 + r; const uint8_t* wq = &sWq[cur][ncol * KT_half]; float sf = c_w4a16_ue4m3[sSFB[cur][ncol * GM_KSUB + ksub]] * alpha; - __half2_raw h0 = __nv_cvt_fp4x2_to_halfraw2( - static_cast<__nv_fp4x2_storage_t>(wq[byte0]), __NV_E2M1); - __half2_raw h1 = __nv_cvt_fp4x2_to_halfraw2( - static_cast<__nv_fp4x2_storage_t>(wq[byte0 + 4]), __NV_E2M1); + __half2_raw h0 = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(wq[byte0])); + __half2_raw h1 = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(wq[byte0 + 4])); float2 f0 = __half22float2(*reinterpret_cast(&h0)); float2 f1 = __half22float2(*reinterpret_cast(&h1)); bb0[jb] = pack_bf16x2(f0.x * sf, f0.y * sf); diff --git a/csrc/kernels/w4a16_matvec_sm120.cu b/csrc/kernels/w4a16_matvec_sm120.cu index ddc2f89b..32c386db 100644 --- a/csrc/kernels/w4a16_matvec_sm120.cu +++ b/csrc/kernels/w4a16_matvec_sm120.cu @@ -6,7 +6,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -38,9 +38,8 @@ __device__ __forceinline__ float blockdot(uint64_t b_pack, float acc = 0.0f; #pragma unroll for (int j = 0; j < 8; ++j) { - const __nv_fp4x2_storage_t bb = - static_cast<__nv_fp4x2_storage_t>(b_pack >> (j * 8)); - const __half2_raw wr = __nv_cvt_fp4x2_to_halfraw2(bb, __NV_E2M1); + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); const float2 wf = __half22float2(*reinterpret_cast(&wr)); const float2 xf = __bfloat1622float2(xb2[j]); acc = fmaf(wf.x, xf.x, acc); diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cu b/csrc/kernels/w4a16_mrows_edge_sm120.cu new file mode 100644 index 00000000..67f99878 --- /dev/null +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cu @@ -0,0 +1,258 @@ +#include "w4a16_mrows_edge_sm120.cuh" + +#include "fp4_e2m1_compat.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +// Same shape constants as the M=1 entry this extends; see its file for why +// each is what it is. The padded shared stride is what keeps the 16-element +// blocks off each other's banks. +constexpr int kWarps = 2; +constexpr int kThreads = kWarps * 32; +// Loads in flight per row: the same build-time constant the single-row entry +// uses, so the verify and the step it stands in for stay on the same tuning. +// See w4a16_edge_sm120.cu for why it is set per architecture. +#ifndef FLASHRT_W4A16_EDGE_UNROLL +#define FLASHRT_W4A16_EDGE_UNROLL 4 +#endif +constexpr int kUnroll = FLASHRT_W4A16_EDGE_UNROLL; +constexpr int kBlockSlots = 24; +constexpr int kBlockInt4 = kBlockSlots / 8; +constexpr int kRowsDense = 2; +constexpr int kRowsSmall = 8; +constexpr int kMaxM = 8; + +static_assert(32 % kRowsDense == 0 && 32 % kRowsSmall == 0, + "rows per warp must divide the 32-row scale group"); + +__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { + const int e = (v >> 3) & 0xF; + const int m = v & 0x7; + if (e == 0) return ldexpf(static_cast(m) / 8.0f, -6); + return ldexpf(1.0f + static_cast(m) / 8.0f, e - 7); +} + +__device__ __forceinline__ int sf_off(int rb_ncs, int row_inner, int k_block) { + return (rb_ncs + (k_block >> 2)) * 512 + row_inner + (k_block & 3); +} + +// One packed block against M activation rows. The weight byte pair is decoded +// once and used M times, which is the whole point: the decode is the same work +// the M=1 kernel does, and the extra rows cost only shared reads and fmas. +template +__device__ __forceinline__ void blockdot_m( + uint64_t b_pack, const __nv_bfloat162* x0, size_t x_row_slots, + float (&acc)[M]) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); + const float2 wf = __half22float2(*reinterpret_cast(&wr)); +#pragma unroll + for (int m = 0; m < M; ++m) { + const float2 xf = __bfloat1622float2( + x0[m * (x_row_slots >> 1) + j]); + acc[m] = fmaf(wf.x, xf.x, acc[m]); + acc[m] = fmaf(wf.y, xf.y, acc[m]); + } + } +} + +__device__ __forceinline__ void stage_padded_row( + const __nv_bfloat16* __restrict__ x, __nv_bfloat16* x_sh, int K) { + const int4* x_i4 = reinterpret_cast(x); + int4* sh_i4 = reinterpret_cast(x_sh); + const int n_i4 = K >> 3; + for (int j = threadIdx.x; j < n_i4; j += kThreads) + sh_i4[(j >> 1) * kBlockInt4 + (j & 1)] = x_i4[j]; +} + +// The K loop, R output rows by M activation rows. Identical in structure to +// the M=1 version: same lane-to-block mapping, same unroll, same order of +// accumulation per (output row, activation row), same final shuffle. +template +__device__ __forceinline__ void row_dot_m( + const uint64_t* __restrict__ w_row0, size_t row_stride_u64, + const uint8_t* __restrict__ SFB, const __nv_bfloat16* x_sh, + size_t x_row_slots, int K_BLOCKS, int rb_ncs, int row_inner, int lane, + float (&acc)[R][M]) { +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] = 0.0f; + + int kb = lane; + const int step = 32 * kUnroll; + for (; kb + 32 * (kUnroll - 1) < K_BLOCKS; kb += step) { + uint64_t wv[R][kUnroll]; + float sf[R][kUnroll]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + wv[r][u] = w_row0[r * row_stride_u64 + kb + 32 * u]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + sf[r][u] = ue4m3_to_float(__ldg( + SFB + sf_off(rb_ncs, row_inner + 16 * r, kb + 32 * u))); +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) { + float part[M]; +#pragma unroll + for (int m = 0; m < M; ++m) part[m] = 0.0f; + blockdot_m( + wv[r][u], + reinterpret_cast( + x_sh + (size_t)(kb + 32 * u) * kBlockSlots), + x_row_slots, part); +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] += part[m] * sf[r][u]; + } + } + for (; kb < K_BLOCKS; kb += 32) { + uint64_t wv[R]; + float sf[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wv[r] = w_row0[r * row_stride_u64 + kb]; +#pragma unroll + for (int r = 0; r < R; ++r) + sf[r] = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner + 16 * r, kb))); +#pragma unroll + for (int r = 0; r < R; ++r) { + float part[M]; +#pragma unroll + for (int m = 0; m < M; ++m) part[m] = 0.0f; + blockdot_m( + wv[r], + reinterpret_cast( + x_sh + (size_t)kb * kBlockSlots), + x_row_slots, part); +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] += part[m] * sf[r]; + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[r][m] += __shfl_xor_sync(0xffffffff, acc[r][m], off); +} + +template +__global__ void w4a16_mrows_edge_kernel( + const __nv_bfloat16* __restrict__ x, + const uint8_t* __restrict__ W, + const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ out, + float alpha, int N, int K, int n_col_super) { + extern __shared__ __nv_bfloat16 x_sh[]; + const size_t row_slots = (size_t)(K >> 4) * kBlockSlots; +#pragma unroll + for (int m = 0; m < M; ++m) + stage_padded_row(x + (size_t)m * K, x_sh + m * row_slots, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; + + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R][M]; + row_dot_m( + reinterpret_cast(W + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB, x_sh, row_slots, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { +#pragma unroll + for (int r = 0; r < R; ++r) { + if (row0 + r >= N) continue; +#pragma unroll + for (int m = 0; m < M; ++m) + out[(size_t)m * N + row0 + r] = __float2bfloat16(acc[r][m] * alpha); + } + } +} + +inline size_t smem_bytes(int K, int M) { + return (size_t)(K >> 4) * kBlockSlots * sizeof(__nv_bfloat16) * M; +} + +inline int rows_per_warp(int N, int K) { + const int r = (K >= 2048) ? kRowsDense : kRowsSmall; + return (N >= r * kWarps) ? r : 1; +} + +#define FLASHRT_MROWS_LAUNCH(R, M) \ + do { \ + const size_t sb = smem_bytes(K, M); \ + if (sb > 48 * 1024) { \ + cudaFuncSetAttribute(w4a16_mrows_edge_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + static_cast(sb)); \ + } \ + w4a16_mrows_edge_kernel \ + <<>>( \ + reinterpret_cast(x_bf16), \ + reinterpret_cast(W_packed), \ + reinterpret_cast(SFB), \ + reinterpret_cast<__nv_bfloat16*>(out), alpha, N, K, n_col_super); \ + } while (0) + +#define FLASHRT_MROWS_BY_M(R) \ + do { \ + switch (M) { \ + case 1: FLASHRT_MROWS_LAUNCH(R, 1); break; \ + case 2: FLASHRT_MROWS_LAUNCH(R, 2); break; \ + case 3: FLASHRT_MROWS_LAUNCH(R, 3); break; \ + case 4: FLASHRT_MROWS_LAUNCH(R, 4); break; \ + case 5: FLASHRT_MROWS_LAUNCH(R, 5); break; \ + case 6: FLASHRT_MROWS_LAUNCH(R, 6); break; \ + case 7: FLASHRT_MROWS_LAUNCH(R, 7); break; \ + default: FLASHRT_MROWS_LAUNCH(R, 8); break; \ + } \ + } while (0) + +} // namespace + +int w4a16_mrows_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int M, + int N, + int K, + float alpha, + cudaStream_t stream) { + if (!x_bf16 || !W_packed || !SFB || !out) return 1; + if (N <= 0 || K <= 0 || (K & 15) != 0) return 2; + if (M <= 0 || M > kMaxM) return 3; + + const int n_col_super = ((K >> 4) + 3) / 4; + const int R = rows_per_warp(N, K); + if (R == kRowsDense) { + FLASHRT_MROWS_BY_M(kRowsDense); + } else if (R == kRowsSmall) { + FLASHRT_MROWS_BY_M(kRowsSmall); + } else { + FLASHRT_MROWS_BY_M(1); + } + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cuh b/csrc/kernels/w4a16_mrows_edge_sm120.cuh new file mode 100644 index 00000000..cb3ba4bf --- /dev/null +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cuh @@ -0,0 +1,49 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// A few rows of activation against the same 4-bit weight the decode GEMV reads. +// +// A speculative verify has to be the same function as the decode step it +// verifies, or the tokens it keeps are not the ones plain greedy emits. Two +// things were in the way. The verify read the dense weights at BF16 while +// decode reads them at 4 bits -- four times the traffic, and a different +// answer: measured logit cosine 0.988 between the two forwards. And the +// general W4A16 GEMM, which does read 4 bits, is not the same arithmetic +// either, and at these shapes it is 7.5 to 9.3 times off the GEMV: 250 us +// against 33 for an 8192x2048 projection, and flat in M, so it is not reading +// the weight at bandwidth at all. +// +// This is the decode GEMV with M rows of activation. The weight stream, the +// lane-to-block mapping, the unroll and the reduction order are unchanged -- +// only the number of activation rows staged in shared memory and the number of +// accumulators a warp carries. The weight is what costs, and it is read once +// regardless of M, so a window of four verifies for what one costs. +// +// Because each output row accumulates in exactly the order the GEMV uses, the +// result at M=1 is bit-identical to it, and rows of a larger M agree with what +// the GEMV would have produced for each row on its own. +// +// x (M, K) bf16, row-major +// W (N, K/2) NVFP4 e2m1 nibbles +// SFB swizzled UE4M3 block scales, as bf16_weight_to_nvfp4_swizzled writes +// out (M, N) bf16 +// alpha weight per-tensor global scale +// +// K must be a multiple of 16, M at most 8. Returns 0 on success. +int w4a16_mrows_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int M, + int N, + int K, + float alpha, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index d2f31235..9a981d30 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -1,27 +1,37 @@ # Qwen3.6-35B-A3B text inference FlashRT runs the language backbone from the official -`Qwen/Qwen3.6-35B-A3B` BF16 checkpoint on an RTX 5090. The checkpoint uses the -same `qwen3_5_moe` text architecture as Nex-N2-mini, so both models share the -same weight loader, prefill, attention, MoE, recurrent-state, and CUDA Graph -decode implementation. +`Qwen/Qwen3.6-35B-A3B` BF16 checkpoint on an RTX 5090 (SM120) and on Jetson AGX +Thor (SM110). The checkpoint uses the same `qwen3_5_moe` text architecture as +Nex-N2-mini, so both models share the same weight loader, prefill, attention, +MoE, recurrent-state, and CUDA Graph decode implementation, and both +architectures run the same frontend -- what differs is which kernel tiers the +build has. -This entry is text-only. It does not load the vision tower and it validates but -does not execute the checkpoint's MTP head. Image/video input and speculative -decode are not part of this interface. +This entry is text-only: it does not load the vision tower, and image or video +input is not part of this interface. + +It does execute the checkpoint's MTP draft head, but only when asked. Pass +`load_mtp=True` to the constructor and call `generate_spec()`; see *Speculative +decode* below. `generate()` never reads the head, and without `load_mtp=True` +the loader does not read it either. ## Requirements | | | |---|---| | Checkpoint | `Qwen/Qwen3.6-35B-A3B` BF16 safetensors | -| Hardware | RTX 5090 / SM120 | -| GPU memory | 32 GB | +| Hardware | RTX 5090 / SM120, Jetson AGX Thor / SM110 | +| GPU memory | 32 GB (SM120); unified memory on Thor | | Framework | PyTorch | | Runtime quantization | NVFP4 | -| Build flags | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` | -Configure and build the gated `qwen3_5_moe` kernels: +Configure and build the gated `qwen3_5_moe` kernels. **The two targets take +different flags** -- `FLASHRT_ENABLE_QWEN35MOE` turns on the block-scaled 4-bit +MMA tier as well, which sm_110 refuses at configure time, so Thor names the two +tiers it can compile: + +RTX 5090 (SM120): ```bash cmake -S . -B build \ @@ -31,14 +41,136 @@ cmake --build build -j pip install -e ".[torch]" ``` +Jetson AGX Thor (SM110): + +```bash +cmake -S . -B build \ + -DGPU_ARCH=110 \ + -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON \ + -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON \ + -DFLASHRT_ENABLE_THOR_FA2=ON +cmake --build build -j 2 +pip install -e ".[torch]" +``` + +`FLASHRT_ENABLE_THOR_FA2` is what builds the vendored FA2 kernels on sm_110; +see *Attention differs by target* below for what it buys and why it is opt-in. +Without it the model still runs, and produces the same tokens, but a long +prefill is far slower. + +### Kernel tiers + +`FLASHRT_ENABLE_QWEN35MOE=ON` is a convenience switch for all three tiers +below. Targets that cannot run a tier can select the remainder explicitly. + +| Flag | Kernels | Requires | +|---|---|---| +| `FLASHRT_ENABLE_QWEN35MOE_CORE` | QKV layout/split, bf16 matvec, router top-k, SiLU/sigmoid fusion, GDN recurrence, weighted-sum reducer, bf16 GEMM | SM80 and newer | +| `FLASHRT_ENABLE_QWEN35MOE_W4A16` | weight-only 4-bit matvec, grouped matvec, GEMM | SM80 and newer; hardware operand conversion from SM89 | +| `FLASHRT_ENABLE_QWEN35MOE_W4A4` | block-scaled 4-bit MMA: grouped GEMV, M16/M64/block-tile MMA | sm_120a / sm_121a | + +The upper tiers depend on the core tier, so enabling either turns it on. SM120 +runs all three; sm_110 runs the first two, and the block-scaled tier refuses to +configure there rather than building kernels that fail at run time. + +**These three tiers are not the whole dependency set.** Walking every `fvk` +call the pipeline makes and resolving each to the preprocessor guard active +where it is defined gives 32 kernels across seven gates: + +| gate | kernels | +|---|---:| +| `FLASHRT_HAVE_QWEN36_KERNELS` | 12 | +| `FLASHRT_HAVE_QWEN35MOE_CORE` | 10 | +| `FLASHRT_HAVE_QWEN35MOE_W4A16` | 3 | +| `ENABLE_CUTLASS_SM120_NVFP4_W4A16` | 2 | +| `FLASHRT_HAVE_QWEN35MOE_W4A4` | 2 | +| `FLASHRT_HAVE_NVFP4_SWIZZLE` | 1 | +| ungated | 1 | + +The twelve under `FLASHRT_HAVE_QWEN36_KERNELS` are the linear-attention path: +causal convolution and its update, the gated-DeltaNet recurrence, the WY chunk +stack, the fused RMSNorm-gated-SiLU, partial RoPE and argmax. They are shared +with the rest of the Qwen3.6 family and are gated on `NOT FLASHRT_SLIM_BUILD`, +not on architecture — so `-DFLASHRT_SLIM_BUILD=ON` removes them and the +frontend then refuses to start, naming what is missing. Do not use a slim build +for this model. + +Selecting tiers by reading the source's own grouping is therefore not enough to +know what a target needs; the call sites are what decide. + +Two further kernels are **optional** and are not part of that required set, +because the frontend resolves each through `getattr` and falls back to the +kernel it replaces when a build does not carry it: + +| symbol | gate | replaces | why | +|---|---|---|---| +| `gated_deltanet_recurrent_edge_qwen36_bf16` | `FLASHRT_HAVE_QWEN36_KERNELS` | `gated_deltanet_recurrent_qwen36_bf16` | same arithmetic without the local-memory round trip for the state column | +| `moe_router_topk_warp_sm120_bf16` | `FLASHRT_HAVE_QWEN35MOE_CORE` | `moe_router_topk_sm120_bf16` | same selection in one warp instead of `k` rounds of block-wide barriers | +| `moe_shared_gate_combine_edge_bf16` | `FLASHRT_HAVE_QWEN35MOE_CORE` | a five-launch tensor chain | one kernel, same arithmetic in the same order, rounded once at the store | +| `moe_grouped_gemm_nvfp4_sm100_bf16out` | `FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100` (sm_110 + `_W4A16`) | the per-expert GEMM loop, or the grouped GEMV | every routed expert of a layer in one launch, with the per-group shapes read from device memory | + +All produce output identical to the path they stand in for, so the fallback is +a performance difference and never a numerical one. The edge recurrence is +shape-specialized to a head dim of 128 and raises for anything else rather than +leaving the output buffer undefined. + +Which of each pair actually runs is a `KernelPolicy` field, not the answer to +"is the symbol there" — see *Runtime controls*. + +### Attention differs by target + +The ten full-attention layers do not use the same kernel everywhere: + +| target | attention | how it is built | +|---|---|---| +| SM120 / SM89 / SM87 | vendored FA2 | automatic: the SM80-family source, which `__CUDA_ARCH__ >= 800` admits | +| Thor SM110 | vendored FA2, or the decomposed reference | opt-in: `-DFLASHRT_ENABLE_THOR_FA2=ON` | +| Thor SM110 | FA4 | its SM100-class CuTe-DSL kernel needs Blackwell tensor memory; ships as the `thor-fa4` pip extra, not compiled into `flash_rt_kernels` | + +FA2 was originally excluded from sm_110 on the grounds that Thor has its own +attention path and FA2 would add about 10 MB of `.so` for nothing. That holds +for the models that use the decomposed path, and it does not hold for a long +prefill of this one: the decomposed path materialises an `(S * heads, S_kv)` +score buffer -- 3.4 GB per layer at ten thousand tokens -- and this model needs +one instantiation (bf16, head_dim 256), not the twelve the size estimate +assumed. Enabling it takes the chunking penalty of a chunked prefill from 64% +to 4%. + +It stays opt-in because every other Thor model still uses its own attention +path and would only be paying the compile time and the binary size. A build +without it runs this model correctly and emits the same tokens; only a long +prefill is slower. Treat a missing FA2 as a signal to fall back, not as a build +error. + +At the *decode* shape the two are the same answer to bf16 precision -- measured +against an fp32 reference, 2.0e-3 relative for both at kv=64, 2.2e-3 against +2.1e-3 at kv=2048 -- so decode on sm_110 keeps the reference path the golden +fixture was recorded through, and takes FA2 only for prefill. +`FLASHRT_NEXN2_DECODE_FA2=1` overrides that. + +The attention backend probes its kernel at construction: it runs one case +through the same launch the hot path uses and compares against +`scaled_dot_product_attention`, falling back if they disagree. That is not +belt-and-braces. Three times in this work a kernel compiled, linked and loaded +while being unable to run — the block-scaled 4-bit tier substitutes an invalid +control path off its own architecture, the lm_head kernel is simply absent +outside GPU_ARCH 120/121, and the vendored FA2 on an SM110 part printed a +complaint and returned without writing its output. That last one still produced +15 of 16 reference tokens, because ten of forty layers contributing nothing is +survivable for a residual stream — which is precisely why a symbol check is not +a capability check. + +`_W4A4` refuses to configure on a target without block-scaled MMA. CUTLASS +still compiles those translation units elsewhere, but substitutes +`CUTE_INVALID_CONTROL_PATH` for the MMA, so the build would succeed and then +fail at run time. The explicit gate turns that into a configure-time error. + ## Usage ```python -from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, -) +from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend -frontend = Qwen36MoeTextFrontendRtx( +frontend = Qwen36MoeTextFrontend( "/models/Qwen3.6-35B-A3B", device="cuda:0", max_seq=4096, @@ -89,7 +221,7 @@ validation can be run without loading model weights: ```bash PYTHONPATH=. python - <<'PY' -from flash_rt.frontends.torch.qwen36_moe_rtx import ( +from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) print(validate_qwen36_moe_checkpoint("/models/Qwen3.6-35B-A3B")) @@ -104,12 +236,55 @@ The shared architecture uses: `8192`; `0` disables chunking. - `FLASHRT_QWEN35MOE_GRAPH_CACHE_MAX` — decode CUDA Graph LRU capacity, default `256`. +- `FLASHRT_QWEN35MOE_SPEC_GRAPH_CACHE_MAX` — speculative-window CUDA Graph LRU + capacity, default `16`. Bounded separately and far lower than the decode + cache: a speculative graph covers `k+1` positions through the whole stack, so + its memory pool is several times a decode step's. The constructor argument + `spec_graph_cache_max` sets it per frontend. The older `FLASHRT_NEXN2_PREFILL_CHUNK` and `FLASHRT_NEXN2_GRAPH_CACHE_MAX` names remain compatible aliases. +Which of several interchangeable kernels each step calls is a +`KernelPolicy` (`flash_rt.frontends.torch._nexn2_rtx_forward`), not a symbol +lookup: every field selects between implementations checked against each other +with `torch.equal`, so a field decides speed and never output. The environment +variables above and `NEXN2_WY_GDN`, `NEXN2_ROUTE_KERNEL`, +`NEXN2_DENSE_CUBLASLT`, `FLASHRT_QWEN35MOE_W4A16_EDGE` and +`FLASHRT_QWEN35MOE_VERIFY_K_ROWS` are its defaults. A policy must not be +changed between a CUDA graph capture and its replay. + ## Validation +### Build and symbol matrix + +A build with every `qwen3_5_moe` option off must compile the same sources and +export the same symbols it did before the tiers existed. That is a property of +the gates, so it is checked by reading them: + +```bash +python scripts/qwen35moe_build_matrix.py # print sources + symbols per tier +python scripts/qwen35moe_build_matrix.py --check # exit 1 if any tier leaks +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + pytest -q -p no:cacheprovider tests/test_qwen35moe_build_matrix.py +``` + +The five configurations behind it, each configure-only: + +| configuration | flags | result | +|---|---|---| +| baseline SM120 | `-DGPU_ARCH=120` | `FA2 ENABLED`; no `qwen3_5_moe` source or symbol | +| baseline SM110 | `-DGPU_ARCH=110` | `FA2 DISABLED`; no `qwen3_5_moe` source or symbol | +| SM110 supported | `-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON -DFLASHRT_ENABLE_THOR_FA2=ON` | core + weight-only tiers, grouped MoE GEMM, FA2 at `hdim={256} x dtype={bf16}` | +| SM120 supported | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` | all three tiers | +| SM110 block-scaled | `-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_W4A4=ON` | **configure fails**, naming the two tiers that do apply | + +The last row is the point of the explicit gate: CUTLASS would otherwise compile +those translation units on sm_110 with the MMA replaced by an invalid control +path, and the failure would arrive at run time instead. + +### Tests + The repository smoke test is checkpoint-independent: ```bash @@ -117,6 +292,18 @@ PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ pytest -q -p no:cacheprovider tests/test_qwen36_moe_smoke.py ``` +Speculative decode has its own file. The constructor contract and the graph +cache's eviction policy run anywhere; the equivalence tests -- K=1 and K=2 +against plain greedy, the window's logits and recurrent, conv and KV state +against the decode steps they stand in for, the rejected-tail rewind, and the +boundary token counts -- need a GPU and a checkpoint and skip without them: + +```bash +FLASHRT_QWEN36_MOE_CKPT_DIR=/models/Qwen3.6-35B-A3B \ +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ +pytest -q -p no:cacheprovider tests/test_qwen36_moe_spec_decode.py +``` + Set `FLASHRT_QWEN36_MOE_CKPT_DIR` to include the official checkpoint contract test. Performance and precision numbers must be measured on Qwen3.6 weights; Nex-N2-mini measurements are not interchangeable even though the compute @@ -146,9 +333,43 @@ official BF16 checkpoint: | 64-token prompt, 32-token eager decode | 48.14 tok/s | | 64-token prompt, 32-token warm CUDA Graph decode | 195.49 tok/s | +These are the original first-light figures, measured with a 32-token +generation. The current SM120 reference for this path is the same-shape +measurement at `P=64, N=64`, warm CUDA Graph: + +| Measurement | Result | +|---|---:| +| Prefill | 40.42 ms | +| Decode, median of 8 runs | 257.95 tok/s | +| Decode, range across 8 runs | 256.90–258.22 tok/s | +| Repeated sequences identical | 8 / 8 | + +Quote that one, not the 32-token row above: the two use different generation +lengths and are not interchangeable. + +The decode work described under *Speculative decode* was tuned and measured on +Thor, and its one tuning constant is scoped to that architecture, so the SM120 +kernels are byte-identical to the ones these numbers were taken on. + The eager, first-capture, and warm-graph runs produced the same 32 token IDs. These numbers are a first-light correctness run, not a context-length sweep. +Reproduce with: + +```bash +PYTHONPATH=. python benchmarks/qwen36_moe_edge_decode.py \ + --checkpoint /path/to/Qwen3.6-35B-A3B \ + --prompt-tokens 64 --max-new-tokens 32 +``` + +The benchmark refuses to report throughput if the eager and captured paths +disagree on any token, because a rate for a path that emits different text is +not a rate for the same work. + +The table above predates the decode work described under *Speculative decode* +below; the correctness gate (`tests/test_qwen36_moe_gpu.py`) passes on the +current tree, but the SM120 latency figures have not been re-measured since. + Four chat prompts from 12 to 45 tokens were also compared with the official Transformers BF16 implementation: @@ -162,12 +383,147 @@ Transformers BF16 implementation: The logit cosine is lower than the Nex-N2-mini measurement, but the tested greedy sequences were token-exact for 16 generated tokens on all four prompts. +## Jetson AGX Thor numbers + +Measured on Jetson AGX Thor (sm_110), unified memory, the same BF16 checkpoint +with runtime NVFP4 conversion, against vLLM 0.26.0 on the same part with the +same pre-tokenized prompts and the same protocol on both sides: TTFT is the +wall time of a one-token generate, decode is the rest of a 64-token generate +with that TTFT subtracted, best of three after a warm-up. One prompt length per +process on both sides. The vision tower is off on both, since this frontend is +text-only. + +| prompt | TTFT | vLLM TTFT | | +|---:|---:|---:|---| +| 20 | **89.5 ms** | 102.3 ms | +14% | +| 256 | **104.7 ms** | 214.5 ms | +105% | +| 512 | **144.5 ms** | 251.1 ms | +74% | +| 1024 | **216.0 ms** | 319.4 ms | +48% | +| 2048 | **379.6 ms** | 495.0 ms | +30% | +| 4096 | **748.6 ms** | 867.4 ms | +16% | +| 10240 | **1890.7 ms** | 2144.5 ms | +13% | +| 32768 | **7207.5 ms** | 7231.8 ms | +0.3% | + +Where the prefill is dominated by per-layer work -- projections, routing, the +linear-attention scan -- the kernels win, and win by more the shorter the +prompt. Where it is dominated by attention's O(S^2) the lead narrows, because +attention is the one component still reached through torch. + +Decode. Each row carries the tree it was taken on, because the decode round +below moved the step 15.3% and the vLLM sweep predates it: + +| measurement | tree | decode | vLLM decode | | +|---|---|---:|---:|---| +| @1024, from the sweep above | before the decode round | 87.1 tok/s | 31.6 tok/s | 2.8x | +| @2048 | before the decode round | 86.3 tok/s | 31.5 tok/s | 2.7x | +| @4096 | before the decode round | 85.1 tok/s | 31.2 tok/s | 2.7x | +| captured steady step @20 | before the decode round | 89.0 tok/s | | | +| captured steady step @20 | **current** | **102.6 tok/s** | | | + +The two protocols agree on the same tree -- the sweep reads 87.1 at 1024, the +captured step 89.0 at 20, a 2% spread -- so the distance from 87 to 102.6 is +the decode round, not the way it was timed. vLLM's side is a different binary +and nothing here changes it, which makes 2.7-2.8x a lower bound on the current +ratio; the current tree has not been swept at 1024-4096, so no figure is quoted +for it. + +Context reaches 128 K on this board, at 2470 tok/s of prefill. It could not +before FA2 was available here: the default configuration chunks past 8192 +tokens and every chunk asked for a non-square causal window, which had no fused +backend, so the scores were materialised. + +### The decode round + +Five changes, each bit-identical to what it replaced, measured in the captured +steady step at a 20-token prompt. + +| | step | tok/s | +|---|---:|---:| +| round start | 11.238 ms | 89.0 | +| gating constants derived once, not per step | 11.059 ms | 90.4 | +| MoE tail fused into one kernel | 10.827 ms | 92.4 | +| spill-free gated-DeltaNet recurrence | 10.379 ms | 96.4 | +| single-warp router top-8 | 10.297 ms | 97.1 | +| grouped GEMV `kUnroll` 4 -> 2 | **9.743 ms** | **102.6** | + +The last row is a build-time constant and is bit-identical by construction: the +main loop advances by `32*kUnroll` and the tail takes the remainder, so a lane +visits the same k-blocks in the same order for any value. The sweep is not +monotone -- 1: 99.7, **2: 102.6**, 3: 96.5, 4: 96.3 -- so two is a genuine +optimum, and it is scoped to sm_110 in CMake because it was measured on a +20-SM part. + +## Speculative decode + +The MTP head ships with the checkpoint and is loaded on request. It is a +DeepSeek-V3-style single module: it reads the pre-final-norm hidden state of the +previous position and the token emitted at this one, and predicts the next. +Drafts are chained, so acceptance decays with each additional draft. + +The window is verified through the decode kernels at `K+1` rows, over the +weights the decode step caches, so a verified row is the decode step it stands +in for -- bit for bit, not approximately. That is what allows the emitted text +to be plain greedy's, and it is checked directly: logits rows, per-token +recurrent and conv snapshots, and the KV rows written are all compared with +`torch.equal` against a decode step run over the same tokens. + +Plain and speculative are measured in the same process, so every ratio is +paired with a baseline from its own run. The absolute rates move a few percent +with what else the board is doing; the ratio is the stable part. Every row +emitted plain greedy's sequence token for token, with the 16-token golden +fixture passing on the same build. + +| tree | plain | K=2 | | +|---|---:|---:|---:| +| after the fused MoE tail | 91.18 | 97.34 | 1.07x | +| after the single-warp router | 91.98 | 98.57 | 1.07x | +| after `kUnroll` 4 -> 2 | 96.64 | 100.66 | 1.04x | +| current | **100.35** | **106.74** | **1.06x** | + +K=1 on the current tree reads 105.22 against the same 100.35. + +`K=2` is the operating point. Above it the window costs more than the extra +accepted tokens return: each additional verified row re-reads the routed +experts, which do not amortise across a window the way the dense weights do, +and each additional draft pays a full-vocabulary projection. + +Acceptance rises with context -- 2.60 tokens kept per window at a 20-token +prompt, 2.72 at 512, 2.74 at 2048 -- while the ratio does not, because the +verify runs `K+1` separate single-query attention passes. That is the price of +keeping the window bit-exact, and it is the largest remaining lever on this +path. + +For scale, vLLM 0.26.0 supports `Qwen3_5MoeMTP` -- the same head -- and with +`num_speculative_tokens=2` reaches 55.00 tok/s at a 20-token prompt against its +own 31.59, a 1.74x gain. The larger gain rests on a step three times heavier: a +fixed per-draft cost is proportionally three times cheaper against 31.7 ms than +against 9.7 ms. Plain greedy decoding here is faster than that speculative +figure by 1.8x, with no speculation at all. + +Enable it with `load_mtp=True` on the constructor; the window width is the `k` +argument to `generate_spec`: + +```python +frontend = Qwen36MoeTextFrontend( + "/models/Qwen3.6-35B-A3B", + device="cuda:0", + max_seq=2048, + load_mtp=True, + spec_graph_cache_max=16, +) +frontend.set_prompt("Explain why deterministic reductions matter.") +token_ids = frontend.generate_spec(max_new_tokens=128, k=2) +print(frontend.tokenizer.decode(token_ids)) +``` + +`FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0` falls back to verifying through the prefill +forward, which is slower and produces the same tokens. + ## Limitations - Text only; the vision tower is not loaded. - The kernelized runtime NVFP4 path is required. -- Greedy decode only. -- The MTP tensors are validated but not loaded, so speculative decode is not - enabled. +- Greedy decode only. Speculative decode is greedy as well: it emits the + sequence plain greedy decoding would emit, token for token, or it is a bug. - Only the BF16 source checkpoint with runtime NVFP4 conversion is supported. -- SM120 only. +- Sampling, batching, and beam search are not implemented. diff --git a/flash_rt/api.py b/flash_rt/api.py index 13d4e20e..a29b163b 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -625,8 +625,8 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, raise NotImplementedError( "config='qwen36_moe' is a text LLM and is not served through " "load_model's VLA wrapper. Construct it directly:\n" - " from flash_rt.frontends.torch.qwen36_moe_rtx import " - "Qwen36MoeTextFrontendRtx\n" + " from flash_rt.frontends.torch.qwen36_moe import " + "Qwen36MoeTextFrontend\n" "See docs/qwen36_moe_usage.md.") from flash_rt.hardware import detect_arch, resolve_pipeline_class diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 7ff6f8f0..261a1afa 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -29,7 +29,8 @@ from flash_rt.frontends.torch._nexn2_rtx_forward import ( CONV, HD, HID, HK, HV, INTER, KD, KS, NKV, NQ, NV, ROPE, TOPK, VD, - _quant_act, build_rope_tables, nexn2_forward_nvfp4, + _quant_act, _w4a16_mrows, build_rope_tables, kernel_policy, + moe_grouped_w4a16, nexn2_forward_nvfp4, set_spec_verify, w4a16_matvec, ) from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 @@ -47,6 +48,22 @@ def _qwen35moe_env(name: str, default: str) -> str: _BATCHED_PREFILL_MIN_S = 8 +def _cache_put(cache, key, value, cap): + """Insert ``value`` and evict least-recently-used down to ``cap``. + + Both graph caches go through this. A captured graph owns its memory pool, + so an unbounded cache leaks device memory across a long generation -- one + graph per absolute position. ``cap <= 0`` disables the bound. + + Returns the value, so a caller can insert and use in one expression. + """ + cache[key] = value + if cap > 0: + while len(cache) > cap: + cache.popitem(last=False) # evict LRU + return value + + def _cs(): """Current CUDA stream handle. Inside torch.cuda.graph capture this is the capture stream; eager, the default stream. fvk calls MUST use it -- @@ -114,7 +131,7 @@ def _w4a16_mv(x1k, w_bf16, ld, key, fvk, device): ld[key + '_w4a16_a'] = float(og.item()) xc = x1k.contiguous() y = torch.empty(1, n, dtype=torch.bfloat16, device=device) - fvk.w4a16_matvec_sm120_bf16( + w4a16_matvec(fvk)( xc.data_ptr(), ld[pk].data_ptr(), ld[key + '_w4a16_sf'].data_ptr(), y.data_ptr(), n, k, ld[key + '_w4a16_a'], _cs()) return y @@ -177,7 +194,8 @@ def _proj_mma(x2d, ld, base, n, fvk, device, state=None): class Nexn2DecodeState: """Persistent decode state: GDN recurrent/conv caches, KV cache, RoPE.""" - def __init__(self, handles, max_seq, device): + def __init__(self, handles, max_seq, device, *, + spec_graph_cache_max=None): self.handles = handles self.device = device self.max_seq = int(max_seq) @@ -208,8 +226,18 @@ def __init__(self, handles, max_seq, device): torch.zeros(1, CONV, KS - 1, dtype=bf16, device=device) for _ in range(nlin)] - # Full-attn KV cache. - self.attn = RtxFlashAttnBackendNexn2(max_seq=self.max_seq, max_q_seq=1) + # Full-attn KV cache. A loaded draft head is one more full-attention + # layer and takes the slot after the model's own. + self.mtp = p.get('mtp') + self.mtp_rank = nfull if self.mtp is not None else None + self.attn = RtxFlashAttnBackendNexn2( + max_seq=self.max_seq, max_q_seq=1, + num_full_layers=nfull + (1 if self.mtp is not None else 0)) + # The pre-final-norm hidden state of the last step, which is what the + # draft head reads. Written every step whether or not one is loaded: + # a 4 KB device copy, and making it conditional would put a Python + # branch inside the captured region. + self.last_hidden = torch.zeros(HID, dtype=bf16, device=device) # RoPE tables for the whole window. theta = float(p['rope_theta']) @@ -273,6 +301,79 @@ def __init__(self, handles, max_seq, device): # chunk (64). 0 disables (always single-pass). self.prefill_chunk = int( _qwen35moe_env("PREFILL_CHUNK", "8192")) + # Optional eager-only traces used to size edge expert caches and to + # score expert quantization against real activations. Keep these + # disabled during CUDA Graph capture. + self.router_trace = None + self.moe_input_trace = None + self._active_layer = -1 + # Per-token recurrent/conv snapshots for a speculative window, sized + # on first use. Only allocated when speculation runs: 30 layers of + # (NV, HK, HV) bf16 per window slot. + self.spec_states = None + self.spec_conv = None + # Set only around a verify block. Prefill runs the same layer code and + # would otherwise pay for -- and overrun -- snapshots it never uses. + self.spec_capture = False + # One captured graph per (pos, window): the KV slots, attention length + # and RoPE slice are baked per position exactly as the decode graph's + # are, and each owns a memory pool, so this is LRU-bounded the same way. + # + # It is NOT bounded at the same number. A speculative graph covers k+1 + # positions through the whole stack, so its pool is several times a + # decode step's, and the decode cap of 256 is sized for a step. Holding + # 256 of these alongside the model is more than a 32 GB board has at a + # 2048-token context -- measured there, it is what runs it out of + # memory. Sixteen keeps the windows a generation actually revisits + # (recapture costs two warmup runs) and bounds the pools at something + # the smallest supported board carries. + self._spec_graphs = collections.OrderedDict() + self.spec_graph_cache_max = int( + spec_graph_cache_max if spec_graph_cache_max is not None + else _qwen35moe_env("SPEC_GRAPH_CACHE_MAX", "16")) + # Its own memory pool, not the decode graphs'. The two are replayed + # interleaved -- a window, then whatever the caller does next -- and + # sharing a pool between graphs used that way is the case the runtime + # does not promise to handle. Measured: with the pool shared, a 64-token + # speculative run ran at half the rate of a 32-token one on identical + # code, the cost growing with the number of live graphs. + self._spec_pool = torch.cuda.graph_pool_handle() + self._spec_tokens = None + self._spec_argmax = None + # Which half of the draft head's fc input carries the hidden state. + # The checkpoint does not say and fc is square in the concatenated + # width, so it was settled by measuring acceptance both ways -- and the + # answer is the embedding first. Over 48 decoded tokens: + # + # cat[embed, hidden] first draft 0.896, chained 0.646, 0.417 + # cat[hidden, embed] 0.000, 0.000, 0.000 + # + # The wrong half drafts noise, so nothing is ever accepted and every + # window pays for a verify that keeps one token. It agrees with the + # reference implementation of this head, which concatenates the + # embedding first as well. + self.mtp_hidden_first = ( + _qwen35moe_env("MTP_HIDDEN_FIRST", "0") != "0") + # Set to an ExpertCache to read the routed experts from storage. Only + # meaningful when the loader skipped them; see _moe_experts_streamed. + self.expert_cache = None + self._scratch = None + self._hadamard = None + + def _streamed_scratch(self, device): + """The two decode buffers a streamed expert is unpacked into. + + Allocated once and reused: 4 MiB for gate_up and 2 MiB for down, which + would otherwise be allocated 8 times per layer per token. + """ + if self._scratch is None: + self._scratch = { + 'gate_up': torch.empty( + 2 * INTER, HID, dtype=torch.bfloat16, device=device), + 'down': torch.empty( + HID, INTER, dtype=torch.bfloat16, device=device), + } + return self._scratch def reset(self): for s in self.lin_state: @@ -282,6 +383,57 @@ def reset(self): self.attn.reset_cache() +def router_topk(fvk): + """The router top-k entry this build should call. + + The warp variant returns identical indices and values -- argmax under a + total order picks one element regardless of the reduction tree, checked + over 800 inputs of which 397 had a tie inside the top-8 -- without the + block kernel's 24 barriers. + """ + if kernel_policy().warp_router_topk: + fn = getattr(fvk, 'moe_router_topk_warp_sm120_bf16', None) + if fn is not None: + return fn + return fvk.moe_router_topk_sm120_bf16 + + +def gdn_recurrent(fvk): + """The single-token GDN recurrence entry this build should call. + + The edge variant is the same arithmetic in the same order -- checked + exactly, over chained steps so the state drift is exercised too -- and it is + 1.70x the shipped one. The shipped one holds the thread's whole state column + in a 128-float array that cannot live in registers, so it is in local + memory and walked five times; ncu measures 39 registers per thread for a + 128-float array. 51% of bandwidth against 87%. + """ + if kernel_policy().gdn_recurrent_edge: + fn = getattr(fvk, 'gated_deltanet_recurrent_edge_qwen36_bf16', None) + if fn is not None: + return fn + return fvk.gated_deltanet_recurrent_qwen36_bf16 + + +def _gdn_gate_consts(ld, device): + """The gating kernel's two constant inputs, derived once per layer. + + ``A_log`` and ``dt_bias`` are weights, so -exp(A_log) and the fp32 bias are + the same on every step -- but deriving them per call put four elementwise + launches per GDN layer inside the captured region, thirty layers of them, + recomputing values identical to the previous replay's. Each is a couple of + microseconds of dispatch quantum for no arithmetic anyone reads. + + Same expressions in the same order, so the bytes handed to the kernel are + the bytes it was getting before. + """ + if 'gdn_neg_exp_a' not in ld: + ld['gdn_neg_exp_a'] = ( + -ld['A_log_t'].float().exp()).float().contiguous() + ld['gdn_dt_bias_f'] = ld['dt_bias_t'].float().contiguous() + return ld['gdn_neg_exp_a'], ld['gdn_dt_bias_f'] + + def _decode_gdn(h, ld, state, lin_rank, fvk, device): """GDN layer at one token, updating recurrent + conv state in place.""" eps = state.eps @@ -289,7 +441,7 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): Wz = ld['in_proj_z_w_t'] Wb, Wa = ld['in_proj_b_w_t'], ld['in_proj_a_w_t'] convw = ld['conv1d_w_t'].reshape(CONV, KS).contiguous() - A_log, dtb = ld['A_log_t'].float(), ld['dt_bias_t'].float() + neg, dtb_c = _gdn_gate_consts(ld, device) nw = ld['gdn_norm_w_t'] s = _cs() @@ -326,8 +478,6 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): conv_out.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), 1, s) - neg = (-A_log.exp()).float().contiguous() - dtb_c = dtb.contiguous() g_out = torch.empty(1, NV, dtype=torch.bfloat16, device=device) bo = torch.empty(1, NV, dtype=torch.bfloat16, device=device) fvk.qwen36_gdn_gating_bf16( @@ -340,7 +490,7 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): gt = g_out.reshape(NV).contiguous() bt = bo.reshape(NV).contiguous() core = torch.empty(NV, HV, dtype=torch.bfloat16, device=device) - fvk.gated_deltanet_recurrent_qwen36_bf16( + gdn_recurrent(fvk)( qt.data_ptr(), kt.data_ptr(), vt.data_ptr(), gt.data_ptr(), bt.data_ptr(), state.lin_state[lin_rank].data_ptr(), core.data_ptr(), 1, NV, HK, HV, True, s) @@ -407,6 +557,109 @@ def _decode_full(h, ld, state, full_rank, pos, fvk, device): 1, 1, HID) +def _hadamard16(device): + """The block-16 transform, built once. Symmetric and its own inverse.""" + m = torch.ones(1, 1, dtype=torch.float32, device=device) + for _ in range(4): + m = torch.cat((torch.cat((m, m), 1), torch.cat((m, -m), 1)), 0) + return m / 4.0 + + +def _rotate16(x, h): + """Apply the transform along the last dimension, in blocks of 16.""" + shape = x.shape + return (x.reshape(-1, 16).float() @ h).reshape(shape).to(x.dtype) + + +def _moe_experts_streamed(x, idx, state, fvk, device, s): + """The routed experts' outputs, read from storage instead of memory. + + Returns only ``d_dn`` -- the per-slot expert outputs. The weighted sum, the + shared expert and its gate are identical to the resident path and stay + there; replacing the whole layer here is how an earlier version silently + dropped the shared expert from every layer. + + Reachable only when the loader was told to stream, in which case the + per-layer stacked tensors were never allocated. Each block is decoded to + bf16 and multiplied with the shared bf16 GEMV, because the block-scaled + 4-bit GEMMs read neither this codebook nor this scale layout. + + When the bundle was written with the transform applied, the stored weight is + H*W, so the activation entering each GEMM has to be rotated the same way or + the products are wrong -- while staying finite and plausible, which is + exactly how it goes unnoticed. + """ + cache = state.expert_cache + layer = state._active_layer + experts = [int(value) for value in idx.cpu().tolist()] + cache.get_many(layer, experts) + + rotated = bool(cache.manifest.get('rht')) + if rotated and state._hadamard is None: + state._hadamard = _hadamard16(device) + h16 = state._hadamard + + scratch = state._streamed_scratch(device) + d_gu = torch.empty(TOPK, 2 * INTER, dtype=torch.bfloat16, device=device) + d_dn = torch.empty(TOPK, HID, dtype=torch.bfloat16, device=device) + xc = (_rotate16(x, h16) if rotated else x).contiguous() + + for slot, expert in enumerate(experts): + parts = cache.components(layer, expert) + gu_alpha, dn_alpha = parts['global_scales'].tolist() + rc = fvk.qwen35moe_e0m3_dequant_bf16( + parts['gate_up_weight'].data_ptr(), + parts['gate_up_scale'].data_ptr(), + scratch['gate_up'].data_ptr(), + 2 * INTER, HID, cache.group_size, gu_alpha, s) + if rc: + raise RuntimeError(f'gate_up decode failed with {rc}') + fvk.bf16_matvec_sm120_bf16( + xc.data_ptr(), scratch['gate_up'].data_ptr(), + d_gu[slot].data_ptr(), 2 * INTER, HID, s) + + gated = _silu_mul( + d_gu[slot:slot + 1, :INTER], d_gu[slot:slot + 1, INTER:], + fvk, device) + if rotated: + gated = _rotate16(gated, h16) + gated = gated.contiguous() + rc = fvk.qwen35moe_e0m3_dequant_bf16( + parts['down_weight'].data_ptr(), + parts['down_scale'].data_ptr(), + scratch['down'].data_ptr(), + HID, INTER, cache.group_size, dn_alpha, s) + if rc: + raise RuntimeError(f'down decode failed with {rc}') + fvk.bf16_matvec_sm120_bf16( + gated.data_ptr(), scratch['down'].data_ptr(), + d_dn[slot].data_ptr(), HID, INTER, s) + return d_dn + + +def _shared_combine(routed, shared, glog, rows, fvk, device): + """out = routed(fp32) + shared(bf16) * sigmoid(gate), in one kernel. + + The tensor-op form is a cast, a sigmoid, a broadcast multiply, an add and a + cast: five launches a layer, forty layers, in a step that is 99% kernel + time and where a launch costs its dispatch quantum whether or not it + computes much. The kernel does the same arithmetic in the same order and + rounds once at the store, so it stands in for the chain rather than + approximating it -- checked bit for bit at the shapes and scales decode and + the window issue, because the fixture and the speculative verify both rest + on it. + """ + if (kernel_policy().fused_shared_combine + and hasattr(fvk, 'moe_shared_gate_combine_edge_bf16')): + out = torch.empty(rows, HID, dtype=torch.bfloat16, device=device) + fvk.moe_shared_gate_combine_edge_bf16( + routed.data_ptr(), shared.data_ptr(), glog.data_ptr(), + out.data_ptr(), rows, HID, _cs()) + return out + sgate = torch.sigmoid(glog.float()).reshape(rows, 1) + return (routed + shared.float() * sgate).to(torch.bfloat16) + + def _moe_layer_decode(h, ld, state, fvk, device): """M=1 fine-grained MoE via the grouped GEMV kernel: the 8 routed experts run in one launch each for gate_up (shared act) and down (per-slot act), @@ -440,38 +693,59 @@ def _moe_layer_decode(h, ld, state, fvk, device): lr = logit_raw.reshape(-1).contiguous() idx = torch.empty(TOPK, dtype=torch.int32, device=device) topv = torch.empty(TOPK, dtype=torch.float32, device=device) - fvk.moe_router_topk_sm120_bf16(lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), - lr.numel(), TOPK, s) + # idx and topv come from torch.empty, so an unchecked failure here leaves + # uninitialised memory to be used as expert indices -- which reaches a file + # offset before anything notices. + rc = router_topk(fvk)( + lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), lr.numel(), TOPK, s) + if rc: + raise RuntimeError( + f'router top-k failed with {rc} for {lr.numel()} experts, k={TOPK}') tw_row = F.softmax(topv, -1) # (TOPK,) device - - if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer - ld['experts_gate_up_alpha_dev'] = \ - ld['experts_gate_up_alpha_t'].to(device).contiguous() - ld['experts_down_alpha_dev'] = \ - ld['experts_down_alpha_t'].to(device).contiguous() - gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] - dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] - gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] - n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] # 1024 / HID - - # gate_up: shared BF16 activation, grouped W4A16 over the 8 experts. BF16 - # activation -> no activation quant, higher cos than the W4A4 mma, and - # faster at this scale (6.2 vs 8.2 us standalone). - xc = x.contiguous() - d_gu = torch.empty(TOPK, n_gu, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( - xc.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), - idx.data_ptr(), d_gu.data_ptr(), TOPK, n_gu, HID, - 0, gu_p[0].numel(), gu_s[0].numel(), s) - - # down: silu(gate)*up (BF16, fused) then grouped W4A16 (per-slot activation). - g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] - inter = _silu_mul(g_, u_, fvk, device).contiguous() - d_dn = torch.empty(TOPK, n_dn, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( - inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), - idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, - INTER, dn_p[0].numel(), dn_s[0].numel(), s) + if state.router_trace is not None: + state.router_trace[state._active_layer].append( + tuple(int(v) for v in idx.cpu().tolist())) + if state.moe_input_trace is not None: + state.moe_input_trace[state._active_layer].append( + x.detach().to("cpu", copy=True)) + + # Streaming replaces only the routed experts' own GEMVs. Everything after + # this -- the weighted sum, the shared expert, its gate -- is identical, and + # returning early from here is how an earlier version silently dropped the + # shared expert from every layer. + if ld.get('experts_streamed'): + d_dn = _moe_experts_streamed(x, idx, state, fvk, device, s) + n_dn = HID + else: + if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + gu_a = ld['experts_gate_up_alpha_dev'] + dn_a = ld['experts_down_alpha_dev'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] # 1024 / HID + + # gate_up: shared BF16 activation, grouped W4A16 over the 8 experts. + # BF16 activation -> no activation quant, higher cos than the W4A4 mma, + # and faster at this scale (6.2 vs 8.2 us standalone). + xc = x.contiguous() + d_gu = torch.empty(TOPK, n_gu, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + xc.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), + idx.data_ptr(), d_gu.data_ptr(), TOPK, n_gu, HID, + 0, gu_p[0].numel(), gu_s[0].numel(), s) + + # down: silu(gate)*up (BF16, fused) then grouped W4A16 (per-slot act). + g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] + inter = _silu_mul(g_, u_, fvk, device).contiguous() + d_dn = torch.empty(TOPK, n_dn, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), + idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, + INTER, dn_p[0].numel(), dn_s[0].numel(), s) # Fixed-order weighted sum. The generic torch matmul may choose a # reduction whose accumulation order changes between launches, which can # flip a later greedy decision when two logits are nearly tied. @@ -501,10 +775,11 @@ def _moe_layer_decode(h, ld, state, fvk, device): si = _silu_mul(sg, su, fvk, device) shared = _proj_mma(si, ld, 'shared_down_proj', HID, fvk, device, state) # shared-expert scalar gate: N=1 GEMV via the bf16 matvec kernel (was a - # torch matmul -- the last fp32 matmul in the captured decode step). - sgate = torch.sigmoid( - _bf16_mv(x, ld['shared_gate_w_t'], fvk, device).float()) - return (out + shared.float() * sgate).reshape(1, 1, HID).to(torch.bfloat16) + # torch matmul -- the last fp32 matmul in the captured decode step). The + # sigmoid, the broadcast multiply, the add and the cast are one kernel. + glog = _bf16_mv(x, ld['shared_gate_w_t'], fvk, device) + return _shared_combine(out, shared, glog, 1, fvk, device).reshape( + 1, 1, HID) def decode_step(state, token_id, pos, fvk, device): @@ -528,13 +803,54 @@ def decode_step(state, token_id, pos, fvk, device): h = res + attn res = h n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + state._active_layer = L h = res + _moe_layer_decode(n, ld, state, fvk, device) + # The pre-final-norm hidden state is what a DeepSeek-V3-style draft head + # consumes. Keeping it in a fixed buffer costs one 4 KB device copy and + # survives graph capture, unlike reading it out per step. + state.last_hidden.copy_(h.reshape(HID)) h = _rms_fvk(h, p['final_norm_w_t'], fvk, device, state.eps) # lm_head as NVFP4 W4A16: 4x less weight read (1GB -> 0.25GB) via the # hand-tuned mma (3.1x the bf16 GEMV; the CUTLASS widen is M=1-broken). # The weight is quantised once during the eager seed (cached on p), so # the captured graph only runs the activation quant + fp4 GEMM. + return _lm_head(state, h, fvk, device) + + +def _ensure_lm_head_nvfp4(state, fvk, device): + """Quantise the lm_head to swizzled NVFP4 once, on the handles. + + Both the single-row decode head and the M-row verify read this one copy, + so the verify cannot drift from decode by having been handed a second + quantisation of the same weight. The .item() lands here, on the first + eager call, and never inside a captured region. + """ + p = state.handles.ptrs + if 'lm_head_packed_t' in p: + return + w = p['lm_head_w_t'].contiguous() + nn, kk = w.shape + packed = torch.empty(nn, kk // 2, dtype=torch.uint8, device=device) + sf = torch.zeros(_sf_swz_bytes(nn, kk), dtype=torch.uint8, device=device) + scr = torch.zeros(1, dtype=torch.float32, device=device) + og = torch.zeros(1, dtype=torch.float32, device=device) + fvk.bf16_weight_to_nvfp4_swizzled( + w.data_ptr(), packed.data_ptr(), sf.data_ptr(), + scr.data_ptr(), og.data_ptr(), nn, kk, 0) + torch.cuda.synchronize() + p['lm_head_packed_t'] = packed + p['lm_head_sf_t'] = sf + p['lm_head_alpha'] = float(og.item()) + + +def _lm_head(state, h, fvk, device): + """Project a hidden state to logits over the full vocabulary. + + Taken out of decode_step so the speculative draft head, which ends the + same way, does not carry a second copy of the quantise-once bookkeeping. + """ + p = state.handles.ptrs vocab = p['vocab_size'] logits = torch.empty(1, vocab, dtype=torch.bfloat16, device=device) if not state.lm_head_nvfp4: @@ -542,28 +858,584 @@ def decode_step(state, token_id, pos, fvk, device): h.reshape(1, HID).contiguous().data_ptr(), p['lm_head_w_t'].data_ptr(), logits.data_ptr(), vocab, HID, _cs()) return logits - if 'lm_head_packed_t' not in p: - w = p['lm_head_w_t'].contiguous() - nn, kk = w.shape - packed = torch.empty(nn, kk // 2, dtype=torch.uint8, device=device) - sf = torch.zeros(_sf_swz_bytes(nn, kk), dtype=torch.uint8, device=device) - scr = torch.zeros(1, dtype=torch.float32, device=device) - og = torch.zeros(1, dtype=torch.float32, device=device) - fvk.bf16_weight_to_nvfp4_swizzled( - w.data_ptr(), packed.data_ptr(), sf.data_ptr(), - scr.data_ptr(), og.data_ptr(), nn, kk, 0) - torch.cuda.synchronize() - p['lm_head_packed_t'] = packed - p['lm_head_sf_t'] = sf - p['lm_head_alpha'] = float(og.item()) - xp, xsf = _quant_act(h.reshape(1, HID), fvk, device, _cs()) - fvk.fp4_w4a4_mma_sm120_full_n_bf16out( - xp.data_ptr(), p['lm_head_packed_t'].data_ptr(), logits.data_ptr(), - vocab, HID, xsf.data_ptr(), p['lm_head_sf_t'].data_ptr(), - p['lm_head_alpha'], _cs()) + _ensure_lm_head_nvfp4(state, fvk, device) + if hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out'): + xp, xsf = _quant_act(h.reshape(1, HID), fvk, device, _cs()) + fvk.fp4_w4a4_mma_sm120_full_n_bf16out( + xp.data_ptr(), p['lm_head_packed_t'].data_ptr(), + logits.data_ptr(), vocab, HID, xsf.data_ptr(), + p['lm_head_sf_t'].data_ptr(), p['lm_head_alpha'], _cs()) + return logits + # That kernel is built only for GPU_ARCH 120/121, so on every other target + # this path had no implementation at all. The W4A16 matvec reads the same + # swizzled weight and the same scale factors, leaves the activation in + # bf16 -- so it also skips the activation quantisation and its error -- and + # lives in a tier that builds wherever the core does. + w4a16_matvec(fvk)( + h.reshape(1, HID).contiguous().data_ptr(), + p['lm_head_packed_t'].data_ptr(), p['lm_head_sf_t'].data_ptr(), + logits.data_ptr(), vocab, HID, p['lm_head_alpha'], _cs()) return logits +def mtp_draft(state, token_id, pos, fvk, device, *, hidden=None): + """Draft the token after next with the MTP head. + + A DeepSeek-V3 single-module head: it sees the main model's last hidden + state for position p-1 and the token emitted at p, and predicts p+1. The + layer under it is an ordinary full-attention layer with its own MoE, so it + runs through the same per-layer code as the model -- which is the point of + loading it through the same loader. + + ``hidden`` defaults to the buffer the last decode step wrote. The head + carries its own KV at the same absolute positions as the model, so calling + this advances that cache and nothing else. + """ + p = state.handles.ptrs + mtp = state.mtp + if mtp is None: + raise RuntimeError( + 'no MTP head is loaded; build the frontend with speculation ' + 'enabled so the loader reads it') + ld = mtp['layer'] + h_prev = state.last_hidden if hidden is None else hidden + + # The draft head runs on its BF16 weights, not on the runtime W4A16 the + # model's own projections take. A draft is one layer -- its weights are a + # rounding error against the window's traffic -- while its accuracy is the + # whole point, since a rejected draft costs a verified position. The + # sibling frontends keep the head BF16 for the same reason; this path had + # been quantising it along with everything else. + was_w4a16, state.dense_w4a16 = state.dense_w4a16, False + try: + return _mtp_draft_bf16(state, mtp, ld, p, token_id, h_prev, pos, + fvk, device) + finally: + state.dense_w4a16 = was_w4a16 + + +def _mtp_draft_bf16(state, mtp, ld, p, token_id, h_prev, pos, fvk, device): + e = F.embedding(token_id.view(1, 1), p['embed_w_t']).reshape(1, HID) + hn = _rms_fvk(h_prev.reshape(1, HID), mtp['pre_h_w_t'], fvk, device, + state.eps) + en = _rms_fvk(e, mtp['pre_e_w_t'], fvk, device, state.eps) + # Which half goes first is a checkpoint convention, not something the + # shapes pin down -- fc is square in the concatenated width. It is + # measured, not assumed: the wrong order drafts noise. + cat = (torch.cat([hn, en], -1) if state.mtp_hidden_first + else torch.cat([en, hn], -1)) + h = _dense_mv(cat, mtp['fc_w_t'], mtp, 'fc_w_t', state, fvk, device) + + res = h + n = _rms_fvk(h, ld['input_norm_w_t'], fvk, device, state.eps) + h = res + _decode_full(n, ld, state, state.mtp_rank, pos, fvk, device) + res = h + n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + prev_layer, state._active_layer = state._active_layer, None + try: + h = res + _moe_layer_decode(n, ld, state, fvk, device) + finally: + state._active_layer = prev_layer + # Return the state before the head's own final norm as well: chaining a + # second draft means feeding the head what the model would have fed it, + # and that is a pre-final-norm hidden state. Handing it the model's stale + # one instead costs real acceptance -- measured 0.208 against 0.539 on the + # second draft. + return _lm_head(state, _rms_fvk(h, mtp['norm_w_t'], fvk, device, + state.eps), fvk, device), h.reshape(HID) + + +def _ensure_spec_buffers(state, window, device): + """Allocate what a window of `window` tokens needs.""" + if state._spec_tokens is None or state._spec_tokens.numel() < window: + state._spec_tokens = torch.zeros(window, dtype=torch.long, + device=device) + state._spec_argmax = torch.zeros(window, dtype=torch.long, + device=device) + have = (state.spec_states is not None + and len(state.spec_states[0]) >= window) + if have: + return + state.spec_states = [ + [torch.empty(NV, HK, HV, dtype=torch.bfloat16, device=device) + for _ in range(window)] + for _ in range(state.n_lin)] + state.spec_conv = [ + [torch.empty(1, CONV, KS - 1, dtype=torch.bfloat16, device=device) + for _ in range(window)] + for _ in range(state.n_lin)] + + +def _rewind_to(state, kept): + """Put the recurrent and conv states where `kept` tokens of the window end. + + The KV cache needs nothing: it is written by absolute position, so the + rejected tail is simply overwritten by whatever comes next. The recurrent + state is the opposite -- it has already absorbed the whole window -- which + is what the per-token snapshots are for. + """ + for rank in range(state.n_lin): + state.lin_state[rank].copy_(state.spec_states[rank][kept - 1]) + state.lin_conv_state[rank].copy_(state.spec_conv[rank][kept - 1]) + + +def _verify_dense(x2d, w_bf16, ld, key, fvk, device): + """A window's rows against the 4-bit weight the decode GEMV reads. + + Same tensor under the same cache key, and the M-row form of that GEMV, + whose per-row accumulation order is the GEMV's -- so row t of the result + equals what the decode step at that position would have computed, bit for + bit, while the weight crosses the bus once for the whole window. + """ + return _w4a16_mrows(x2d, w_bf16, ld, key, fvk, device) + + +def _verify_gdn(h, ld, state, lin_rank, w, fvk, device): + """The GDN layer over a window of w tokens, snapshotting per token. + + The projections and the elementwise stages run at w rows; the two stages + that carry state -- the causal conv and the recurrence -- run a token at a + time through the very kernels the decode step calls, because a window is + accepted up to a prefix and the state at that prefix has to be the state + decode would have been in. The sequential-scan variant would do both in one + launch, but it is cos 0.99999 against the per-token kernel rather than + equal to it, and this layer is ~6% of a step: not worth paying for in + tokens that diverge. + """ + eps = state.eps + convw = ld['conv1d_w_t'].reshape(CONV, KS).contiguous() + neg, dtb_c = _gdn_gate_consts(ld, device) + nw = ld['gdn_norm_w_t'] + s = _cs() + x = h.reshape(w, HID) + + if 'in_proj_fused_w' not in ld: + ld['in_proj_fused_w'] = torch.cat( + [ld['in_proj_qkv_w_t'], ld['in_proj_z_w_t'], + ld['in_proj_a_w_t'], ld['in_proj_b_w_t']], 0).contiguous() + fused = _verify_dense(x, ld['in_proj_fused_w'], ld, 'in_proj_fused', + fvk, device) + mixed = fused[:, :KD * 2 + VD].contiguous() + z = fused[:, KD * 2 + VD:KD * 2 + VD + NV * HV].reshape( + w * NV, HV).contiguous() + a = fused[:, -2 * NV:-NV].contiguous() + b = fused[:, -NV:].contiguous() + + conv_out = torch.empty(w, CONV, dtype=torch.bfloat16, device=device) + st_in = state.lin_conv_state[lin_rank] + for t in range(w): + st_out = state.spec_conv[lin_rank][t] + fvk.causal_conv1d_qwen36_update_inout_bf16( + mixed[t].data_ptr(), convw.data_ptr(), 0, + conv_out[t].data_ptr(), st_in.data_ptr(), st_out.data_ptr(), + 1, CONV, KS, True, s) + st_in = st_out + state.lin_conv_state[lin_rank].copy_(st_in) + + qb = torch.empty(w, NV, HK, dtype=torch.bfloat16, device=device) + kb = torch.empty(w, NV, HK, dtype=torch.bfloat16, device=device) + vb = torch.empty(w, NV, HV, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_lin_split_qkv_broadcast_bf16( + conv_out.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), + w, s) + + g_out = torch.empty(w, NV, dtype=torch.bfloat16, device=device) + bo = torch.empty(w, NV, dtype=torch.bfloat16, device=device) + fvk.qwen36_gdn_gating_bf16( + a.data_ptr(), b.data_ptr(), neg.data_ptr(), dtb_c.data_ptr(), + g_out.data_ptr(), bo.data_ptr(), w, NV, s) + + core = torch.empty(w, NV, HV, dtype=torch.bfloat16, device=device) + lin_state = state.lin_state[lin_rank] + for t in range(w): + qt, kt, vt = qb[t], kb[t], vb[t] + gt, bt = g_out[t], bo[t] + gdn_recurrent(fvk)( + qt.data_ptr(), kt.data_ptr(), vt.data_ptr(), gt.data_ptr(), + bt.data_ptr(), lin_state.data_ptr(), core[t].data_ptr(), + 1, NV, HK, HV, True, s) + state.spec_states[lin_rank][t].copy_(lin_state) + + nf = torch.empty(w * NV, HV, dtype=torch.bfloat16, device=device) + fvk.rms_norm_gated_silu_qwen36_bf16( + core.reshape(w * NV, HV).data_ptr(), z.data_ptr(), nw.data_ptr(), + nf.data_ptr(), w * NV, HV, eps, s) + out = _verify_dense(nf.reshape(w, VD), ld['out_proj_w_t'], ld, + 'out_proj_w_t', fvk, device) + return out.reshape(1, w, HID) + + +def _verify_full(h, ld, state, full_rank, pos, w, fvk, device): + """The full-attention layer over a window of w tokens. + + Projections, norms and rope run at w rows. The attention itself runs a + token at a time at q_seq=1 against [0..pos+t], which is the call the decode + step makes: a batched q_seq=w call would have to carry a bottom-right + causal mask and would reduce over a different tiling, and this is the one + place where the two would stop being the same function. The KV it reads is + small next to the weights the window is here to amortise. + """ + eps = state.eps + s = _cs() + qnw, knw = ld['q_norm_w_t'], ld['k_norm_w_t'] + x2 = h.reshape(w, HID) + + nqg = NQ * 2 * HD + if 'qkv_fused_w' not in ld: + ld['qkv_fused_w'] = torch.cat( + [ld['q_proj_w_t'], ld['k_proj_w_t'], ld['v_proj_w_t']], + 0).contiguous() + fused = _verify_dense(x2, ld['qkv_fused_w'], ld, 'qkv_fused', fvk, device) + qg = fused[:, :nqg].contiguous() + kk = fused[:, nqg:nqg + NKV * HD].reshape(w * NKV, HD).contiguous() + v = fused[:, nqg + NKV * HD:].reshape(w, NKV, HD).contiguous() + + q_pre = torch.empty(w, NQ, HD, dtype=torch.bfloat16, device=device) + gate = torch.empty(w, NQ * HD, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_split_q_gate_bf16( + qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), w, s) + q = _rms_fvk(q_pre.reshape(w * NQ, HD), qnw, fvk, device, eps) + kn = _rms_fvk(kk, knw, fvk, device, eps) + + ct = state.rope_cos[pos:pos + w].contiguous() + st = state.rope_sin[pos:pos + w].contiguous() + qin = q.reshape(w, NQ, HD).contiguous() + kin = kn.reshape(w, NKV, HD).contiguous() + qo = torch.empty(w, NQ, HD, dtype=torch.bfloat16, device=device) + ko = torch.empty(w, NKV, HD, dtype=torch.bfloat16, device=device) + fvk.qwen36_partial_rope_qk_bf16( + qin.data_ptr(), kin.data_ptr(), ct.data_ptr(), st.data_ptr(), + qo.data_ptr(), ko.data_ptr(), w, NQ, NKV, HD, ROPE, s) + + attn = state.attn + at = torch.empty(w, NQ * HD, dtype=torch.bfloat16, device=device) + for t in range(w): + attn.Q_buf[:, :1].copy_(qo[t].reshape(1, 1, NQ, HD)) + attn.K_cache[full_rank, pos + t:pos + t + 1].copy_( + ko[t].reshape(1, NKV, HD)) + attn.V_cache[full_rank, pos + t:pos + t + 1].copy_( + v[t].reshape(1, NKV, HD)) + attn.run('full', layer_idx=full_rank, q_seq=1, kv_seq=pos + t + 1, + stream=s, softmax_scale=float(HD) ** -0.5) + at[t].copy_(attn.O_buf[:, :1].reshape(NQ * HD)) + at = _sigmoid_mul(at, gate, fvk, device) + out = _verify_dense(at, ld['o_proj_w_t'], ld, 'o_proj_w_t', fvk, device) + return out.reshape(1, w, HID) + + +def _verify_moe(h, ld, state, w, fvk, device): + """The MoE layer over a window of w tokens. + + The routed experts are the one part of a window that does not amortise: + w tokens pick up to w*TOPK distinct experts out of 256, so the weight + traffic here scales with the window where everything else is read once. + They still go through one grouped launch rather than w of them -- the + kernel already takes the slot count, and a slot is an independent GEMV, so + w*TOPK slots compute exactly what w separate TOPK-slot launches would. + """ + s = _cs() + x = h.reshape(w, HID) + ne = ld['router_w_t'].shape[0] + + if 'router_shared_fused_w' not in ld: + ld['router_shared_fused_w'] = torch.cat( + [ld['router_w_t'], ld['shared_gate_proj_w_t'], + ld['shared_up_proj_w_t']], 0).contiguous() + rs = _verify_dense(x, ld['router_shared_fused_w'], ld, + 'router_shared_fused', fvk, device) + logit_raw = rs[:, :ne].contiguous() + sg, su = rs[:, ne:ne + INTER], rs[:, ne + INTER:] + + # Top-8 a row at a time through the decode router. It is a single-block + # kernel, so w launches is w small launches -- and the selected set has to + # be the set decode selects, ties included, or the window keeps a token + # from a different mixture. + idx = torch.empty(w, TOPK, dtype=torch.int32, device=device) + topv = torch.empty(w, TOPK, dtype=torch.float32, device=device) + for t in range(w): + rc = router_topk(fvk)( + logit_raw[t].data_ptr(), idx[t].data_ptr(), topv[t].data_ptr(), + ne, TOPK, s) + if rc: + raise RuntimeError( + f'router top-k failed with {rc} for {ne} experts, k={TOPK}') + tw = F.softmax(topv, -1) + + if 'experts_gate_up_alpha_dev' not in ld: + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + + slots = w * TOPK + eidx = idx.reshape(-1).contiguous() + # One activation row per slot: the grouped kernel indexes A by slot, and + # the decode call gets the same effect from a zero stride over its single + # row. The copy is w*TOPK*HID bf16 -- tens of KB. + xrep = x.repeat_interleave(TOPK, 0).contiguous() + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + xrep.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), + eidx.data_ptr(), d_gu.data_ptr(), slots, n_gu, HID, + HID, gu_p[0].numel(), gu_s[0].numel(), s) + + g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] + inter = _silu_mul(g_, u_, fvk, device).contiguous() + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), + eidx.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, + INTER, dn_p[0].numel(), dn_s[0].numel(), s) + + rk = f'verify_topk_rows_{w}' + if rk not in ld: + ld[rk] = torch.arange(slots, dtype=torch.int32, device=device) + twf = tw.reshape(-1).contiguous() + out = torch.empty(w, n_dn, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), ld[rk].data_ptr(), twf.data_ptr(), + out.data_ptr(), w, TOPK, n_dn, n_dn, s) + + si = _silu_mul(sg.contiguous(), su.contiguous(), fvk, device) + shared = _verify_dense(si, ld['shared_down_proj_w_t'], ld, + 'shared_down_proj_w_t', fvk, device) + # The scalar gate is an N=1 GEMV over a 4 KB weight, so a row at a time + # costs nothing and is the decode kernel's own arithmetic. + xc = x.contiguous() + gsc = torch.empty(w, 1, dtype=torch.bfloat16, device=device) + for t in range(w): + xr = xc[t] + fvk.bf16_matvec_sm120_bf16( + xr.data_ptr(), ld['shared_gate_w_t'].data_ptr(), + gsc[t].data_ptr(), 1, HID, s) + return _shared_combine(out, shared, gsc, w, fvk, device).reshape( + 1, w, HID) + + +def _verify_block_K(state, toks, pos, w, fvk, device): + """Run a window of w tokens through the decode kernels, at w rows. + + This is what makes the verify the same function as the steps it verifies. + Every stage is the kernel decode calls, at w rows instead of one, over the + weights decode caches; the two state-carrying stages and the attention run + per token for the reasons given above. So the window's row t is the decode + step at pos+t, and the largest weights cross the bus once instead of w + times. + + Returns (logits (w, vocab), hidden (w, HID) pre-final-norm). + """ + p = state.handles.ptrs + layers = p['layers'] + h = F.embedding(toks.view(1, w), p['embed_w_t']) + + for L in range(state.num_layers): + ld = layers[L] + res = h + n = _rms_fvk(h, ld['input_norm_w_t'], fvk, device, state.eps) + if state.types[L] == 'linear_attention': + attn = _verify_gdn(n, ld, state, state._lin_rank[L], w, + fvk, device) + else: + attn = _verify_full(n, ld, state, state._full_rank[L], pos, w, + fvk, device) + h = res + attn + res = h + n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + state._active_layer = L + h = res + _verify_moe(n, ld, state, w, fvk, device) + + hidden = h.reshape(w, HID) + hn = _rms_fvk(h, p['final_norm_w_t'], fvk, device, state.eps) + vocab = p['vocab_size'] + _ensure_lm_head_nvfp4(state, fvk, device) + logits = torch.empty(w, vocab, dtype=torch.bfloat16, device=device) + hc = hn.reshape(w, HID).contiguous() + rc = fvk.w4a16_mrows_edge_sm120_bf16( + hc.data_ptr(), p['lm_head_packed_t'].data_ptr(), + p['lm_head_sf_t'].data_ptr(), logits.data_ptr(), + w, vocab, HID, p['lm_head_alpha'], _cs()) + if rc: + raise RuntimeError(f'M-row lm_head failed with {rc} at M={w}') + return logits, hidden + + +def _verify_block_usable(state) -> bool: + """Can the window run on the decode kernels? + + The M-row GEMV stages a window's activations in shared memory, so it has a + width limit; and the whole point is that the window reads what decode + reads, which is only true where decode takes the W4A16 dense path over + BF16-scope weights. Anywhere else the prefill forward is still the answer. + """ + # gdn_in_proj_w4a16 is gated separately from the rest of the dense path, so + # with it off decode reads the GDN in_proj at BF16 while the window reads + # it at four bits -- a different function in thirty of the forty layers, + # which is exactly the thing this block exists to rule out. + if not kernel_policy().verify_k_rows or not state.dense_w4a16: + return False + if not state.gdn_in_proj_w4a16: + return False + # The window fuses the router with the shared gate/up and reads every + # projection at four bits, which is what decode does only when the loader + # kept these BF16. One NVFP4 site among them and decode takes the W4A4 mma + # instead, so ask about each of the three the window assumes. + ld = state.handles.ptrs['layers'][0] + return (ld.get('router_packed') is None + and ld.get('shared_gate_proj_packed') is None + and ld.get('out_proj_packed') is None + and not ld.get('experts_streamed')) + + +def _spec_block(state, pos, k, fvk, device): + """The whole window as one dependency chain: k drafts, then the verify. + + Written to be capturable end to end. Each draft's token is chosen on the + device -- ``qwen36_argmax_bf16`` writes it straight into the token buffer + the next draft reads -- so the chain never leaves the GPU, and the only + host decision left is how much of the window to keep. + """ + vocab = state.handles.ptrs['vocab_size'] + toks = state._spec_tokens + window = k + 1 + + hidden = state.last_hidden + for j in range(k): + d_logits, hidden = mtp_draft(state, toks[j:j + 1], pos + j, fvk, + device, hidden=hidden) + fvk.qwen36_argmax_bf16(d_logits.data_ptr(), + toks[j + 1:j + 2].data_ptr(), 1, vocab, _cs()) + + if _verify_block_usable(state): + logits, hid = _verify_block_K(state, toks[:window], pos, window, + fvk, device) + else: + state.spec_capture = True + set_spec_verify(True) + try: + logits, hid = nexn2_forward_nvfp4( + state.handles, toks[:window].view(1, window), fvk, device, + cap=state, pos_offset=pos, last_logits_only=False, + return_hidden=True) + finally: + state.spec_capture = False + set_spec_verify(False) + logits = logits.reshape(window, -1) + fvk.qwen36_argmax_bf16(logits.data_ptr(), state._spec_argmax.data_ptr(), + window, vocab, _cs()) + return hid + + +def _ensure_spec_graph(state, pos, k, fvk, device): + """Capture the draft-and-verify window at ``pos``, or return the cached one. + + Everything the block mutates is snapshotted and restored around the warmup + and capture runs -- the recurrent and conv states, the KV rows the window + writes across every rank including the draft head's, and the drafted token + slots -- so a later replay advances from the true pre-window state rather + than from whatever the capture left behind. + """ + key = (pos, k) + cached = state._spec_graphs.get(key) + if cached is not None: + state._spec_graphs.move_to_end(key) + return cached + + window = k + 1 + snap_lin = [t.clone() for t in state.lin_state] + snap_conv = [t.clone() for t in state.lin_conv_state] + snap_k = state.attn.K_cache[:, pos:pos + window].clone() + snap_v = state.attn.V_cache[:, pos:pos + window].clone() + snap_tok = state._spec_tokens.clone() + + def _restore(): + for i, t in enumerate(state.lin_state): + t.copy_(snap_lin[i]) + for i, t in enumerate(state.lin_conv_state): + t.copy_(snap_conv[i]) + state.attn.K_cache[:, pos:pos + window].copy_(snap_k) + state.attn.V_cache[:, pos:pos + window].copy_(snap_v) + state._spec_tokens.copy_(snap_tok) + + with torch.no_grad(): # settle allocator, kernel order, and the + for _ in range(2): # weight quantisation the draft does lazily + _spec_block(state, pos, k, fvk, device) + _restore() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=state._graph_stream, + pool=state._spec_pool), torch.no_grad(): + hid = _spec_block(state, pos, k, fvk, device) + with torch.no_grad(): + _restore() + + return _cache_put(state._spec_graphs, key, (g, hid), + state.spec_graph_cache_max) + + +def spec_decode_step(state, token_id, pos, k, fvk, device): + """One speculative step: draft k tokens, verify k+1 positions, keep a prefix. + + Returns (tokens, next_pos) where `tokens` are the ids actually emitted -- + between 1 and k+1 of them. + + Verifying the window is one batched forward over k+1 positions, which reads + the dense weights once instead of k+1 times; that, and nothing about the + drafts being good, is where the time comes from. A draft is kept only where + the model's own argmax agrees with it, so the emitted sequence is what that + verifier's greedy decode would have produced. + """ + window = k + 1 + _ensure_spec_buffers(state, window, device) + state._spec_tokens[0].copy_(token_id.view(1)[0]) + + g, hid = _ensure_spec_graph(state, pos, k, fvk, device) + g.replay() + + # One D2H for the whole decision: the drafted ids and what the model said + # at each position. Everything before this stayed on the device. + drafted = state._spec_tokens[:window].tolist() + argmax = state._spec_argmax[:window].tolist() + kept = 1 + for j in range(k): + if argmax[j] != drafted[j + 1]: + break + kept += 1 + + if kept < window: + _rewind_to(state, kept) + # The draft head reads the pre-final-norm hidden state of the last emitted + # position. Without this the next window would draft off a stale one. + state.last_hidden.copy_(hid[kept - 1]) + tokens = drafted[1:kept] + [argmax[kept - 1]] + return tokens, pos + kept + + +def generate_greedy_spec(state, input_ids, max_new_tokens, k, fvk, device): + """Greedy decode through the draft-and-verify step. + + Emits exactly what generate_greedy would; the tokens are a check on the + machinery, not an approximation of it. + """ + logits = seed_prefill(state, input_ids, fvk, device) + pos = input_ids.view(-1).shape[0] + nxt = logits[0].argmax().view(1) + out = [] + state.spec_windows = 0 + state.spec_kept = 0 + while len(out) < max_new_tokens: + tokens, pos = spec_decode_step(state, nxt, pos, k, fvk, device) + emitted = [int(nxt)] + tokens[:-1] + state.spec_windows += 1 + state.spec_kept += len(emitted) + out.extend(emitted) + nxt = torch.tensor([tokens[-1]], dtype=torch.long, device=device) + return out[:max_new_tokens] + + def seed_prefill(state, input_ids, fvk, device): """Run the decode step over prompt tokens 0..S-1, building all state. @@ -594,9 +1466,15 @@ def seed_prefill_batched(state, input_ids, fvk, device): return seed_prefill_chunked(state, input_ids, fvk, device, state.prefill_chunk) state.reset() - logits = nexn2_forward_nvfp4( + logits, hidden = nexn2_forward_nvfp4( state.handles, input_ids.view(1, -1), fvk, device, cap=state, - last_logits_only=True) + last_logits_only=True, return_hidden=True) + # The last prompt position's pre-final-norm hidden state, which is what a + # draft head reads. The per-token path writes it every step; this one has + # to do it explicitly, and without it the first window drafts off whatever + # the previous generation left behind -- so how much of that window is kept + # depends on what ran before it, and the run stops being reproducible. + state.last_hidden.copy_(hidden[-1]) return logits # already (1, vocab): only the seeding logit @@ -614,9 +1492,11 @@ def seed_prefill_chunked(state, input_ids, fvk, device, block): logits = None for b0 in range(0, S, block): b1 = min(b0 + block, S) - logits = nexn2_forward_nvfp4( + logits, hidden = nexn2_forward_nvfp4( state.handles, ids[:, b0:b1], fvk, device, cap=state, - pos_offset=b0, last_logits_only=True, compute_logits=(b1 == S)) + pos_offset=b0, last_logits_only=True, compute_logits=(b1 == S), + return_hidden=True) + state.last_hidden.copy_(hidden[-1]) # see seed_prefill_batched return logits @@ -691,11 +1571,7 @@ def _restore(): with torch.no_grad(): _restore() - state._graphs[pos] = (g, out) - cap = state.graph_cache_max - if cap > 0 and len(state._graphs) > cap: - state._graphs.popitem(last=False) # evict LRU - return state._graphs[pos] + return _cache_put(state._graphs, pos, (g, out), state.graph_cache_max) def generate_greedy_graph(state, input_ids, max_new_tokens, fvk, device): diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 31389d3d..7f9cb7b2 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1,7 +1,10 @@ """FlashRT -- Nex-N2-mini (qwen3_5_moe) kernelized NVFP4 forward. -Production prefill forward (S>1) that drives the fvk SM120 kernels off the +Production prefill forward (S>1) that drives the gated fvk kernels off the pre-quantized :class:`WeightHandles` produced by ``extract_weights_nexn2_nvfp4``. +The ``sm120`` in several kernel names is where each was written, not where it +runs: this path also serves Qwen3.6 on Jetson AGX Thor, and which of several +interchangeable kernels each step calls is a :class:`KernelPolicy` below. Every heavy op runs on a FlashRT kernel -- no ``torch`` matmul, no ``F.scaled_dot_product_attention``, no host-side sync in the hot path -- so the prefill is fully on-device and bit-reproducible (it seeds the decode state). @@ -32,6 +35,19 @@ import torch import torch.nn.functional as F + +def _cs(): + """Current CUDA stream handle. + + Inside a graph capture this is the capture stream; eager, the default one. + A hard-coded 0 is not the same thing: during capture it names a stream that + is not being captured, which is an illegal access at capture_end. This + forward was written for prefill, which is never captured -- a speculative + verify block is. + """ + return torch.cuda.current_stream().cuda_stream + + from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes # Static Nex-N2-mini dims (config.json:text_config). Kept module-local so @@ -68,10 +84,29 @@ def _rms_k(x, w, fvk, device, eps): x2 = x.reshape(-1, dim).contiguous() out = torch.empty(x2.shape[0], dim, dtype=torch.bfloat16, device=device) fvk.rms_norm(x2.data_ptr(), w.data_ptr(), out.data_ptr(), - x2.shape[0], dim, eps, 0) + x2.shape[0], dim, eps, _cs()) return out.reshape(shp) +def _add_rms_k(h, x, w, fvk, device, eps): + """h += x in place, and return rmsnorm(h, w). + + The residual add and the norm that always follows it were a tensor add and + a separate kernel: two passes over (S, HID) and two launches per half + layer, 160 of each per forward. Same weight and eps convention as the + plain norm, so this is the two of them and not a third behaviour. + + Both must be bf16 and contiguous -- the kernel writes through `h`. + """ + dim = h.shape[-1] + h2 = h.reshape(-1, dim) + out = torch.empty(h2.shape[0], dim, dtype=torch.bfloat16, device=device) + fvk.residual_add_rms_norm( + h2.data_ptr(), x.reshape(-1, dim).data_ptr(), w.data_ptr(), + out.data_ptr(), h2.shape[0], dim, eps, _cs()) + return out.reshape(h.shape) + + def _proj(x2d, ld, base, n, fvk, device): """y = x @ w.T for one projection, dispatching on the loader's scope. @@ -84,8 +119,16 @@ def _proj(x2d, ld, base, n, fvk, device): return _nvfp4_gemm(x2d, ld[base + '_packed'], ld[base + '_sf'], ld[base + '_alpha'], n, fvk, device) w = ld[base + '_w_t'] - if (_DENSE_W4A16 and x2d.shape[0] >= 64 and (w.shape[0] % 64) == 0 - and (x2d.shape[1] % 64) == 0): + # A speculative verify block is a handful of rows, so the M>=64 heuristic + # below would send it to the BF16 GEMM -- four times the weight bytes, on + # the pass whose whole purpose is to read the weights once. It wants the + # 4-bit weight for the same reason decode does: at this M the cost is + # traffic, not throughput. + if (_SPEC_VERIFY and x2d.shape[0] <= 8 + and (w.shape[1] % 16) == 0): + return _w4a16_mrows(x2d, w, ld, base + '_w_t', fvk, device) + if (_DENSE_W4A16 and x2d.shape[0] >= 64 + and (w.shape[0] % 64) == 0 and (x2d.shape[1] % 64) == 0): return _gemm_w4a16(x2d, w, ld, base + '_w_t', fvk, device) if (_DENSE_W16A16 and x2d.shape[0] >= _DENSE_BF16_MIN_M and (x2d.shape[1] % 64) == 0): @@ -100,17 +143,176 @@ def _proj(x2d, ld, base, n, fvk, device): return (x2d.float() @ w.float().T).to(torch.bfloat16) +# cuBLASLt for the same product, where the build has it. It is 4-8x the +# hand-written kernel at every shape prefill issues -- 140 ms of a 1024-token +# prefill against 22.5 -- and it is a drop-in in the strongest sense: bitwise +# identical output at every shape checked, and bit-reproducible across repeated +# launches. +# +# The determinism caveat elsewhere in this file is about torch.matmul, whose +# split-K reduction order can vary and flip a near-tie argmax. It does not apply +# to this entry point, which was measured rather than assumed. Set False to +# force the hand-written kernel. +import os as _os_early + + +class KernelPolicy: + """Which implementation each interchangeable step of this model calls. + + The forward and decode paths have, at several steps, more than one kernel + that computes the same thing: a fused form and the chain it replaces, a + warp-per-row form and a block-per-row one, an "edge" shared-memory layout + and the original. Each pair was checked against the other with + ``torch.equal`` -- not a tolerance -- so which one runs decides speed and + cannot decide output. + + They are gathered here rather than decided at each call site by asking the + module what symbols it happens to export. A kernel appearing in a build is + not a reason to change what an already-validated model path does; a caller + saying so is. The frontend owns one of these and can hand a different one + down, and the environment variables that predate it remain the defaults, so + an existing configuration behaves exactly as it did. + + Fields are read at call time. A policy must therefore not be changed + between a CUDA graph capture and its replay -- the replay repeats whichever + branch the capture took, so the two would disagree. + """ + + __slots__ = ('dense_cublaslt', 'cublaslt_max_algos', 'wy_gdn', + 'edge_w4a16', 'route_kernel', 'fused_shared_combine', + 'warp_router_topk', 'gdn_recurrent_edge', 'verify_k_rows') + + def __init__(self, *, + dense_cublaslt=None, + cublaslt_max_algos=1, + wy_gdn=None, + edge_w4a16=None, + route_kernel=None, + fused_shared_combine=True, + warp_router_topk=True, + gdn_recurrent_edge=True, + verify_k_rows=None): + env = _os_early.environ.get + + def _flag(value, name, default='1'): + if value is not None: + return bool(value) + return env(name, default) != '0' + + # cuBLASLt for the dense bf16 GEMMs; the in-house kernel is also + # deterministic but 66% slower at 2048 (693 against 418 ms). + self.dense_cublaslt = _flag(dense_cublaslt, 'NEXN2_DENSE_CUBLASLT') + # How many cuBLASLt candidates the first call for a shape times. 1 + # takes the heuristic's own pick; see _gemm_w16a16. + self.cublaslt_max_algos = int(cublaslt_max_algos) + # WY chunked gated-delta scan for the GDN prefill instead of the + # sequential scan (11x at S=2048). + self.wy_gdn = _flag(wy_gdn, 'NEXN2_WY_GDN') + # The "edge" shared-memory layout of the two weight-only 4-bit GEMVs. + self.edge_w4a16 = _flag( + edge_w4a16, 'FLASHRT_QWEN35MOE_W4A16_EDGE') + # The five-kernel routing producer instead of the tensor chain. + self.route_kernel = _flag(route_kernel, 'NEXN2_ROUTE_KERNEL') + # Decode-side fusions, each bit-identical to the chain it replaces. + self.fused_shared_combine = bool(fused_shared_combine) + self.warp_router_topk = bool(warp_router_topk) + self.gdn_recurrent_edge = bool(gdn_recurrent_edge) + # Run a speculative verify window through the decode kernels at k+1 + # rows rather than through the prefill forward. Two names, as the rest + # of this model's variables have: the generic one and the one it + # shipped under. + if verify_k_rows is not None: + self.verify_k_rows = bool(verify_k_rows) + else: + self.verify_k_rows = env( + 'FLASHRT_QWEN35MOE_VERIFY_K_ROWS', + env('FLASHRT_NEXN2_VERIFY_K_ROWS', '1')) != '0' + + def __repr__(self) -> str: # pragma: no cover + fields = ', '.join( + f'{name}={getattr(self, name)!r}' for name in self.__slots__) + return f'KernelPolicy({fields})' + + +_POLICY = KernelPolicy() + + +def kernel_policy(): + """The policy the forward and decode paths are currently reading.""" + return _POLICY + + +def set_kernel_policy(policy): + """Install ``policy``; returns the one it replaced. + + Not to be called between a CUDA graph capture and its replay. + """ + global _POLICY + if not isinstance(policy, KernelPolicy): + raise TypeError( + f'expected a KernelPolicy, got {type(policy).__name__}') + previous, _POLICY = _POLICY, policy + return previous + + +# The cuBLASLt wrapper picks its algorithm by *timing* eight candidates at +# first use. Timing is noisy, so different processes pick different algorithms, +# and different algorithms reduce in different orders -- which makes the model +# itself non-deterministic across processes. Measured on one binary: the golden +# prefix came out 16/16 three times and 14/16 three times, flipping between +# exactly two token streams. Asking for one candidate takes the heuristic's own +# choice instead, and five of five processes then agree. +# +# It is not a speed trade worth making either way round: at 1024 tokens the +# timed pick is worth 1.6% warm (213.6 against 217.1 ms) and costs 25% of the +# cold time (about 1020 against 770 ms), because the timing loop runs inside +# the first call. Determinism and a faster first token for 1.6% of the warm +# path. +# +# Requested per call through KernelPolicy.cublaslt_max_algos, not by setting +# the kernel's environment variable: that +# variable is process-global and shared with every other frontend, so setting +# it here would decide the algorithm for a model loaded later in the same +# process that never asked. The kernel caches its plan per +# (M, N, K, max_algos), so this choice stays with these call sites. + +# Set once, on the first call: a build predating the max_algos argument still +# links and still runs, one autotune behaviour older. +_CUBLASLT_TAKES_ALGOS = None + + def _gemm_w16a16(x2d, w, fvk, device): """y = x @ w.T via the deterministic bf16-act x bf16-weight tensor-core GEMM (fp32 register accumulate). Matches the fp32 path's argmax (cos 1.0) and is bit-identical run-to-run, at ~1.75x the fp32/TF32 op.""" + global _CUBLASLT_TAKES_ALGOS m, k = x2d.shape n = w.shape[0] xc = x2d.contiguous() wc = w.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + policy = kernel_policy() + if policy.dense_cublaslt and hasattr(fvk, 'bf16_matmul_cublaslt_bf16'): + algos = policy.cublaslt_max_algos + if _CUBLASLT_TAKES_ALGOS is None: + try: + fvk.bf16_matmul_cublaslt_bf16( + xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, + _cs(), algos) + _CUBLASLT_TAKES_ALGOS = True + return y + except TypeError: + _CUBLASLT_TAKES_ALGOS = False + if _CUBLASLT_TAKES_ALGOS: + fvk.bf16_matmul_cublaslt_bf16( + xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, _cs(), + algos) + else: + fvk.bf16_matmul_cublaslt_bf16(xc.data_ptr(), wc.data_ptr(), + y.data_ptr(), m, n, k, _cs()) + return y fvk.w16a16_gemm_sm120_bf16(xc.data_ptr(), wc.data_ptr(), y.data_ptr(), - m, n, k, 1.0, 0) + m, n, k, 1.0, _cs()) return y @@ -119,8 +321,41 @@ def _gemm_w16a16(x2d, w, fvk, device): # fp4 *weight* (not the activation), so W4A16 lands at the same ~0.987 as W4A4 # while being slower than the CUTLASS W4A4 -- dominated. Default OFF; the 0.994 # path needs a bf16-*weight* GEMM (repurpose this kernel's 2.18x structure). +# +# Retested on a part with no W4A4 at all, where it might have been expected to +# win on traffic: at S=1024 it moves TTFT 1335.0 -> 1346.6 ms, i.e. nothing. +# The 178 ms those GEMMs cost is not what bounds this prefill. _DENSE_W4A16 = False +# Set only around a speculative verify block; see _proj and the lm_head below. +# +# Off, on evidence, twice over. Routing the verify block through _gemm_w4a16 +# loses whether or not the window is captured -- 36.15 to 33.38 eager, and +# 39.77 to 20.84 captured -- and it also reintroduces the divergence from plain +# greedy that a BF16 verify does not have. +# +# Both readings say the same thing: this is the wrong path, not the wrong idea. +# _gemm_w4a16 quantises the weight through its own helper into its own cache, +# so it is neither the tensor the decode GEMV reads nor a kernel shaped for +# three rows. What the verify wants is a small-M GEMM over the *same* packed +# weights the decode path already caches -- then it reads a quarter of the +# bytes and differs from decode only by reduction order. Until that exists, +# BF16 is both faster and the one that agrees with plain greedy token for +# token. +# On: the verify runs the dense projections through the M-row form of the +# decode GEMV, over the tensor the decode path caches. Off, it reads the same +# weights at BF16 -- four times the bytes, and a different answer from the step +# it is verifying (measured logit cosine 0.988 against decode, which is what +# made the emitted text diverge from plain greedy). +_SPEC_VERIFY_W4A16 = _os_early.environ.get( + 'NEXN2_SPEC_VERIFY_W4A16', '1') != '0' +_SPEC_VERIFY = False + + +def set_spec_verify(on: bool) -> None: + global _SPEC_VERIFY + _SPEC_VERIFY = bool(on) and _SPEC_VERIFY_W4A16 + # BF16 tensor-core dense projections (vs the default fp32/TF32 matmul). The # experts-scope q/k/v/o/out/shared/router projections dominate the prefill # profile as fp32 GEMMs; bf16 inputs with fp32 accumulate roughly halve that @@ -142,6 +377,42 @@ def _gemm_w16a16(x2d, w, fvk, device): _DENSE_W16A16 = True +def _w4a16_mrows(x2d, w, ld, key, fvk, device): + """A few rows against the 4-bit weight the *decode* path caches. + + This is what makes a speculative verify the same function as the step it + verifies. It reads the identical packed tensor under the identical cache + keys the decode GEMV uses -- not a second copy quantised by a different + helper -- and runs the M-row form of that GEMV, whose per-row accumulation + order is the GEMV's. So a verified row equals the decode row it stands in + for, bit for bit, and the window reads the weight once rather than once per + token and at a quarter of the bytes the BF16 path reads. + """ + n, k = w.shape + pk = key + '_w4a16_p' + if pk not in ld: + packed = torch.empty(n, k // 2, dtype=torch.uint8, device=device) + sf = torch.zeros(_sf_swz_bytes(n, k), dtype=torch.uint8, device=device) + scr = torch.zeros(1, dtype=torch.float32, device=device) + og = torch.zeros(1, dtype=torch.float32, device=device) + fvk.bf16_weight_to_nvfp4_swizzled( + w.contiguous().data_ptr(), packed.data_ptr(), sf.data_ptr(), + scr.data_ptr(), og.data_ptr(), n, k, _cs()) + torch.cuda.synchronize() + ld[pk] = packed + ld[key + '_w4a16_sf'] = sf + ld[key + '_w4a16_a'] = float(og.item()) + m = x2d.shape[0] + xc = x2d.contiguous() + y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + rc = fvk.w4a16_mrows_edge_sm120_bf16( + xc.data_ptr(), ld[pk].data_ptr(), ld[key + '_w4a16_sf'].data_ptr(), + y.data_ptr(), m, n, k, ld[key + '_w4a16_a'], _cs()) + if rc: + raise RuntimeError(f'M-row W4A16 failed with {rc} at M={m}') + return y + + def _gemm_w4a16(x2d, w, ld, key, fvk, device): """y = x @ w.T via the bf16-act x fp4-weight tensor-core GEMM. Weight quantised to NVFP4 once (cached); activation stays BF16 (precise).""" @@ -151,7 +422,7 @@ def _gemm_w4a16(x2d, w, ld, key, fvk, device): xc = x2d.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) fvk.w4a16_gemm_sm120_bf16(xc.data_ptr(), p.data_ptr(), s.data_ptr(), - y.data_ptr(), m, n, k, a, 0) + y.data_ptr(), m, n, k, a, _cs()) return y @@ -175,7 +446,7 @@ def _wquant(w, ld, key, fvk, device): og = torch.zeros(1, dtype=torch.float32, device=device) fvk.bf16_weight_to_nvfp4_swizzled( w.contiguous().data_ptr(), p.data_ptr(), s.data_ptr(), - scr.data_ptr(), og.data_ptr(), nn, kk, 0) + scr.data_ptr(), og.data_ptr(), nn, kk, _cs()) torch.cuda.synchronize() ld[pk] = p ld[key + '_w4s'] = s @@ -195,7 +466,7 @@ def _gemm_fp4(x2d, w, ld, key, fvk, device, xp=None, xsf=None): y = torch.empty(m, n, dtype=torch.bfloat16, device=device) fvk.fp4_w4a16_gemm_sm120_bf16out( xp.data_ptr(), p.data_ptr(), y.data_ptr(), m, n, k, - xsf.data_ptr(), s.data_ptr(), a, 0) + xsf.data_ptr(), s.data_ptr(), a, _cs()) return y @@ -215,9 +486,16 @@ def _quant_act(x2d, fvk, device, stream=0): def _nvfp4_gemm_preq(xp, xsf, wp_ptr, wsf_ptr, alpha, m, n, k, fvk, device, - stream=0): - """y = x @ w.T from a pre-quantised activation (xp, xsf).""" - y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + stream=0, out=None): + """y = x @ w.T from a pre-quantised activation (xp, xsf). + + ``out`` lets a caller point the result at a slice of a buffer it already + owns, which is what the per-expert loop wants: it writes 256 blocks into + one matrix, and allocating each of them separately costs more in Python + than the GEMM costs on the device. + """ + y = torch.empty(m, n, dtype=torch.bfloat16, device=device) if out is None \ + else out fvk.fp4_w4a16_gemm_sm120_bf16out( xp.data_ptr(), wp_ptr, y.data_ptr(), m, n, k, xsf.data_ptr(), wsf_ptr, alpha, stream) @@ -242,7 +520,7 @@ def _silu_mul(g, u, fvk, device): gc = g.reshape(-1).contiguous() uc = u.reshape(-1).contiguous() out = torch.empty(n, dtype=torch.bfloat16, device=device) - fvk.silu_mul_sm120_bf16(gc.data_ptr(), uc.data_ptr(), out.data_ptr(), n, 0) + fvk.silu_mul_sm120_bf16(gc.data_ptr(), uc.data_ptr(), out.data_ptr(), n, _cs()) return out.reshape(g.shape) @@ -253,75 +531,57 @@ def _silu_mul(g, u, fvk, device): # the inter-chunk state recurrence is sequential -> 11x faster at S=2048, # bit-exact (out cos 0.99998, state cos 0.99997 vs the seq-scan). Default on. import os as _os -_USE_WY_GDN = _os.environ.get('NEXN2_WY_GDN', '1') != '0' _WY_MIN_S = 64 # below this the seq-scan's lower fixed overhead wins -def _wy_pack_t(x, ch=64): - """(S, H, D) -> (chunks, H, ch, D): x_pack[ci, h, i, d] = x[ci*ch+i, h, d] - (zero-padded last chunk). The packed chunk-major layout the mma kernels read.""" - s, hh, d = x.shape - pad = (-s) % ch - if pad: - x = F.pad(x, (0, 0, 0, 0, 0, pad)) - return x.reshape(-1, ch, hh, d).permute(0, 2, 1, 3).contiguous() - - -def _wy_l2(x): - """l2norm over the last dim, eps inside rsqrt (matches the seq-scan kEps).""" - xf = x.float() - return (xf * torch.rsqrt((xf * xf).sum(-1, keepdim=True) + 1e-6)).to( - torch.bfloat16) - - -def _wy_gcumsum(g, ch=64): - """(S, NV) -> (S, NV) per-chunk cumulative sum of the (log-space) gate.""" - s = g.shape[0] - pad = (-s) % ch - gp = F.pad(g, (0, 0, 0, pad)) if pad else g - return torch.cumsum(gp.float().reshape(-1, ch, g.shape[1]), 1).reshape( - -1, g.shape[1])[:s].to(torch.bfloat16) - - -def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): - """WY chunked scan. q16/k16 (S,16,128) raw post-conv, v (S,32,128), - g/beta (S,32). Returns core (S,32,128) + final state (32,128,128). +def _gdn_wy_chunk(qb, kb, v, g, beta, fvk, device, init_state=None): + """WY chunked scan. qb/kb (S,32,128) raw post-conv with q/k already + GQA-broadcast across the 32 v-head slots, v (S,32,128), g/beta (S,32). + Returns core (S,32,128) + final state (32,128,128). ``init_state`` (NV,HK,HV) is the recurrent state to continue from -- the chunk_h kernel reads it as h0[0] and writes the post-block state back, so a chunked prefill carries it across blocks (probe-verified bit-exact: whole vs two state-carried halves match at cos 1.0). Defaults to zeros. - Pipeline (FLA chunked delta rule, all add-only existing kernels): - l2norm + per-chunk g-cumsum (torch glue) -> kkt -> solve_tril(+pack) -> - recompute_wu -> chunk_h (inter-chunk state) -> output_o.""" - S = q16.shape[0] + Pipeline (FLA chunked delta rule, kernels throughout): norm+pack_q+cumsum + -> kkt -> solve_tril(+pack) -> recompute_wu -> chunk_h (inter-chunk state) + -> pack_v -> output_o. + """ + S = qb.shape[0] chunks = (S + 63) // 64 CH, QKG = 64, NV // NK - q_l2 = _wy_l2(q16) - k_l2 = _wy_l2(k16).contiguous() - gc = _wy_gcumsum(g).contiguous() - betac = beta.contiguous() - vc = v.contiguous() + + # l2norm of q and k, the GQA broadcast of q into the 32 v-head slots, the + # chunk-major packing of q, and the per-chunk gate cumulative sum, in one + # kernel. The broadcast never materialises and q is normalised straight + # into its packed slots, so the only q traffic is the packed write. + k_l2 = torch.empty(S, NK, HK, dtype=torch.bfloat16, device=device) + q_pack = torch.empty(chunks, NV, CH, HK, dtype=torch.bfloat16, + device=device) + gc = torch.empty(S, NV, dtype=torch.bfloat16, device=device) + fvk.gdn_wy_norm_pack_q_cumsum_edge_bf16( + qb.data_ptr(), kb.data_ptr(), g.data_ptr(), k_l2.data_ptr(), + q_pack.data_ptr(), gc.data_ptr(), S, NK, NV, HK, QKG, _cs()) k_pack = torch.empty(chunks, NK, CH, HK, dtype=torch.bfloat16, device=device) kkt_base = torch.empty(chunks, NK, CH, CH, dtype=torch.float32, device=device) A = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) fvk.linear_attn_gdn_wy_kkt_b64_bf16_cublaslt( - k_l2.data_ptr(), betac.data_ptr(), gc.data_ptr(), k_pack.data_ptr(), - kkt_base.data_ptr(), A.data_ptr(), S, NK, NV, HK, QKG, 0) + k_l2.data_ptr(), beta.data_ptr(), gc.data_ptr(), k_pack.data_ptr(), + kkt_base.data_ptr(), A.data_ptr(), S, NK, NV, HK, QKG, _cs()) Ai = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) Ai_pack = torch.empty(chunks, NV, CH, CH, dtype=torch.bfloat16, device=device) fvk.linear_attn_gdn_wy_solve_tril_b64_f32_parallel_pack( - A.data_ptr(), Ai.data_ptr(), Ai_pack.data_ptr(), S, NV, 0) + A.data_ptr(), Ai.data_ptr(), Ai_pack.data_ptr(), S, NV, _cs()) w_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) u_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) fvk.linear_attn_gdn_wy_recompute_wu_b64_bf16_mma_fla( - k_l2.data_ptr(), vc.data_ptr(), betac.data_ptr(), gc.data_ptr(), + k_l2.data_ptr(), v.data_ptr(), beta.data_ptr(), gc.data_ptr(), Ai_pack.data_ptr(), w_pack.data_ptr(), u_pack.data_ptr(), - S, NK, NV, HK, QKG, 0) + S, NK, NV, HK, QKG, _cs()) state = (init_state.clone() if init_state is not None else torch.zeros(NV, HK, HV, dtype=torch.bfloat16, device=device)) @@ -330,16 +590,19 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): fvk.linear_attn_gdn_wy_chunk_h_b64_bf16_mma_fla( k_l2.data_ptr(), w_pack.data_ptr(), u_pack.data_ptr(), gc.data_ptr(), state.data_ptr(), h0.data_ptr(), v_new.data_ptr(), 0, 0, - S, NK, NV, HK, QKG, 0) - - q_pack = _wy_pack_t(q_l2.repeat_interleave(QKG, 1)) - k_pack_hv = _wy_pack_t(k_l2.repeat_interleave(QKG, 1)) - v_pack = _wy_pack_t(v_new) + S, NK, NV, HK, QKG, _cs()) + + # v is the only side still needing a packed copy; the raw-K output_o does + # the GQA expansion of k in-kernel, so k never gets a 32-head buffer. + v_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, + device=device) + fvk.gdn_wy_pack_v_edge_bf16( + v_new.data_ptr(), v_pack.data_ptr(), S, NV, HV, _cs()) core = torch.empty(S, NV, HV, dtype=torch.bfloat16, device=device) - fvk.linear_attn_gdn_wy_output_o_b64_bf16_mma_fla( - q_pack.data_ptr(), k_pack_hv.data_ptr(), v_pack.data_ptr(), + fvk.linear_attn_gdn_wy_output_o_b64_bf16_mma_fla_rawk( + q_pack.data_ptr(), k_l2.data_ptr(), v_pack.data_ptr(), h0.data_ptr(), gc.data_ptr(), core.data_ptr(), - S, NV, HV, float(HV ** -0.5), 0) + S, NK, NV, HV, QKG, float(HV ** -0.5), _cs()) return core, state @@ -385,22 +648,39 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, # Same (B, S, conv_dim) layout the decode update kernel uses; no bias. For a # chunked block, prepend the previous block's last KS-1 inputs (conv_hist) # so the block's first outputs see the right history, then drop them. + # The row-blocked entry walks several tokens per thread with the window in + # registers, so each input is read once instead of once per output that + # needs it. Bit-identical to the per-token entry (probe: exact at every + # length measured) and 3.3x, which puts it within 1.3x of what its traffic + # implies rather than 4.2x off. It also lifts the gridDim.y ceiling that + # capped a single launch at 65535 tokens. + _conv = getattr(fvk, 'causal_conv1d_qwen36_rows_bf16', + fvk.causal_conv1d_qwen36_bf16) convw_k = convw.reshape(CONV, KS).contiguous() - if conv_hist is not None: + xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) + _hist_conv = getattr(fvk, 'causal_conv1d_qwen36_rows_hist_bf16', None) + if conv_hist is not None and _hist_conv is not None: + # The conv reads the previous block's trailing inputs where it needs + # them. Prepending them to the activations instead meant concatenating + # and then slicing the whole block back off -- two copies of it per + # layer, 691 ms of a 32768-token prefill, to supply three tokens. + # conv_hist is already (1, CONV, KS-1), newest last, which is the + # layout the kernel reads, so the transpose goes with the copy. + _hist_conv(mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, + conv_hist.contiguous().data_ptr(), xc.data_ptr(), + B, S, CONV, KS, True, _cs()) + elif conv_hist is not None: hist = conv_hist[0].transpose(0, 1).reshape(1, KS - 1, CONV) mixed_ext = torch.cat( [hist.to(mixed.dtype), mixed], dim=1).contiguous() Se = mixed_ext.shape[1] xc_ext = torch.empty(B, Se, CONV, dtype=torch.bfloat16, device=device) - fvk.causal_conv1d_qwen36_bf16( - mixed_ext.data_ptr(), convw_k.data_ptr(), 0, - xc_ext.data_ptr(), B, Se, CONV, KS, True, 0) + _conv(mixed_ext.data_ptr(), convw_k.data_ptr(), 0, + xc_ext.data_ptr(), B, Se, CONV, KS, True, _cs()) xc = xc_ext[:, KS - 1:, :].contiguous() else: - xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) - fvk.causal_conv1d_qwen36_bf16( - mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, - xc.data_ptr(), B, S, CONV, KS, True, 0) + _conv(mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, + xc.data_ptr(), B, S, CONV, KS, True, _cs()) # split conv output + broadcast q/k 16 -> 32 heads in one fvk kernel. xc_bf = xc.reshape(B * S, CONV).contiguous() qb = torch.empty(B, S, NV, HK, dtype=torch.bfloat16, device=device) @@ -408,7 +688,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, vb = torch.empty(B, S, NV, HV, dtype=torch.bfloat16, device=device) fvk.qwen35moe_lin_split_qkv_broadcast_bf16( xc_bf.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), - B * S, 0) + B * S, _cs()) neg = (-A_log.exp()).float().contiguous() dtb_c = dtb.contiguous() @@ -418,16 +698,16 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, bo = torch.empty(B, S, NV, dtype=torch.bfloat16, device=device) fvk.qwen36_gdn_gating_bf16( a_bf.data_ptr(), b_bf.data_ptr(), neg.data_ptr(), dtb_c.data_ptr(), - g_out.data_ptr(), bo.data_ptr(), B * S, NV, 0) + g_out.data_ptr(), bo.data_ptr(), B * S, NV, _cs()) - if _USE_WY_GDN and S >= _WY_MIN_S: + if kernel_policy().wy_gdn and S >= _WY_MIN_S: # WY chunked delta-rule scan: 11x faster than the seq-scan at S=2048, - # bit-exact. qb/kb are the 16->32 broadcast heads (src_h = h//2), so the - # 16 unique K-heads are the even slots; the WY kernels re-expand by GQA. - q16 = qb.reshape(S, NV, HK)[:, 0::2, :] - k16 = kb.reshape(S, NV, HK)[:, 0::2, :] + # bit-exact. qb/kb carry the 16->32 broadcast heads (src_h = h//2); the + # front kernel reads the group leaders and re-expands where it packs, + # so no strided slice is taken here. core, state = _gdn_wy_chunk( - q16, k16, vb.reshape(S, NV, HV), g_out.reshape(S, NV), + qb.reshape(S, NV, HK), kb.reshape(S, NV, HK), + vb.reshape(S, NV, HV), g_out.reshape(S, NV), bo.reshape(S, NV), fvk, device, init_state=init_state) core = core.reshape(B, S, NV, HV) else: @@ -445,10 +725,18 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, vb.reshape(S, NV, HV).contiguous().data_ptr(), g_out.reshape(S, NV).contiguous().data_ptr(), bo.reshape(S, NV).contiguous().data_ptr(), - state.data_ptr(), core.data_ptr(), S, NV, HK, True, 0) + state.data_ptr(), core.data_ptr(), S, NV, HK, True, _cs()) core = core.reshape(B, S, NV, HV) if cap is not None: + # The per-step snapshots go FIRST. `init_state` and `conv_hist` are the + # very tensors the two copies below overwrite -- the caller passes + # cap.lin_state[rank] / cap.lin_conv_state[rank] in directly -- so + # replaying the scan after the copies would start it from the state the + # block ended at, and every rewind would restore a fabricated one. + if getattr(cap, 'spec_capture', False): + _capture_per_token_state(cap, rank, S, init_state, conv_hist, + mixed, qb, kb, vb, g_out, bo, fvk, device) # GDN recurrent final state = `state` after the S-step scan; conv state # = the last KS-1 `mixed` inputs (channel-major, newest at index -1), # matching the causal_conv1d_update rolling buffer (1, CONV, KS-1). @@ -461,7 +749,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, nf = torch.empty_like(cf) fvk.rms_norm_gated_silu_qwen36_bf16( cf.data_ptr(), zf.data_ptr(), nw.data_ptr(), nf.data_ptr(), - cf.shape[0], HV, eps, 0) + cf.shape[0], HV, eps, _cs()) out = _proj(nf.reshape(B * S, VD), ld, 'out_proj', HID, fvk, device) return out.reshape(B, S, HID) @@ -475,15 +763,146 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, # the pre-existing flash_rt_fa2.so (already a hard dep of the decode backend), # so this adds no new csrc. _FA2_MOD = None +_FA2_USABLE = None _NUM_SMS = None def _get_fa2(): + """The vendored FA2 module, or None where the target does not build it. + + Thor is such a target: its arch list omits FA2 because it uses FA4. Absence + is a fallback, not an error -- ``_sdpa_causal_attn`` computes the same + thing -- so this returns None rather than raising, the way the decode + attention backend already treats it. + """ global _FA2_MOD if _FA2_MOD is None: - from flash_rt import flash_rt_fa2 as _m - _FA2_MOD = _m - return _FA2_MOD + try: + from flash_rt import flash_rt_fa2 as _m + except ImportError: + _FA2_MOD = False + else: + _FA2_MOD = _m + return _FA2_MOD or None + + +def _fa2_usable(device): + """Does the vendored kernel actually compute here? + + Importing it and finding its symbols proves neither: its arch handling can + leave a build that links, loads, prints a complaint and returns without + writing the output, which downstream looks like wrong attention rather than + a failure. The decode backend probes for the same reason. One launch, once. + """ + global _FA2_USABLE + if _FA2_USABLE is not None: + return _FA2_USABLE + if _get_fa2() is None: + _FA2_USABLE = False + return False + g = torch.Generator(device=device).manual_seed(1) + q = torch.randn(1, 8, NQ, HD, generator=g, device=device, + dtype=torch.bfloat16) + k = torch.randn(1, 8, NKV, HD, generator=g, device=device, + dtype=torch.bfloat16) + v = torch.randn_like(k) + try: + produced = _fa2_causal_attn(q, k, v, device, _probe=True).float() + torch.cuda.synchronize(device) + except Exception: # noqa: BLE001 + _FA2_USABLE = False + return False + expected = _sdpa_causal_attn(q, k, v, device).float() + _FA2_USABLE = bool( + torch.isfinite(produced).all() + and ((produced - expected).norm() + / expected.norm().clamp_min(1e-6)).item() < 0.05) + return _FA2_USABLE + + +_FLEX_CACHE = {} + + +def _flex_causal(sq, sk, device): + """A bottom-right-causal flex_attention closure for this block shape. + + Block masks are built per (Sq, Sk) and cached, because a chunked prefill + revisits the same shapes as it walks the prompt. Returns None where flex is + unavailable, leaving the explicit-mask path to handle it. + """ + key = (sq, sk, str(device)) + got = _FLEX_CACHE.get(key) + if got is not None: + return got + if key in _FLEX_CACHE: + return None + try: + from torch.nn.attention.flex_attention import ( + create_block_mask, flex_attention, + ) + + off = sk - sq + + def mask_mod(b, h, q_idx, kv_idx): + return kv_idx <= q_idx + off + + block_mask = create_block_mask(mask_mod, 1, NQ, sq, sk, device=device) + + def run(q, k, v): + return flex_attention(q, k, v, block_mask=block_mask, + scale=float(HD) ** -0.5, enable_gqa=True) + + _FLEX_CACHE[key] = run + return run + except Exception: # noqa: BLE001 + _FLEX_CACHE[key] = None + return None + + +def _sdpa_causal_attn(qf, kf, vf, device): + """Reference causal GQA attention, for a build without the FA2 kernel. + + FA2 causal aligns bottom-right -- query i attends to keys [0, Sk-Sq+i] -- + which is exactly a chunked block's absolute causal window. torch's + ``is_causal=True`` aligns top-left, and the two only agree when Sq == Sk, + so the mask is built explicitly rather than left to a flag whose convention + differs where it matters. + """ + import torch.nn.functional as F + + Sq, Sk = qf.shape[1], kf.shape[1] + q = qf.transpose(1, 2) # (1, NQ, Sq, HD) + k, v = kf.transpose(1, 2), vf.transpose(1, 2) # (1, NKV, Sk, HD) + # An explicit mask forces the math backend, which materialises the scores + # and runs them through SIMT fp32 GEMMs -- 131.7 ms of a 2048-token prefill + # in twenty launches, growing as S^2. When the block is square the two + # causal conventions coincide, so say is_causal and let the fused backend + # take it; the mask is only needed when Sq < Sk, which is chunked prefill. + if Sq == Sk: + return F.scaled_dot_product_attention( + q, k, v, is_causal=True, scale=float(HD) ** -0.5, enable_gqa=True + ).transpose(1, 2).contiguous() + # A chunked block's window is bottom-right causal, which is_causal does not + # mean (measured: cos 0.24 against this mask, i.e. it silently truncates the + # history) and which a boolean mask only expresses by materialising the + # scores. flex_attention states it as a predicate and skips fully-masked + # blocks -- numerically right, cos 0.999997 -- but it compiles per shape, + # and a chunked prefill hands it a new (Sq, Sk) for every chunk: measured + # 10240 tokens 4071 -> 4330 ms, 16384 tokens 13858. Left out on that + # evidence; it becomes the right answer once the chunk shapes are fixed and + # warmed, which is where the long-context work goes next. + qi = torch.arange(Sk - Sq, Sk, device=device).unsqueeze(1) + mask = torch.arange(Sk, device=device).unsqueeze(0) <= qi + try: + o = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, scale=float(HD) ** -0.5, enable_gqa=True) + except TypeError: # torch without native GQA + groups = NQ // NKV + o = F.scaled_dot_product_attention( + q, k.repeat_interleave(groups, dim=1), + v.repeat_interleave(groups, dim=1), + attn_mask=mask, scale=float(HD) ** -0.5) + return o.transpose(1, 2).contiguous() def _num_sms(): @@ -494,13 +913,18 @@ def _num_sms(): return _NUM_SMS -def _fa2_causal_attn(qf, kf, vf, device): +def _fa2_causal_attn(qf, kf, vf, device, *, _probe=False): """Causal GQA attention via the vendored FA2 kernel (bf16, native GQA -- no KV repeat). qf (1,Sq,NQ,HD), kf/vf (1,Sk,NKV,HD). Returns (1,Sq,NQ,HD). Sk may exceed Sq (chunked prefill: a block of Sq queries against the Sk accumulated KV); FA2 causal uses bottom-right alignment, so query i attends to keys [0, Sk-Sq+i] -- exactly the block's absolute causal window. splitkv - off (large-q parallelism).""" + off (large-q parallelism). + + Falls back to the reference where the kernel is absent or refuses. ``_probe`` + forces the kernel, since the probe is what decides that question.""" + if not _probe and not _fa2_usable(device): + return _sdpa_causal_attn(qf, kf, vf, device) Sq = qf.shape[1] Sk = kf.shape[1] qc, kc, vc = qf.contiguous(), kf.contiguous(), vf.contiguous() @@ -512,7 +936,7 @@ def _fa2_causal_attn(qf, kf, vf, device): batch=1, seqlen_q=Sq, seqlen_k=Sk, num_heads_q=NQ, num_heads_kv=NKV, head_dim=HD, q_strides=qc.stride()[:3], k_strides=kc.stride()[:3], v_strides=vc.stride()[:3], o_strides=o.stride()[:3], - softmax_scale=float(HD) ** -0.5, num_sms=_num_sms(), stream=0) + softmax_scale=float(HD) ** -0.5, num_sms=_num_sms(), stream=_cs()) return o @@ -535,7 +959,7 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, q_pre = torch.empty(B * S, NQ, HD, dtype=torch.bfloat16, device=device) gate = torch.empty(B * S, NQ * HD, dtype=torch.bfloat16, device=device) fvk.qwen35moe_split_q_gate_bf16( - qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), B * S, 0) + qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), B * S, _cs()) q = q_pre.view(B, S, NQ, HD) gate = gate.view(B, S, NQ * HD) q = _rms_k(q.to(torch.bfloat16), qnw, fvk, device, eps) @@ -550,7 +974,7 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, ctc, stc = ct.contiguous(), st.contiguous() fvk.qwen36_partial_rope_qk_bf16( qin.data_ptr(), kin.data_ptr(), ctc.data_ptr(), stc.data_ptr(), - qo.data_ptr(), ko.data_ptr(), S, NQ, NKV, HD, ROPE, 0) + qo.data_ptr(), ko.data_ptr(), S, NQ, NKV, HD, ROPE, _cs()) # Causal GQA attention via the vendored FA2 kernel (native GQA: KV stays at # NKV=2, no repeat_interleave; layout is FA2's (B,S,H,HD), no transpose). @@ -576,13 +1000,52 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, gc = gate.reshape(-1).contiguous() ato = torch.empty_like(atc) fvk.sigmoid_mul_sm120_bf16(atc.data_ptr(), gc.data_ptr(), - ato.data_ptr(), atc.numel(), 0) + ato.data_ptr(), atc.numel(), _cs()) at = ato.reshape(B * S, NQ * HD) return _proj(at, ld, 'o_proj', HID, fvk, device).reshape(B, S, HID) +# The two weight-only 4-bit GEMVs each have an "edge" variant that is bitwise +# identical -- it differs only in a shared-memory layout and in decoding the +# UE4M3 scale byte arithmetically rather than through a constant-memory lookup +# whose index differs per lane. Because the outputs are identical to the bit, +# choosing between them is purely a performance decision and cannot move a +# token, so preferring the variant needs no accuracy argument. Set +# KernelPolicy.edge_w4a16 to False (or FLASHRT_QWEN35MOE_W4A16_EDGE=0) to force +# the original. + + +def w4a16_matvec(fvk): + """The dense 4-bit GEMV entry point this build should call.""" + if kernel_policy().edge_w4a16: + fn = getattr(fvk, 'w4a16_matvec_edge_sm120_bf16', None) + if fn is not None: + return fn + return fvk.w4a16_matvec_sm120_bf16 + + +def moe_grouped_w4a16(fvk): + """The grouped per-slot 4-bit GEMV entry point this build should call.""" + if kernel_policy().edge_w4a16: + fn = getattr(fvk, 'moe_grouped_w4a16_edge_sm120_bf16', None) + if fn is not None: + return fn + return fvk.moe_grouped_w4a16_sm120_bf16 + + # Grouped MoE for prefill (on by default); set False to use the per-expert loop. _USE_GROUPED_MOE = True +# One GEMM per expert, for a build without the block-scaled MMA tiles. Reads +# each expert's weight once instead of once per token that routed to it. +# +# It only pays once the tokens per expert are worth a launch. Each expert costs +# about five launches (quantise, two GEMMs, the gate), so with 256 of them a +# layer that is thousands of launches whichever way; below the threshold the +# grouped GEMV's two launches win even though it re-reads the weight. Measured +# at S=256 -- eight tokens an expert -- the per-expert path takes TTFT from 585 +# to 923 ms, while at S=1024 it takes it from 2237 to 1354. +_USE_PER_EXPERT_GEMM = True +_PER_EXPERT_MIN_M = 16 # mean tokens per expert, = S * TOPK / 256 # M=16 tensor-core mma MoE: tokens are sorted into 16-row expert tiles and the # SM120 block-scaled mma runs each expert once at full M-utilisation -- ~5.6x # the SIMT grouped W4A16 at large S (the compute wall). W4A4 (FP4 activation), @@ -596,6 +1059,48 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, _USE_BT_MOE = True +def _capture_per_token_state(cap, rank, S, init_state, conv_hist, mixed, + qb, kb, vb, g_out, bo, fvk, device): + """Record what the recurrent state would be after each token of a block. + + A verified speculative window is accepted up to some prefix, and the layer + that has to be rewound is this one: the KV cache is a cursor, but the + recurrent and conv states are not -- they have already absorbed every token + of the block, including the rejected tail. + + Rather than re-deriving them afterwards, run the scan a token at a time and + keep each intermediate. The block is a handful of tokens, so this is a few + extra launches and a few MB of state copies; recovering the state any other + way means either re-running the block's projections or reconstructing the + recurrence from saved inputs, both of which cost more than they save at + this length. + """ + state = (init_state.clone() if init_state is not None + else torch.zeros(NV, HK, HV, dtype=torch.bfloat16, device=device)) + q3 = qb.reshape(S, NV, HK).contiguous() + k3 = kb.reshape(S, NV, HK).contiguous() + v3 = vb.reshape(S, NV, HV).contiguous() + g2 = g_out.reshape(S, NV).contiguous() + b2 = bo.reshape(S, NV).contiguous() + core1 = torch.empty(1, NV, HV, dtype=torch.bfloat16, device=device) + for t in range(S): + fvk.gdn_recurrent_seq_sm120_bf16( + q3[t:t + 1].data_ptr(), k3[t:t + 1].data_ptr(), + v3[t:t + 1].data_ptr(), g2[t:t + 1].data_ptr(), + b2[t:t + 1].data_ptr(), state.data_ptr(), core1.data_ptr(), + 1, NV, HK, True, _cs()) + cap.spec_states[rank][t].copy_(state) + + # Conv state after t+1 tokens: the last KS-1 entries of the block's inputs + # preceded by whatever history the block started from. + prev = (conv_hist[0] if conv_hist is not None + else torch.zeros(mixed.shape[-1], KS - 1, + dtype=mixed.dtype, device=device)) + hist = torch.cat([prev, mixed[0].transpose(0, 1)], dim=1) + for t in range(S): + cap.spec_conv[rank][t].copy_(hist[:, t + 1:t + KS].unsqueeze(0)) + + def _moe_experts_m16(x, ti, tw, ld, fvk, device): """Routed experts via the M=16 tensor-core block-scaled mma. Sort the S*TOPK assignments by expert, pack into zero-padded 16-row tiles, quant once @@ -615,10 +1120,11 @@ def _moe_experts_m16(x, ti, tw, ld, fvk, device): exp_flat = ti.reshape(-1).to(torch.int32) tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) - order = exp_flat.argsort() + # Stable, so equal-expert ties keep token order and the rows packed into + # each quantisation tile are the same run to run. + order = exp_flat.argsort(stable=True) se = exp_flat[order].long() stok = tok_flat[order] - sw = tw.reshape(-1)[order] counts = torch.bincount(se, minlength=E) tile_counts = (counts + 15) // 16 tile_off = torch.cumsum(tile_counts, 0) - tile_counts @@ -636,16 +1142,26 @@ def _moe_experts_m16(x, ti, tw, ld, fvk, device): fvk.moe_m16_mma_sm120_bf16( ap.data_ptr(), gu_p.data_ptr(), asf.data_ptr(), gu_s.data_ptr(), d_gu.data_ptr(), gu_a.data_ptr(), tile_expert.data_ptr(), - total_tiles, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), 0) + total_tiles, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), _cs()) inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device).contiguous() ip, isf = _quant_act(inter, fvk, device) d_dn = torch.empty(total_tiles * 16, n_dn, dtype=torch.bfloat16, device=device) fvk.moe_m16_mma_sm120_bf16( ip.data_ptr(), dn_p.data_ptr(), isf.data_ptr(), dn_s.data_ptr(), d_dn.data_ptr(), dn_a.data_ptr(), tile_expert.data_ptr(), - total_tiles, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), 0) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn[tiled_row].float() * sw.unsqueeze(-1)) + total_tiles, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), _cs()) + # Deterministic unpermute, as the grouped-GEMM path does: one kernel sums + # each token's TOPK rows in a fixed order. Slot i of the token-major + # routing sits at sorted position inv[i], whose output row is tiled_row of + # that position. + inv = torch.empty(S * TOPK, dtype=torch.long, device=device) + inv[order] = torch.arange(S * TOPK, device=device) + rows = tiled_row[inv].to(torch.int32).contiguous() + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), rows.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -709,14 +1225,14 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): fvk.moe_blocktile_mma_sm120_bf16( ap.data_ptr(), gu_p.data_ptr(), asf.data_ptr(), gu_s.data_ptr(), d_gu.data_ptr(), gu_a.data_ptr(), tile_expert.data_ptr(), - MAX_TILES, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), 0) + MAX_TILES, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), _cs()) inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device).contiguous() ip, isf = _quant_act(inter, fvk, device) d_dn = torch.empty(MAX_TILES * 64, n_dn, dtype=torch.bfloat16, device=device) fvk.moe_blocktile_mma_sm120_bf16( ip.data_ptr(), dn_p.data_ptr(), isf.data_ptr(), dn_s.data_ptr(), d_dn.data_ptr(), dn_a.data_ptr(), tile_expert.data_ptr(), - MAX_TILES, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), 0) + MAX_TILES, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), _cs()) # Deterministic unpermute via the fused gather-weighted-sum kernel: invert # the routing permutation (inv: orig slot -> sorted position, a 131 KB int # scatter) to get each token's TOPK d_dn rows, then one kernel computes @@ -730,7 +1246,321 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): out = torch.empty(S, HID, dtype=torch.float32, device=device) fvk.moe_weighted_sum_sm120_bf16( d_dn.data_ptr(), rows.data_ptr(), twc.data_ptr(), out.data_ptr(), - S, TOPK, n_dn, n_dn, 0) + S, TOPK, n_dn, n_dn, _cs()) + return out + + +_GROUPED_SCRATCH = {} +_ROUTE_CONST = {} +_ROUTE_BUF = {} +# Off puts the routing back on the tensor chain, which is how the kernel's +# output is A/B'd against it end to end rather than only in a probe. + + +def _route_constants(S, device): + """The parts of the routing permutation that depend only on the shape. + + Each layer routes differently, but the token index per slot and the slot + index itself do not change -- they are a function of S alone, and every + layer was rebuilding both. Forty layers of arange + repeat_interleave is + launches and traffic spent to recompute a constant. + """ + key = (S, str(device)) + got = _ROUTE_CONST.get(key) + if got is None: + tok_flat = torch.arange( + S, device=device).repeat_interleave(TOPK).contiguous() + slot_ix = torch.arange(S * TOPK, device=device) + got = (tok_flat, slot_ix) + _ROUTE_CONST[key] = got + return got + + +def _route_buffers(S, fvk, device): + """The routing kernel's outputs, allocated once per prompt length. + + Their sizes depend only on S, so the forty layers of a prefill write + through the same buffers at the same addresses -- which is what lets the + call sit inside a captured region, and incidentally saves forty rounds of + allocation per forward. + """ + key = (S, str(device)) + got = _ROUTE_BUF.get(key) + if got is None: + slots = S * TOPK + ws_bytes = int(fvk.moe_route_prefill_workspace_bytes( + S, TOPK, _N_EXPERTS)) + got = { + 'ti': torch.empty(S, TOPK, dtype=torch.int32, device=device), + 'tw': torch.empty(S, TOPK, dtype=torch.float32, device=device), + 'se': torch.empty(slots, dtype=torch.int32, device=device), + # int64: the activation quantiser reads this gather index + # as a long, and int32 there is an illegal access. + 'stok': torch.empty(slots, dtype=torch.int64, device=device), + 'inv': torch.empty(slots, dtype=torch.int32, device=device), + 'group_off': torch.empty(_N_EXPERTS + 1, dtype=torch.int32, + device=device), + 'ws': torch.empty(ws_bytes, dtype=torch.uint8, device=device), + 'ws_bytes': ws_bytes, + 'sfa_off': {}, + } + _ROUTE_BUF[key] = got + return got + + +def _route_prefill(logits, fvk, device): + """Softmax, top-k, and the permutation the grouped GEMM reads, in kernels. + + Replaces softmax + top-k + a renormalising divide + a stable argsort + two + gathers + a bincount + a cumulative sum + a scatter: ten tensor ops a + layer, of which the top-k alone was 25 ms of a 2048-token prefill. + + Returns None where the kernel is absent, so the tensor chain stays the + fallback rather than this being a hard dependency. + """ + if (not kernel_policy().route_kernel + or not hasattr(fvk, 'moe_route_prefill_bf16')): + return None + S = logits.shape[0] + b = _route_buffers(S, fvk, device) + rc = fvk.moe_route_prefill_bf16( + logits.data_ptr(), b['ti'].data_ptr(), b['tw'].data_ptr(), + b['se'].data_ptr(), b['stok'].data_ptr(), b['inv'].data_ptr(), + b['group_off'].data_ptr(), b['ws'].data_ptr(), b['ws_bytes'], + S, _N_EXPERTS, TOPK, _cs()) + if rc: + raise RuntimeError(f'prefill routing failed with {rc}') + return b + + +def _route_sfa_off(route, k, fvk, device): + """Per-expert scale-factor byte offsets for one projection's K.""" + n_col = ((k // 16) + 3) // 4 + off = route['sfa_off'].get(k) + if off is None: + off = torch.empty(_N_EXPERTS, dtype=torch.int32, device=device) + route['sfa_off'][k] = off + fvk.moe_route_sfa_offsets( + route['group_off'].data_ptr(), off.data_ptr(), _N_EXPERTS, n_col, + _cs()) + return off, n_col + + +def _grouped_scratch(fvk, device): + """One scratch buffer per device for the grouped GEMM's descriptor arrays. + + Its size depends only on the expert count, not on how the routing falls, so + it is allocated once and never resized -- which is also what lets the call + sit inside a captured region. + """ + key = str(device) + got = _GROUPED_SCRATCH.get(key) + if got is None: + nbytes = int(fvk.moe_grouped_gemm_nvfp4_sm100_scratch_bytes(_N_EXPERTS)) + got = (torch.empty(nbytes, dtype=torch.uint8, device=device), nbytes) + _GROUPED_SCRATCH[key] = got + return got + + +def _sf_layout(counts, k, device): + """Per-group scale-factor byte offsets, and a host-known bound on the total. + + The block-scaled layout blocks rows by 128, so a group of c rows needs + ceil(c/128) super-blocks. Summing that needs the counts, which live on the + device -- but the sum is bounded by (experts + slots/128) super-blocks + whatever the routing does, and that bound follows from the shapes alone. + Sizing the buffer from the bound rather than the sum is what keeps this free + of a host read. + """ + n_col = ((k // 16) + 3) // 4 + per_group = ((counts + 127) // 128) * (n_col * 512) + off = torch.zeros(_N_EXPERTS, dtype=torch.int32, device=device) + off[1:] = per_group.cumsum(0)[:-1].to(torch.int32) + return off, n_col + + +def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device, route=None): + """Every routed expert of the layer in two GEMM launches. + + The per-expert loop below reads each weight once, which is the right amount, + but pays a launch and a host iteration per expert -- and the host iteration + is fatal twice over: it dominated the time at S=1024, and it makes the layer + impossible to capture. A grouped GEMM takes the per-group shapes from device + memory, so the launch geometry depends only on the expert count and the + routing never reaches the host. + + Measured against the loop it replaces, at the shapes prefill issues: 6.0x on + gate_up and 14.5x on down at S=1024, 512 launches down to 2, output bitwise + identical. + """ + S = x.shape[0] + slots = S * TOPK + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + if 'experts_gate_up_alpha_dev' not in ld: + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] + scratch, scratch_bytes = _grouped_scratch(fvk, device) + + if route is not None: + se, stok, group_off, order, counts = ( + route['se'], route['stok'], route['group_off'], None, None) + else: + tok_flat, slot_ix = _route_constants(S, device) + exp_flat = ti.reshape(-1).to(torch.int32) + order = exp_flat.argsort(stable=True) + se = exp_flat[order].contiguous() + stok = tok_flat[order] + + counts = torch.bincount(se, minlength=_N_EXPERTS) + group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, + device=device) + group_off[1:] = counts.cumsum(0).to(torch.int32) + + def project(A, k, n, w_p, w_s, alpha, out, gate=False, perm=None): + if route is not None: + sfa_off, n_col = _route_sfa_off(route, k, fvk, device) + else: + sfa_off, n_col = _sf_layout(counts, k, device) + bound = (_N_EXPERTS + slots // 128 + 1) * n_col * 512 + packed = torch.empty(slots, k // 2, dtype=torch.uint8, device=device) + sfa = torch.empty(bound, dtype=torch.uint8, device=device) + if gate: + # A is the merged (slots, 2k) gate/up output: gate it and quantise + # in one pass rather than slicing two strided halves out of it, + # copying both, gating into a third buffer and reading that back. + # Warp-per-row: a lane owns one scale-factor group and keeps + # it in registers, so there is no shared memory and no barrier. + # Byte-identical to the block-per-row form and 2.7x at the shape + # prefill issues, which puts it at 1.04x of its traffic bound. + _sq = getattr(fvk, 'moe_grouped_silu_quant_nvfp4_warp_bf16', + fvk.moe_grouped_silu_quant_nvfp4_bf16) + rc = _sq( + A.data_ptr(), se.data_ptr(), group_off.data_ptr(), + sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + slots, k, _cs()) + else: + rc = fvk.moe_grouped_quant_nvfp4_bf16( + A.data_ptr(), se.data_ptr(), group_off.data_ptr(), + sfa_off.data_ptr(), 0 if perm is None else perm.data_ptr(), + packed.data_ptr(), sfa.data_ptr(), slots, k, _cs()) + if rc: + raise RuntimeError(f'grouped activation quant failed with {rc}') + rc = fvk.moe_grouped_gemm_nvfp4_sm100_bf16out( + packed.data_ptr(), sfa.data_ptr(), w_p.data_ptr(), + w_s.data_ptr(), alpha.data_ptr(), out.data_ptr(), + group_off.data_ptr(), sfa_off.data_ptr(), + _N_EXPERTS, n, k, w_p[0].numel(), w_s[0].numel(), + scratch.data_ptr(), scratch_bytes, _cs()) + if rc: + raise RuntimeError(f'grouped MoE GEMM failed with {rc}') + + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + project(x, HID, n_gu, gu_p, gu_s, gu_a, d_gu, perm=stok) + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) + project(d_gu, INTER, n_dn, dn_p, dn_s, dn_a, d_dn, gate=True) + + # Deterministic unpermute, the same one the block-tile path uses: invert the + # routing permutation and let one kernel sum each token's TOPK rows in fixed + # order. index_add_ was 37.8 ms of a 1024-token prefill and reduces through + # atomics, so its order varies -- which prefill cannot afford, since it + # seeds a decode that has to be reproducible. + # rows[i] is which sorted row holds slot i, which is exactly the inverse + # permutation -- gathering arange through it, as the tiled path has to, + # would just reproduce it. + if route is not None: + inv, twc = route['inv'], route['tw'] + else: + inv = torch.empty(slots, dtype=torch.int32, device=device) + inv[order] = slot_ix.to(torch.int32) + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) + return out + + +def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): + """Routed experts as one GEMM per expert over the tokens that chose it. + + The grouped GEMV below issues one GEMV per (token, expert) slot, so an + expert's weight is re-read once per token that routed to it -- 8192 reads + of a 1.18 MB weight per layer at S=1024, which is 9.7 GB of traffic a layer + even before the down projection. Sorting by expert makes those reads hit L2 + rather than DRAM, which is why it works at all, but it is still bounded by + L2 bandwidth and it dominates prefill: measured 74.6% of a 1024-token + prefill, 1762 ms of 2361. + + Grouping the tokens instead turns each expert into a single M-row GEMM that + reads its weight once. The block-scaled 4-bit MMA tile does the same thing + and better, but it is a build tier that is not present everywhere; this path + needs only the NVFP4 W4A16 GEMM, which is. + + The count per expert is data-dependent, so this reads it to the host -- one + sync per layer, which prefill can afford and a captured decode could not. + """ + S = x.shape[0] + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + if 'experts_gate_up_alpha_list' not in ld: + ld['experts_gate_up_alpha_list'] = ld['experts_gate_up_alpha_t'].tolist() + ld['experts_down_alpha_list'] = ld['experts_down_alpha_t'].tolist() + gu_a = ld['experts_gate_up_alpha_list'] + dn_a = ld['experts_down_alpha_list'] + + exp_flat = ti.reshape(-1).to(torch.int32) + tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) + # Stable, so equal-expert ties keep token order and the rows packed into + # each quantisation tile are the same run to run. + order = exp_flat.argsort(stable=True) + se = exp_flat[order] + stok = tok_flat[order] + + counts = torch.bincount(se, minlength=_N_EXPERTS).tolist() + slots = S * TOPK + A = x[stok].contiguous() # (slots, HID) bf16 + # One buffer per projection, written in place by each expert's GEMM. The + # activation is a slot-major matrix throughout, so the gate is one launch + # over all of it rather than one per expert -- 256 launches a layer and two + # slice copies each, for an op that does not care where the rows came from. + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) + + off = 0 + bounds = [] + for e, cnt in enumerate(counts): + if cnt == 0: + continue + bounds.append((e, off, cnt)) + xp, xsf = _quant_act(A[off:off + cnt], fvk, device, _cs()) + _nvfp4_gemm_preq(xp, xsf, gu_p[e].data_ptr(), gu_s[e].data_ptr(), + gu_a[e], cnt, n_gu, HID, fvk, device, _cs(), + out=d_gu[off:off + cnt]) + off += cnt + + inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device) + for e, off_e, cnt in bounds: + xp, xsf = _quant_act(inter[off_e:off_e + cnt].contiguous(), fvk, + device, _cs()) + _nvfp4_gemm_preq(xp, xsf, dn_p[e].data_ptr(), dn_s[e].data_ptr(), + dn_a[e], cnt, n_dn, INTER, fvk, device, _cs(), + out=d_dn[off_e:off_e + cnt]) + + # Deterministic unpermute, as the grouped-GEMM path does; index_add_ + # reduces through atomics, so its order varies run to run. + inv = torch.empty(S * TOPK, dtype=torch.int32, device=device) + inv[order] = torch.arange(S * TOPK, dtype=torch.int32, device=device) + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -754,26 +1584,36 @@ def _moe_experts_grouped(x, ti, tw, ld, fvk, device): slots = S * TOPK exp_flat = ti.reshape(-1).to(torch.int32) tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) - order = exp_flat.argsort() + # Stable, so equal-expert ties keep token order run to run. + order = exp_flat.argsort(stable=True) se = exp_flat[order].contiguous() stok = tok_flat[order] - sw = tw.reshape(-1)[order] A = x[stok].contiguous() # (slots, HID) bf16 d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( A.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), se.data_ptr(), d_gu.data_ptr(), slots, n_gu, HID, - HID, gu_p[0].numel(), gu_s[0].numel(), 0) + HID, gu_p[0].numel(), gu_s[0].numel(), _cs()) g, u = d_gu[:, :INTER], d_gu[:, INTER:] inter = _silu_mul(g, u, fvk, device).contiguous() d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), se.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, - INTER, dn_p[0].numel(), dn_s[0].numel(), 0) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + INTER, dn_p[0].numel(), dn_s[0].numel(), _cs()) + # Deterministic unpermute, as the grouped-GEMM path does: invert the + # routing permutation and let one kernel sum each token's TOPK rows in a + # fixed order. index_add_ reduces through atomics, so eight fp32 addends + # land in whatever order the blocks retire -- and a prefill cannot afford + # that, because it seeds a decode that has to be reproducible. + inv = torch.empty(slots, dtype=torch.int32, device=device) + inv[order] = torch.arange(slots, dtype=torch.int32, device=device) + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -791,15 +1631,43 @@ def _moe_layer(h, ld, fvk, device): # Router GEMM via the deterministic w16a16 kernel (bf16 weight, fp32 # accumulate) instead of the fp32 upcast matmul. bf16 logits match the - # bf16 reference router; softmax/topk stay (already CUDA ops). - logit = F.softmax(_gemm_w16a16(x, rw, fvk, device).float(), -1) - tw, ti = torch.topk(logit, TOPK, -1) - tw = tw / tw.sum(-1, keepdim=True) + # bf16 reference router. + lg = _gemm_w16a16(x, rw, fvk, device) + + # The block-scaled 4-bit MMA tiles are a build tier, not a given: a target + # whose toolchain has no block-scaled mma builds the weight-only tier + # instead. Ask the module what it has rather than assuming, so the tile + # choice degrades to the grouped GEMV instead of raising mid-prefill. + big = x.shape[0] >= _M16_MIN_S + use_bt = (_USE_BT_MOE and big + and hasattr(fvk, 'moe_blocktile_mma_sm120_bf16')) + use_m16 = (not use_bt and _USE_M16_MOE and big + and hasattr(fvk, 'moe_m16_mma_sm120_bf16')) + grouped = (not use_bt and not use_m16 and big + and hasattr(fvk, 'moe_grouped_gemm_nvfp4_sm100_bf16out')) + # Only the grouped path reads the kernel's permutation, and the tiled + # paths index with the tensor top-k's own indices, so the routing is not + # computed twice for a path that will not use it. + route = _route_prefill(lg, fvk, device) if grouped else None + if route is not None: + ti, tw = route['ti'], route['tw'] + else: + logit = F.softmax(lg.float(), -1) + tw, ti = torch.topk(logit, TOPK, -1) + tw = tw / tw.sum(-1, keepdim=True) - if _USE_BT_MOE and x.shape[0] >= _M16_MIN_S: + if use_bt: out = _moe_experts_bt(x, ti, tw, ld, fvk, device) - elif _USE_M16_MOE and x.shape[0] >= _M16_MIN_S: + elif use_m16: out = _moe_experts_m16(x, ti, tw, ld, fvk, device) + elif grouped: + # No threshold: the grouped path wins at every prefill length measured, + # because it does not pay per expert for anything. + out = _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device, route) + elif (big and _USE_PER_EXPERT_GEMM + and x.shape[0] * TOPK >= _PER_EXPERT_MIN_M * _N_EXPERTS + and hasattr(fvk, 'fp4_w4a16_gemm_sm120_bf16out')): + out = _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device) elif _USE_GROUPED_MOE: out = _moe_experts_grouped(x, ti, tw, ld, fvk, device) else: @@ -825,9 +1693,18 @@ def _moe_layer(h, ld, fvk, device): su = _proj(x, ld, 'shared_up_proj', INTER, fvk, device) si = _silu_mul(sg, su, fvk, device) shared = _proj(si, ld, 'shared_down_proj', HID, fvk, device) - # shared-expert scalar gate: GEMM (N=1) via w16a16, then sigmoid. - sgate = torch.sigmoid( - _gemm_w16a16(x, ld['shared_gate_w_t'], fvk, device).float()) + # shared-expert scalar gate: GEMM (N=1) via w16a16. The sigmoid, the + # broadcast multiply, the add onto the routed sum and the cast are one + # kernel; the routed sum stays fp32 until the single rounding at its store. + glog = _gemm_w16a16(x, ld['shared_gate_w_t'], fvk, device) + if hasattr(fvk, 'moe_shared_gate_combine_edge_bf16'): + comb = torch.empty(x.shape[0], HID, dtype=torch.bfloat16, + device=device) + fvk.moe_shared_gate_combine_edge_bf16( + out.data_ptr(), shared.data_ptr(), glog.data_ptr(), + comb.data_ptr(), x.shape[0], HID, _cs()) + return comb.reshape(B, S, HID) + sgate = torch.sigmoid(glog.float()) return (out + shared.float() * sgate).reshape(B, S, HID).to(torch.bfloat16) @@ -874,10 +1751,14 @@ def nexn2_forward_nvfp4(handles, input_ids, fvk, device, cap=None, ct, st = ct_full[pos_offset:], st_full[pos_offset:] chunked = pos_offset > 0 lin_rank = full_rank = 0 + # Every residual add is immediately followed by the norm of what it + # produced, so the two run as one kernel that updates the residual stream + # in place -- which means the loop carries the *normed* tensor across each + # boundary and takes the first norm before entering it. + h = h.contiguous() + n = _rms_k(h, layers[0]['input_norm_w_t'], fvk, device, eps) for L in range(p['num_layers']): ld = layers[L] - res = h - n = _rms_k(h, ld['input_norm_w_t'], fvk, device, eps) if types[L] == 'linear_attention': init_s = cap.lin_state[lin_rank] if chunked else None conv_h = cap.lin_conv_state[lin_rank] if chunked else None @@ -888,21 +1769,31 @@ def nexn2_forward_nvfp4(handles, input_ids, fvk, device, cap=None, attn = _full_attn_layer(n, ld, ct, st, fvk, device, eps, cap, full_rank, pos_offset=pos_offset) full_rank += 1 - h = res + attn - res = h - n = _rms_k(h, ld['post_norm_w_t'], fvk, device, eps) - h = res + _moe_layer(n, ld, fvk, device) + n = _add_rms_k(h, attn, ld['post_norm_w_t'], fvk, device, eps) + moe = _moe_layer(n, ld, fvk, device) + # The norm after the last layer's residual is the final norm, and + # between layers it is the next layer's input norm -- one call either + # way, so the boundary is a choice of weight rather than a branch. + nxt = (layers[L + 1]['input_norm_w_t'] if L + 1 < p['num_layers'] + else p['final_norm_w_t']) + n = _add_rms_k(h, moe, nxt, fvk, device, eps) hidden = h[0] # (S, HID) residual stream, pre-final-norm if not compute_logits: return (None, hidden) if return_hidden else None - h = _rms_k(h, p['final_norm_w_t'], fvk, device, eps) + h = n # already the final norm, see above # lm_head via w16a16 (bf16 weight, fp32 accumulate): reads the ~1GB weight # as bf16 (no fp32 widen), same argmax. logits returned bf16. Slice to the # last position first when only the seeding logit is needed (avoids the # (S, vocab) materialisation that dominates long-context prefill memory). h_lm = h[0][-1:].contiguous() if last_logits_only else h[0] - logits = _gemm_w16a16(h_lm, p['lm_head_w_t'], fvk, device) + if _SPEC_VERIFY: + # The lm_head is the single largest weight; at BF16 it is a gigabyte a + # verify, which on its own outweighs what the window saves. + logits = _gemm_w4a16(h_lm, p['lm_head_w_t'], p, 'lm_head_w_t', + fvk, device) + else: + logits = _gemm_w16a16(h_lm, p['lm_head_w_t'], fvk, device) if return_hidden: return logits, hidden return logits diff --git a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py index 985162aa..7fc977bb 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py @@ -168,7 +168,8 @@ def _bf16_from_ckpt(handles, out_dict, name, key, handles_d, wmap, device, def _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, - n_experts: int, *, quantize_shared: bool = True) -> None: + n_experts: int, *, quantize_shared: bool = True, + stream_experts: bool = False) -> None: """Load one layer's MoE block: router (BF16) + experts + shared expert.""" # Router gate (BF16) and shared-expert sigmoid gate (BF16). _bf16_from_ckpt(handles, ld, 'router_w', lp + 'mlp.gate.weight', @@ -189,6 +190,20 @@ def _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, # Routed experts: packed 3D tensors (E, out, in). Quantize each expert # into a contiguous slice of a per-layer stacked NVFP4 buffer so the # downstream grouped GEMM sees one contiguous weight per projection. + if stream_experts: + # The routed experts are read from storage at decode time, so the + # stacked per-layer tensors are never built. Skipping them is the + # point: they are 16.9 GiB of the resident footprint, and a cache + # added on top of them would cost memory rather than save it. The + # shapes are still checked, because a bundle is generated against them. + for name in ('mlp.experts.gate_up_proj', 'mlp.experts.down_proj'): + if not _has(wmap, lp + name): + raise ValueError( + f'{lp}{name} is absent, so a streamed expert bundle ' + 'cannot correspond to this checkpoint') + ld['experts_streamed'] = True + return + gate_up = _get(handles_d, wmap, lp + 'mlp.experts.gate_up_proj') down = _get(handles_d, wmap, lp + 'mlp.experts.down_proj') e_gu, n_gu, k_gu = gate_up.shape # (E, 2*inter, hidden) @@ -244,6 +259,8 @@ def extract_weights_nexn2_nvfp4( fvk, device: str = 'cuda:0', quant_scope: str = 'experts', + stream_experts: bool = False, + load_mtp: bool = False, ) -> WeightHandles: """Build :class:`WeightHandles` from a Nex-N2-mini BF16 ckpt directory. @@ -253,6 +270,12 @@ def extract_weights_nexn2_nvfp4( * ``'experts'``: only the storage-dominant routed experts go NVFP4; full-attn / out_proj / shared stay BF16. ~21 GB; E2E cos ~0.99 -- the precision-per-VRAM baseline until the Step-3 W4A16 mixed kernel. + + stream_experts: skip the routed experts entirely, leaving the decode path + to read them from a prepared bundle. They are 16.9 GiB of the resident + footprint, so this is the difference between a model that fits a small + device and one that does not; a cache added without skipping them would + only add to the total. The decode path must then be given an ExpertCache. """ if quant_scope not in ('full', 'experts'): raise ValueError( @@ -288,10 +311,13 @@ def extract_weights_nexn2_nvfp4( handles_d, wmap, device) # ── Per-layer ── - per_layer: list = [None] * num_layers - for i in range(num_layers): - lp = f'model.language_model.layers.{i}.' - ltype = layer_types[i] + def _load_layer(lp: str, ltype: str, *, streamed: bool = None) -> dict: + """Build one layer's weight dict from its checkpoint prefix. + + Taken out of the loop so the MTP head can use it: its layer lives under + a different prefix but has exactly a full-attention layer's keys, and + loading it a second way would be a second thing to keep correct. + """ ld: dict = {'type': ltype, 'quant_format': 'nvfp4'} _bf16_from_ckpt(handles, ld, 'input_norm_w', lp + 'input_layernorm.weight', @@ -331,12 +357,19 @@ def extract_weights_nexn2_nvfp4( _proj_load(handles, ld, 'out_proj', gp + 'out_proj.weight', handles_d, wmap, fvk, device, quantize=quant_main) else: - raise ValueError(f'layer {i}: unknown layer_type {ltype!r}') + raise ValueError(f'{lp}: unknown layer_type {ltype!r}') # Every layer has a MoE FFN (mlp_only_layers is empty). _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, n_experts, - quantize_shared=quant_main) - per_layer[i] = ld + quantize_shared=quant_main, + stream_experts=(stream_experts if streamed is None + else streamed)) + return ld + + per_layer: list = [None] * num_layers + for i in range(num_layers): + per_layer[i] = _load_layer( + f'model.language_model.layers.{i}.', layer_types[i]) handles.ptrs['layers'] = per_layer handles.ptrs['vocab_size'] = vocab @@ -354,6 +387,30 @@ def extract_weights_nexn2_nvfp4( handles.ptrs['quant_format'] = 'nvfp4' handles.ptrs['quant_scope'] = quant_scope handles.ptrs['ckpt_dir'] = ckpt_dir - handles.ptrs['mtp'] = None # MTP weights not in the base ckpt + # ── Multi-token-prediction head ── + # + # One full-attention layer plus its own 256-expert MoE, under `mtp.`, with + # four head-level tensors around it. It drafts the token after next from + # the main model's last hidden state and the token just emitted, which is + # only useful with a verifier, so it is opt-in: it costs another layer's + # worth of weights and a KV slot. + handles.ptrs['mtp'] = None + if load_mtp: + if not _has(wmap, 'mtp.fc.weight'): + raise RuntimeError( + f'{ckpt_dir} has no MTP head (mtp.fc.weight is absent), so ' + 'speculative drafting cannot be built from it.') + # A bundle holds the model's own layers, so the head's experts have + # nowhere to stream from and stay resident whatever the model does. + mtp: dict = {'layer': _load_layer('mtp.layers.0.', 'full_attention', + streamed=False)} + _bf16_from_ckpt(handles, mtp, 'fc_w', 'mtp.fc.weight', + handles_d, wmap, device) + for name, key in (('norm_w', 'mtp.norm.weight'), + ('pre_h_w', 'mtp.pre_fc_norm_hidden.weight'), + ('pre_e_w', 'mtp.pre_fc_norm_embedding.weight')): + _bf16_from_ckpt(handles, mtp, name, key, + handles_d, wmap, device, fold_one=True) + handles.ptrs['mtp'] = mtp handles.ptrs['dflash'] = None return handles diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index b1791552..987ae483 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -34,27 +34,68 @@ ) +# The tier combination each target can build. FLASHRT_ENABLE_QWEN35MOE turns on +# all three tiers including the block-scaled 4-bit MMA one, which needs +# sm_120a/sm_121a; recommending it on a target whose toolchain refuses it sends +# the reader to a configure error. So the advice is keyed by the device in +# front of them. +_TIER_ADVICE = { + (11, 0): ("-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON " + "-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON"), + (12, 0): "-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON", + (12, 1): "-DGPU_ARCH=121 -DFLASHRT_ENABLE_QWEN35MOE=ON", +} + + +def _build_advice() -> str: + """The configure flags for the device this process is actually running.""" + try: + import torch + + cap = torch.cuda.get_device_capability() + except Exception: # pragma: no cover + return ("-DFLASHRT_ENABLE_QWEN35MOE=ON on sm_120a/sm_121a, or " + "-DFLASHRT_ENABLE_QWEN35MOE_CORE=ON " + "-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON elsewhere") + return _TIER_ADVICE.get( + cap, + f"the tiers sm_{cap[0]}{cap[1]} can compile " + "(FLASHRT_ENABLE_QWEN35MOE_CORE / _W4A16; _W4A4 needs sm_120a)") + + def _require_kernels( fvk, *, model_label: str = "Nex-N2", - usage_doc: str = "docs/nexn2_usage.md") -> None: + usage_doc: str = "docs/nexn2_usage.md", + required=None, require_fa2: bool = True) -> None: """Raise a clear RuntimeError if the gated qwen3_5_moe kernels or the FA2 - module are missing (build was not configured with - -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent).""" - missing = [s for s in _REQUIRED_FVK if not hasattr(fvk, s)] + module are missing (the build did not enable the qwen3_5_moe tiers, or + flash_rt_fa2 is absent). + + ``required`` lets a configuration that calls fewer kernels say so. A list + demanding more than a path uses turns a working build into a refusal; one + demanding less lets a missing symbol surface mid-forward. Both are wrong, + so the list belongs with whatever decides which kernels get called. + """ + missing = [s for s in (required or _REQUIRED_FVK) if not hasattr(fvk, s)] if missing: raise RuntimeError( - f"{model_label} kernelized path needs the qwen3_5_moe SM120 " - "kernels, which " - "are absent from flash_rt_kernels (missing: " - f"{', '.join(missing)}). Rebuild on an SM120 toolchain with " - f"-DFLASHRT_ENABLE_QWEN35MOE=ON. See {usage_doc}.") + f"{model_label} kernelized path needs the gated qwen3_5_moe " + "kernels, which are absent from flash_rt_kernels (missing: " + f"{', '.join(missing)}). Reconfigure with {_build_advice()}. " + f"See {usage_doc}.") + if not require_fa2: + # The attention backend probes its kernel and falls back to a + # reference implementation, so a target that builds no FA2 still runs + # -- more slowly on a long prompt, and never differently. + return try: from flash_rt import flash_rt_fa2 as _fa2 except Exception as e: # pragma: no cover raise RuntimeError( f"{model_label} full attention needs the vendored FA2 module " "(flash_rt_fa2), which failed to import. Build with FA2 enabled " - "(ENABLE_FA2, auto-on for SM120).") from e + "(automatic on sm_80/86/87/89/120/121; on Thor sm_110 it is " + "opt-in with -DFLASHRT_ENABLE_THOR_FA2=ON).") from e fa2_missing = [s for s in ('fwd_bf16', 'fwd_bf16_causal') if not hasattr(_fa2, s)] if fa2_missing: # pragma: no cover @@ -67,6 +108,14 @@ def _require_kernels( class Nexn2TorchFrontendRtx: """Nex-N2-mini inference frontend (PyTorch + RTX SM120).""" + # Kernels this configuration calls; a subclass whose path calls fewer + # narrows it. See _require_kernels. + _REQUIRED_KERNELS = _REQUIRED_FVK + + # Whether the vendored FA2 module must be present. A subclass whose + # attention backend can fall back sets this False. + _REQUIRE_FA2 = True + _MODEL_LABEL = "Nex-N2" _USAGE_DOC = "docs/nexn2_usage.md" @@ -108,6 +157,13 @@ def __init__(self, checkpoint_path: str, *, self._quant_format = quant self._kernelized = bool(kernelized) self._quant_scope = quant_scope + # Set by a subclass that streams the routed experts from a bundle + # instead of holding them; see _nexn2_rtx_decode._moe_experts_streamed. + self._stream_experts = getattr(self, '_stream_experts', False) + # Read by the loader before any weight is touched, like the above: + # the draft head is another layer's worth of weights and is only + # useful with a verifier, so nothing loads it unless asked. + self._load_mtp = getattr(self, '_load_mtp', False) self._tokenizer = None self._prompt_ids = None self._pipeline: Nexn2Pipeline | None = None @@ -147,8 +203,6 @@ def _build_kernelized_nvfp4(self) -> None: to NVFP4 (GDN in_proj / norms / router kept BF16) and frees the BF16 source as it goes, fitting in ~22 GB. """ - from transformers import AutoTokenizer - from flash_rt import flash_rt_kernels as fvk from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import ( extract_weights_nexn2_nvfp4, @@ -159,22 +213,50 @@ def _build_kernelized_nvfp4(self) -> None: fvk, model_label=self._MODEL_LABEL, usage_doc=self._USAGE_DOC, + required=self._REQUIRED_KERNELS, + require_fa2=self._REQUIRE_FA2, ) - self._tokenizer = AutoTokenizer.from_pretrained(self.checkpoint_path) self._fvk = fvk self._weights = extract_weights_nexn2_nvfp4( self.checkpoint_path, fvk, device=self.device, - quant_scope=self._quant_scope) + quant_scope=self._quant_scope, + stream_experts=self._stream_experts, + load_mtp=self._load_mtp) @property def tokenizer(self): - """The HF tokenizer loaded from the checkpoint.""" + """The checkpoint's tokenizer, loaded when something asks for it. + + Loading it eagerly would make ``transformers`` a hard requirement of + the runtime, which it is not: a caller that supplies token ids through + :meth:`set_prompt_ids` never needs one. That matters for a deployment + target where the dependency may be absent or unwelcome, and it keeps + the kernel and weight paths testable without it. + """ + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained( + self.checkpoint_path) return self._tokenizer + def set_prompt_ids(self, token_ids) -> None: + """Set the prompt from token ids, requiring no tokenizer.""" + import torch + + ids = torch.as_tensor( + token_ids, dtype=torch.long, device=self.device).reshape(1, -1) + if ids.shape[1] == 0: + raise ValueError('token_ids is empty') + # Matches set_prompt: the decode state is not discarded, because + # seed_prefill resets the recurrent and KV caches itself and + # reallocating them per prompt would be waste. + self._prompt_ids = ids + def set_prompt(self, text: str) -> None: """Tokenize ``text`` for the next ``infer()`` / ``generate()`` call.""" - enc = self._tokenizer(text, return_tensors='pt') + enc = self.tokenizer(text, return_tensors='pt') self._prompt_ids = enc['input_ids'].to(self.device) def infer(self): diff --git a/flash_rt/frontends/torch/qwen36_moe.py b/flash_rt/frontends/torch/qwen36_moe.py new file mode 100644 index 00000000..8684224d --- /dev/null +++ b/flash_rt/frontends/torch/qwen36_moe.py @@ -0,0 +1,426 @@ +"""Qwen3.6-35B-A3B text inference. + +The language backbone is the same ``qwen3_5_moe`` architecture used by +Nex-N2-mini, so this frontend reuses that implementation. Nothing in it is +specific to one GPU: it is registered for RTX SM120 and for Jetson AGX Thor +(SM110), and each target's build tiers decide which kernels the shared forward +and decode paths resolve to. (The base class keeps its ``Rtx`` name, which +predates the Thor path; see :mod:`flash_rt.frontends.torch.nexn2_rtx`.) + +The official Qwen3.6 checkpoint also contains a vision tower and an MTP draft +head. The vision tower is validated but not executed here. The draft head is +loaded only when ``load_mtp=True`` is passed, which is what +:meth:`Qwen36MoeTextFrontend.generate_spec` needs; ``generate`` never uses it. + +See ``docs/qwen36_moe_usage.md`` for the per-architecture build commands. +""" + +from __future__ import annotations + +import json +import os +from contextlib import ExitStack +from typing import Any + +from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx + + +_EXPECTED_LAYER_TYPES = tuple( + "full_attention" if (i + 1) % 4 == 0 else "linear_attention" + for i in range(40) +) + +_EXPECTED_TEXT_CONFIG = { + "model_type": "qwen3_5_moe_text", + "num_hidden_layers": 40, + "hidden_size": 2048, + "vocab_size": 248320, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "attention_bias": False, + "attn_output_gate": True, + "hidden_act": "silu", + "num_experts": 256, + "num_experts_per_tok": 8, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "mamba_ssm_dtype": "float32", + "full_attention_interval": 4, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-6, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": False, + "tie_word_embeddings": False, +} + +_EXPECTED_ROPE_PARAMETERS = { + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], +} + + +def _expected_mlp_shapes(prefix: str) -> dict[str, tuple[int, ...]]: + return { + prefix + "mlp.gate.weight": (256, 2048), + prefix + "mlp.experts.gate_up_proj": (256, 1024, 2048), + prefix + "mlp.experts.down_proj": (256, 2048, 512), + prefix + "mlp.shared_expert.gate_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.up_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.down_proj.weight": (2048, 512), + prefix + "mlp.shared_expert_gate.weight": (1, 2048), + } + + +def _expected_attention_shapes( + prefix: str, layer_type: str) -> dict[str, tuple[int, ...]]: + if layer_type == "full_attention": + return { + prefix + "self_attn.q_proj.weight": (8192, 2048), + prefix + "self_attn.k_proj.weight": (512, 2048), + prefix + "self_attn.v_proj.weight": (512, 2048), + prefix + "self_attn.o_proj.weight": (2048, 4096), + prefix + "self_attn.q_norm.weight": (256,), + prefix + "self_attn.k_norm.weight": (256,), + } + return { + prefix + "linear_attn.in_proj_qkv.weight": (8192, 2048), + prefix + "linear_attn.in_proj_z.weight": (4096, 2048), + prefix + "linear_attn.in_proj_a.weight": (32, 2048), + prefix + "linear_attn.in_proj_b.weight": (32, 2048), + prefix + "linear_attn.conv1d.weight": (8192, 1, 4), + prefix + "linear_attn.A_log": (32,), + prefix + "linear_attn.dt_bias": (32,), + prefix + "linear_attn.norm.weight": (128,), + prefix + "linear_attn.out_proj.weight": (2048, 4096), + } + + +def _expected_text_shapes( + layer_types: tuple[str, ...]) -> dict[str, tuple[int, ...]]: + shapes = { + "lm_head.weight": (248320, 2048), + "model.language_model.embed_tokens.weight": (248320, 2048), + "model.language_model.norm.weight": (2048,), + } + for i, layer_type in enumerate(layer_types): + prefix = f"model.language_model.layers.{i}." + shapes[prefix + "input_layernorm.weight"] = (2048,) + shapes[prefix + "post_attention_layernorm.weight"] = (2048,) + shapes.update(_expected_mlp_shapes(prefix)) + shapes.update(_expected_attention_shapes(prefix, layer_type)) + return shapes + + +def _expected_mtp_shapes() -> dict[str, tuple[int, ...]]: + prefix = "mtp.layers.0." + shapes = { + "mtp.fc.weight": (2048, 4096), + "mtp.norm.weight": (2048,), + "mtp.pre_fc_norm_embedding.weight": (2048,), + "mtp.pre_fc_norm_hidden.weight": (2048,), + prefix + "input_layernorm.weight": (2048,), + prefix + "post_attention_layernorm.weight": (2048,), + } + shapes.update(_expected_attention_shapes(prefix, "full_attention")) + shapes.update(_expected_mlp_shapes(prefix)) + return shapes + + +_MTP_KEYS = set(_expected_mtp_shapes()) + + +def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: + return set(_expected_text_shapes(layer_types)) + + +def _read_tensor_shapes( + checkpoint_path: str, + weight_map: dict[str, str], + tensor_names: set[str], +) -> dict[str, tuple[int, ...]]: + from safetensors import safe_open + + shapes = {} + with ExitStack() as stack: + readers = { + shard: stack.enter_context( + safe_open( + os.path.join(checkpoint_path, shard), + framework="pt", + device="cpu", + ) + ) + for shard in set(weight_map[name] for name in tensor_names) + } + for name in tensor_names: + shapes[name] = tuple( + readers[weight_map[name]].get_slice(name).get_shape()) + return shapes + + +def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: + """Validate the official BF16 checkpoint before allocating GPU weights.""" + checkpoint_path = os.path.abspath(os.fspath(checkpoint_path)) + config_path = os.path.join(checkpoint_path, "config.json") + index_path = os.path.join( + checkpoint_path, "model.safetensors.index.json") + + for path in (config_path, index_path): + if not os.path.isfile(path): + raise FileNotFoundError( + f"Qwen3.6-35B-A3B checkpoint is missing {path!r}") + + with open(config_path, "r", encoding="utf-8") as f: + config = json.load(f) + if config.get("model_type") != "qwen3_5_moe": + raise ValueError( + "Qwen3.6-35B-A3B text frontend requires " + "model_type='qwen3_5_moe'; " + f"got {config.get('model_type')!r} in {config_path}") + + text_config = config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError(f"missing text_config object in {config_path}") + + mismatches = [] + for name, expected in _EXPECTED_TEXT_CONFIG.items(): + actual = text_config.get(name) + if actual != expected: + mismatches.append(f"{name}={actual!r} (expected {expected!r})") + rope_parameters = text_config.get("rope_parameters") + if not isinstance(rope_parameters, dict): + mismatches.append( + "rope_parameters is missing or is not an object") + else: + for name, expected in _EXPECTED_ROPE_PARAMETERS.items(): + actual = rope_parameters.get(name) + if actual != expected: + mismatches.append( + f"rope_parameters.{name}={actual!r} " + f"(expected {expected!r})") + layer_types = tuple(text_config.get("layer_types") or ()) + if layer_types != _EXPECTED_LAYER_TYPES: + mismatches.append( + "layer_types does not match the 30-linear/10-full attention " + "qwen3_5_moe schedule") + if mismatches: + raise ValueError( + "checkpoint is not compatible with the Qwen3.6-35B-A3B " + "SM120 text pipeline: " + "; ".join(mismatches)) + + with open(index_path, "r", encoding="utf-8") as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"missing weight_map object in {index_path}") + + required_text = _required_text_keys(layer_types) + missing_text = sorted(required_text.difference(weight_map)) + if missing_text: + preview = ", ".join(missing_text[:8]) + if len(missing_text) > 8: + preview += f", ... ({len(missing_text)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing text backbone tensors: " + + preview) + + missing_mtp = sorted(_MTP_KEYS.difference(weight_map)) + if missing_mtp: + preview = ", ".join(missing_mtp[:8]) + if len(missing_mtp) > 8: + preview += f", ... ({len(missing_mtp)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing its MTP tensor group: " + + preview) + + shards = sorted(set(weight_map.values())) + bad_shards = [] + for shard in shards: + path = os.path.join(checkpoint_path, shard) + if not os.path.isfile(path) or os.path.getsize(path) == 0: + bad_shards.append(shard) + if bad_shards: + preview = ", ".join(bad_shards[:8]) + if len(bad_shards) > 8: + preview += f", ... ({len(bad_shards)} missing or empty)" + raise FileNotFoundError( + "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " + + preview) + + expected_shapes = _expected_text_shapes(layer_types) + expected_shapes.update(_expected_mtp_shapes()) + actual_shapes = _read_tensor_shapes( + checkpoint_path, weight_map, set(expected_shapes)) + shape_mismatches = [ + f"{name}={actual_shapes[name]!r} (expected {expected!r})" + for name, expected in sorted(expected_shapes.items()) + if actual_shapes[name] != expected + ] + if shape_mismatches: + preview = "; ".join(shape_mismatches[:8]) + if len(shape_mismatches) > 8: + preview += f"; ... ({len(shape_mismatches)} mismatched)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint tensor shape mismatches: " + + preview) + + return { + "checkpoint_path": checkpoint_path, + "text_tensor_count": len(required_text), + "mtp_tensor_count": len(_MTP_KEYS), + "vision_tensor_count": sum( + ".visual." in name for name in weight_map), + "tensor_count": len(weight_map), + "shard_count": len(shards), + } + + +class Qwen36MoeTextFrontend(Nexn2TorchFrontendRtx): + """Qwen3.6-35B-A3B text-only frontend (RTX SM120 and Jetson AGX Thor).""" + + _MODEL_LABEL = "Qwen3.6-35B-A3B text" + _USAGE_DOC = "docs/qwen36_moe_usage.md" + + # The block-scaled 4-bit MMA tier is a build tier, and the prefill now picks + # its MoE tile from what the module actually has: without the tier it uses + # the weight-only grouped GEMV, which is slower on a long prompt and + # otherwise identical. So the tier is a performance requirement, not a + # correctness one, and demanding it here would refuse a build -- a Jetson + # one, whose toolchain has no block-scaled mma -- that runs this correctly. + _REQUIRED_KERNELS = tuple( + name for name in Nexn2TorchFrontendRtx._REQUIRED_KERNELS + if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) + ) + ('moe_grouped_w4a16_sm120_bf16',) + + # Same for FA2: prefill and decode both probe the vendored kernel and fall + # back to a reference, so a target that builds FA4 instead still runs. + _REQUIRE_FA2 = False + + def __init__(self, checkpoint_path: str, *, + device: str = "cuda:0", + max_seq: int = 2048, + quant: str = "nvfp4", + kernelized: bool = True, + quant_scope: str = "experts", + load_mtp: bool = False, + spec_graph_cache_max: int | None = None) -> None: + """Construct the frontend. + + Args beyond the shared ones (see + :class:`~flash_rt.frontends.torch.nexn2_rtx.Nexn2TorchFrontendRtx`): + + load_mtp: read the checkpoint's MTP draft head, which + :meth:`generate_spec` needs. Off by default: it is a full extra + transformer layer plus its KV, and ``generate`` never reads it. + spec_graph_cache_max: how many captured speculative windows to keep. + Each owns a CUDA graph memory pool, and a speculative window is + larger than a decode step, so this is bounded separately from the + decode graph cache and much lower. ``None`` takes the runtime + default, which is sized for the smallest board this path runs on. + """ + if quant != "nvfp4": + raise NotImplementedError( + f"quant={quant!r} is not implemented; only 'nvfp4' is " + "supported") + if not kernelized: + raise NotImplementedError( + "Qwen3.6-35B-A3B text only supports kernelized=True with " + "runtime NVFP4 conversion") + if spec_graph_cache_max is not None: + spec_graph_cache_max = int(spec_graph_cache_max) + if spec_graph_cache_max < 1: + raise ValueError( + "spec_graph_cache_max must be at least 1 (got " + f"{spec_graph_cache_max}); a speculative step captures a " + "graph per position, so a cache that holds none would " + "recapture every step") + # Read by the loader, before any weight is touched, through the base + # class -- hence set before super().__init__. + self._load_mtp = bool(load_mtp) + self._spec_graph_cache_max = spec_graph_cache_max + contract = validate_qwen36_moe_checkpoint(checkpoint_path) + super().__init__( + checkpoint_path, + device=device, + max_seq=max_seq, + quant=quant, + kernelized=kernelized, + quant_scope=quant_scope, + ) + self._checkpoint_contract = contract + + @property + def load_mtp(self) -> bool: + """Whether the MTP draft head was loaded (see ``generate_spec``).""" + return self._load_mtp + + def _decode_state_or_new(self): + from flash_rt.frontends.torch._nexn2_rtx_decode import Nexn2DecodeState + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device, + spec_graph_cache_max=self._spec_graph_cache_max) + return self._decode_state + + def generate_spec(self, max_new_tokens: int, *, k: int = 2): + """Greedy decode through draft-and-verify with the MTP head. + + Emits exactly what ``generate`` emits: a draft is kept only where the + model's own argmax agrees with it, so this is a speed change and + nothing else. Requires ``load_mtp=True`` at construction. + """ + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate_spec()") + if not self._load_mtp: + raise RuntimeError( + "speculative decoding needs the MTP draft head, which this " + "frontend did not load. Construct it with load_mtp=True.") + if int(k) < 1: + raise ValueError( + f"k must be at least 1 (got {k}); k is how many tokens the " + "draft head proposes per window") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + generate_greedy_spec, + ) + + return generate_greedy_spec( + self._decode_state_or_new(), self._prompt_ids, max_new_tokens, + int(k), self._fvk, self.device) + + def generate(self, max_new_tokens: int, *, do_sample: bool = False): + """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate()") + if do_sample: + raise NotImplementedError( + "the Qwen3.6-35B-A3B kernelized path supports greedy " + "decoding only") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + generate_greedy_graph, + ) + + return generate_greedy_graph( + self._decode_state_or_new(), + self._prompt_ids, + max_new_tokens, + self._fvk, + self.device, + ) + + +# The name this frontend shipped under while it was registered for RTX only. +# Kept so an existing import keeps working; new code should use the class +# above, which is what both architectures resolve to. +Qwen36MoeTextFrontendRtx = Qwen36MoeTextFrontend diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py index cf4b3e85..ffd0ed2a 100644 --- a/flash_rt/frontends/torch/qwen36_moe_rtx.py +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -1,333 +1,20 @@ -"""Qwen3.6-35B-A3B text inference on RTX SM120. +"""Compatibility import path for the Qwen3.6-35B-A3B text frontend. -The language backbone is the same ``qwen3_5_moe`` architecture used by -Nex-N2-mini, so this frontend reuses that implementation. The official -Qwen3.6 checkpoint also contains a vision tower and an MTP head; this entry -validates those weights but intentionally exposes only text prefill and greedy -decode. Vision and speculative decoding are separate integration surfaces. +The frontend moved to :mod:`flash_rt.frontends.torch.qwen36_moe` when it stopped +being RTX-only. This module re-exports the public names so an existing import +keeps working; new code should import from the module above. """ from __future__ import annotations -import json -import os -from contextlib import ExitStack -from typing import Any - -from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx - - -_EXPECTED_LAYER_TYPES = tuple( - "full_attention" if (i + 1) % 4 == 0 else "linear_attention" - for i in range(40) +from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, + Qwen36MoeTextFrontendRtx, + validate_qwen36_moe_checkpoint, ) -_EXPECTED_TEXT_CONFIG = { - "model_type": "qwen3_5_moe_text", - "num_hidden_layers": 40, - "hidden_size": 2048, - "vocab_size": 248320, - "num_attention_heads": 16, - "num_key_value_heads": 2, - "head_dim": 256, - "attention_bias": False, - "attn_output_gate": True, - "hidden_act": "silu", - "num_experts": 256, - "num_experts_per_tok": 8, - "moe_intermediate_size": 512, - "shared_expert_intermediate_size": 512, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_key_head_dim": 128, - "linear_value_head_dim": 128, - "linear_conv_kernel_dim": 4, - "mamba_ssm_dtype": "float32", - "full_attention_interval": 4, - "partial_rotary_factor": 0.25, - "rms_norm_eps": 1e-6, - "mtp_num_hidden_layers": 1, - "mtp_use_dedicated_embeddings": False, - "tie_word_embeddings": False, -} - -_EXPECTED_ROPE_PARAMETERS = { - "rope_type": "default", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25, - "mrope_interleaved": True, - "mrope_section": [11, 11, 10], -} - - -def _expected_mlp_shapes(prefix: str) -> dict[str, tuple[int, ...]]: - return { - prefix + "mlp.gate.weight": (256, 2048), - prefix + "mlp.experts.gate_up_proj": (256, 1024, 2048), - prefix + "mlp.experts.down_proj": (256, 2048, 512), - prefix + "mlp.shared_expert.gate_proj.weight": (512, 2048), - prefix + "mlp.shared_expert.up_proj.weight": (512, 2048), - prefix + "mlp.shared_expert.down_proj.weight": (2048, 512), - prefix + "mlp.shared_expert_gate.weight": (1, 2048), - } - - -def _expected_attention_shapes( - prefix: str, layer_type: str) -> dict[str, tuple[int, ...]]: - if layer_type == "full_attention": - return { - prefix + "self_attn.q_proj.weight": (8192, 2048), - prefix + "self_attn.k_proj.weight": (512, 2048), - prefix + "self_attn.v_proj.weight": (512, 2048), - prefix + "self_attn.o_proj.weight": (2048, 4096), - prefix + "self_attn.q_norm.weight": (256,), - prefix + "self_attn.k_norm.weight": (256,), - } - return { - prefix + "linear_attn.in_proj_qkv.weight": (8192, 2048), - prefix + "linear_attn.in_proj_z.weight": (4096, 2048), - prefix + "linear_attn.in_proj_a.weight": (32, 2048), - prefix + "linear_attn.in_proj_b.weight": (32, 2048), - prefix + "linear_attn.conv1d.weight": (8192, 1, 4), - prefix + "linear_attn.A_log": (32,), - prefix + "linear_attn.dt_bias": (32,), - prefix + "linear_attn.norm.weight": (128,), - prefix + "linear_attn.out_proj.weight": (2048, 4096), - } - - -def _expected_text_shapes( - layer_types: tuple[str, ...]) -> dict[str, tuple[int, ...]]: - shapes = { - "lm_head.weight": (248320, 2048), - "model.language_model.embed_tokens.weight": (248320, 2048), - "model.language_model.norm.weight": (2048,), - } - for i, layer_type in enumerate(layer_types): - prefix = f"model.language_model.layers.{i}." - shapes[prefix + "input_layernorm.weight"] = (2048,) - shapes[prefix + "post_attention_layernorm.weight"] = (2048,) - shapes.update(_expected_mlp_shapes(prefix)) - shapes.update(_expected_attention_shapes(prefix, layer_type)) - return shapes - - -def _expected_mtp_shapes() -> dict[str, tuple[int, ...]]: - prefix = "mtp.layers.0." - shapes = { - "mtp.fc.weight": (2048, 4096), - "mtp.norm.weight": (2048,), - "mtp.pre_fc_norm_embedding.weight": (2048,), - "mtp.pre_fc_norm_hidden.weight": (2048,), - prefix + "input_layernorm.weight": (2048,), - prefix + "post_attention_layernorm.weight": (2048,), - } - shapes.update(_expected_attention_shapes(prefix, "full_attention")) - shapes.update(_expected_mlp_shapes(prefix)) - return shapes - - -_MTP_KEYS = set(_expected_mtp_shapes()) - - -def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: - return set(_expected_text_shapes(layer_types)) - - -def _read_tensor_shapes( - checkpoint_path: str, - weight_map: dict[str, str], - tensor_names: set[str], -) -> dict[str, tuple[int, ...]]: - from safetensors import safe_open - - shapes = {} - with ExitStack() as stack: - readers = { - shard: stack.enter_context( - safe_open( - os.path.join(checkpoint_path, shard), - framework="pt", - device="cpu", - ) - ) - for shard in set(weight_map[name] for name in tensor_names) - } - for name in tensor_names: - shapes[name] = tuple( - readers[weight_map[name]].get_slice(name).get_shape()) - return shapes - - -def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: - """Validate the official BF16 checkpoint before allocating GPU weights.""" - checkpoint_path = os.path.abspath(os.fspath(checkpoint_path)) - config_path = os.path.join(checkpoint_path, "config.json") - index_path = os.path.join( - checkpoint_path, "model.safetensors.index.json") - - for path in (config_path, index_path): - if not os.path.isfile(path): - raise FileNotFoundError( - f"Qwen3.6-35B-A3B checkpoint is missing {path!r}") - - with open(config_path, "r", encoding="utf-8") as f: - config = json.load(f) - if config.get("model_type") != "qwen3_5_moe": - raise ValueError( - "Qwen3.6-35B-A3B text frontend requires " - "model_type='qwen3_5_moe'; " - f"got {config.get('model_type')!r} in {config_path}") - - text_config = config.get("text_config") - if not isinstance(text_config, dict): - raise ValueError(f"missing text_config object in {config_path}") - - mismatches = [] - for name, expected in _EXPECTED_TEXT_CONFIG.items(): - actual = text_config.get(name) - if actual != expected: - mismatches.append(f"{name}={actual!r} (expected {expected!r})") - rope_parameters = text_config.get("rope_parameters") - if not isinstance(rope_parameters, dict): - mismatches.append( - "rope_parameters is missing or is not an object") - else: - for name, expected in _EXPECTED_ROPE_PARAMETERS.items(): - actual = rope_parameters.get(name) - if actual != expected: - mismatches.append( - f"rope_parameters.{name}={actual!r} " - f"(expected {expected!r})") - layer_types = tuple(text_config.get("layer_types") or ()) - if layer_types != _EXPECTED_LAYER_TYPES: - mismatches.append( - "layer_types does not match the 30-linear/10-full attention " - "qwen3_5_moe schedule") - if mismatches: - raise ValueError( - "checkpoint is not compatible with the Qwen3.6-35B-A3B " - "SM120 text pipeline: " + "; ".join(mismatches)) - - with open(index_path, "r", encoding="utf-8") as f: - index = json.load(f) - weight_map = index.get("weight_map") - if not isinstance(weight_map, dict): - raise ValueError(f"missing weight_map object in {index_path}") - - required_text = _required_text_keys(layer_types) - missing_text = sorted(required_text.difference(weight_map)) - if missing_text: - preview = ", ".join(missing_text[:8]) - if len(missing_text) > 8: - preview += f", ... ({len(missing_text)} missing)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint is missing text backbone tensors: " - + preview) - - missing_mtp = sorted(_MTP_KEYS.difference(weight_map)) - if missing_mtp: - preview = ", ".join(missing_mtp[:8]) - if len(missing_mtp) > 8: - preview += f", ... ({len(missing_mtp)} missing)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint is missing its MTP tensor group: " - + preview) - - shards = sorted(set(weight_map.values())) - bad_shards = [] - for shard in shards: - path = os.path.join(checkpoint_path, shard) - if not os.path.isfile(path) or os.path.getsize(path) == 0: - bad_shards.append(shard) - if bad_shards: - preview = ", ".join(bad_shards[:8]) - if len(bad_shards) > 8: - preview += f", ... ({len(bad_shards)} missing or empty)" - raise FileNotFoundError( - "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " - + preview) - - expected_shapes = _expected_text_shapes(layer_types) - expected_shapes.update(_expected_mtp_shapes()) - actual_shapes = _read_tensor_shapes( - checkpoint_path, weight_map, set(expected_shapes)) - shape_mismatches = [ - f"{name}={actual_shapes[name]!r} (expected {expected!r})" - for name, expected in sorted(expected_shapes.items()) - if actual_shapes[name] != expected - ] - if shape_mismatches: - preview = "; ".join(shape_mismatches[:8]) - if len(shape_mismatches) > 8: - preview += f"; ... ({len(shape_mismatches)} mismatched)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint tensor shape mismatches: " - + preview) - - return { - "checkpoint_path": checkpoint_path, - "text_tensor_count": len(required_text), - "mtp_tensor_count": len(_MTP_KEYS), - "vision_tensor_count": sum( - ".visual." in name for name in weight_map), - "tensor_count": len(weight_map), - "shard_count": len(shards), - } - - -class Qwen36MoeTextFrontendRtx(Nexn2TorchFrontendRtx): - """Qwen3.6-35B-A3B text-only frontend for RTX SM120.""" - - _MODEL_LABEL = "Qwen3.6-35B-A3B text" - _USAGE_DOC = "docs/qwen36_moe_usage.md" - - def __init__(self, checkpoint_path: str, *, - device: str = "cuda:0", - max_seq: int = 2048, - quant: str = "nvfp4", - kernelized: bool = True, - quant_scope: str = "experts") -> None: - if quant != "nvfp4": - raise NotImplementedError( - f"quant={quant!r} is not implemented; only 'nvfp4' is " - "supported") - if not kernelized: - raise NotImplementedError( - "Qwen3.6-35B-A3B text only supports kernelized=True with " - "runtime NVFP4 conversion") - contract = validate_qwen36_moe_checkpoint(checkpoint_path) - super().__init__( - checkpoint_path, - device=device, - max_seq=max_seq, - quant=quant, - kernelized=kernelized, - quant_scope=quant_scope, - ) - self._checkpoint_contract = contract - - def generate(self, max_new_tokens: int, *, do_sample: bool = False): - """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" - if self._prompt_ids is None: - raise ValueError("call set_prompt(...) before generate()") - if do_sample: - raise NotImplementedError( - "the Qwen3.6-35B-A3B kernelized path supports greedy " - "decoding only") - - from flash_rt.frontends.torch._nexn2_rtx_decode import ( - Nexn2DecodeState, - generate_greedy_graph, - ) - - if self._decode_state is None: - self._decode_state = Nexn2DecodeState( - self._weights, self._user_max_seq, self.device) - return generate_greedy_graph( - self._decode_state, - self._prompt_ids, - max_new_tokens, - self._fvk, - self.device, - ) +__all__ = [ + "Qwen36MoeTextFrontend", + "Qwen36MoeTextFrontendRtx", + "validate_qwen36_moe_checkpoint", +] diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 5d231ead..4048dda5 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -165,17 +165,26 @@ def detect_arch() -> str: # ── Nex-N2-mini / Qwen3.6-35B-A3B (qwen3_5_moe) ── # Text LLM, not a VLA: GDN linear-attn + full-attn-every-4th + 256-expert - # NVFP4 MoE. RTX 5090 (SM120) only, and requires the gated kernel build - # (-DFLASHRT_ENABLE_QWEN35MOE=ON). Registered here for discoverability / - # resolve_pipeline_class, but the frontend exposes an LLM surface - # (infer()->logits, generate_greedy) rather than the VLA predict(images) - # API, so these are used via direct frontend construction rather than - # load_model's VLAModel wrapper. + # NVFP4 MoE. Registered here for discoverability / resolve_pipeline_class, + # but the frontend exposes an LLM surface (infer()->logits, + # generate_greedy) rather than the VLA predict(images) API, so these are + # used via direct frontend construction rather than load_model's VLAModel + # wrapper. + # + # Nex-N2 is RTX 5090 (SM120) and needs the full gated kernel build + # (-DFLASHRT_ENABLE_QWEN35MOE=ON). ("nexn2", "torch", "rtx_sm120"): ("flash_rt.frontends.torch.nexn2_rtx", "Nexn2TorchFrontendRtx"), + # Qwen3.6 runs the same frontend on RTX SM120 and on Jetson AGX Thor + # (SM110). The two differ only in which kernel tiers the build has: + # SM120 takes the whole switch, Thor takes the two tiers its toolchain can + # compile. See docs/qwen36_moe_usage.md for the exact command per target. ("qwen36_moe", "torch", "rtx_sm120"): - ("flash_rt.frontends.torch.qwen36_moe_rtx", - "Qwen36MoeTextFrontendRtx"), + ("flash_rt.frontends.torch.qwen36_moe", + "Qwen36MoeTextFrontend"), + ("qwen36_moe", "torch", "thor"): + ("flash_rt.frontends.torch.qwen36_moe", + "Qwen36MoeTextFrontend"), # ── Pi0-FAST ── (SM120 runtime fork inside pipeline, no AttentionBackend protocol.) ("pi0fast", "torch", "thor"): diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 863bea4d..3c24f354 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -46,12 +46,20 @@ class RtxFlashAttnBackendNexn2: NUM_KV_HEADS = 2 HEAD_DIM = 256 - def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None): + def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None, + num_full_layers: int | None = None, + use_fa2: bool | None = None): import torch self._torch = torch bf16 = dtype if dtype is not None else torch.bfloat16 d = "cuda" + # The model's ten full-attention layers by default. A speculative + # draft head is one more full-attention layer carrying its own KV, so + # it asks for an extra slot rather than owning a second cache. + self.NUM_FULL_LAYERS = ( + int(num_full_layers) if num_full_layers is not None + else type(self).NUM_FULL_LAYERS) self._max_seq = int(max_seq) self._max_q_seq = int(max_q_seq) @@ -95,12 +103,129 @@ def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None): dtype=torch.float32, device=d, ) - from flash_rt import flash_rt_fa2 as _fa2 - self._fa2 = _fa2 - self._fa2_fwd = _fa2.fwd_bf16 + # A target may not build FA2 at all. Absence is a fallback, not an + # error, because the reference path below computes the same thing. + # + # Prefill and decode want different answers about FA2. Prefill gains + # 20x on a chunked block's non-square window. At the decode shape the + # two are the same answer to bf16 precision -- measured against an + # fp32 reference, 2.0e-3 relative for both at kv=64, 2.2e-3 against + # 2.1e-3 at kv=2048 -- so taking it there buys about 1% of a step and + # moves two of the sixteen golden tokens, because a bf16-level + # difference in one step flips a later one. That is the fixture losing + # its meaning in exchange for 1%, which is the wrong trade. + # + # So the caller says. ``use_fa2=None`` keeps what each target already + # validated: on the arch that has always had FA2 in decode, keep it; + # on sm_110, where FA2 has only just started building and the fixture + # was recorded through the reference path, decline it. + # FLASHRT_NEXN2_DECODE_FA2 overrides either way. + import os as _os + if use_fa2 is None: + _cap = torch.cuda.get_device_capability() + _default = "0" if _cap == (11, 0) else "1" + want_fa2 = _os.environ.get( + "FLASHRT_NEXN2_DECODE_FA2", _default) != "0" + else: + want_fa2 = bool(use_fa2) + try: + from flash_rt import flash_rt_fa2 as _fa2 + except ImportError: + self._fa2 = None + self._fa2_fwd = None + else: + self._fa2 = _fa2 if want_fa2 else None + self._fa2_fwd = _fa2.fwd_bf16 if want_fa2 else None self._num_sms = torch.cuda.get_device_properties( torch.cuda.current_device() ).multi_processor_count + self._fa2_usable = ( + self._fa2_fwd is not None and self._probe_fa2()) + + def _probe_fa2(self) -> bool: + """Does the vendored kernel actually compute on this device? + + Importing it and finding its symbols proves neither. Its own arch + handling can leave a build that links, loads, prints a complaint to + stdout and returns without writing the output -- which downstream looks + like plausible-but-wrong attention rather than a failure. Measured on + an SM110 part: the module imported, every symbol was present, and the + kernel refused at run time. + + So run one small case against a reference and compare. The cost is one + launch at construction. + """ + # Imported here, not at module scope, for the same reason as F: this + # module is written to import without torch present. The body had + # never run on a target that builds FA2, so the missing name sat + # unnoticed until this arch started building one. + import torch + import torch.nn.functional as F + + q_seq, kv_seq = 1, 8 + generator = torch.Generator(device=self.Q_buf.device).manual_seed(1) + q = torch.randn( + 1, q_seq, self.NUM_Q_HEADS, self.HEAD_DIM, generator=generator, + device=self.Q_buf.device, dtype=torch.bfloat16) + k = torch.randn( + 1, kv_seq, self.NUM_KV_HEADS, self.HEAD_DIM, generator=generator, + device=self.Q_buf.device, dtype=torch.bfloat16) + v = torch.randn_like(k) + self.Q_buf[:, :q_seq].copy_(q) + self.K_cache[0:1, :kv_seq].copy_(k) + self.V_cache[0:1, :kv_seq].copy_(v) + self.O_buf[:, :q_seq].zero_() + try: + self._launch_fa2(0, q_seq, kv_seq, 0, + 1.0 / (self.HEAD_DIM ** 0.5)) + torch.cuda.synchronize() + except Exception: # noqa: BLE001 + return False + produced = self.O_buf[:, :q_seq].float().clone() + + groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS + kr = k.repeat_interleave(groups, dim=2) + vr = v.repeat_interleave(groups, dim=2) + expected = F.scaled_dot_product_attention( + q.transpose(1, 2).float(), kr.transpose(1, 2).float(), + vr.transpose(1, 2).float()).transpose(1, 2) + if not torch.isfinite(produced).all(): + return False + reference = expected.norm().clamp_min(1e-6) + return bool( + ((produced - expected).norm() / reference).item() < 0.05) + + def _sdpa(self, layer_idx: int, q_seq: int, kv_seq: int, + softmax_scale: float) -> None: + """Reference attention, for a device the vendored kernel refuses.""" + import torch.nn.functional as F + + # BF16 throughout. The cache is bf16, so upcasting it materialises the + # whole history in fp32 every step -- hundreds of MB of temporaries per + # token at a long context, which is what made decode fall from 26 to 11 + # tok/s between 4k and 10k. SDPA accumulates in fp32 regardless. + q = self.Q_buf[:, :q_seq].transpose(1, 2) + k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] + v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] + # Broadcasting the KV to the query head count materialises it: at + # decode that is 8x2xkv_seqx256 floats twice per layer, ~48 MB a step + # across the ten full-attention layers, purely to be read once. Native + # GQA does the same thing without the copy. fp32 either way, so the + # numerics are untouched -- this path seeds a token-exact decode. + groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS + try: + out = F.scaled_dot_product_attention( + q, k.transpose(1, 2), v.transpose(1, 2), + is_causal=q_seq > 1, scale=softmax_scale, enable_gqa=True, + ).transpose(1, 2) + except TypeError: # torch without native GQA + out = F.scaled_dot_product_attention( + q, + k.repeat_interleave(groups, dim=2).transpose(1, 2), + v.repeat_interleave(groups, dim=2).transpose(1, 2), + is_causal=q_seq > 1, scale=softmax_scale, + ).transpose(1, 2) + self.O_buf[:, :q_seq].copy_(out.to(self.O_buf.dtype)) # ── Layer cache pointer math ── @@ -175,6 +300,21 @@ def run(self, site: str, layer_idx: int, q_seq: int, if softmax_scale is None: softmax_scale = 1.0 / (self.HEAD_DIM ** 0.5) + if not self._fa2_usable: + self._sdpa(layer_idx, q_seq, kv_seq, softmax_scale) + return o.data_ptr() + + self._launch_fa2(layer_idx, q_seq, kv_seq, stream, softmax_scale) + return o.data_ptr() + + def _launch_fa2(self, layer_idx: int, q_seq: int, kv_seq: int, + stream: int, softmax_scale: float) -> None: + """One vendored-FA2 launch. Shared with the construction-time probe so + the probe exercises the same call the hot path makes.""" + q = self.Q_buf[:, :q_seq] + k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] + v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] + o = self.O_buf[:, :q_seq] self._fa2_fwd( Q=q.data_ptr(), K=k.data_ptr(), V=v.data_ptr(), O=o.data_ptr(), softmax_lse=self.lse_buf.data_ptr(), @@ -192,7 +332,6 @@ def run(self, site: str, layer_idx: int, q_seq: int, num_sms=self._num_sms, stream=stream, ) - return o.data_ptr() def make_nexn2_attention_spec(*, max_seq: int, max_q_seq: int = 1) -> dict: diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md new file mode 100644 index 00000000..ca1ae248 --- /dev/null +++ b/qwen36_moe_edge/README.md @@ -0,0 +1,153 @@ +# Qwen3.6-MoE edge experiments + +This directory contains checkpoint-independent development utilities for a +memory-constrained Qwen3.6-35B-A3B runtime. It is not a production frontend. + +The intended runtime layout follows the MiniMax-M3 Spark prototype: + +- non-routed weights remain resident in a mixed-precision format; +- each routed expert is stored as one fixed-size block; +- a bounded per-layer LRU holds hot expert blocks; +- misses are read from local storage into reusable staging buffers. + +Inspect projected checkpoint sizes and sampled expert quality: + +```bash +PYTHONPATH=. python qwen36_moe_edge/probe.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --mode memory \ + --group-size 16 + +PYTHONPATH=. python qwen36_moe_edge/probe.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --mode quality \ + --group-size 16 +``` + +Check that the gated kernels compute the same thing on another architecture. +Compiling for a target says nothing about what it computes there, so record a +reference where the kernels are known good and replay it elsewhere: + +```bash +# On a known-good target: +PYTHONPATH=. python qwen36_moe_edge/kernel_parity.py \ + --output parity_sm120.json + +# On the target under test, with the reference alongside it: +PYTHONPATH=. python qwen36_moe_edge/kernel_parity.py \ + --output parity_sm110.json \ + --reference parity_sm120.json +``` + +No checkpoint is involved: shapes come from the Qwen3.6 geometry and inputs +from a fixed generator. Inputs are stored in the reference and replayed rather +than regenerated, because CUDA RNG is not bit-reproducible across +architectures — the Philox thread mapping follows occupancy, so regenerating on +the target compares kernels on different data and reads as a kernel failure. +Divergence appears only past the first launch block, which is why small tensors +appear to agree and large ones do not. + +Each case also checks its kernel against a Torch expression on the local +device, so a genuine kernel fault is distinguishable from a harness or input +problem: a broken kernel fails its local check first. + +Score the quantization schemes against the activations the router actually +sends each expert, and optionally save the references a device can check +itself against: + +```bash +PYTHONPATH=. python qwen36_moe_edge/expert_quality.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --prompt "Explain edge mixture-of-experts inference. " \ + --prompt-tokens 32 \ + --new-tokens 32 \ + --output qwen36_expert_quality.json \ + --golden qwen36_expert_golden.safetensors +``` + +Prefer this over `probe.py --mode quality` when deciding what to generate. +The probe uses `torch.randn` activations and quantizes the activations too; +neither matches the runtime, where the activation is 4 KiB against a 1.7 MiB +weight block and so is left in BF16. Random activations also hide errors that +real inputs expose, because a scale calibrated against noise is not the scale +real inputs need. + +Generate fixed-size routed-expert blocks for a layer range: + +```bash +PYTHONPATH=. python qwen36_moe_edge/quantize_experts.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --output /models/Qwen3.6-35B-A3B-INT8E \ + --format int8 \ + --layers 0:40 + +PYTHONPATH=. python qwen36_moe_edge/quantize_experts.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --output /models/Qwen3.6-35B-A3B-INT4E-RHT16 \ + --format int4-rht \ + --group-size 16 \ + --layers 0:40 +``` + +INT8 uses symmetric per-output-channel FP16 scales. INT4 follows the Thor +Pi0.5 numerical contract: sign-magnitude values, one UE4M3 scale per 16 K +values, and two values per byte with the low nibble first. `int4-rht` applies +the same orthonormal H16/4 transform to every K block that the runtime applies +to activations. Scale bytes in these edge block files are linear; a loader +must convert them to the SM1xx SFB tile-interleaved layout before calling the +native block-scaled MMA kernels. + +Each block carries a trailing pad so its offset and length are multiples of +`BLOCK_ALIGNMENT`. On a device whose memory holds only a fraction of the +experts, the expert stream cannot go through the page cache — it would compete +with the resident weights for the same physical memory — so the reader has to +use `O_DIRECT`, which requires aligned offsets and lengths. The INT4 group-16 +payload is already a multiple of 4096; the INT8 payload is 3,151,872 bytes and +takes 2048 bytes of pad. `manifest.json` records `block_bytes`, +`block_alignment`, and the padding entry in `block_sizes`. + +An SM120 machine can collect real router selections for cache sizing: + +```bash +PYTHONPATH=. python qwen36_moe_edge/route_trace.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --prompt "Explain edge mixture-of-experts inference. " \ + --prompt-tokens 32 \ + --new-tokens 64 \ + --quotas 16,27,32,43,64 \ + --output qwen36_moe_route_trace.json +``` + +Tracing deliberately uses eager per-token prefill. It must not be enabled +during CUDA Graph capture. + +Each quota is scored under three policies: + +- `single_lru` — one per-layer LRU behind both prefill and decode. Prefill + touches every expert in a layer, so this measures what survives prompt + churn. +- `two_tier` — a per-layer warm set pinned from prompt-phase selection counts + plus an evictable ring sized by `--stream-fraction`. Prefill cannot displace + the warm set. +- `two_tier_oracle_warm` — the same split with the warm set chosen from the + decode phase. Not implementable; it bounds what a better warm-set heuristic + could add. + +On Qwen3.6-35B-A3B the plain LRU wins from 16 slots per layer up, and its +margin grows with prompt length — at 43 slots per layer, 0.745 against 0.731 +for a 32-token prompt and 0.711 against 0.664 for a 128-token prompt. Once a +per-layer quota exists, recency predicts this router's next selections better +than prompt-phase frequency, and a longer prompt makes the frequency estimate +more diffuse rather than more reliable. The oracle variant stays ahead of both +(0.776 at 43 slots for the 128-token prompt), so pinning is sound and the +prompt-derived choice of what to pin is what falls short. Treat the policy as +something to measure per checkpoint, not to assume. + +Capacity dominates policy either way: going from 43 to 64 slots per layer cuts +read volume by a third, while any policy change at a fixed quota moves it by a +few percent. + +`--block-bytes` (default: the INT4 group-16 block) and `--bandwidths` turn +misses per token into a read volume and the token rate each storage bandwidth +would allow, which is the number that decides whether a memory budget is +viable. diff --git a/qwen36_moe_edge/__init__.py b/qwen36_moe_edge/__init__.py new file mode 100644 index 00000000..b423fee7 --- /dev/null +++ b/qwen36_moe_edge/__init__.py @@ -0,0 +1 @@ +"""Development utilities for memory-constrained Qwen3.6-MoE inference.""" diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py new file mode 100644 index 00000000..a9c8f9f8 --- /dev/null +++ b/qwen36_moe_edge/expert_cache.py @@ -0,0 +1,427 @@ +#!/usr/bin/env python3 +"""A bounded, streaming cache for routed-expert blocks. + +The target device holds only a fraction of the experts, so the cache is the +runtime: what it can hold and how fast it can refill decide both token rate and +time to first token. Three properties follow from that and are enforced here +rather than left to convention. + +**The budget is a hard limit, not a projection.** On a unified-memory device the +weights, the cache, the staging buffers and the operating system all draw on the +same physical memory, so a runtime that merely intends to stay small is not +measurable. Construction computes its own footprint and refuses to allocate if +it would exceed the budget. + +**Reads bypass the page cache.** Streaming tens of GiB of blocks through the +page cache would make it compete with the resident weights for that same +memory. Reads use ``O_DIRECT``, which is why the bundle pads each block to a +4096-byte boundary. + +**Misses are fetched concurrently.** A single reader leaves a large part of an +NVMe device idle; measured on one, four readers were worth 1.7x over one and +eight saturated it. ``get_many`` issues a layer's misses together. + +A per-layer quota of at least ``num_experts_per_token`` also makes a class of +bug structurally impossible: the experts one token needs cannot evict each +other, so a caller may hold several pointers from the same ``get_many`` at once. +""" + +from __future__ import annotations + +import json +import os +import queue +from collections import Counter, OrderedDict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path + +import torch + + +@dataclass +class CacheConfig: + """Sizing and placement. ``budget_bytes`` of 0 disables the check.""" + + bundle: Path + slots_per_layer: int + staging_buffers: int = 4 + budget_bytes: int = 0 + reserve_bytes: int = 0 + resident_bytes: int = 0 + device: str = "cuda:0" + experts_per_token: int = 8 + read_chunk: int = 1 << 28 + # Bypass the page cache. Off only to demonstrate what happens when it is + # not bypassed, or on a filesystem without O_DIRECT. + direct: bool = True + metadata: dict = field(default_factory=dict) + + +class CacheBudgetError(RuntimeError): + """The requested cache does not fit the budget it was given.""" + + +def _load_manifest(bundle: Path) -> dict: + path = bundle / "manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"expert bundle is missing {path}") + with path.open(encoding="utf-8") as f: + manifest = json.load(f) + for key in ("block_bytes", "block_alignment", "num_layers", + "num_experts", "block_sizes"): + if key not in manifest: + raise ValueError(f"{path} has no {key!r}") + block_bytes = int(manifest["block_bytes"]) + alignment = int(manifest["block_alignment"]) + if block_bytes % alignment: + raise ValueError( + f"{path}: block_bytes {block_bytes} is not a multiple of " + f"{alignment}, so direct reads of it are impossible") + return manifest + + +class ExpertCache: + """Per-layer LRU over fixed-size blocks read straight from storage.""" + + def __init__(self, config: CacheConfig): + self.config = config + self.manifest = _load_manifest(config.bundle) + self.block_bytes = int(self.manifest["block_bytes"]) + self.alignment = int(self.manifest["block_alignment"]) + self.num_layers = int(self.manifest["num_layers"]) + self.num_experts = int(self.manifest["num_experts"]) + # None in the manifest means INT8, whose scales are per output channel + # and need no group. + self.group_size = int(self.manifest.get("group_size") or 0) + + if config.slots_per_layer < config.experts_per_token: + raise ValueError( + f"slots_per_layer={config.slots_per_layer} is below " + f"experts_per_token={config.experts_per_token}; one token's " + "experts would evict each other and a caller could not hold " + "their pointers at once") + + self.footprint = self.plan(config, self.manifest) + if config.budget_bytes: + total = self.footprint["projected_bytes"] + if total > config.budget_bytes: + raise CacheBudgetError( + f"cache needs {total / 2**30:.3f} GiB " + f"(slots {self.footprint['slot_bytes'] / 2**30:.3f} + " + f"staging {self.footprint['staging_bytes'] / 2**30:.3f} + " + f"resident {config.resident_bytes / 2**30:.3f} + " + f"reserve {config.reserve_bytes / 2**30:.3f}) but the " + f"budget is {config.budget_bytes / 2**30:.3f} GiB. " + f"Reduce slots_per_layer below " + f"{self.max_slots_per_layer(config, self.manifest)}.") + + self._total_slots = config.slots_per_layer * self.num_layers + self.slots = torch.empty( + self._total_slots, self.block_bytes, + dtype=torch.uint8, device=config.device) + self._staging = [ + torch.empty(self.block_bytes, dtype=torch.uint8).pin_memory() + for _ in range(config.staging_buffers) + ] + for buffer in self._staging: + if buffer.data_ptr() % self.alignment: + raise RuntimeError( + "a pinned staging buffer is not " + f"{self.alignment}-byte aligned, which direct reads " + "require") + self._pool = ThreadPoolExecutor(max_workers=config.staging_buffers) + # Buffers are taken from here and returned, so a task owns one for as + # long as it runs. Indexing by task number would let task N and task + # N + len(staging) share a buffer: the pool bounds how many run at once, + # not the order they finish in. + self._available: queue.Queue = queue.Queue() + for index in range(len(self._staging)): + self._available.put(index) + + # Per-layer LRU of expert -> slot index, and that layer's free slots. + self._lru: list[OrderedDict[int, int]] = [ + OrderedDict() for _ in range(self.num_layers)] + self._free: list[list[int]] = [ + list(range(layer * config.slots_per_layer, + (layer + 1) * config.slots_per_layer)) + for layer in range(self.num_layers) + ] + self._fds: dict[int, int] = {} + self._global_scales: dict[int, torch.Tensor] = {} + self.hits = 0 + self.misses = 0 + self.bytes_read = 0 + + # ── sizing, answerable before anything is allocated ── + + @staticmethod + def plan(config: CacheConfig, manifest: dict) -> dict[str, int]: + """What this configuration would occupy.""" + block_bytes = int(manifest["block_bytes"]) + slot_bytes = ( + config.slots_per_layer * int(manifest["num_layers"]) * block_bytes) + staging_bytes = config.staging_buffers * block_bytes + return { + "block_bytes": block_bytes, + "slots_per_layer": config.slots_per_layer, + "slot_bytes": slot_bytes, + "staging_bytes": staging_bytes, + "resident_bytes": config.resident_bytes, + "reserve_bytes": config.reserve_bytes, + "projected_bytes": ( + slot_bytes + staging_bytes + + config.resident_bytes + config.reserve_bytes), + } + + @staticmethod + def max_slots_per_layer(config: CacheConfig, manifest: dict) -> int: + """Largest per-layer quota that fits ``config.budget_bytes``.""" + if not config.budget_bytes: + return int(manifest["num_experts"]) + block_bytes = int(manifest["block_bytes"]) + available = ( + config.budget_bytes - config.resident_bytes + - config.reserve_bytes - config.staging_buffers * block_bytes) + if available <= 0: + return 0 + return min( + int(manifest["num_experts"]), + available // (block_bytes * int(manifest["num_layers"]))) + + # ── reading ── + + def _fd(self, layer: int) -> int: + if layer not in self._fds: + path = self.config.bundle / f"experts_layer_{layer:02d}.bin" + expected = self.num_experts * self.block_bytes + size = path.stat().st_size + if size != expected: + raise ValueError( + f"{path} is {size} bytes; expected {expected} " + f"({self.num_experts} x {self.block_bytes})") + flags = os.O_RDONLY + if self.config.direct: + flags |= getattr(os, "O_DIRECT", 0) + self._fds[layer] = os.open(path, flags) + return self._fds[layer] + + def _fetch(self, layer: int, expert: int, slot: int) -> None: + if not 0 <= expert < self.num_experts: + raise ValueError( + f"expert {expert} is outside 0..{self.num_experts - 1}; a " + "negative or oversized index becomes an invalid file offset") + buffer = self._available.get() + try: + staging = self._staging[buffer] + view = memoryview(staging.numpy()) + fd = self._fd(layer) + base = expert * self.block_bytes + offset = 0 + while offset < self.block_bytes: + length = min(self.config.read_chunk, self.block_bytes - offset) + try: + read = os.preadv( + fd, [view[offset:offset + length]], base + offset) + except OSError as error: + # A direct read rejects a misaligned offset, length or + # buffer with the same EINVAL, which says nothing about + # which of the three it was. + raise OSError( + f"{error.strerror} reading layer {layer} expert " + f"{expert}: offset {base + offset} aligned=" + f"{(base + offset) % self.alignment == 0}, length " + f"{length} aligned={length % self.alignment == 0}, " + f"buffer {staging.data_ptr():#x} aligned=" + f"{staging.data_ptr() % self.alignment == 0}") from error + if read <= 0: + raise IOError( + f"short read of layer {layer} expert {expert} at " + f"{offset}/{self.block_bytes}") + offset += read + self.slots[slot].copy_(staging) + self.bytes_read += self.block_bytes + finally: + self._available.put(buffer) + + def _claim(self, layer: int, expert: int) -> int: + """A slot for an expert not currently held, evicting if necessary.""" + free = self._free[layer] + if free: + slot = free.pop() + else: + slot = self._lru[layer].popitem(last=False)[1] + self._lru[layer][expert] = slot + return slot + + def get_many(self, layer: int, experts) -> list[int]: + """Device pointers for several experts of one layer, misses in parallel. + + With ``slots_per_layer >= experts_per_token`` none of the returned + pointers can be invalidated by the others. + """ + wanted = list(dict.fromkeys(int(expert) for expert in experts)) + out_of_range = [ + expert for expert in wanted + if not 0 <= expert < self.num_experts + ] + if out_of_range: + raise ValueError( + f"layer {layer} was asked for experts {out_of_range} outside " + f"0..{self.num_experts - 1}; the full request was {wanted}") + if len(wanted) > self.config.slots_per_layer: + raise ValueError( + f"asked for {len(wanted)} experts of layer {layer} but the " + f"quota is {self.config.slots_per_layer}") + pending = [] + for expert in wanted: + slot = self._lru[layer].get(expert) + if slot is not None: + self._lru[layer].move_to_end(expert) + self.hits += 1 + continue + self.misses += 1 + pending.append((expert, self._claim(layer, expert))) + if pending: + futures = [ + self._pool.submit(self._fetch, layer, expert, slot) + for expert, slot in pending + ] + for future in futures: + future.result() + torch.cuda.synchronize(self.config.device) + return [ + int(self.slots[self._lru[layer][expert]].data_ptr()) + for expert in wanted + ] + + def get(self, layer: int, expert: int) -> int: + return self.get_many(layer, (expert,))[0] + + def components(self, layer: int, expert: int) -> dict[str, torch.Tensor]: + """The block's four parts as views over its slot, plus its scales. + + Views, not copies: the caller reads them where the block already lies. + The manifest's ``block_layout`` gives the order, so a consumer never + reproduces the offset arithmetic and cannot drift from the writer. + """ + self.get(layer, expert) + raw = self.slots[self._lru[layer][expert]] + sizes = self.manifest["block_sizes"] + offset = 0 + parts = {} + for name in self.manifest["block_layout"]: + length = int(sizes[name]) + if name != "padding": + parts[name] = raw[offset:offset + length] + offset += length + parts["global_scales"] = self.global_scales(layer)[expert] + return parts + + def global_scales(self, layer: int) -> torch.Tensor: + """This layer's per-expert (gate_up, down) scales, read once. + + They live beside the blocks rather than inside them so a block stays + exactly ``block_bytes`` and aligned; the kernel takes them as its GEMM + alpha. + """ + cached = self._global_scales.get(layer) + if cached is None: + name = self.manifest.get( + "global_scales", "global_scales_layer_NN.bin") + path = self.config.bundle / name.replace( + "NN", f"{layer:02d}") + expected = self.num_experts * 2 * 4 + size = path.stat().st_size + if size != expected: + raise ValueError( + f"{path} is {size} bytes; expected {expected} " + f"({self.num_experts} experts x 2 x float32)") + cached = torch.frombuffer( + bytearray(path.read_bytes()), dtype=torch.float32 + ).view(self.num_experts, 2) + self._global_scales[layer] = cached + return cached + + # ── startup ── + + def warm(self, frequency: list[Counter]) -> int: + """Preload each layer's most frequently selected experts. + + Entries stay evictable. Measured on held-out prompts, a set built from + unrelated traffic removes about a quarter of decode misses and half of + the cold prefill read at this quota; pinning it instead costs more + adaptivity than it gains. + """ + if len(frequency) != self.num_layers: + raise ValueError( + f"frequency has {len(frequency)} layers, expected " + f"{self.num_layers}") + loaded = 0 + for layer in range(self.num_layers): + experts = [ + expert for expert, _ in + frequency[layer].most_common(self.config.slots_per_layer) + ] + for start in range(0, len(experts), self.config.slots_per_layer): + chunk = experts[start:start + self.config.slots_per_layer] + self.get_many(layer, chunk) + loaded += len(chunk) + self.hits = 0 + self.misses = 0 + self.bytes_read = 0 + return loaded + + # ── reporting ── + + def stats(self) -> dict[str, float]: + requests = self.hits + self.misses + report = dict(self.footprint) + report.update({ + "hits": self.hits, + "misses": self.misses, + "hit_rate": self.hits / requests if requests else 0.0, + "bytes_read": self.bytes_read, + "resident_experts": sum(len(lru) for lru in self._lru), + }) + if torch.cuda.is_available(): + free, total = torch.cuda.mem_get_info(self.config.device) + report.update({ + "device_free_bytes": free, + "device_total_bytes": total, + "torch_allocated_bytes": torch.cuda.memory_allocated( + self.config.device), + "torch_peak_allocated_bytes": torch.cuda.max_memory_allocated( + self.config.device), + "torch_reserved_bytes": torch.cuda.memory_reserved( + self.config.device), + "torch_peak_reserved_bytes": torch.cuda.max_memory_reserved( + self.config.device), + }) + return report + + def close(self) -> None: + """Release everything, including the slots. + + The slot array is the largest single allocation a runtime makes, so a + close that only dropped file descriptors would leak the entire cache + on any reconfiguration -- on a device where the budget is the whole + point, that is not a detail. After this the cache is unusable. + """ + self._pool.shutdown(wait=True) + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + for lru in self._lru: + lru.clear() + self._free = [[] for _ in range(self.num_layers)] + self.slots = None + self._staging = [] + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def __enter__(self) -> "ExpertCache": + return self + + def __exit__(self, *_) -> None: + self.close() diff --git a/qwen36_moe_edge/expert_quality.py b/qwen36_moe_edge/expert_quality.py new file mode 100644 index 00000000..987f400f --- /dev/null +++ b/qwen36_moe_edge/expert_quality.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +"""Score routed-expert quantization against real routed activations. + +``probe.py --mode quality`` samples experts with ``torch.randn`` activations +and quantizes the activations too. Neither matches how the edge runtime will +use these weights: + +- The activations an expert actually sees are the post-norm hidden states of + tokens the router sent to *that* expert. Gaussian noise has none of their + structure, and a scale calibrated against it hides errors that real inputs + expose. +- At M=1 the activation is 4 KiB against a 1.7 MiB weight block, so quantizing + it buys no bandwidth. The expert path is weight-only, W4A16 or W8A16. + +This tool captures the real activations from a forward pass, replays each +sampled expert in BF16 for the reference, and scores the weight-only +reconstructions against it. It can also write a small bundle so a device can +check its own dequantization and kernel against the same references without +loading the source checkpoint. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +import torch +import torch.nn.functional as F + +from qwen36_moe_edge.quantize_experts import ( + E4M3_MAX, + HIDDEN, + INTERMEDIATE, + NUM_LAYERS, + CheckpointReader, + _int4_weight, + _int8_weight, + _rht16, + dequantize_int4, +) + + +# nvfp4_e2m1 is the control: the shipped SM120 runtime uses that format for +# these same experts and reproduces greedy tokens exactly, so it is the bar an +# alternative 4-bit format has to match. Scoring a format without it invites +# reading a metric artefact as a defect. +SCHEMES = ("w8a16", "w4a16", "w4a16_rht16", "nvfp4_e2m1") + +# The sixteen E2M1 magnitudes, for nearest-value rounding. +_E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def expert_forward( + activation: torch.Tensor, + gate_up: torch.Tensor, + down: torch.Tensor) -> torch.Tensor: + """One routed expert: gate_up, SwiGLU, down.""" + projected = activation @ gate_up.T + hidden = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + return hidden @ down.T + + +def _reconstruct_e2m1( + weight: torch.Tensor, group_size: int) -> torch.Tensor: + """Two-level E2M1, matching what the shipped NVFP4 expert path stores.""" + rows, columns = weight.shape + grouped = weight.float().reshape(rows, columns // group_size, group_size) + amax = grouped.abs().amax(dim=2).clamp_min(1e-12) + peak = max(_E2M1_MAGNITUDES) + global_scale = max(float(amax.max()) / (E4M3_MAX * peak), 1e-12) + scale = (amax / peak / global_scale).to(torch.float8_e4m3fn) + effective = (scale.float() * global_scale).unsqueeze(-1) + normalized = grouped / effective.clamp_min(1e-30) + codebook = torch.tensor( + _E2M1_MAGNITUDES, dtype=torch.float32, device=weight.device) + nearest = codebook[ + (normalized.abs().unsqueeze(-1) - codebook).abs().argmin(-1)] + return ( + torch.sign(normalized) * nearest * effective + ).reshape(rows, columns) + + +def _reconstruct( + weight: torch.Tensor, + *, + scheme: str, + group_size: int) -> torch.Tensor: + """Quantize a weight and dequantize it, as the runtime's kernel will.""" + columns = weight.shape[1] + if scheme == "w8a16": + quantized, scale = _int8_weight(weight) + return quantized.float() * scale.float()[:, None] + if scheme == "nvfp4_e2m1": + return _reconstruct_e2m1(weight, group_size) + source = _rht16(weight) if scheme == "w4a16_rht16" else weight + packed, scale, global_scale = _int4_weight(source, group_size) + return dequantize_int4( + packed, scale, columns, group_size, global_scale) + + +def score_expert( + activation: torch.Tensor, + gate_up: torch.Tensor, + down: torch.Tensor, + *, + scheme: str, + group_size: int) -> dict[str, float]: + """Cosine and relative L2 of a weight-only scheme against BF16.""" + reference = expert_forward(activation, gate_up, down) + + gate_up_q = _reconstruct( + gate_up, scheme=scheme, group_size=group_size) + down_q = _reconstruct(down, scheme=scheme, group_size=group_size) + if scheme == "w4a16_rht16": + # The transform is orthonormal, so rotating both sides of each dot + # product leaves it unchanged. The runtime rotates activations the + # same way before calling the kernel. + projected = _rht16(activation) @ gate_up_q.T + hidden = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + output = _rht16(hidden) @ down_q.T + else: + output = expert_forward(activation, gate_up_q, down_q) + + difference = (output - reference).flatten() + return { + "cosine": F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item(), + "relative_l2": ( + difference.norm() / reference.flatten().norm().clamp_min(1e-12) + ).item(), + # Absolute error and reference magnitude, so results can be pooled + # across experts. Per-expert relative L2 alone is misleading here: + # output norms span three orders of magnitude across experts, and the + # router weights the small ones down before summing, so an expert with + # a near-zero output shows a huge relative error that contributes + # almost nothing to the layer. + "absolute_l2": difference.norm().item(), + "reference_l2": reference.flatten().norm().item(), + } + + +def collect_activations( + checkpoint: str, + *, + prompt: str, + prompt_tokens: int, + new_tokens: int, + max_seq: int, + device: str) -> tuple[list[list[list[int]]], list[list[torch.Tensor]]]: + """Run one eager forward, returning per-layer selections and MoE inputs.""" + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, + ) + + frontend = Qwen36MoeTextFrontend( + checkpoint, device=device, max_seq=max_seq, quant_scope="experts") + input_ids = frontend.tokenizer( + prompt, return_tensors="pt", add_special_tokens=False, + ).input_ids[:, :prompt_tokens].to(device) + if input_ids.shape[1] != prompt_tokens: + raise ValueError("the supplied prompt is shorter than prompt_tokens") + + state = Nexn2DecodeState(frontend._weights, max_seq, device) + state.batched_prefill = False + state.router_trace = {layer: [] for layer in range(state.num_layers)} + state.moe_input_trace = {layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generate_greedy( + state, input_ids, new_tokens, frontend._fvk, device) + + selections = [ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ] + activations = [ + list(state.moe_input_trace[layer]) + for layer in range(state.num_layers) + ] + return selections, activations + + +def _sampled_pairs( + selections: list[list[list[int]]], + *, + layers: tuple[int, ...], + experts_per_layer: int) -> list[tuple[int, int, int]]: + """Pick (layer, expert, step) triples the router actually produced.""" + pairs = [] + for layer in layers: + seen: dict[int, int] = {} + for step, experts in enumerate(selections[layer]): + for expert in experts: + seen.setdefault(expert, step) + for expert, step in list(seen.items())[:experts_per_layer]: + pairs.append((layer, expert, step)) + return pairs + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--golden", type=Path) + parser.add_argument("--prompt", required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=32) + parser.add_argument("--max-seq", type=int, default=256) + parser.add_argument( + "--layers", default="0,1,3,19,20,39", + help="layers to sample; the default spans both attention kinds and " + "the first, middle and last MoE blocks") + parser.add_argument("--experts-per-layer", type=int, default=4) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + layers = tuple(int(value) for value in args.layers.split(",")) + for layer in layers: + if not 0 <= layer < NUM_LAYERS: + parser.error(f"layer {layer} is outside 0..{NUM_LAYERS - 1}") + + selections, activations = collect_activations( + str(args.checkpoint), + prompt=args.prompt, + prompt_tokens=args.prompt_tokens, + new_tokens=args.new_tokens, + max_seq=args.max_seq, + device=args.device, + ) + pairs = _sampled_pairs( + selections, layers=layers, experts_per_layer=args.experts_per_layer) + print(f"scoring {len(pairs)} routed (layer, expert) pairs", flush=True) + + reader = CheckpointReader(args.checkpoint) + metrics = ("cosine", "relative_l2", "absolute_l2", "reference_l2") + scores: dict[str, list[float]] = { + f"{scheme}.{metric}": [] + for scheme in SCHEMES for metric in metrics + } + records = [] + golden: dict[str, torch.Tensor] = {} + for layer, expert, step in pairs: + activation = activations[layer][step].to( + device=args.device, dtype=torch.float32) + gate_up = reader.expert(layer, "gate_up_proj", expert).to( + device=args.device, dtype=torch.float32) + down = reader.expert(layer, "down_proj", expert).to( + device=args.device, dtype=torch.float32) + + record = {"layer": layer, "expert": expert, "step": step} + for scheme in SCHEMES: + values = score_expert( + activation, gate_up, down, + scheme=scheme, group_size=args.group_size) + record[scheme] = values + for metric, value in values.items(): + scores[f"{scheme}.{metric}"].append(value) + records.append(record) + + if args.golden is not None: + key = f"layer{layer:02d}.expert{expert:03d}" + golden[f"{key}.activation"] = ( + activation.to(torch.bfloat16).cpu()) + golden[f"{key}.reference"] = expert_forward( + activation, gate_up, down).to(torch.bfloat16).cpu() + + summary = { + name: { + "min": min(values), + "mean": statistics.mean(values), + "max": max(values), + } + for name, values in scores.items() + } + # Pooled relative L2: total error energy over total reference energy. This + # weights each expert by how much signal it carries, which is what the + # router does downstream, so it is the figure to judge a format on. + for scheme in SCHEMES: + error = sum( + value ** 2 for value in scores[f"{scheme}.absolute_l2"]) + signal = sum( + value ** 2 for value in scores[f"{scheme}.reference_l2"]) + summary[f"{scheme}.pooled_relative_l2"] = ( + error ** 0.5 / max(signal ** 0.5, 1e-12)) + result = { + "prompt_tokens": args.prompt_tokens, + "new_tokens": args.new_tokens, + "group_size": args.group_size, + "layers": list(layers), + "pair_count": len(pairs), + "summary": summary, + "records": records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(result, f, indent=2) + f.write("\n") + + if args.golden is not None: + from safetensors.torch import save_file + + args.golden.parent.mkdir(parents=True, exist_ok=True) + save_file(golden, str(args.golden), metadata={ + "hidden_size": str(HIDDEN), + "intermediate_size": str(INTERMEDIATE), + "pairs": ",".join( + f"{layer}:{expert}" for layer, expert, _ in pairs), + }) + print(f"wrote {len(golden) // 2} reference pairs to {args.golden}") + + print(f"\n{'scheme':<14} {'pooled relL2':>13} {'cos mean':>10} " + f"{'cos min':>10} {'per-expert relL2 mean':>22}") + for scheme in SCHEMES: + cosine = summary[f"{scheme}.cosine"] + l2 = summary[f"{scheme}.relative_l2"] + pooled = summary[f"{scheme}.pooled_relative_l2"] + print(f"{scheme:<14} {pooled:>13.5f} {cosine['mean']:>10.6f} " + f"{cosine['min']:>10.6f} {l2['mean']:>22.5f}") + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/kernel_parity.py b/qwen36_moe_edge/kernel_parity.py new file mode 100644 index 00000000..0673efc7 --- /dev/null +++ b/qwen36_moe_edge/kernel_parity.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +"""Cross-architecture parity for the qwen3_5_moe core and W4A16 kernels. + +The tier split was verified by compiling for sm_87 and sm_110, which proves +nothing about what the kernels compute there. This exercises each binding with +seeded inputs and records its output, so the same script run on another target +can be diffed against a reference produced on sm_120a. + +Deliberately no model and no checkpoint: shapes come from the Qwen3.6 geometry +and values from a fixed generator, so any machine can run it. + +Inputs are generated on the CPU and stored alongside the outputs. A comparison +run loads them from the reference rather than regenerating: CUDA RNG is not +bit-reproducible across architectures -- the Philox thread mapping follows +occupancy -- so regenerating on the target would compare kernels on different +data and read as a kernel failure. Divergence appears only past the first +launch block, which is why small tensors matched and large ones did not. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +HIDDEN = 2048 +INTERMEDIATE = 512 +TOPK = 8 +NUM_EXPERTS = 256 +# Linear-attention geometry. The split kernel broadcasts q and k from the 16 +# stored key heads to all 32 value heads, so every one of its three outputs is +# NV * HK wide -- not the 16-head width the stored layout suggests. Sizing +# q/k at 2048 makes the kernel write past them into whatever the allocator +# placed next, which shows up as a corrupted third output. +NV = 32 +HK = 128 +HV = 128 + + +# Inputs recorded by the reference run and replayed by comparison runs. +_INPUTS: dict[str, torch.Tensor] = {} +_REPLAY: dict[str, torch.Tensor] | None = None +_CPU_GEN = torch.Generator().manual_seed(20260730) +_CASE = "" + + +def _record(name: str, tensor: torch.Tensor, device) -> torch.Tensor: + """Return the replayed input if one exists, else keep what we generated.""" + key = f"{_CASE}.in.{name}" + if _REPLAY is not None: + if key not in _REPLAY: + raise KeyError(f"reference has no input {key}") + tensor = _REPLAY[key].to(dtype=tensor.dtype) + _INPUTS[key] = tensor.detach().cpu() + return tensor.to(device) + + +def _bf16(name, shape, device, scale=1.0): + """A bfloat16 input, generated on the CPU so it is machine-independent.""" + values = torch.randn(*shape, generator=_CPU_GEN, dtype=torch.float32) + return _record(name, (values * scale).to(torch.bfloat16), device) + + +def case_bf16_matvec(fvk, device): + x = _bf16("x", (1, HIDDEN), device) + w = _bf16("w", (HIDDEN, HIDDEN), device, 0.02) + out = torch.zeros(1, HIDDEN, dtype=torch.bfloat16, device=device) + rc = fvk.bf16_matvec_sm120_bf16( + x.data_ptr(), w.data_ptr(), out.data_ptr(), HIDDEN, HIDDEN, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": (x.float() @ w.float().T)} + + +def case_router_topk(fvk, device): + logits = _bf16("logits", (NUM_EXPERTS,), device, 4.0).contiguous() + idx = torch.empty(TOPK, dtype=torch.int32, device=device) + val = torch.empty(TOPK, dtype=torch.float32, device=device) + rc = fvk.moe_router_topk_sm120_bf16( + logits.data_ptr(), idx.data_ptr(), val.data_ptr(), + NUM_EXPERTS, TOPK, 0) + torch.cuda.synchronize() + reference = torch.topk(logits.float(), TOPK) + return {"rc": rc, "idx": idx, "val": val, + "torch_idx": reference.indices.to(torch.int32), + "torch_val": reference.values} + + +def case_silu_mul(fvk, device): + n = 4096 + g = _bf16("g", (n,), device) + u = _bf16("u", (n,), device) + out = torch.zeros(n, dtype=torch.bfloat16, device=device) + rc = fvk.silu_mul_sm120_bf16( + g.data_ptr(), u.data_ptr(), out.data_ptr(), n, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, + "torch": torch.nn.functional.silu(g.float()) * u.float()} + + +def case_sigmoid_mul(fvk, device): + n = 4096 + x = _bf16("x", (n,), device) + gate = _bf16("gate", (n,), device) + out = torch.zeros(n, dtype=torch.bfloat16, device=device) + rc = fvk.sigmoid_mul_sm120_bf16( + x.data_ptr(), gate.data_ptr(), out.data_ptr(), n, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, + "torch": x.float() * torch.sigmoid(gate.float())} + + +def case_weighted_sum(fvk, device): + d_dn = _bf16("d_dn", (TOPK, HIDDEN), device) + rows = torch.arange(TOPK, dtype=torch.int32, device=device) + weights = _record("weights", torch.softmax( + torch.randn(TOPK, generator=_CPU_GEN), -1), device) + # The reducer writes float32 and expects a flat buffer: see + # tests/test_qwen36_moe_gpu.py and the decode call site. Handing it a + # bfloat16 destination yields NaN, not a wrong-but-plausible answer. + out = torch.zeros(HIDDEN, dtype=torch.float32, device=device) + rc = fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), rows.data_ptr(), weights.contiguous().data_ptr(), + out.data_ptr(), 1, TOPK, HIDDEN, HIDDEN, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": weights.float() @ d_dn.float()} + + +def case_w16a16_gemm(fvk, device): + m, n, k = 16, HIDDEN, HIDDEN + x = _bf16("x", (m, k), device) + w = _bf16("w", (n, k), device, 0.02) + out = torch.zeros(m, n, dtype=torch.bfloat16, device=device) + rc = fvk.w16a16_gemm_sm120_bf16( + x.data_ptr(), w.data_ptr(), out.data_ptr(), m, n, k, 1.0, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": x.float() @ w.float().T} + + +def case_lin_split_qkv(fvk, device): + S = 4 + conv_out = _bf16("conv_out", (S, 8192), device).contiguous() + q32 = torch.zeros(S, NV, HK, dtype=torch.bfloat16, device=device) + k32 = torch.zeros(S, NV, HK, dtype=torch.bfloat16, device=device) + v32 = torch.zeros(S, NV, HV, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_lin_split_qkv_broadcast_bf16( + conv_out.data_ptr(), q32.data_ptr(), k32.data_ptr(), v32.data_ptr(), + S, 0) + torch.cuda.synchronize() + return {"q32": q32, "k32": k32, "v32": v32} + + +def case_e0m3_dequant(fvk, device): + """The streamed-block decode: sign-magnitude 4-bit plus a two-level scale.""" + from qwen36_moe_edge.quantize_experts import _int4_weight, dequantize_int4 + + rows, cols, group = 2 * INTERMEDIATE, HIDDEN, 16 + weight = _bf16("weight", (rows, cols), "cpu", 0.02).float() + packed, scale, global_scale = _int4_weight(weight, group) + packed = _record("packed", packed, device).contiguous() + scale = _record("scale", scale, device).contiguous() + out = torch.zeros(rows, cols, dtype=torch.bfloat16, device=device) + rc = fvk.qwen35moe_e0m3_dequant_bf16( + packed.data_ptr(), scale.data_ptr(), out.data_ptr(), + rows, cols, group, float(global_scale), 0) + torch.cuda.synchronize() + reference = dequantize_int4( + packed.cpu(), scale.cpu(), cols, group, global_scale) + return {"rc": rc, "out": out, "torch": reference.to(device)} + + +def case_split_q_gate(fvk, device): + S = 4 + q_proj = _bf16("q_proj", (S, 8192), device).contiguous() + q_pre = torch.zeros(S, 4096, dtype=torch.bfloat16, device=device) + gate = torch.zeros(S, 4096, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_split_q_gate_bf16( + q_proj.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), S, 0) + torch.cuda.synchronize() + return {"q_pre": q_pre, "gate": gate} + + +CASES = { + "bf16_matvec": case_bf16_matvec, + "moe_router_topk": case_router_topk, + "silu_mul": case_silu_mul, + "sigmoid_mul": case_sigmoid_mul, + "moe_weighted_sum": case_weighted_sum, + "w16a16_gemm": case_w16a16_gemm, + "lin_split_qkv": case_lin_split_qkv, + "split_q_gate": case_split_q_gate, + "e0m3_dequant": case_e0m3_dequant, +} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + from flash_rt import flash_rt_kernels as fvk + + capability = torch.cuda.get_device_capability() + print(f"device: {torch.cuda.get_device_name(0)} sm_{capability[0]}" + f"{capability[1]}") + + global _REPLAY, _CASE + if args.reference and args.reference.with_suffix(".pt").is_file(): + _REPLAY = torch.load( + args.reference.with_suffix(".pt"), weights_only=True) + print(f"replaying inputs from {args.reference.with_suffix('.pt')}") + + outputs, report = {}, {} + for name, case in CASES.items(): + _CASE = name + if not hasattr(fvk, _binding_of(name)): + report[name] = {"status": "binding absent"} + print(f"{name:<20} binding absent") + continue + try: + result = case(fvk, args.device) + except Exception as error: # noqa: BLE001 + report[name] = {"status": f"raised {type(error).__name__}: {error}"} + print(f"{name:<20} RAISED {error}") + continue + entry = {"status": "ok"} + if "rc" in result: + entry["rc"] = int(result.pop("rc")) + for key in list(result): + if key.startswith("torch"): + continue + outputs[f"{name}.{key}"] = result[key].detach().float().cpu() + # Local agreement with torch, where a reference was computed. + for key in ("out", "val"): + if key in result and "torch" in result: + entry["torch_cosine"] = _cosine( + result[key], result["torch"]) + if "torch_idx" in result: + # Compare as a set: the top-8 of this input contains an exact tie + # (two logits at 8.9375), and the kernel and torch.topk are free to + # break it differently while both being right. + entry["topk_index_set_match"] = int( + set(result["idx"].tolist()) + == set(result["torch_idx"].tolist())) + entry["topk_value_match"] = int(torch.allclose( + result["val"].cpu(), result["torch_val"].cpu())) + report[name] = entry + print(f"{name:<20} {entry}") + + args.output.parent.mkdir(parents=True, exist_ok=True) + torch.save({**_INPUTS, **outputs}, args.output.with_suffix(".pt")) + with args.output.open("w", encoding="utf-8") as f: + json.dump({"capability": list(capability), "cases": report}, f, + indent=2) + f.write("\n") + + if args.reference and args.reference.with_suffix(".pt").is_file(): + print("\n--- against reference ---") + expected = _REPLAY + worst = 1.0 + for key in sorted(k for k in set(expected) & set(outputs) + if ".in." not in k): + cosine = _cosine(outputs[key], expected[key]) + exact = torch.equal(outputs[key], expected[key]) + worst = min(worst, cosine) + print(f"{key:<34} cos={cosine:.8f} bitwise={'yes' if exact else 'no'}") + compared = {k for k in set(expected) | set(outputs) if ".in." not in k} + missing = sorted(compared - (set(expected) & set(outputs))) + if missing: + print(f"outputs present on only one side: {missing}") + print(f"worst cosine: {worst:.8f}") + + +def _binding_of(name: str) -> str: + return { + "lin_split_qkv": "qwen35moe_lin_split_qkv_broadcast_bf16", + "split_q_gate": "qwen35moe_split_q_gate_bf16", + "e0m3_dequant": "qwen35moe_e0m3_dequant_bf16", + }.get(name, f"{name}_sm120_bf16") + + +def _cosine(a: torch.Tensor, b: torch.Tensor) -> float: + x = a.detach().float().flatten().cpu() + y = b.detach().float().flatten().cpu() + return torch.nn.functional.cosine_similarity(x, y, dim=0).item() + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/probe.py b/qwen36_moe_edge/probe.py new file mode 100644 index 00000000..660b09d4 --- /dev/null +++ b/qwen36_moe_edge/probe.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 +"""Memory and quantization probes for the Qwen3.6-MoE edge path.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from qwen36_moe_edge.quantize_experts import ( + HIDDEN, + INTERMEDIATE, + NUM_EXPERTS, + NUM_LAYERS, + CheckpointReader, + _hadamard16, + _int4_weight, + _int8_weight, + _layout, + dequantize_int4, +) + + +def _numel(shape: tuple[int, ...]) -> int: + result = 1 + for value in shape: + result *= value + return result + + +def _category(name: str) -> str: + if ".mlp.experts." in name: + return "routed_experts" + if name in ( + "lm_head.weight", + "model.language_model.embed_tokens.weight", + ): + return "embed_lm_head" + if ".linear_attn." in name and not name.endswith("norm.weight"): + return "gdn_weights" + dense_markers = ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.out_proj.weight", + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ".mlp.shared_expert.down_proj.weight", + ) + if any(marker in name for marker in dense_markers): + return "other_dense" + return "norm_router_misc" + + +def _quantized_bytes( + shape: tuple[int, ...], *, bits: int, group_size: int +) -> int: + elements = _numel(shape) + if len(shape) < 2: + return elements * 2 + rows = _numel(shape[:-1]) + columns = shape[-1] + if bits == 8: + return elements + rows * 2 + return (elements + 1) // 2 + rows * ( + (columns + group_size - 1) // group_size) + + +def memory_probe( + checkpoint: Path, + group_size: int, + budget_gib: float, + runtime_reserve_gib: float) -> None: + with (checkpoint / "model.safetensors.index.json").open( + encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + readers = { + shard: safe_open( + checkpoint / shard, framework="pt", device="cpu") + for shard in set(weight_map.values()) + } + categories: dict[str, list[tuple[int, tuple[int, ...]]]] = {} + for name, shard in weight_map.items(): + if ".visual." in name or name.startswith("mtp."): + continue + shape = tuple(readers[shard].get_slice(name).get_shape()) + categories.setdefault(_category(name), []).append( + (_numel(shape), shape)) + + print("category bf16 GiB int8 GiB int4 GiB") + for category, tensors in sorted(categories.items()): + bf16 = sum(elements * 2 for elements, _ in tensors) + int8 = sum( + _quantized_bytes(shape, bits=8, group_size=group_size) + for _, shape in tensors + ) + int4 = sum( + _quantized_bytes(shape, bits=4, group_size=group_size) + for _, shape in tensors + ) + print( + f"{category:22s} {bf16 / 2**30:8.3f} " + f"{int8 / 2**30:10.3f} {int4 / 2**30:10.3f}" + ) + for quant_format in ("int8", "int4"): + layout = _layout(quant_format, group_size) + print( + f"{quant_format} expert block: " + f"{sum(layout.values()) / 2**20:.4f} MiB" + ) + + int8_resident = sum( + _quantized_bytes(shape, bits=8, group_size=group_size) + for category, tensors in categories.items() + if category != "routed_experts" + for _, shape in tensors + ) + int4_resident = sum( + ( + elements * 2 + if category == "gdn_weights" + else _quantized_bytes( + shape, bits=4, group_size=group_size) + ) + for category, tensors in categories.items() + if category != "routed_experts" + for elements, shape in tensors + ) + budget_bytes = int(budget_gib * 2**30) + reserve_bytes = int(runtime_reserve_gib * 2**30) + print( + f"\n{budget_gib:.2f} GiB budget, " + f"{runtime_reserve_gib:.2f} GiB runtime reserve" + ) + for quant_format, resident in ( + ("int8", int8_resident), + ("int4-mixed", int4_resident), + ): + block_format = "int8" if quant_format == "int8" else "int4" + block_bytes = sum(_layout(block_format, group_size).values()) + available = max(0, budget_bytes - reserve_bytes - resident) + quota = available // block_bytes // NUM_LAYERS + cache_bytes = quota * NUM_LAYERS * block_bytes + projected = resident + cache_bytes + reserve_bytes + print( + f"{quant_format:10s}: resident={resident / 2**30:.3f} GiB, " + f"quota={quota} experts/layer, " + f"cache={cache_bytes / 2**30:.3f} GiB, " + f"projected={projected / 2**30:.3f} GiB" + ) + + +def quality_probe( + checkpoint: Path, + *, + layers: int, + group_size: int, + device: str, +) -> None: + reader = CheckpointReader(checkpoint) + generator = torch.Generator(device=device).manual_seed(2026) + scores = { + "w8a16": [], + "w8a8": [], + "int4_w4a4": [], + "int4_rht16_w4a4": [], + } + selected_layers = torch.linspace( + 0, NUM_LAYERS - 1, steps=layers).round().int().tolist() + for layer in selected_layers: + expert = (layer * 37 + 11) % NUM_EXPERTS + gate_up = reader.expert( + layer, "gate_up_proj", expert + ).to(device=device, dtype=torch.float32) + down = reader.expert( + layer, "down_proj", expert + ).to(device=device, dtype=torch.float32) + activation = torch.randn( + 16, HIDDEN, generator=generator, device=device) + projected = activation @ gate_up.T + reference = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) @ down.T + + gu8, gu8_scale = _int8_weight(gate_up) + dn8, dn8_scale = _int8_weight(down) + gu8 = gu8.float() * gu8_scale.float()[:, None] + dn8 = dn8.float() * dn8_scale.float()[:, None] + for mode in ("w8a16", "w8a8"): + current = activation + if mode == "w8a8": + scale = ( + current.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) + / 127.0 + ) + current = ( + current / scale + ).round().clamp(-127, 127) * scale + projected = current @ gu8.T + current = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + if mode == "w8a8": + scale = ( + current.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) + / 127.0 + ) + current = ( + current / scale + ).round().clamp(-127, 127) * scale + output = current @ dn8.T + scores[mode].append(F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item()) + + for mode, use_rht, current_group in ( + ("int4_w4a4", False, group_size), + ("int4_rht16_w4a4", True, 16), + ): + transform = _hadamard16(gate_up.device) + gu_source = gate_up + current = activation + if use_rht: + gu_source = ( + gate_up.reshape(-1, 16) @ transform + ).reshape_as(gate_up) + current = ( + activation.reshape(-1, 16) @ transform + ).reshape_as(activation) + gu4, gu4_scale, gu4_alpha = _int4_weight( + gu_source, current_group) + gu4 = dequantize_int4( + gu4, gu4_scale, HIDDEN, current_group, gu4_alpha) + current4, current4_scale, current4_alpha = _int4_weight( + current, current_group) + current = dequantize_int4( + current4, current4_scale, HIDDEN, current_group, + current4_alpha) + projected = current @ gu4.T + current = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + dn_source = down + if use_rht: + dn_source = ( + down.reshape(-1, 16) @ transform + ).reshape_as(down) + current = ( + current.reshape(-1, 16) @ transform + ).reshape_as(current) + dn4, dn4_scale, dn4_alpha = _int4_weight( + dn_source, current_group) + dn4 = dequantize_int4( + dn4, dn4_scale, INTERMEDIATE, current_group, dn4_alpha) + current4, current4_scale, current4_alpha = _int4_weight( + current, current_group) + current = dequantize_int4( + current4, current4_scale, INTERMEDIATE, current_group, + current4_alpha) + output = current @ dn4.T + scores[mode].append(F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item()) + + for mode, values in scores.items(): + print( + f"{mode}: min={min(values):.6f} " + f"mean={statistics.mean(values):.6f}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--mode", choices=("memory", "quality"), required=True) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--sample-layers", type=int, default=20) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--budget-gib", type=float, default=7.0) + parser.add_argument("--runtime-reserve-gib", type=float, default=1.5) + args = parser.parse_args() + if args.mode == "memory": + memory_probe( + args.checkpoint, + args.group_size, + args.budget_gib, + args.runtime_reserve_gib, + ) + else: + quality_probe( + args.checkpoint, + layers=args.sample_layers, + group_size=args.group_size, + device=args.device, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/quantize_experts.py b/qwen36_moe_edge/quantize_experts.py new file mode 100644 index 00000000..c70e8de4 --- /dev/null +++ b/qwen36_moe_edge/quantize_experts.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +"""Write Qwen3.6 routed experts as fixed-size INT8 or INT4 blocks.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +import torch +from safetensors import safe_open + + +NUM_LAYERS = 40 +NUM_EXPERTS = 256 +HIDDEN = 2048 +INTERMEDIATE = 512 + +# Expert blocks are padded so that every block's offset and length are a +# multiple of the logical block size. A device whose memory holds only a +# fraction of the experts cannot afford to stream them through the page +# cache -- the cache competes with the resident weights for the same +# physical memory -- so the reader has to use O_DIRECT, which requires +# aligned offsets and lengths. +BLOCK_ALIGNMENT = 4096 + +# Largest magnitude an e4m3 scale byte can hold. The per-group scale is +# expressed as a fraction of a per-tensor global scale so that it lands in +# e4m3's normal range: with one level, every scale in this checkpoint falls +# into e4m3's subnormal range, where the format keeps about three bits and +# carries ~18 % relative error, which swamps the 4-bit value grid entirely. +E4M3_MAX = 448.0 + + +class CheckpointReader: + def __init__(self, checkpoint: Path): + index_path = checkpoint / "model.safetensors.index.json" + with index_path.open(encoding="utf-8") as f: + self.weight_map = json.load(f)["weight_map"] + self.readers = { + shard: safe_open( + checkpoint / shard, framework="pt", device="cpu") + for shard in set(self.weight_map.values()) + } + + def expert(self, layer: int, name: str, expert: int) -> torch.Tensor: + key = ( + f"model.language_model.layers.{layer}." + f"mlp.experts.{name}" + ) + shard = self.weight_map[key] + return self.readers[shard].get_slice(key)[expert] + + +def _int8_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = ( + weight.float().abs().amax(dim=1).clamp_min(1e-8) / 127.0 + ).to(torch.float16) + quantized = ( + weight.float() / scale.float()[:, None] + ).round().clamp(-127, 127).to(torch.int8) + return quantized, scale + + +def _hadamard16(device: torch.device) -> torch.Tensor: + matrix = torch.ones(1, 1, dtype=torch.float32, device=device) + for _ in range(4): + matrix = torch.cat(( + torch.cat((matrix, matrix), dim=1), + torch.cat((matrix, -matrix), dim=1), + ), dim=0) + return matrix / 4.0 + + +def _rht16(weight: torch.Tensor) -> torch.Tensor: + rows, columns = weight.shape + return ( + weight.float().reshape(rows, columns // 16, 16) + @ _hadamard16(weight.device) + ).reshape(rows, columns) + + +def _int4_weight( + weight: torch.Tensor, group_size: int +) -> tuple[torch.Tensor, torch.Tensor, float]: + """Pack to 4-bit with a two-level scale. + + Returns the packed values, the per-group e4m3 scale bytes, and the global + scale the kernel applies as its GEMM alpha. The effective scale of a group + is ``global_scale * e4m3(scale_byte)``. + """ + rows, columns = weight.shape + if columns % group_size: + raise ValueError( + f"K={columns} is not divisible by group_size={group_size}") + grouped = weight.float().reshape(rows, columns // group_size, group_size) + amax = grouped.abs().amax(dim=2).clamp_min(1e-12) + global_scale = max(float(amax.max()) / (E4M3_MAX * 7.0), 1e-12) + scale = (amax / 7.0 / global_scale).to(torch.float8_e4m3fn) + effective = scale.float() * global_scale + values = ( + grouped / effective.unsqueeze(-1).clamp_min(1e-30) + ).round().clamp(-7, 7).to(torch.int8).reshape(rows, columns) + magnitude = values.abs().to(torch.uint8) + code = magnitude | ((values < 0).to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return (packed.contiguous(), scale.view(torch.uint8).contiguous(), + global_scale) + + +def dequantize_int4( + packed: torch.Tensor, + scale: torch.Tensor, + columns: int, + group_size: int, + global_scale: float = 1.0, +) -> torch.Tensor: + """Inverse of :func:`_int4_weight`, for scoring and reference paths.""" + low = packed & 0x0F + high = (packed >> 4) & 0x0F + low = (low & 0x07).to(torch.int8) * torch.where( + (low & 0x08) != 0, -1, 1).to(torch.int8) + high = (high & 0x07).to(torch.int8) * torch.where( + (high & 0x08) != 0, -1, 1).to(torch.int8) + values = torch.stack((low, high), dim=-1).flatten(1) + rows = values.shape[0] + scale_float = scale.view(torch.float8_e4m3fn).float() * global_scale + return ( + values.float().reshape(rows, columns // group_size, group_size) + * scale_float.unsqueeze(-1) + ).reshape(rows, columns) + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + return tensor.contiguous().cpu().numpy().tobytes() + + +def quantize_expert( + gate_up: torch.Tensor, + down: torch.Tensor, + *, + quant_format: str, + group_size: int, + device: str, +) -> tuple[bytes, tuple[float, float]]: + """Return the fixed-size block and its (gate_up, down) global scales. + + INT8 uses per-output-channel scales that need no second level, so its + global scales are both 1.0. + """ + if quant_format not in ("int8", "int4", "int4-rht"): + raise ValueError(f"unsupported quantization format: {quant_format}") + if quant_format == "int4-rht" and group_size != 16: + raise ValueError("int4-rht requires group_size=16") + gate_up = gate_up.to(device=device, dtype=torch.float32) + down = down.to(device=device, dtype=torch.float32) + if quant_format == "int8": + gu_weight, gu_scale = _int8_weight(gate_up) + dn_weight, dn_scale = _int8_weight(down) + alphas = (1.0, 1.0) + else: + if quant_format == "int4-rht": + gate_up = _rht16(gate_up) + down = _rht16(down) + gu_weight, gu_scale, gu_alpha = _int4_weight(gate_up, group_size) + dn_weight, dn_scale, dn_alpha = _int4_weight(down, group_size) + alphas = (gu_alpha, dn_alpha) + return b"".join(( + _tensor_bytes(gu_weight), + _tensor_bytes(gu_scale), + _tensor_bytes(dn_weight), + _tensor_bytes(dn_scale), + bytes(_layout(quant_format, group_size)["padding"]), + )), alphas + + +def _parse_layers(value: str) -> range: + start, stop = (int(part) for part in value.split(":", 1)) + if start < 0 or stop > NUM_LAYERS or start >= stop: + raise argparse.ArgumentTypeError( + f"layers must satisfy 0 <= start < stop <= {NUM_LAYERS}") + return range(start, stop) + + +def _layout(quant_format: str, group_size: int) -> dict[str, int]: + if quant_format == "int8": + gu_weight = 2 * INTERMEDIATE * HIDDEN + gu_scale = 2 * INTERMEDIATE * 2 + dn_weight = HIDDEN * INTERMEDIATE + dn_scale = HIDDEN * 2 + else: + gu_weight = 2 * INTERMEDIATE * HIDDEN // 2 + gu_scale = 2 * INTERMEDIATE * (HIDDEN // group_size) + dn_weight = HIDDEN * INTERMEDIATE // 2 + dn_scale = HIDDEN * (INTERMEDIATE // group_size) + layout = { + "gate_up_weight": gu_weight, + "gate_up_scale": gu_scale, + "down_weight": dn_weight, + "down_scale": dn_scale, + } + # Trailing pad keeps every expert's offset and length aligned. The INT4 + # group-16 payload happens to be a multiple already; the INT8 payload is + # 3,151,872 B, which is 769.5 blocks. + layout["padding"] = -sum(layout.values()) % BLOCK_ALIGNMENT + return layout + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--format", choices=("int8", "int4", "int4-rht"), required=True) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--layers", type=_parse_layers, default=range(40)) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + if args.format != "int8" and args.group_size not in (16, 32, 64, 128): + parser.error("--group-size must be one of 16, 32, 64, or 128") + if args.format == "int4-rht" and args.group_size != 16: + parser.error("--format int4-rht requires --group-size 16") + + args.output.mkdir(parents=True, exist_ok=True) + layout = _layout(args.format, args.group_size) + block_bytes = sum(layout.values()) + manifest = { + "format": f"flashrt-qwen36-moe-{args.format}-experts-v2", + "group_size": args.group_size if args.format != "int8" else None, + "rht": args.format == "int4-rht", + "num_layers": NUM_LAYERS, + "num_experts": NUM_EXPERTS, + "hidden_size": HIDDEN, + "intermediate_size": INTERMEDIATE, + "block_layout": list(layout), + "block_sizes": layout, + "block_bytes": block_bytes, + "block_alignment": BLOCK_ALIGNMENT, + # Per-expert global scales, one pair per expert. Kept out of the + # blocks so a block stays exactly block_bytes and 4096-aligned; + # they are a few bytes each and belong with the resident weights, + # where the kernel reads them as its GEMM alpha. + "global_scales": "global_scales_layer_NN.bin", + "global_scales_dtype": "float32", + "global_scales_layout": ["gate_up", "down"], + } + with (args.output / "manifest.json").open("w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2) + f.write("\n") + + reader = CheckpointReader(args.checkpoint) + for layer in args.layers: + output_path = args.output / f"experts_layer_{layer:02d}.bin" + scale_check = args.output / f"global_scales_layer_{layer:02d}.bin" + expected_bytes = NUM_EXPERTS * block_bytes + if (output_path.is_file() + and output_path.stat().st_size == expected_bytes + and scale_check.is_file() + and scale_check.stat().st_size == NUM_EXPERTS * 2 * 4): + print(f"layer {layer}: already complete") + continue + temporary_path = output_path.with_suffix(".bin.tmp") + started = time.perf_counter() + alphas = [] + with temporary_path.open("wb") as f: + for expert in range(NUM_EXPERTS): + gate_up = reader.expert( + layer, "gate_up_proj", expert) + down = reader.expert(layer, "down_proj", expert) + block, expert_alphas = quantize_expert( + gate_up, + down, + quant_format=args.format, + group_size=args.group_size, + device=args.device, + ) + if len(block) != block_bytes: + raise RuntimeError( + f"expert block is {len(block)} bytes; " + f"expected {block_bytes}") + f.write(block) + alphas.extend(expert_alphas) + os.replace(temporary_path, output_path) + scale_path = args.output / f"global_scales_layer_{layer:02d}.bin" + temporary_scale_path = scale_path.with_suffix(".bin.tmp") + with temporary_scale_path.open("wb") as f: + f.write(_tensor_bytes( + torch.tensor(alphas, dtype=torch.float32))) + os.replace(temporary_scale_path, scale_path) + elapsed = time.perf_counter() - started + gib = expected_bytes / 2**30 + print( + f"layer {layer}: {gib:.3f} GiB in {elapsed:.2f}s " + f"({gib / elapsed:.2f} GiB/s)", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py new file mode 100644 index 00000000..7731d05d --- /dev/null +++ b/qwen36_moe_edge/route_trace.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +"""Collect router selections and simulate bounded per-layer expert caches. + +Two cache policies are simulated from the same trace: + +``simulate_lru`` + One LRU per layer holding every access. Prefill touches far more experts + than the cache can hold, so by the time decode starts the LRU contains + whatever the end of the prompt happened to use. + +``simulate_two_tier`` + A per-layer warm set chosen from prompt-phase selection counts and never + evicted, plus a small LRU ring for everything else. Prefill cannot + displace the warm set, so the decode hit rate follows warm-set coverage + rather than what the end of the prompt happened to leave behind. + +Which is better is a property of the checkpoint, not a foregone conclusion. +On Qwen3.6-35B-A3B, with a per-layer quota already in place, the plain LRU +wins at every quota from 16 slots up, and its margin *grows* with prompt +length: at 43 slots per layer it reaches 0.745 against 0.731 for a 32-token +prompt and 0.711 against 0.664 for a 128-token prompt. Recency predicts this +router's next selections better than prompt-phase frequency does, and a longer +prompt makes the frequency estimate more diffuse rather than more reliable. + +The oracle variant is consistently best, so pinning itself is not the problem +— the prompt-derived choice of what to pin is. Measure before committing a +runtime to either policy. + +The reported miss count per token, multiplied by the expert block size, is the +per-token read volume a streaming runtime has to sustain. +""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter, OrderedDict +from pathlib import Path + +import torch + +from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, +) +from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, +) + + +# INT4 group-16 routed-expert block, matching quantize_experts._layout. +DEFAULT_BLOCK_BYTES = 1769472 + + +def simulate_lru( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quota: int) -> dict[str, float]: + """Single per-layer LRU shared by prefill and decode.""" + prompt_accesses = prompt_misses = 0 + decode_accesses = decode_misses = 0 + for layer_trace in trace: + cache: OrderedDict[int, None] = OrderedDict() + for step, experts in enumerate(layer_trace): + prompt = step < prompt_tokens + for expert in experts: + if prompt: + prompt_accesses += 1 + else: + decode_accesses += 1 + if expert in cache: + cache.move_to_end(expert) + continue + if prompt: + prompt_misses += 1 + else: + decode_misses += 1 + if len(cache) >= quota: + cache.popitem(last=False) + cache[expert] = None + decode_steps = len(trace[0]) - prompt_tokens + return { + "prompt_hit_rate": 1.0 - prompt_misses / prompt_accesses, + "decode_hit_rate": 1.0 - decode_misses / decode_accesses, + "decode_misses_per_token": decode_misses / decode_steps, + } + + +def simulate_two_tier( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + pinned: int, + stream: int, + warm_from: str = "prompt") -> dict[str, float]: + """Pinned per-layer warm set plus a per-layer LRU ring. + + ``warm_from="prompt"`` selects the warm set from prompt-phase selection + counts, which is what a runtime can actually do. ``warm_from="decode"`` + selects it from the decode phase instead; that is not implementable, but + it bounds how much a better warm-set heuristic could win. + """ + if warm_from not in ("prompt", "decode"): + raise ValueError(f"unsupported warm_from: {warm_from!r}") + decode_accesses = decode_misses = warm_hits = 0 + for layer_trace in trace: + source = ( + layer_trace[:prompt_tokens] if warm_from == "prompt" + else layer_trace[prompt_tokens:] + ) + counts: Counter[int] = Counter() + for experts in source: + counts.update(experts) + warm = {expert for expert, _ in counts.most_common(pinned)} + ring: OrderedDict[int, None] = OrderedDict() + for experts in layer_trace[prompt_tokens:]: + for expert in experts: + decode_accesses += 1 + if expert in warm: + warm_hits += 1 + continue + if expert in ring: + ring.move_to_end(expert) + continue + decode_misses += 1 + if stream: + if len(ring) >= stream: + ring.popitem(last=False) + ring[expert] = None + decode_steps = len(trace[0]) - prompt_tokens + return { + "decode_hit_rate": 1.0 - decode_misses / decode_accesses, + "warm_hit_rate": warm_hits / decode_accesses, + "decode_misses_per_token": decode_misses / decode_steps, + } + + +def global_frequency( + trace: list[list[list[int]]]) -> list[Counter[int]]: + """Per-layer selection counts over a whole trace. + + Intended to be built from traces other than the one being evaluated: a set + derived from the trace it is scored on is an oracle, not a predictor. + """ + return [ + Counter(expert for step in layer_trace for expert in step) + for layer_trace in trace + ] + + +def simulate_warm_lru( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quota: int, + preload: list[Counter[int]] | None = None, + window: int = 1) -> dict[str, float]: + """Per-layer LRU, optionally warm-started, over ``window`` tokens at a time. + + ``preload`` fills each layer's cache with its most frequent experts before + decode, which is what a runtime can do at startup from offline statistics. + Entries are evictable: an earlier experiment showed that pinning a + prompt-derived set costs more adaptivity than it gains. + + ``window`` groups that many decode steps into one request, as a + multi-token verification step would. Note that this does not reduce reads: + an LRU already captures the reuse that a window's union would. + """ + if window < 1: + raise ValueError(f"window must be at least 1, got {window}") + misses = accesses = tokens = 0 + for index, layer_trace in enumerate(trace): + cache: OrderedDict[int, None] = OrderedDict() + if preload is not None: + for expert, _ in preload[index].most_common(quota): + cache[expert] = None + steps = layer_trace[prompt_tokens:] + for start in range(0, len(steps) - window + 1, window): + # Order-preserving dedupe, matching what the cache does. Within one + # request the insertion order decides which entry becomes the + # oldest, so it changes what a later eviction picks: iterating a set + # instead put this 0.25 % away from the measured cache. + requested = dict.fromkeys( + expert + for step in steps[start:start + window] + for expert in step + ) + if index == 0: + tokens += window + for expert in requested: + accesses += 1 + if expert in cache: + cache.move_to_end(expert) + continue + misses += 1 + if len(cache) >= quota: + cache.popitem(last=False) + cache[expert] = None + return { + "distinct_hit_rate": 1.0 - misses / max(accesses, 1), + "decode_misses_per_token": misses / max(tokens, 1), + } + + +def cold_prefill_blocks( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + resident: list[set[int]]) -> int: + """Blocks prefill must read before the first token can be emitted. + + Prefill routes every prompt token independently, so a layer's cost is the + union of its tokens' selections. Whatever is already resident is free; the + rest sets the floor on time to first token. + """ + return sum( + len({ + expert + for step in layer_trace[:prompt_tokens] + for expert in step + } - resident[index]) + for index, layer_trace in enumerate(trace) + ) + + +def read_volume( + misses_per_token: float, + *, + block_bytes: int, + bandwidths: tuple[float, ...]) -> dict[str, float]: + """Per-token read volume and the tok/s each bandwidth would allow.""" + per_token = misses_per_token * block_bytes + result = {"mb_per_token": per_token / 1e6} + for bandwidth in bandwidths: + result[f"tok_s_at_{bandwidth:g}gbps"] = ( + bandwidth * 1e9 / per_token if per_token else float("inf")) + return result + + +_POLICIES = ("single_lru", "two_tier", "two_tier_oracle_warm") + + +def summarize( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quotas: tuple[int, ...], + stream_fraction: float, + block_bytes: int, + bandwidths: tuple[float, ...]) -> dict[str, dict]: + """Compare both policies across per-layer quotas.""" + summary: dict[str, dict] = {} + for quota in quotas: + stream = max(1, int(round(quota * stream_fraction))) + pinned = max(0, quota - stream) + entry = { + "quota": quota, + "pinned": pinned, + "stream": stream, + "single_lru": simulate_lru( + trace, prompt_tokens=prompt_tokens, quota=quota), + "two_tier": simulate_two_tier( + trace, prompt_tokens=prompt_tokens, + pinned=pinned, stream=stream), + "two_tier_oracle_warm": simulate_two_tier( + trace, prompt_tokens=prompt_tokens, + pinned=pinned, stream=stream, warm_from="decode"), + } + for policy in _POLICIES: + entry[policy].update(read_volume( + entry[policy]["decode_misses_per_token"], + block_bytes=block_bytes, + bandwidths=bandwidths, + )) + summary[str(quota)] = entry + return summary + + +def format_summary( + summary: dict[str, dict], + *, + quotas: tuple[int, ...], + bandwidths: tuple[float, ...]) -> str: + header = f"{'quota':>6} {'pin/str':>8} {'policy':<22}" + header += f" {'hit':>7} {'miss/tok':>9} {'MB/tok':>8}" + for bandwidth in bandwidths: + header += f" {f'{bandwidth:g}GB/s':>9}" + lines = [header] + for quota in quotas: + entry = summary[str(quota)] + split = f"{entry['pinned']}/{entry['stream']}" + for policy in _POLICIES: + values = entry[policy] + line = f"{quota:>6} {split:>8} {policy:<22}" + line += f" {values['decode_hit_rate']:>7.4f}" + line += f" {values['decode_misses_per_token']:>9.2f}" + line += f" {values['mb_per_token']:>8.1f}" + for bandwidth in bandwidths: + line += f" {values[f'tok_s_at_{bandwidth:g}gbps']:>9.2f}" + lines.append(line) + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=64) + parser.add_argument("--max-seq", type=int, default=128) + parser.add_argument("--quotas", default="8,16,24,27,32,43,64") + parser.add_argument( + "--stream-fraction", type=float, default=0.25, + help="share of each layer's quota held as an evictable LRU ring") + parser.add_argument("--block-bytes", type=int, default=DEFAULT_BLOCK_BYTES) + parser.add_argument( + "--bandwidths", default="1.0,1.5,2.0", + help="storage read bandwidths in GB/s to project tok/s for") + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + frontend = Qwen36MoeTextFrontend( + args.checkpoint, + device=args.device, + max_seq=args.max_seq, + quant_scope="experts", + ) + input_ids = frontend.tokenizer( + args.prompt, + return_tensors="pt", + add_special_tokens=False, + ).input_ids[:, :args.prompt_tokens].to(args.device) + if input_ids.shape[1] != args.prompt_tokens: + parser.error("the supplied prompt is shorter than --prompt-tokens") + + state = Nexn2DecodeState( + frontend._weights, args.max_seq, args.device) + state.batched_prefill = False + state.router_trace = { + layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generated = generate_greedy( + state, + input_ids, + args.new_tokens, + frontend._fvk, + args.device, + ) + + trace = [ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ] + quotas = tuple(int(value) for value in args.quotas.split(",")) + bandwidths = tuple( + float(value) for value in args.bandwidths.split(",")) + summary = summarize( + trace, + prompt_tokens=args.prompt_tokens, + quotas=quotas, + stream_fraction=args.stream_fraction, + block_bytes=args.block_bytes, + bandwidths=bandwidths, + ) + result = { + "prompt_tokens": args.prompt_tokens, + "block_bytes": args.block_bytes, + "stream_fraction": args.stream_fraction, + "generated_tokens": generated, + "trace": trace, + "cache": summary, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(result, f) + f.write("\n") + + print(format_summary( + summary, quotas=quotas, bandwidths=bandwidths)) + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/streaming_frontend.py b/qwen36_moe_edge/streaming_frontend.py new file mode 100644 index 00000000..f7dc4e0b --- /dev/null +++ b/qwen36_moe_edge/streaming_frontend.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Qwen3.6-35B-A3B with the routed experts read from storage. + +The shipped frontend holds every expert in memory, which is 16.9 GiB of a +21.4 GiB footprint. This one skips them at load and serves each token's top-k +from a bounded cache over a prepared bundle, so what stays resident is the +non-expert weights plus however many slots the budget affords. + +It is the same pipeline otherwise: same attention, same recurrence, same +router, same reducer. Only where the expert weights come from changes, which is +why a token-level comparison against the ordinary frontend is meaningful. + +Greedy decode only, and not for CUDA Graph capture: a miss issues host reads, +which a captured graph cannot replay. +""" + +from __future__ import annotations + +from pathlib import Path + +from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend + +from qwen36_moe_edge.expert_cache import CacheConfig, ExpertCache + + +class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontend): + """Routed experts streamed from a bundle rather than held in memory.""" + + _MODEL_LABEL = "Qwen3.6-35B-A3B text, streamed experts" + + # The block-scaled 4-bit MMA kernels are absent from this list because this + # path never calls them: they serve the batched prefill, and streaming runs + # prefill through the per-token loop instead, since a miss issues host reads. + # Demanding them would refuse a build that can run this perfectly well -- + # which is what happened on the first attempt, on a target where the tier is + # correctly not built at all. + _REQUIRED_KERNELS = tuple( + name for name in Qwen36MoeTextFrontend._REQUIRED_KERNELS + if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) + ) + ('qwen35moe_e0m3_dequant_bf16', 'bf16_matvec_sm120_bf16') + + # The attention backend probes its kernel and falls back, so this runs on a + # target that builds no FA2. Thor is one: it uses FA4, whose SM100-class + # kernel needs Blackwell tensor memory that Orin's SM87 does not have -- + # so the two targets take different attention paths by design. + _REQUIRE_FA2 = False + + def __init__(self, checkpoint_path: str, bundle: str | Path, *, + slots_per_layer: int, + device: str = "cuda:0", + max_seq: int = 2048, + staging_buffers: int = 4, + budget_bytes: int = 0, + reserve_bytes: int = 0, + warm_frequency=None) -> None: + # Read by the loader through the base class, before any weight is + # touched, so the expert tensors are never built. + self._stream_experts = True + super().__init__( + checkpoint_path, device=device, max_seq=max_seq, + quant_scope="experts") + + resident = 0 + try: + import torch + + resident = int(torch.cuda.memory_allocated(device)) + except Exception: # noqa: BLE001 + pass + self.cache = ExpertCache(CacheConfig( + bundle=Path(bundle), + slots_per_layer=slots_per_layer, + staging_buffers=staging_buffers, + budget_bytes=budget_bytes, + reserve_bytes=reserve_bytes, + # Measured, not assumed: what the weights actually took. + resident_bytes=resident, + device=device, + )) + if warm_frequency is not None: + self.cache.warm(warm_frequency) + + def generate(self, max_new_tokens: int, *, do_sample: bool = False): + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate()") + if do_sample: + raise NotImplementedError("greedy decoding only") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device) + state = self._decode_state + state.expert_cache = self.cache + # A miss reads from storage on the host, which a captured graph cannot + # replay, so this path stays eager. + state.batched_prefill = False + return generate_greedy( + state, self._prompt_ids, max_new_tokens, self._fvk, self.device) + + def close(self) -> None: + self.cache.close() diff --git a/qwen36_moe_edge/warm_start_validation.py b/qwen36_moe_edge/warm_start_validation.py new file mode 100644 index 00000000..2d04826d --- /dev/null +++ b/qwen36_moe_edge/warm_start_validation.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +"""Validate a startup expert set on prompts it was not derived from. + +Filling each layer's cache at startup with its most frequently selected +experts looks strongly positive when the frequencies come from the same trace +being scored -- but that is an oracle, not a predictor. A deployment builds the +set offline, from other traffic, and then meets an unseen prompt. + +This runs several unrelated prompts through one model load and reports +leave-one-out results: for each prompt, the startup set is built from the +*other* prompts' traces only. The gap between that and the oracle is what the +heuristic actually costs. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import Counter +from pathlib import Path + +import torch + +from qwen36_moe_edge.route_trace import ( + DEFAULT_BLOCK_BYTES, + cold_prefill_blocks, + global_frequency, + simulate_warm_lru, +) + + +def collect_traces( + checkpoint: str, + prompts: list[str], + *, + prompt_tokens: int, + new_tokens: int, + max_seq: int, + device: str) -> list[list[list[list[int]]]]: + """Trace every prompt's router selections in one model load.""" + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, + ) + + frontend = Qwen36MoeTextFrontend( + checkpoint, device=device, max_seq=max_seq, quant_scope="experts") + state = Nexn2DecodeState(frontend._weights, max_seq, device) + state.batched_prefill = False + + traces = [] + for index, prompt in enumerate(prompts): + input_ids = frontend.tokenizer( + prompt, return_tensors="pt", add_special_tokens=False, + ).input_ids[:, :prompt_tokens].to(device) + if input_ids.shape[1] != prompt_tokens: + raise ValueError( + f"prompt {index} is shorter than {prompt_tokens} tokens") + state.router_trace = { + layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generate_greedy( + state, input_ids, new_tokens, frontend._fvk, device) + traces.append([ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ]) + print(f"traced prompt {index + 1}/{len(prompts)}", flush=True) + return traces + + +def _merge(frequencies: list[list[Counter[int]]]) -> list[Counter[int]]: + """Sum per-layer selection counts across several traces.""" + layers = len(frequencies[0]) + merged = [Counter() for _ in range(layers)] + for frequency in frequencies: + for layer in range(layers): + merged[layer].update(frequency[layer]) + return merged + + +def leave_one_out( + traces: list[list[list[list[int]]]], + *, + prompt_tokens: int, + quota: int, + block_bytes: int) -> list[dict[str, float]]: + """Score each prompt against a set built from the other prompts only.""" + frequencies = [global_frequency(trace) for trace in traces] + results = [] + for index, trace in enumerate(traces): + others = [f for position, f in enumerate(frequencies) + if position != index] + held_out = _merge(others) if others else None + oracle = frequencies[index] + + cold = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota) + warm = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota, + preload=held_out) + best = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota, preload=oracle) + + def resident(frequency): + if frequency is None: + return [set() for _ in trace] + return [ + {expert for expert, _ in frequency[layer].most_common(quota)} + for layer in range(len(trace)) + ] + + results.append({ + "prompt": index, + "cold_misses_per_token": cold["decode_misses_per_token"], + "held_out_misses_per_token": warm["decode_misses_per_token"], + "oracle_misses_per_token": best["decode_misses_per_token"], + "cold_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(None)) * block_bytes / 2 ** 30, + "held_out_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(held_out)) * block_bytes / 2 ** 30, + "oracle_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(oracle)) * block_bytes / 2 ** 30, + }) + return results + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--prompts-file", type=Path, required=True, + help="one prompt per line; blank lines ignored") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=32) + parser.add_argument("--max-seq", type=int, default=128) + parser.add_argument("--quotas", default="43,57,64") + parser.add_argument("--block-bytes", type=int, default=DEFAULT_BLOCK_BYTES) + parser.add_argument( + "--save-traces", type=Path, + help="write the per-prompt traces, so a warm set built from some of " + "them can be replayed against another on real hardware") + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + prompts = [ + line.strip() + for line in args.prompts_file.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + if len(prompts) < 3: + parser.error("leave-one-out needs at least three prompts") + + traces = collect_traces( + args.checkpoint, prompts, + prompt_tokens=args.prompt_tokens, + new_tokens=args.new_tokens, + max_seq=args.max_seq, + device=args.device, + ) + + if args.save_traces is not None: + args.save_traces.parent.mkdir(parents=True, exist_ok=True) + with args.save_traces.open("w", encoding="utf-8") as f: + json.dump({"prompt_tokens": args.prompt_tokens, + "traces": traces}, f) + f.write("\n") + print(f"wrote {len(traces)} traces to {args.save_traces}") + + quotas = tuple(int(value) for value in args.quotas.split(",")) + report = {"prompt_count": len(prompts), "quotas": list(quotas), + "prompt_tokens": args.prompt_tokens, "by_quota": {}} + for quota in quotas: + rows = leave_one_out( + traces, prompt_tokens=args.prompt_tokens, quota=quota, + block_bytes=args.block_bytes) + report["by_quota"][str(quota)] = rows + + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(report, f, indent=2) + f.write("\n") + + print(f"\n{'quota':>6} {'metric':<22} {'cold':>9} {'held-out':>9} " + f"{'oracle':>9} {'held-out win':>13}") + for quota in quotas: + rows = report["by_quota"][str(quota)] + for label, keys in ( + ("decode miss/token", ( + "cold_misses_per_token", "held_out_misses_per_token", + "oracle_misses_per_token")), + ("cold prefill GiB", ( + "cold_prefill_gib", "held_out_prefill_gib", + "oracle_prefill_gib")), + ): + cold, held, oracle = ( + statistics.mean(row[key] for row in rows) for key in keys) + win = (1.0 - held / cold) * 100.0 if cold else 0.0 + print(f"{quota:>6} {label:<22} {cold:>9.2f} {held:>9.2f} " + f"{oracle:>9.2f} {win:>12.1f}%") + + +if __name__ == "__main__": + main() diff --git a/scripts/qwen35moe_build_matrix.py b/scripts/qwen35moe_build_matrix.py new file mode 100644 index 00000000..0143f236 --- /dev/null +++ b/scripts/qwen35moe_build_matrix.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""What each qwen3_5_moe build tier adds, and what a build without them has. + +The claim these tiers make is that a build with every qwen3_5_moe option off +compiles the same sources and exports the same symbols it did before they +existed. That is a property of the gates, so it is checked by reading them: +which translation units CMake adds under each tier, and which ``m.def`` names +sit inside the matching preprocessor guard in the bindings. + +Reading the gates rather than building has a specific limit and a specific +advantage. It cannot catch a kernel that fails to compile -- only a build does +that, and the configure matrix printed at the end is how to run one. It can +catch the thing a single build cannot: a source or a symbol that leaks into a +configuration nobody built. + + python scripts/qwen35moe_build_matrix.py # print the matrix + python scripts/qwen35moe_build_matrix.py --check # exit 1 on a leak + +``tests/test_qwen35moe_build_matrix.py`` runs the same checks. +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent + +# CMake option -> the compile definition it sets -> what the bindings guard on. +TIERS = { + "FLASHRT_ENABLE_QWEN35MOE_CORE": "FLASHRT_HAVE_QWEN35MOE_CORE", + "FLASHRT_ENABLE_QWEN35MOE_W4A16": "FLASHRT_HAVE_QWEN35MOE_W4A16", + "FLASHRT_ENABLE_QWEN35MOE_W4A4": "FLASHRT_HAVE_QWEN35MOE_W4A4", +} + +# Gates that are not tiers but are still model-specific: the grouped MoE GEMM +# object, which only the weight-only tier on sm_110 builds. +EXTRA_GATES = ("FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100",) + +# Sources gated somewhere other than a tier, with the gate that owns each. +# Checked by name because that is the whole point: the grouped MoE GEMM used to +# be a second source in an object library every Thor build compiles. +ELSEWHERE = { + "csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu": + "qwen35moe_nvfp4_grouped_sm100_obj", +} + + +def _cmake_text() -> str: + return (ROOT / "CMakeLists.txt").read_text(encoding="utf-8") + + +def _bindings_text() -> str: + return (ROOT / "csrc" / "bindings.cpp").read_text(encoding="utf-8") + + +def tier_sources() -> dict[str, list[str]]: + """Sources CMake adds inside each ``if()`` block.""" + text = _cmake_text() + out: dict[str, list[str]] = {} + for option in TIERS: + # The block runs from `if(