liangsu9988 commited on about 13 hours ago

Commit

9022070

verified ·

1 Parent(s): e877859

Publish PI0.5 Thor kernels from 503100c (torch213-cxx11-cu130-aarch64-linux)

Browse files

Files changed (42) hide show

CARD.md +28 -0
README.md +135 -0
SYNC.md +80 -0
VALIDATION.md +155 -0
benchmarks/README.md +11 -0
benchmarks/RESULTS.md +144 -0
benchmarks/benchmark.py +22 -2
benchmarks/benchmark_bias.py +241 -0
build.toml +75 -0
build/torch213-cxx11-cu130-aarch64-linux/__init__.py +93 -2
build/torch213-cxx11-cu130-aarch64-linux/{fp8_gemm_source_test.abi3.so → _fp8_gemm_cuda_503100c.abi3.so} +2 -2
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +5 -9
csrc/README.md +7 -0
csrc/cublaslt_fp8_bias_sm110.cu +179 -0
csrc/cublaslt_fp8_bias_sm110.cuh +21 -0
csrc/cutlass_sm110_fp8_gemm.cu +176 -0
csrc/cutlass_sm110_fp8_gemm.cuh +31 -0
csrc/cutlass_sm120_block128_fp8_gemm.cu +273 -0
csrc/cutlass_sm120_block128_fp8_gemm.cuh +51 -0
csrc/fp8_block128_gemm_mma_sm89.cu +318 -0
csrc/fp8_block128_gemm_mma_sm89.cuh +102 -0
csrc/fp8_bs_gemm_device.cuh +1327 -0
csrc/fp8_gemv_m1_sm120.cu +159 -0
csrc/fp8_gemv_m1_sm120.cuh +31 -0
csrc/fp8_gemv_m1_sm89.cu +188 -0
csrc/fp8_gemv_m1_sm89.cuh +38 -0
csrc/fp8_smallM_handtuned_ldmatrix_sm120.cu +377 -0
csrc/fp8_smallM_handtuned_ldmatrix_sm120.cuh +68 -0
csrc/fp8_smallM_handtuned_sm120.cu +338 -0
csrc/fp8_smallM_handtuned_sm120.cuh +89 -0
csrc/gemm_types_sm110.h +362 -0
examples/README.md +3 -0
flake.lock +118 -0
flake.nix +20 -0
scripts/README.md +3 -0
tests/README.md +9 -0
tests/test_fp8_gemm.py +698 -0
torch-ext/README.md +6 -0
torch-ext/fp8_gemm/__init__.py +345 -0
torch-ext/torch_binding.cpp +577 -0
torch-ext/torch_binding.h +54 -0

CARD.md ADDED Viewed

	@@ -0,0 +1,28 @@

+# flashrt/fp8-gemm
+FlashRT native CUDA FP8 GEMV/GEMM kernels for low-latency transformer and
+diffuser linear layers.
+The block-128 scaled API supports Ada `sm_89` and Blackwell `sm_120a`.
+The per-tensor API supports Blackwell `sm_110a` (Jetson AGX Thor) and
+`sm_120a`. SM110 uses the production FlashRT Sq/T1/Wide CUTLASS family and has
+been swept across PI0.5, GROOT, Cosmos Edge, and LingBot VLA projection shapes.
+## Functions
+- `fp8_linear_bf16(input, weight, alpha=1.0, out=None, variant=0)`
+- `fp8_linear_residual_bf16(input, weight, residual, alpha=1.0, variant=0)`
+- `fp8_linear_bias_bf16(input, weight, bias, alpha=1.0, out=None)`
+- `fp8_linear_bias_residual_bf16(input, weight, bias, residual, alpha=1.0)`
+- `fp8_linear_bias_gelu_bf16(input, weight, bias, alpha=1.0, out=None)`
+- `fp8_blockwise_linear_bf16(input, weight, input_scale, weight_scale, out=None)`
+- `select_fp8_linear_tile(m, n, k, variant=0)`
+On SM110, keep `variant=0` for the tuned public dispatcher. Variants `1`, `2`,
+and `3` force Sq, T1, and Wide respectively for diagnostics.
+SM110 also provides BF16 bias, in-place bias+residual, and tanh-GELU+bias
+epilogues for SigLIP-style projection and MLP shapes. The validated large-M
+plain GEMM band is `M=65..1024`.
+See the repository README for shape contracts, validation status, and examples.

README.md ADDED Viewed

	@@ -0,0 +1,135 @@

+# fp8-gemm
+FlashRT native CUDA FP8 GEMV/GEMM kernels for low-latency transformer and
+diffuser linear layers on NVIDIA Ada SM89 and Blackwell SM110/SM120 GPUs.
+This package exposes the hand-tuned FP8 E4M3 decode and small-M kernels as
+Tensor APIs for Hugging Face Kernel Hub. It is intended for model runtimes that
+already hold activations and weights in FP8 and want a low-overhead BF16 output
+linear path.
+## Available Functions
+- `fp8_linear_bf16(input, weight, alpha=1.0, out=None, variant=0)`
+- `fp8_linear_residual_bf16(input, weight, residual, alpha=1.0, variant=0)`
+- `fp8_linear_bias_bf16(input, weight, bias, alpha=1.0, out=None)`
+- `fp8_linear_bias_residual_bf16(input, weight, bias, residual, alpha=1.0)`
+- `fp8_linear_bias_gelu_bf16(input, weight, bias, alpha=1.0, out=None)`
+- `fp8_blockwise_linear_bf16(input, weight, input_scale, weight_scale, out=None)`
+- `fp8_blockwise_swiglu_quantize_fp8(input, gate_up_weight, input_scale, gate_up_weight_scale, output=None, output_scale=None)`
+- `select_fp8_linear_tile(m, n, k, variant=0)`
+Tensor contract:
+- `input`: `torch.float8_e4m3fn`, shape `(M, K)`, contiguous CUDA tensor.
+- `weight`: `torch.float8_e4m3fn`, shape `(N, K)`, contiguous CUDA tensor.
+- `out`: `torch.bfloat16`, shape `(M, N)`.
+- `residual`: `torch.bfloat16`, shape `(1, N)` or `(N,)`, only supported for
+  the `M=1` decode GEMV path.
+- `K % 16 == 0`; SM120 additionally requires `K % 32 == 0`.
+- On SM120, `M == 1` uses dedicated GEMV and `2 <= M <= 64` uses small-M
+  GEMM tiles.
+- On SM110 (Jetson AGX Thor), the per-tensor API uses the production FlashRT
+  CUTLASS Sq/T1/Wide family and supports the validated model-shape matrix from
+  decode through large vision/backbone rows. The large-M production band is
+  validated from `M=65` through `M=1024`, including PI0.5 prefill QKV, O,
+  gate/up, and down projections at `M=712..970`. `N` and `K` must be divisible
+  by 16.
+- The three BF16 bias APIs are SM110-only. They accept BF16 `(N,)` bias and
+  preserve the same row-major FP8 `(M,K)` input and `(N,K)` weight contract.
+  The residual API updates a BF16 `(M,N)` tensor in place. The GELU API uses
+  the tanh approximation.
+- SM110 `variant=0` is the production auto dispatcher. Diagnostic variants are
+  `1=Sq`, `2=T1`, and `3=Wide`; they are correctness-tested but should not be
+  pinned by model integrations without a shape-specific benchmark.
+- The per-tensor kernels use Blackwell FP8 MMA instructions and are not valid
+  for SM89. SM89 support is provided by the blockwise API below.
+- `alpha` is a host float. For per-tensor FP8 quantization, pass
+  `float(input_scale * weight_scale)` from your static calibration metadata.
+The blockwise API uses a separate contract:
+- `input`: FP8 E4M3 `(M, K)`.
+- `weight`: FP8 E4M3 `(N, K)`.
+- `input_scale`: FP32 `(M, K / 128)`.
+- `weight_scale`: FP32 `(N / 128, K / 128)`.
+- `N` and `K` must be divisible by 128; `M` is unrestricted.
+- Output is BF16 `(M, N)`.
+- On SM89, the blockwise API dispatches to the production FlashRT native
+  `mma.sync.aligned.m16n8k32` GEMM/GEMV implementation.
+- On SM120, it dispatches to the production FlashRT CUTLASS block-scaled
+  implementation.
+- SM110 is intentionally not claimed by the blockwise API; use the per-tensor
+  static-scale path there. Other architectures are rejected explicitly.
+The fused SM89 producer accepts FP8 `(M,K)` input, FP8 `(2*N,K)` gate/up
+weight, block-128 FP32 scales, and returns FP8 `(M,N)` plus FP32 `(M,N/128)`
+output scales. Its public range is `1 <= M <= 256` with `N` and `K` divisible
+by 128. It is rejected explicitly on non-SM89 GPUs.
+## Minimal Usage
+```python
+from kernels import get_kernel
+import torch
+ops = get_kernel("flashrt/fp8-gemm", version=1, trust_remote_code=True)
+x = torch.randn((16, 4096), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+w = torch.randn((8192, 4096), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+y = ops.fp8_linear_bf16(x, w, alpha=1.0)
+```
+SM110 bias epilogues:
+```python
+bias = torch.randn((8192,), device="cuda", dtype=torch.bfloat16)
+residual = torch.randn((16, 8192), device="cuda", dtype=torch.bfloat16)
+y = ops.fp8_linear_bias_bf16(x, w, bias, alpha=1.0)
+ops.fp8_linear_bias_residual_bf16(x, w, bias, residual, alpha=1.0)
+y_gelu = ops.fp8_linear_bias_gelu_bf16(x, w, bias, alpha=1.0)
+```
+Warm each distinct SM110 bias shape once before CUDA Graph capture. The
+cuBLASLt fallback lazily creates and caches its descriptor, algorithm, and
+workspace on the first call; replay itself performs no allocation.
+Decode residual path:
+```python
+x = torch.randn((1, 4096), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+w = torch.randn((4096, 4096), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+residual = torch.zeros((1, 4096), device="cuda", dtype=torch.bfloat16)
+ops.fp8_linear_residual_bf16(x, w, residual, alpha=1.0)
+```
+Block-128 scaling:
+```python
+m, k, n = 51, 1536, 1536
+x = torch.randn((m, k), device="cuda").to(torch.float8_e4m3fn)
+w = torch.randn((n, k), device="cuda").to(torch.float8_e4m3fn)
+x_scale = torch.ones((m, k // 128), device="cuda", dtype=torch.float32)
+w_scale = torch.ones((n // 128, k // 128), device="cuda", dtype=torch.float32)
+y = ops.fp8_blockwise_linear_bf16(x, w, x_scale, w_scale)
+```
+## Validation
+```bash
+python fp8-gemm/tests/test_fp8_gemm.py --backend source --mode full
+python fp8-gemm/benchmarks/benchmark.py --backend source --mode headline
+python fp8-gemm/benchmarks/benchmark.py --backend source --mode pi05-prefill
+python fp8-gemm/benchmarks/benchmark_bias.py --backend source
+```
+The SM110 full sweep covers PI0.5, GROOT N1.6/N1.7, Cosmos Edge, and LingBot
+VLA projection families, plus decode, generic small-M, the `M=65` large-M
+boundary, and the three SigLIP bias epilogues. Public
+benchmark tables are only updated after source correctness, installed artifact
+correctness, shape/tile sweeps, `torch.compile(fullgraph=True)`, CUDA Graph
+replay, and parity against the original FlashRT native pointer entry pass.

SYNC.md ADDED Viewed

	@@ -0,0 +1,80 @@

+# Source Sync
+- Upstream FlashRT source: `../official/FlashRT`
+- Initial package date: June 20, 2026
+- SM89 source revision: `70b8eac4b05e9193bd99631cf872c5a971b59f5d`
+- SM110 sync revision: `132049d7c3a3534fb7d35676cd726f39408b1af6`
+Copied source files:
+- `csrc/gemm/fp8_gemv_m1_sm120.cu`
+- `csrc/gemm/fp8_gemv_m1_sm120.cuh`
+- `csrc/gemm/fp8_smallM_handtuned_sm120.cu`
+- `csrc/gemm/fp8_smallM_handtuned_sm120.cuh`
+- `csrc/gemm/fp8_smallM_handtuned_ldmatrix_sm120.cu`
+- `csrc/gemm/fp8_smallM_handtuned_ldmatrix_sm120.cuh`
+- `csrc/gemm/cutlass_sm120_block128_fp8_gemm.cu`
+- `csrc/gemm/cutlass_sm120_block128_fp8_gemm.cuh`
+- `csrc/gemm/fp8_block128_gemm_mma_sm89.cu`
+- `csrc/gemm/fp8_block128_gemm_mma_sm89.cuh`
+- `csrc/gemm/fp8_bs_gemm_device.cuh`
+- `csrc/gemm/fp8_gemv_m1_sm89.cu`
+- `csrc/gemm/fp8_gemv_m1_sm89.cuh`
+- `csrc/gemm/gemm_types_sm100.h`
+- `csrc/gemm/cutlass_sm100.cu`
+The SM110 copies are package-local as `csrc/gemm_types_sm110.h` and
+`csrc/cutlass_sm110_fp8_gemm.cu`. The C declarations in
+`csrc/cutlass_sm110_fp8_gemm.cuh` are packaging glue; the upstream pointer API
+declares them in its aggregate binding instead.
+Local packaging edits:
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/fp8_gemm`.
+- Kept public APIs model-agnostic; no raw pointer or stream arguments.
+- Bound the upstream measured `32x128-w4-s1` fused SwiGLU producer without
+  changing its CUDA tile or arithmetic.
+- Added a Tensor-facing SM110 dispatcher over the upstream BF16-output Sq, T1,
+  and Wide tactics. The public dispatcher and diagnostic variants do not alter
+  the copied GEMM templates or arithmetic.
+- Renamed SM100 source filenames locally to make their SM110 package role
+  explicit; CUTLASS still uses the SM100-family architecture templates when
+  compiling for `sm_110a`.
+- The SM110 build uses `-O3`, `--expt-relaxed-constexpr`, and
+  `--use_fast_math`, matching the validated native path.
+Architecture assumptions:
+- CUDA 12.8+ for SM89/SM120; CUDA 13.0+ for SM110.
+- NVIDIA Ada SM89 for block-128 scaled GEMM/GEMV.
+- NVIDIA Blackwell SM110a for per-tensor Sq/T1/Wide FP8 GEMM with BF16 output.
+- NVIDIA Blackwell SM120a for the original public APIs. The per-tensor FP8 MMA path uses
+  `.kind::f8f6f4` instructions and must be compiled for `sm_120a`, not plain
+  `sm_120`.
+- The SM110 kernel depends on the builder-provided CUTLASS 4.5 package. The
+  package flake is pinned to a builder revision that exports `cutlass_4_5`.
+- The release flake temporarily pins
+  `LiangSu8899/kernels@d720fa9`, based on
+  `huggingface/kernels@e9152aa`. The fork preserves the upstream builder
+  sources and changes only the stale CUTLASS 4.5.2 fixed-output hash from the
+  specified value to the value returned by the upstream archive. Return to an
+  upstream revision after that hash correction lands.
+Runtime constraints:
+- Inputs are FP8 E4M3 tensors with layout `input[M, K]` and `weight[N, K]`.
+- Output is BF16 `out[M, N]`.
+- `K` must be divisible by 32.
+- On SM120, `M` must be `1` or in `2..64`. M=128 remains an internal tuning
+  item because the validated correct SM120 tile is not performance-positive
+  enough for public release.
+- On SM110, `N` and `K` must be divisible by 16. The full-row Sq/T1/Wide path
+  has been validated on `M` from 1 through 1024 across PI0.5, GROOT,
+  Cosmos Edge, and LingBot projection families.
+- `alpha` is a host float scale multiplier, normally
+  `input_scale * weight_scale`.
+- The blockwise path consumes FP32 scales with layouts `(M, K/128)` and
+  `(N/128, K/128)`. It is the same CUTLASS kernel and schedule dispatcher used
+  by the upstream FlashRT pointer API.
+- Blockwise scaling is not exposed on SM110 in this increment.

VALIDATION.md ADDED Viewed

	@@ -0,0 +1,155 @@

+# Validation
+Date: June 20, 2026
+Local environment:
+- GPU: NVIDIA GeForce RTX 5090
+- PyTorch: 2.9.1+cu128
+- CUDA runtime reported by PyTorch: 12.8
+- Source build target: `sm_120a`
+## Source Correctness
+Command:
+```bash
+python fp8-gemm/tests/test_fp8_gemm.py --backend source --mode full
+```
+Result: 14/14 checks passed, plus the blockwise custom op passed
+`torch.compile(fullgraph=True)` with bitwise-equal output to the eager wrapper.
+Covered public v1 rows:
+- M=1 decode GEMV: `K in {512,4096}`, `N in {512,2048,8192}`
+- small-M GEMM: `M in {8,16,32,64}` with representative
+  transformer/diffuser-adjacent `K,N` rows
+- M=1 residual-add GEMV
+- block-128 scaled FP8 GEMM at:
+  - `(M,K,N)=(1,1024,1024)`
+  - `(51,1536,1536)`
+  - `(277,2048,2048)`
+  - `(1024,1152,1152)`
+  - `(2520,3072,3072)`
+  - `(128,4096,12288)`
+Metrics:
+- `max_abs`
+- `mean_abs`
+- `p99_abs`
+- cosine similarity
+- output dtype
+- tolerance
+The blockwise rows use the stricter gate:
+- `max_abs <= 0.0625`
+- `mean_abs <= 0.003`
+- `p99_abs <= 0.015625`
+- cosine similarity `>= 0.9999`
+The release benchmark also compares the Tensor wrapper against an independent
+binding of the original FlashRT pointer API. Matching source code alone is not
+treated as proof of zero wrapper overhead.
+## Source Benchmark
+Command:
+```bash
+python fp8-gemm/benchmarks/benchmark.py \
+  --backend source --mode headline --warmup 20 --iterations 100 --compile-ref
+```
+Result: all public rows passed. Headline rows are recorded in
+`benchmarks/RESULTS.md`.
+## Architecture Scope Boundary
+On SM120, the public per-tensor path supports `M=1` and `2 <= M <= 64`.
+The blockwise path retains its independent unrestricted-M contract.
+On SM110, the public per-tensor path uses the production CUTLASS Sq/T1/Wide
+family and supports the validated model-shape matrix. The current full sweep
+covers the large-M boundary `65`, exact PI0.5 prefill rows `712/768/970`, and
+representative `K,N` rows from PI0.5, GROOT N1.6/N1.7, Cosmos Edge, and
+LingBot VLA. It also gates BF16 bias, in-place bias+residual, and tanh-GELU
+bias epilogues on SigLIP dimensions `1152/3456/4304`. SM110 blockwise scaling
+is not claimed.
+The August 8 Thor source gate passed `39/39` with zero failures. Plain large-M
+GEMM was bitwise equal to the reference, and every PI0.5 prefill auto tile was
+within 2% of the fastest validated package/native tile.
+## Thor SM110 Increment
+Validated August 2, 2026:
+- GPU: NVIDIA Thor, compute capability 11.0;
+- PyTorch: 2.11.0+cu130;
+- CUDA: 13.0;
+- CUTLASS: 4.5.2, matching the current `kernel-builder` `cutlass_4_5`
+  dependency;
+- pinned builder: `e9152aa24e0d99eca255ca9f1beb996de32f9ca4`;
+- source correctness: 23/23;
+- locally installed aarch64 artifact correctness: 23/23;
+- `torch.compile(fullgraph=True)`: exact output parity;
+- CUDA Graph capture/replay: exact output parity;
+- original SM120 source regression on RTX 5090: 14/14.
+The 23 Thor rows include 20 production auto-dispatch checks and three forced
+Sq/T1/Wide diagnostics. Ordinary GEMMs were bitwise equal to the FP32
+accumulation reference after BF16 output conversion. The residual row passed
+with `max_abs=0.0625`, `p99_abs=0.0625`, and cosine `0.9999958` under the
+documented BF16 residual contract.
+Source-to-installed-artifact performance parity passed over 17 public
+auto-dispatch shapes: median artifact/source `0.9986`, p95 `1.0195`, and max
+`1.0244`.
+Comparisons against the original FlashRT pointer entry are reported separately
+in `benchmarks/RESULTS.md`.
+The final clean local artifact was built from
+`d31c69b1cb97ecd703aba01e29f423097f11c86a`. All 17 production rows passed the
+dispatcher gate; the worst auto/fastest-valid-tile paired ratio was `1.0028`.
+Sixteen rows matched the original CUTLASS 4.4.2 native entry within about 1.3%
+in the paired graph comparison. The PI0.5 gate/up row is a documented CUTLASS
+4.5.2 version outlier at `1.128x`; it is not described as native-performance
+parity.
+Before the SM110 update was published, the existing Thor pipeline dependency
+set was cold-loaded from Hub using both `kernels==0.16.0` and
+`kernels==0.12.3`: 20/20 package imports passed for each client. The Thor host
+required `HF_ENDPOINT=https://hf-mirror.com`; direct access to
+`huggingface.co` timed out, so official-endpoint cold loading remains a
+post-publication check on a host with direct Hub access.
+## HF Jobs Publish Status
+`flashrt/fp8-gemm` v1 was built and uploaded through the repository HF Jobs
+workflow.
+- Hub revision checked on June 20, 2026: `166f09be`
+- Uploaded variants:
+  - `torch211-cxx11-cu128-x86_64-linux`
+  - `torch211-cxx11-cu130-x86_64-linux`
+  - `torch212-cxx11-cu130-x86_64-linux`
+  - `torch212-cxx11-cu132-x86_64-linux`
+The existing SM120 Hub variants remain published. The new
+`torch211-cxx11-cu130-aarch64-linux` SM110 artifact is not included in the
+older Hub revision above; it must be published only after the clean-commit
+artifact rebuild and cold-cache checks.
+## SM89 Increment
+The source now also exposes block-128 scaled FP8 GEMM/GEMV and
+`fp8_blockwise_swiglu_quantize_fp8` on SM89. The fused producer performs the
+gate/up GEMMs, SiLU product, and block-128 FP8 requantization in one launch.
+It requires `1<=M<=256`, `N%128==0`, and `K%128==0` and uses the upstream
+measured `32x128-w4-s1` tile. SM120 source regression remains 14/14. SM89
+installed correctness, tile parity, and performance claims remain gated on an
+SM89 release artifact run; source presence alone is not recorded as runtime
+validation.

benchmarks/README.md ADDED Viewed

	@@ -0,0 +1,11 @@

+# Benchmarks
+```bash
+python fp8-gemm/benchmarks/benchmark.py --backend source --mode headline
+```
+The benchmark sweeps dispatcher rows and explicit M=1 GEMV variants.
+On NVIDIA Thor, use `--mode thor-full`. The runner records normal launches,
+CUDA Graph replay, all SM110 Sq/T1/Wide diagnostic tiles, and an optional
+original FlashRT pointer-API comparison when `FLASHRT_NATIVE_ROOT` is set.

benchmarks/RESULTS.md ADDED Viewed

	@@ -0,0 +1,144 @@

+# Benchmark Results: fp8-gemm
+Validated locally on June 20, 2026.
+Environment:
+- GPU: NVIDIA GeForce RTX 5090
+- PyTorch: 2.9.1+cu128
+- CUDA runtime reported by PyTorch: 12.8
+- Build target: `sm_120a`
+- Backend: source extension
+- Benchmark command:
+```bash
+python fp8-gemm/benchmarks/benchmark.py \
+  --backend source --mode headline --warmup 20 --iterations 100 --compile-ref
+```
+Correctness gate:
+```bash
+python fp8-gemm/tests/test_fp8_gemm.py --backend source --mode full
+```
+Result: 8/8 rows passed. Metrics recorded: max absolute error, mean absolute
+error, p99 absolute error, cosine similarity, dtype, and tolerance. Public v1
+SM120 scope is `M=1` decode and `2 <= M <= 64` small-M rows. SM110 uses a
+separate full-row CUTLASS dispatcher described below.
+## Headline Rows
+| Shape | Tile | FlashRT us | Torch eager us | Torch compile us | Speedup vs eager | Speedup vs compile | Max abs | P99 abs | Cosine |
+| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| `M=1,K=4096,N=2048` | `gemv_fp8_m1_w4` | 6.188 | 32.799 | 41.714 | 5.30x | 6.74x | 0.000 | 0.000 | 1.000000 |
+| `M=1,K=4096,N=8192` | `gemv_fp8_m1_w8` | 10.290 | 162.342 | 156.012 | 15.78x | 15.16x | 0.000 | 0.000 | 1.000000 |
+| `M=16,K=4096,N=4096` | `ld_fp8_gemm_16x128x256_w4` | 14.391 | 106.215 | 96.204 | 7.38x | 6.68x | 0.000 | 0.000 | 1.000000 |
+| `M=32,K=4096,N=8192` | `ld_fp8_gemm_32x128x256_w4` | 22.581 | 200.997 | 189.331 | 8.90x | 8.38x | 0.000 | 0.000 | 1.000000 |
+| `M=64,K=512,N=1024` | `ld_fp8_gemm_64x128x256_w4` | 8.259 | 18.085 | 50.002 | 2.19x | 6.05x | 0.000 | 0.000 | 1.000000 |
+## M=1 Variant Sweep
+The dispatcher defaults to `variant=0`. Explicit variants are retained for
+benchmarking and tuning; public callers should use `variant=0` unless they have
+measured their exact shape.
+| Shape | Variant | Tile | FlashRT us | Speedup vs eager | Status |
+| --- | ---: | --- | ---: | ---: | --- |
+| `M=1,K=4096,N=2048` | 0 | `gemv_fp8_m1_w4` | 6.188 | 5.30x | pass |
+| `M=1,K=4096,N=2048` | 4 | `gemv_fp8_m1_w4` | 6.186 | 5.30x | pass |
+| `M=1,K=4096,N=2048` | 8 | `gemv_fp8_m1_w8` | 6.184 | 5.30x | pass |
+| `M=1,K=4096,N=2048` | 16 | `gemv_fp8_m1_w16` | 6.188 | 5.30x | pass |
+| `M=1,K=4096,N=8192` | 0 | `gemv_fp8_m1_w8` | 10.290 | 15.78x | pass |
+| `M=1,K=4096,N=8192` | 4 | `gemv_fp8_m1_w4` | 10.274 | 15.81x | pass |
+| `M=1,K=4096,N=8192` | 8 | `gemv_fp8_m1_w8` | 10.272 | 15.82x | pass |
+| `M=1,K=4096,N=8192` | 16 | `gemv_fp8_m1_w16` | 10.278 | 15.80x | pass |
+## Block-128 Scaled GEMM
+Measured on RTX 5090 against an independent binding of the original FlashRT
+pointer API. The wrapper and native columns execute the same production
+CUTLASS kernel from separate extension modules. PyTorch eager and compile
+dequantize the block-scaled tensors to FP32, run the GEMM, and cast to BF16.
+CUTLASS is already the native implementation, so there is no additional
+contract-equivalent library row.
+| Workload `(M,K,N)` | Native us | Wrapper us | Wrapper/native | Eager us | Compile us | Max abs | P99 abs | Cosine |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| decode `(1,1024,1024)` | 10.287 | 10.276 | 0.999 | 36.256 | 42.242 | 0.000000 | 0.000000 | 1.0000001 |
+| action `(51,1536,1536)` | 14.364 | 14.360 | 1.000 | 55.595 | 51.500 | 0.000061 | 0.000000 | 1.0000001 |
+| GROOT `(277,2048,2048)` | 28.698 | 28.692 | 1.000 | 104.895 | 73.989 | 0.000122 | 0.000000 | 1.0000001 |
+| vision `(1024,1152,1152)` | 18.456 | 18.466 | 1.001 | 101.103 | 79.565 | 0.000122 | 0.000000 | 1.0000000 |
+| video `(2520,3072,3072)` | 114.188 | 114.696 | 1.004 | 1032.687 | 930.219 | 0.000244 | 0.000000 | 1.0000000 |
+| Qwen MLP `(128,4096,12288)` | 51.258 | 51.253 | 1.000 | 892.744 | 390.698 | 0.000244 | 0.000000 | 1.0000000 |
+All wrapper outputs were bitwise equal to the original native entry. The
+PyTorch-reference metrics above use production-scale ranges; the wider full
+correctness sweep remains the release gate.
+## Release Status
+- Source correctness: passed.
+- Source benchmark/tile sweep: passed for v1 public scope.
+- Existing SM120 installed artifacts: published.
+- SM110 local installed artifact: correctness, compile, graph, and parity
+  passed.
+- SM110 Hub artifact: pending clean-commit rebuild and upload.
+## NVIDIA Thor SM110 Results
+Validated August 2, 2026 on NVIDIA Thor with PyTorch 2.11.0+cu130, CUDA 13.0,
+CUTLASS 4.5.2, and a locally installed
+`torch211-cxx11-cu130-aarch64-linux` artifact. Timings are CUDA Graph replay
+latencies. `Native` is the independently loaded original FlashRT pointer API;
+the ratio is installed artifact / native, so values above 1 are slower.
+| Workload `(M,K,N)` | Auto tile | Artifact us | Native us | Artifact/native | Correctness |
+| --- | --- | ---: | ---: | ---: | --- |
+| decode `(1,4096,2048)` | T1 | 17.232 | 17.344 | 0.997 | pass |
+| decode-wide `(1,4096,8192)` | T1 | 50.144 | 50.112 | 1.001 | pass |
+| small-M `(16,4096,4096)` | T1 | 23.952 | 24.304 | 0.997 | pass |
+| small-M `(32,4096,8192)` | T1 | 46.256 | 46.000 | 1.013 | pass |
+| small-M `(64,512,1024)` | T1 | 11.296 | 11.264 | 0.997 | pass |
+| PI0.5 QKV `(51,2048,2560)` | T1 | 14.080 | 14.048 | 1.009 | pass |
+| PI0.5 O `(51,2048,2048)` | T1 | 13.440 | 13.504 | 0.993 | pass |
+| PI0.5 gate/up `(51,2048,16384)` | Wide | 92.496 | 81.312 | 1.128 | pass |
+| PI0.5 down `(51,8192,2048)` | T1 | 23.392 | 23.392 | 1.003 | pass |
+| GROOT DiT QKV `(51,1536,4608)` | T1 | 15.232 | 15.104 | 1.004 | pass |
+| GROOT N1.7 O `(277,2048,2048)` | Wide | 18.704 | 18.704 | 1.002 | pass |
+| GROOT N1.7 gate/up `(277,2048,16384)` | Wide | 186.432 | 189.360 | 0.985 | pass |
+| GROOT N1.7 down `(277,8192,2048)` | Sq | 50.080 | 49.984 | 1.003 | pass |
+| GROOT vision O `(1024,1024,1024)` | Sq | 15.360 | 15.424 | 0.997 | pass |
+| Cosmos Edge action `(64,2048,9216)` | T1 | 25.264 | 25.312 | 0.999 | pass |
+| LingBot vision O `(1024,1280,1280)` | Wide | 17.152 | 17.216 | 0.997 | pass |
+| LingBot action gate/up `(105,2048,16384)` | Wide | 140.080 | 139.984 | 1.002 | pass |
+Each graph ratio is the median of paired, per-launch package/native samples;
+candidate order rotates every round to control Thor DVFS bias. Sixteen rows are
+within about 1.3% of the original FlashRT pointer extension. PI0.5 gate/up is a
+reproducible CUTLASS dependency-version outlier: the Hub-buildable package uses
+CUTLASS 4.5.2 while the original native extension uses 4.4.2, and the paired
+ratio in the final clean artifact run is 1.128. CUTLASS 4.0 was also tested but
+failed at runtime on SM110, so it is not a valid packaging fallback. This row
+is retained explicitly and is not used for a native-parity claim.
+The installed artifact selected the fastest validated tactic on every row;
+worst auto/fastest-valid-tile was 1.0028. Source-to-artifact packaging parity
+passed with median 0.9986, p95 1.0195, and max 1.0244.
+## PI0.5 Thor Prefill and BF16 Bias Update
+Source gate rerun August 8, 2026 on NVIDIA Thor, PyTorch 2.13.0+cu130:
+- correctness: `39/39`, `fail_count=0`;
+- plain FP8 GEMM outputs were bitwise equal to the reference across the
+  `M=65..1024` band;
+- PI0.5 prefill QKV/O/gate-up/down auto dispatch was within 2% of the fastest
+  valid Sq/T1/Wide package tile and the original FlashRT native entry;
+- all 11 BF16 bias, bias+residual, and bias+GELU checks passed;
+- bias-only output was exact; residual p99 was at most one BF16 step; tanh-GELU
+  p99 was at most `3.8147e-6`, with cosine at least `0.999995`.
+The down-projection bias family uses a CUTLASS Wide fused epilogue while other
+SigLIP shapes use the faster cuBLASLt path. The dispatcher preserves the
+public row-major `(N,K)` weight contract.

benchmarks/benchmark.py CHANGED Viewed

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import importlib
 import json
 import os
 import statistics
 import sys
@@ -45,6 +46,10 @@ SHAPES = {
     "cosmos_edge_action": (64, 2048, 9216),
     "lingbot_vision_o": (1024, 1280, 1280),
     "lingbot_action_gate_up": (105, 2048, 16384),
 }
 MODES = {
@@ -58,6 +63,12 @@ MODES = {
         "cosmos_edge_action",
         "lingbot_action_gate_up",
     ],
     "thor-full": list(SHAPES),
 }
@@ -128,7 +139,10 @@ def load_source_ops() -> SourceOps:
     if capability == (11, 0):
         if not (cutlass_include / "cutlass" / "cutlass.h").is_file():
             raise RuntimeError("set CUTLASS_INCLUDE for the SM110 source benchmark")
-        cuda_sources = [str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu")]
         source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY"
         extra_includes = [
             str(cutlass_include),
@@ -179,6 +193,12 @@ def select_tile(m: int, n: int, k: int, variant: int = 0) -> str:
             raise RuntimeError("SM110 variant must be in [0, 3]")
         if variant:
             return forced[variant]
         if n >= 8 * k:
             return "sm110_wide_bf16"
         if m >= 128 and k >= 4 * n:
@@ -309,7 +329,7 @@ def metrics(got, expected):
     return (
         float(diff.max().item()),
         float(diff.mean().item()),
-        float(torch.quantile(diff, 0.99).item()),
         float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()),
     )

 import argparse
 import importlib
 import json
+import math
 import os
 import statistics
 import sys
     "cosmos_edge_action": (64, 2048, 9216),
     "lingbot_vision_o": (1024, 1280, 1280),
     "lingbot_action_gate_up": (105, 2048, 16384),
+    "pi05_prefill_qkv": (712, 2048, 2560),
+    "pi05_prefill_o": (970, 2048, 2048),
+    "pi05_prefill_gate_up": (768, 2048, 32768),
+    "pi05_prefill_down": (768, 16384, 2048),
 }
 MODES = {
         "cosmos_edge_action",
         "lingbot_action_gate_up",
     ],
+    "pi05-prefill": [
+        "pi05_prefill_qkv",
+        "pi05_prefill_o",
+        "pi05_prefill_gate_up",
+        "pi05_prefill_down",
+    ],
     "thor-full": list(SHAPES),
 }
     if capability == (11, 0):
         if not (cutlass_include / "cutlass" / "cutlass.h").is_file():
             raise RuntimeError("set CUTLASS_INCLUDE for the SM110 source benchmark")
+        cuda_sources = [
+            str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu"),
+            str(PACKAGE / "csrc" / "cublaslt_fp8_bias_sm110.cu"),
+        ]
         source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY"
         extra_includes = [
             str(cutlass_include),
             raise RuntimeError("SM110 variant must be in [0, 3]")
         if variant:
             return forced[variant]
+        if m >= 512 and k == 2048 and 2048 <= n <= 2560:
+            return "sm110_sq_bf16"
+        if m >= 512 and n >= 16 * k:
+            return "sm110_t1_bf16"
+        if m >= 512 and k >= 4 * n:
+            return "sm110_wide_bf16"
         if n >= 8 * k:
             return "sm110_wide_bf16"
         if m >= 128 and k >= 4 * n:
     return (
         float(diff.max().item()),
         float(diff.mean().item()),
+        float(diff.kthvalue(max(1, math.ceil(0.99 * diff.numel()))).values.item()),
         float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()),
     )

benchmarks/benchmark_bias.py ADDED Viewed

	@@ -0,0 +1,241 @@

+#!/usr/bin/env python3
+"""Benchmark SM110 BF16-output FP8 GEMM epilogues."""
+from __future__ import annotations
+import argparse
+import importlib
+import json
+import math
+import os
+import statistics
+import sys
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "fp8-gemm"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+SHAPES = {
+    "siglip_qkv": (512, 1152, 3456),
+    "siglip_mlp_up": (768, 1152, 4304),
+    "siglip_mlp_down": (768, 4304, 1152),
+}
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self.ops = getattr(torch.ops, namespace)
+    def fp8_linear_bias_bf16(self, x, w, bias, alpha=1.0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16
+            )
+        self.ops.fp8_linear_bias_bf16(x, w, bias, float(alpha), out)
+        return out
+    def fp8_linear_bias_residual_bf16(
+        self, x, w, bias, residual, alpha=1.0
+    ):
+        self.ops.fp8_linear_bias_residual_bf16(
+            x, w, bias, float(alpha), residual
+        )
+        return residual
+    def fp8_linear_bias_gelu_bf16(self, x, w, bias, alpha=1.0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16
+            )
+        self.ops.fp8_linear_bias_gelu_bf16(x, w, bias, float(alpha), out)
+        return out
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+    cutlass = Path(os.environ["CUTLASS_INCLUDE"])
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "11.0a")
+    namespace = "fp8_gemm_bias_source_bench"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu"),
+            str(PACKAGE / "csrc" / "cublaslt_fp8_bias_sm110.cu"),
+        ],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"),
+            str(REGISTRATION_INCLUDE),
+            str(cutlass),
+            str(cutlass.parent / "tools" / "util" / "include"),
+        ],
+        extra_cflags=[
+            "-O3", "-DNDEBUG", "-DCUDA_KERNEL",
+            "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY",
+        ],
+        extra_cuda_cflags=[
+            "-O3", "-DNDEBUG", "--expt-relaxed-constexpr", "--use_fast_math",
+            "-DCUDA_KERNEL", "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY",
+        ],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("fp8_gemm")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def load_native():
+    root = os.environ.get("FLASHRT_NATIVE_ROOT")
+    if not root:
+        return None
+    sys.path.insert(0, root)
+    try:
+        module = importlib.import_module("flash_rt.flash_rt_kernels")
+        return module.GemmRunner()
+    finally:
+        sys.path.remove(root)
+def measure(fn, warmup: int, iterations: int, rounds: int = 7) -> float:
+    samples = []
+    for _ in range(rounds):
+        for _ in range(warmup):
+            fn()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iterations):
+            fn()
+        end.record()
+        torch.cuda.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0 / iterations)
+    return float(statistics.median(samples))
+def metrics(got, expected):
+    diff = (got.float() - expected.float()).abs().flatten()
+    rank = max(1, math.ceil(0.99 * diff.numel()))
+    return {
+        "max_abs": float(diff.max().item()),
+        "mean_abs": float(diff.mean().item()),
+        "p99_abs": float(diff.kthvalue(rank).values.item()),
+        "cosine": float(torch.nn.functional.cosine_similarity(
+            got.float().flatten(), expected.float().flatten(), dim=0
+        ).item()),
+    }
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=64)
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    if torch.cuda.get_device_capability() != (11, 0):
+        raise SystemExit("SM110 is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    native = load_native()
+    native_stream = int(torch.cuda.current_stream().cuda_stream)
+    rows = []
+    for name, (m, k, n) in SHAPES.items():
+        generator = torch.Generator(device="cuda").manual_seed(m + k + n)
+        x = (torch.randn((m, k), device="cuda", generator=generator) * 0.25).to(
+            torch.float8_e4m3fn
+        )
+        weight = (
+            torch.randn((n, k), device="cuda", generator=generator) * 0.25
+        ).to(torch.float8_e4m3fn)
+        weight_kn = weight.t().contiguous()
+        bias = (torch.randn((n,), device="cuda", generator=generator) * 0.1).to(
+            torch.bfloat16
+        )
+        alpha = 0.75
+        base = x.float() @ weight.float().t() * alpha
+        for epilogue in ("bias", "bias_residual", "bias_gelu"):
+            if epilogue == "bias":
+                out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+                invoke = lambda: ops.fp8_linear_bias_bf16(
+                    x, weight, bias, alpha=alpha, out=out
+                )
+                expected = (base + bias.float()).to(torch.bfloat16)
+            elif epilogue == "bias_residual":
+                initial = (torch.randn((m, n), device="cuda", generator=generator) * 0.1).to(
+                    torch.bfloat16
+                )
+                out = initial.clone()
+                invoke = lambda: ops.fp8_linear_bias_residual_bf16(
+                    x, weight, bias, out, alpha=alpha
+                )
+                expected = (initial.float() + base + bias.float()).to(torch.bfloat16)
+            else:
+                out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+                invoke = lambda: ops.fp8_linear_bias_gelu_bf16(
+                    x, weight, bias, alpha=alpha, out=out
+                )
+                expected = torch.nn.functional.gelu(
+                    base + bias.float(), approximate="tanh"
+                ).to(torch.bfloat16)
+            invoke()
+            torch.cuda.synchronize()
+            accuracy = metrics(out, expected)
+            hub_us = measure(invoke, args.warmup, args.iterations)
+            row = {
+                "shape": name,
+                "M": m,
+                "K": k,
+                "N": n,
+                "epilogue": epilogue,
+                "hub_us": hub_us,
+                **accuracy,
+            }
+            if native is not None and epilogue == "bias":
+                native_out = torch.empty_like(out)
+                native_invoke = lambda: native.fp8_nn_bias_bf16(
+                    x.data_ptr(), weight_kn.data_ptr(), native_out.data_ptr(),
+                    bias.data_ptr(), m, n, k, alpha, native_stream
+                )
+                native_invoke()
+                torch.cuda.synchronize()
+                row["native_us"] = measure(
+                    native_invoke, args.warmup, args.iterations
+                )
+                row["hub_over_native"] = row["hub_us"] / row["native_us"]
+                row["native_metrics"] = metrics(native_out, expected)
+            rows.append(row)
+    payload = {"device": torch.cuda.get_device_name(), "rows": rows}
+    rendered = json.dumps(payload, indent=2, sort_keys=True)
+    print(rendered)
+    if args.json_out:
+        Path(args.json_out).write_text(rendered + "\n")
+    failed = [
+        row for row in rows
+        if row["p99_abs"] > 0.25 or row["mean_abs"] > 0.02 or row["cosine"] < 0.999
+    ]
+    return 1 if failed else 0
+if __name__ == "__main__":
+    raise SystemExit(main())

build.toml ADDED Viewed

	@@ -0,0 +1,75 @@

+[general]
+name = "fp8-gemm"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = ["cuda"]
+[general.cuda]
+minver = "12.8"
+[general.hub]
+repo-id = "flashrt/fp8-gemm"
+[torch]
+include = ["csrc"]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+[kernel.fp8_gemm_sm110]
+backend = "cuda"
+cuda-capabilities = ["11.0a"]
+cuda-flags = [
+    "--expt-relaxed-constexpr",
+    "-O3",
+    "--use_fast_math",
+]
+cuda-minver = "13"
+depends = [
+    "torch",
+    "cutlass_4_5",
+]
+include = ["csrc"]
+src = [
+    "csrc/gemm_types_sm110.h",
+    "csrc/cutlass_sm110_fp8_gemm.cu",
+    "csrc/cutlass_sm110_fp8_gemm.cuh",
+    "csrc/cublaslt_fp8_bias_sm110.cu",
+    "csrc/cublaslt_fp8_bias_sm110.cuh",
+]
+[kernel.fp8_gemm_sm89]
+backend = "cuda"
+cuda-capabilities = ["8.9"]
+cuda-minver = "12.8"
+depends = ["torch"]
+include = ["csrc"]
+src = [
+    "csrc/fp8_block128_gemm_mma_sm89.cu",
+    "csrc/fp8_block128_gemm_mma_sm89.cuh",
+    "csrc/fp8_bs_gemm_device.cuh",
+    "csrc/fp8_gemv_m1_sm89.cu",
+    "csrc/fp8_gemv_m1_sm89.cuh",
+]
+[kernel.fp8_gemm]
+backend = "cuda"
+cuda-capabilities = ["12.0a"]
+cuda-minver = "12.8"
+depends = [
+    "torch",
+    "cutlass_4_0",
+]
+include = ["csrc"]
+src = [
+    "csrc/fp8_gemv_m1_sm120.cu",
+    "csrc/fp8_gemv_m1_sm120.cuh",
+    "csrc/fp8_smallM_handtuned_sm120.cu",
+    "csrc/fp8_smallM_handtuned_sm120.cuh",
+    "csrc/fp8_smallM_handtuned_ldmatrix_sm120.cu",
+    "csrc/fp8_smallM_handtuned_ldmatrix_sm120.cuh",
+    "csrc/cutlass_sm120_block128_fp8_gemm.cu",
+    "csrc/cutlass_sm120_block128_fp8_gemm.cuh",
+]

build/torch213-cxx11-cu130-aarch64-linux/__init__.py CHANGED Viewed

@@ -37,6 +37,34 @@ def _fp8_linear_residual_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_blockwise_linear_bf16"))
 def _fp8_blockwise_linear_bf16_fake(
     input: torch.Tensor,
@@ -95,8 +123,8 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
     variant = int(variant)
     if m <= 0 or n <= 0 or k <= 0:
         raise RuntimeError("m, n, and k must be positive")
-    if k % 32 != 0:
-        raise RuntimeError("k must be divisible by 32")
     capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
     if capability == (11, 0):
         forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"}
@@ -106,6 +134,12 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
             raise RuntimeError("SM110 requires n and k divisible by 16")
         if variant:
             return forced[variant]
         if n >= 8 * k:
             return "sm110_wide_bf16"
         if m >= 128 and k >= 4 * n:
@@ -116,6 +150,8 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
             return "sm110_wide_bf16"
         return "sm110_t1_bf16"
     if m == 1:
         if variant == 4:
             return "gemv_fp8_m1_w4"
         if variant == 8:
@@ -131,6 +167,8 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
         return "gemv_fp8_m1_w16"
     if variant != 0:
         raise RuntimeError("small-M dispatcher currently supports variant=0 only")
     if m <= 16:
         if k % 256 == 0:
             return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
@@ -197,6 +235,56 @@ def fp8_linear_residual_bf16(
     return residual
 def fp8_blockwise_linear_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -248,6 +336,9 @@ def fp8_blockwise_swiglu_quantize_fp8(
 __all__ = [
     "fp8_linear_bf16",
     "fp8_linear_residual_bf16",
     "fp8_blockwise_linear_bf16",
     "fp8_blockwise_swiglu_quantize_fp8",
     "select_fp8_linear_tile",

     return None
+def _check_bias_linear_shapes(input, weight, bias, out) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input and weight K dimensions must match")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias must have shape (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out must have shape (input.shape[0], weight.shape[0])")
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(input, weight, bias, alpha: float, out) -> None:
+    _check_bias_linear_shapes(input, weight, bias, out)
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_residual_bf16"))
+def _fp8_linear_bias_residual_bf16_fake(
+    input, weight, bias, alpha: float, residual
+) -> None:
+    _check_bias_linear_shapes(input, weight, bias, residual)
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_bf16"))
+def _fp8_linear_bias_gelu_bf16_fake(input, weight, bias, alpha: float, out) -> None:
+    _check_bias_linear_shapes(input, weight, bias, out)
 @torch.library.register_fake(add_op_namespace_prefix("fp8_blockwise_linear_bf16"))
 def _fp8_blockwise_linear_bf16_fake(
     input: torch.Tensor,
     variant = int(variant)
     if m <= 0 or n <= 0 or k <= 0:
         raise RuntimeError("m, n, and k must be positive")
+    if k % 16 != 0:
+        raise RuntimeError("k must be divisible by 16")
     capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
     if capability == (11, 0):
         forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"}
             raise RuntimeError("SM110 requires n and k divisible by 16")
         if variant:
             return forced[variant]
+        if m >= 512 and k == 2048 and 2048 <= n <= 2560:
+            return "sm110_sq_bf16"
+        if m >= 512 and n >= 16 * k:
+            return "sm110_t1_bf16"
+        if m >= 512 and k >= 4 * n:
+            return "sm110_wide_bf16"
         if n >= 8 * k:
             return "sm110_wide_bf16"
         if m >= 128 and k >= 4 * n:
             return "sm110_wide_bf16"
         return "sm110_t1_bf16"
     if m == 1:
+        if k % 32:
+            raise RuntimeError("SM120 requires k divisible by 32")
         if variant == 4:
             return "gemv_fp8_m1_w4"
         if variant == 8:
         return "gemv_fp8_m1_w16"
     if variant != 0:
         raise RuntimeError("small-M dispatcher currently supports variant=0 only")
+    if k % 32:
+        raise RuntimeError("SM120 requires k divisible by 32")
     if m <= 16:
         if k % 256 == 0:
             return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
     return residual
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 FP8 linear with fused BF16 bias and BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(input, weight, bias, float(alpha), out)
+    return out
+def fp8_linear_bias_residual_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+) -> torch.Tensor:
+    """SM110 fused ``residual += alpha * input @ weight.T + bias``."""
+    ops.fp8_linear_bias_residual_bf16(
+        input, weight, bias, float(alpha), residual
+    )
+    return residual
+def fp8_linear_bias_gelu_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 FP8 linear with fused BF16 bias and GELU epilogue."""
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_gelu_bf16(input, weight, bias, float(alpha), out)
+    return out
 def fp8_blockwise_linear_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
 __all__ = [
     "fp8_linear_bf16",
     "fp8_linear_residual_bf16",
+    "fp8_linear_bias_bf16",
+    "fp8_linear_bias_residual_bf16",
+    "fp8_linear_bias_gelu_bf16",
     "fp8_blockwise_linear_bf16",
     "fp8_blockwise_swiglu_quantize_fp8",
     "select_fp8_linear_tile",

build/torch213-cxx11-cu130-aarch64-linux/{fp8_gemm_source_test.abi3.so → _fp8_gemm_cuda_503100c.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6da6757d21dbf032ac9c79d54dceffafed47cbf28b392de481d03e6078f198a8
-size 1893488

 version https://git-lfs.github.com/spec/v1
+oid sha256:22cab0b528b64c6915faa2735614461a01a77df8803ba448361d1f3c4e275885
+size 2534560

build/torch213-cxx11-cu130-aarch64-linux/_ops.py CHANGED Viewed

@@ -1,6 +1,6 @@
 import torch
-from . import fp8_gemm_source_test
-ops = torch.ops.fp8_gemm_source_test
 def add_op_namespace_prefix(op_name: str):
-    return f"fp8_gemm_source_test::{op_name}"

 import torch
+from . import _fp8_gemm_cuda_503100c
+ops = torch.ops._fp8_gemm_cuda_503100c
 def add_op_namespace_prefix(op_name: str):
+    return f"_fp8_gemm_cuda_503100c::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "fp8-gemm",
-  "id": "fp8_gemm_source_test",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,20 +13,16 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "hKsdFER2ZmYz36/zz9LYE7ddsK/R1dprHJFx+9yBzks=",
-      "fp8_gemm_source_test.abi3.so": "baZ1fSHb8DKsnHnVTc7/r+1Hy/KLOS3kgdA+YHjxmKg=",
-      "_ops.py": "zi4cl9u1hlmXMSZ8ZPvvrFIj8y1rvp8y3OTUWOXsp84=",
       "fp8_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "456d297",
       "dirty": false
-    },
-    "validation": {
-      "torch": "2.13.0+cu130",
-      "cuda": "13.0"
     }
   }
 }

 {
   "name": "fp8-gemm",
+  "id": "_fp8_gemm_cuda_503100c",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "HsFfsGY7AK/6Fziv4r2VQ+R2kDjS6hxKKgWqlA+zrKY=",
+      "_fp8_gemm_cuda_503100c.abi3.so": "IsqwtSi2TGkV+qJzVhRGGgGnffiAO6RINh0fPE4nWIU=",
+      "_ops.py": "RYYdMXXadIz2MzzsJ5E8N2JwNrfVICVpxfJyn0+g770=",
       "fp8_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
+      "sha": "503100c",
       "dirty": false
     }
   }
 }

csrc/README.md ADDED Viewed

	@@ -0,0 +1,7 @@

+# csrc
+Native FlashRT CUDA sources copied from upstream:
+- `fp8_gemv_m1_sm120.*`
+- `fp8_smallM_handtuned_sm120.*`
+- `fp8_smallM_handtuned_ldmatrix_sm120.*`

csrc/cublaslt_fp8_bias_sm110.cu ADDED Viewed

	@@ -0,0 +1,179 @@

+#include "cublaslt_fp8_bias_sm110.cuh"
+#include <cublasLt.h>
+#include <cstdint>
+#include <mutex>
+#include <unordered_map>
+namespace {
+constexpr size_t kWorkspaceBytes = 32 * 1024 * 1024;
+struct Key {
+  int M;
+  int N;
+  int K;
+  int epilogue;
+  bool operator==(const Key& other) const {
+    return M == other.M && N == other.N && K == other.K &&
+           epilogue == other.epilogue;
+  }
+};
+struct KeyHash {
+  size_t operator()(const Key& key) const {
+    size_t value = std::hash<int>{}(key.M);
+    value ^= std::hash<int>{}(key.N) + 0x9e3779b9 + (value << 6) +
+             (value >> 2);
+    value ^= std::hash<int>{}(key.K) + 0x9e3779b9 + (value << 6) +
+             (value >> 2);
+    value ^= std::hash<int>{}(key.epilogue) + 0x9e3779b9 + (value << 6) +
+             (value >> 2);
+    return value;
+  }
+};
+struct Entry {
+  cublasLtMatmulDesc_t operation = nullptr;
+  cublasLtMatrixLayout_t weight = nullptr;
+  cublasLtMatrixLayout_t input = nullptr;
+  cublasLtMatrixLayout_t output = nullptr;
+  cublasLtMatmulAlgo_t algorithm{};
+};
+cublasLtHandle_t handle = nullptr;
+void* workspace = nullptr;
+std::unordered_map<Key, Entry, KeyHash> cache;
+std::mutex cache_mutex;
+int status_code(cublasStatus_t status) {
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : -1000 - static_cast<int>(status);
+}
+void destroy_entry(Entry& entry) {
+  if (entry.operation) cublasLtMatmulDescDestroy(entry.operation);
+  if (entry.weight) cublasLtMatrixLayoutDestroy(entry.weight);
+  if (entry.input) cublasLtMatrixLayoutDestroy(entry.input);
+  if (entry.output) cublasLtMatrixLayoutDestroy(entry.output);
+  entry = Entry{};
+}
+int ensure_runtime() {
+  if (handle) return 0;
+  cublasStatus_t status = cublasLtCreate(&handle);
+  if (status != CUBLAS_STATUS_SUCCESS) return status_code(status);
+  const cudaError_t cuda_status = cudaMalloc(&workspace, kWorkspaceBytes);
+  return cuda_status == cudaSuccess ? 0 : -2000 - static_cast<int>(cuda_status);
+}
+int create_entry(const Key& key, Entry* entry) {
+  cublasStatus_t status = cublasLtMatmulDescCreate(
+      &entry->operation, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  cublasOperation_t transpose = CUBLAS_OP_T;
+  cublasOperation_t no_transpose = CUBLAS_OP_N;
+  cublasLtEpilogue_t epilogue =
+      key.epilogue == static_cast<int>(FlashRtFp8BiasEpilogue::kBiasGelu)
+          ? CUBLASLT_EPILOGUE_GELU_BIAS
+          : CUBLASLT_EPILOGUE_BIAS;
+  cudaDataType_t bias_type = CUDA_R_16BF;
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulDescSetAttribute(
+        entry->operation, CUBLASLT_MATMUL_DESC_TRANSA, &transpose,
+        sizeof(transpose));
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulDescSetAttribute(
+        entry->operation, CUBLASLT_MATMUL_DESC_TRANSB, &no_transpose,
+        sizeof(no_transpose));
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulDescSetAttribute(
+        entry->operation, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue,
+        sizeof(epilogue));
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulDescSetAttribute(
+        entry->operation, CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, &bias_type,
+        sizeof(bias_type));
+  }
+  // Row-major weight [N,K] is column-major [K,N]. Row-major input [M,K]
+  // is column-major [K,M]. The logical result [N,M] has row-major [M,N]
+  // storage, so no layout conversion or transpose kernel is required.
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatrixLayoutCreate(
+        &entry->weight, CUDA_R_8F_E4M3, key.K, key.N, key.K);
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatrixLayoutCreate(
+        &entry->input, CUDA_R_8F_E4M3, key.K, key.M, key.K);
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatrixLayoutCreate(
+        &entry->output, CUDA_R_16BF, key.N, key.M, key.N);
+  }
+  cublasLtMatmulPreference_t preference = nullptr;
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulPreferenceCreate(&preference);
+  }
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    const size_t workspace_bytes = kWorkspaceBytes;
+    status = cublasLtMatmulPreferenceSetAttribute(
+        preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_bytes, sizeof(workspace_bytes));
+  }
+  cublasLtMatmulHeuristicResult_t results[16]{};
+  int returned = 0;
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    status = cublasLtMatmulAlgoGetHeuristic(
+        handle, entry->operation, entry->weight, entry->input, entry->output,
+        entry->output, preference, 16, results, &returned);
+  }
+  if (preference) cublasLtMatmulPreferenceDestroy(preference);
+  if (status == CUBLAS_STATUS_SUCCESS && returned > 0) {
+    entry->algorithm = results[0].algo;
+    return 0;
+  }
+  destroy_entry(*entry);
+  return status == CUBLAS_STATUS_SUCCESS ? -1100 : status_code(status);
+}
+}  // namespace
+int fp8_linear_bias_sm110_bf16(
+    const void* input_fp8,
+    const void* weight_fp8,
+    const void* bias_bf16,
+    void* out_bf16,
+    int M,
+    int N,
+    int K,
+    float alpha,
+    float beta,
+    FlashRtFp8BiasEpilogue epilogue,
+    cudaStream_t stream) {
+  std::lock_guard<std::mutex> lock(cache_mutex);
+  int rc = ensure_runtime();
+  if (rc != 0) return rc;
+  const Key key{M, N, K, static_cast<int>(epilogue)};
+  auto iterator = cache.find(key);
+  if (iterator == cache.end()) {
+    Entry entry;
+    rc = create_entry(key, &entry);
+    if (rc != 0) return rc;
+    iterator = cache.emplace(key, entry).first;
+  }
+  Entry& entry = iterator->second;
+  cublasStatus_t status = cublasLtMatmulDescSetAttribute(
+      entry.operation, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_bf16,
+      sizeof(bias_bf16));
+  if (status != CUBLAS_STATUS_SUCCESS) return status_code(status);
+  status = cublasLtMatmul(
+      handle, entry.operation, &alpha, weight_fp8, entry.weight, input_fp8,
+      entry.input, &beta, out_bf16, entry.output, out_bf16, entry.output,
+      &entry.algorithm, workspace, kWorkspaceBytes, stream);
+  return status_code(status);
+}

csrc/cublaslt_fp8_bias_sm110.cuh ADDED Viewed

	@@ -0,0 +1,21 @@

+#pragma once
+#include <cuda_runtime.h>
+enum class FlashRtFp8BiasEpilogue : int {
+  kBias = 0,
+  kBiasGelu = 1,
+};
+int fp8_linear_bias_sm110_bf16(
+    const void* input_fp8,
+    const void* weight_fp8,
+    const void* bias_bf16,
+    void* out_bf16,
+    int M,
+    int N,
+    int K,
+    float alpha,
+    float beta,
+    FlashRtFp8BiasEpilogue epilogue,
+    cudaStream_t stream);

csrc/cutlass_sm110_fp8_gemm.cu ADDED Viewed

	@@ -0,0 +1,176 @@

+// CUTLASS FP8 GEMM implementations for SM110a. The package dispatcher uses
+// BF16-output Sq/T1/Wide variants. Weight B is column-major in the CUTLASS
+// problem and is supplied by the public API as a contiguous [N, K] tensor.
+#include "gemm_types_sm110.h"
+#include "cutlass_sm110_fp8_gemm.cuh"
+#include "cutlass/util/device_memory.h"
+#include <cuda_runtime.h>
+#include <cstdio>
+// Generic runner: initialize and run on the caller's stream.
+template <typename GemmOp>
+static int cutlass_run_impl(void* A, void* B, void* D,
+                             int M, int N, int K,
+                             float alpha, float beta,
+                             cudaStream_t stream) {
+    using ElementA = typename GemmOp::ElementA;
+    using ElementB = typename GemmOp::ElementB;
+    using ElementD = typename GemmOp::ElementD;
+    // CUTLASS stride computation
+    auto stride_A = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideA{}, {M, K, 1});
+    auto stride_B = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideB{}, {N, K, 1});
+    auto stride_D = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideD{}, {M, N, 1});
+    typename GemmOp::Arguments args{
+        cutlass::gemm::GemmUniversalMode::kGemm,
+        {M, N, K, 1},  // problem size
+        {(ElementA*)A, stride_A, (ElementB*)B, stride_B},
+        {{alpha, beta}, (ElementD*)D, stride_D, (ElementD*)D, stride_D}
+    };
+    GemmOp gemm;
+    size_t ws_size = GemmOp::get_workspace_size(args);
+    static cutlass::device_memory::allocation<uint8_t> workspace(0);
+    if (ws_size > workspace.size()) {
+        workspace = cutlass::device_memory::allocation<uint8_t>(ws_size);
+    }
+    auto status = gemm.can_implement(args);
+    if (status != cutlass::Status::kSuccess) {
+        fprintf(stderr, "[CUTLASS] cannot implement: M=%d N=%d K=%d\n", M, N, K);
+        return -1;
+    }
+    status = gemm.initialize(args, workspace.get(), stream);
+    if (status != cutlass::Status::kSuccess) {
+        fprintf(stderr, "[CUTLASS] init failed: M=%d N=%d K=%d\n", M, N, K);
+        return -2;
+    }
+    status = gemm.run(stream);
+    if (status != cutlass::Status::kSuccess) {
+        fprintf(stderr, "[CUTLASS] run failed: M=%d N=%d K=%d\n", M, N, K);
+        return -3;
+    }
+    return 0;
+}
+template <typename GemmOp>
+static int cutlass_run_bias_impl(
+    void* A, void* B, void* bias, void* D, int M, int N, int K,
+    float alpha, float beta, cudaStream_t stream) {
+    using ElementA = typename GemmOp::ElementA;
+    using ElementB = typename GemmOp::ElementB;
+    using ElementC = typename GemmOp::ElementC;
+    using ElementD = typename GemmOp::ElementD;
+    using ElementBias = cutlass::bfloat16_t;
+    auto stride_A = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideA{}, {M, K, 1});
+    auto stride_B = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideB{}, {N, K, 1});
+    auto stride_C = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideC{}, {M, N, 1});
+    auto stride_D = cutlass::make_cute_packed_stride(
+        typename GemmOp::GemmKernel::StrideD{}, {M, N, 1});
+    typename GemmOp::Arguments args{
+        cutlass::gemm::GemmUniversalMode::kGemm,
+        {M, N, K, 1},
+        {reinterpret_cast<ElementA*>(A), stride_A,
+         reinterpret_cast<ElementB*>(B), stride_B},
+        {{alpha, beta}, reinterpret_cast<ElementC*>(D), stride_C,
+         reinterpret_cast<ElementD*>(D), stride_D}
+    };
+    args.epilogue.thread.bias_ptr =
+        reinterpret_cast<ElementBias const*>(bias);
+    GemmOp gemm;
+    const size_t ws_size = GemmOp::get_workspace_size(args);
+    static cutlass::device_memory::allocation<uint8_t> workspace(0);
+    if (ws_size > workspace.size()) {
+        workspace = cutlass::device_memory::allocation<uint8_t>(ws_size);
+    }
+    auto status = gemm.can_implement(args);
+    if (status != cutlass::Status::kSuccess) return -11;
+    status = gemm.initialize(args, workspace.get(), stream);
+    if (status != cutlass::Status::kSuccess) return -12;
+    status = gemm.run(stream);
+    return status == cutlass::Status::kSuccess ? 0 : -13;
+}
+// Exported C functions.
+extern "C" {
+int cutlass_fp8_sq(void* A, void* B, void* D, int M, int N, int K,
+                    float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_sq::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_t1(void* A, void* B, void* D, int M, int N, int K,
+                    float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_t1::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_wide(void* A, void* B, void* D, int M, int N, int K,
+                      float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_wide::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_plain(void* A, void* B, void* D, int M, int N, int K,
+                       float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_plain::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_gelu(void* A, void* B, void* D, int M, int N, int K,
+                      float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_gelu::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+// FP32 output variants for models with activations exceeding FP16 range.
+int cutlass_fp8_sq_f32out(void* A, void* B, void* D, int M, int N, int K,
+                           float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_sq_f32out::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_wide_f32out(void* A, void* B, void* D, int M, int N, int K,
+                             float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_wide_f32out::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+// BF16 output variants
+int cutlass_fp8_sq_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                            float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_sq_bf16out::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_wide_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                              float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_wide_bf16out::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_t1_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                            float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_impl<sm100_t1_bf16out::Gemm>(A, B, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_wide_bias_bf16out(
+    void* A, void* B, void* bias, void* D, int M, int N, int K,
+    float alpha, float beta, cudaStream_t stream) {
+    return cutlass_run_bias_impl<sm100_wide_bias_bf16out::Gemm>(
+        A, B, bias, D, M, N, K, alpha, beta, stream);
+}
+int cutlass_fp8_wide_bias_gelu_bf16out(
+    void* A, void* B, void* bias, void* D, int M, int N, int K,
+    float alpha, cudaStream_t stream) {
+    return cutlass_run_bias_impl<sm100_wide_bias_gelu_bf16out::Gemm>(
+        A, B, bias, D, M, N, K, alpha, 0.0f, stream);
+}
+}  // extern "C"

csrc/cutlass_sm110_fp8_gemm.cuh ADDED Viewed

	@@ -0,0 +1,31 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cuda_runtime.h>
+extern "C" {
+int cutlass_fp8_sq(void* A, void* B, void* D, int M, int N, int K,
+                   float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_t1(void* A, void* B, void* D, int M, int N, int K,
+                   float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_wide(void* A, void* B, void* D, int M, int N, int K,
+                     float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_plain(void* A, void* B, void* D, int M, int N, int K,
+                      float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_gelu(void* A, void* B, void* D, int M, int N, int K,
+                     float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_sq_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                           float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_wide_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                             float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_t1_bf16out(void* A, void* B, void* D, int M, int N, int K,
+                           float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_wide_bias_bf16out(
+    void* A, void* B, void* bias, void* D, int M, int N, int K,
+    float alpha, float beta, cudaStream_t stream);
+int cutlass_fp8_wide_bias_gelu_bf16out(
+    void* A, void* B, void* bias, void* D, int M, int N, int K,
+    float alpha, cudaStream_t stream);
+}  // extern "C"

csrc/cutlass_sm120_block128_fp8_gemm.cu ADDED Viewed

	@@ -0,0 +1,273 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Path B implementation: CUTLASS block-128 FP8 GEMM for SM120a.
+// Header: cutlass_sm120_block128_fp8_gemm.cuh.
+//
+// Kernel template ported from CUTLASS 4.x example 87b
+// (third_party/cutlass/examples/87_blackwell_geforce_gemm_blockwise/
+//  87b_blackwell_geforce_fp8_bf16_gemm_groupwise.cu).
+//
+// Two GEMM instantiations are kept live and dispatched by M:
+//   * Pingpong (TileShape 64 x 128 x 128) — M <= 64
+//   * Cooperative (TileShape 128 x 128 x 128) — M > 64
+//
+// Per-shape Arguments + workspace are cached in two thread-safe maps.
+// The kernel itself is fused (no dequant intermediate), removing the
+// 3x memory bandwidth tax of Path D.
+#include "cutlass_sm120_block128_fp8_gemm.cuh"
+#include "cute/tensor.hpp"
+#include "cutlass/cutlass.h"
+#include "cutlass/detail/blockwise_scale_layout.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/util/packed_stride.hpp"
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+namespace flash_rt {
+namespace gemm {
+namespace {
+using namespace cute;
+// ── Element / layout types (match 87b) ───────────────────────────
+using ElementA           = cutlass::float_e4m3_t;
+using LayoutA            = cutlass::layout::RowMajor;
+constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
+using ElementB           = cutlass::float_e4m3_t;
+using LayoutB            = cutlass::layout::ColumnMajor;
+constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
+using ElementC           = cutlass::bfloat16_t;
+using LayoutC            = cutlass::layout::RowMajor;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+using ElementD           = ElementC;
+using LayoutD            = LayoutC;
+constexpr int AlignmentD = AlignmentC;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+// DeepSeek / Qwen3.6 layout: per-token activation, 128x128 weight.
+//
+// majorSFA = majorSFB = K so the SFA tensor is laid out (M, K/128)
+// row-major (the natural ckpt layout produced by HF dynamic FP8 quant)
+// instead of (M, K/128) col-major (the CUTLASS MN-major default).
+// Same for SFB: (N/128, K/128) row-major matches the safetensors
+// weight_scale_inv on disk.
+constexpr int ScaleGranularityM = 1;
+constexpr int ScaleGranularityN = 128;
+constexpr int ScaleGranularityK = 128;
+using ScaleConfig =
+    cutlass::detail::Sm120BlockwiseScaleConfig<ScaleGranularityM,
+                                                ScaleGranularityN,
+                                                ScaleGranularityK,
+                                                cute::UMMA::Major::K,
+                                                cute::UMMA::Major::K>;
+using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
+using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
+// ── Two kernel variants (Pingpong for small M, Cooperative for larger M) ──
+template <class MmaTileShape_, class Schedule_>
+struct GemmInstance {
+  using MmaTileShape  = MmaTileShape_;
+  using ClusterShape  = Shape<_1, _1, _1>;
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+      MmaTileShape, ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator, ElementCompute,
+      ElementC, LayoutC, AlignmentC,
+      ElementD, LayoutD, AlignmentD,
+      cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+      ElementA, cute::tuple<LayoutA, LayoutSFA>, AlignmentA,
+      ElementB, cute::tuple<LayoutB, LayoutSFB>, AlignmentB,
+      ElementAccumulator,
+      MmaTileShape, ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<
+          static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      Schedule_
+    >::CollectiveOp;
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+      Shape<int, int, int, int>,
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      void>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+};
+using PingpongMmaTileShape = Shape<_64, _128, _128>;
+using CooperativeMmaTileShape = Shape<_128, _128, _128>;
+using PingpongGemm =
+    typename GemmInstance<PingpongMmaTileShape,
+        cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::Gemm;
+using CooperativeGemm =
+    typename GemmInstance<CooperativeMmaTileShape,
+        cutlass::gemm::KernelScheduleSm120Blockwise>::Gemm;
+// ── Per-shape workspace cache ───────────────────────────────────
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct ShapeKeyHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (static_cast<size_t>(k.M) * 1315423911u)
+         ^ (static_cast<size_t>(k.N) * 2654435761u)
+         ^ static_cast<size_t>(k.K);
+  }
+};
+struct CachedWorkspace {
+  void*  ptr = nullptr;
+  size_t size = 0;
+};
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache;
+std::mutex g_ws_mu;
+void* get_workspace(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache.find(key);
+  if (it != g_ws_cache.end() && it->second.size >= needed) {
+    return it->second.ptr;
+  }
+  if (it != g_ws_cache.end()) {
+    cudaFree(it->second.ptr);
+    g_ws_cache.erase(it);
+  }
+  CachedWorkspace w;
+  w.size = needed;
+  if (needed > 0) {
+    cudaMalloc(&w.ptr, needed);
+  }
+  g_ws_cache[key] = w;
+  return w.ptr;
+}
+template <class Gemm>
+cutlass::Status run_gemm(
+    const void* A_fp8, const void* B_fp8, void* D_bf16,
+    int M, int N, int K,
+    const float* act_scale, const float* w_scale,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+  StrideA stride_A = cutlass::make_cute_packed_stride(
+      StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(
+      StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(
+      StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(
+      StrideD{}, cute::make_shape(M, N, 1));
+  LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(
+      cute::make_shape(M, N, K, 1));
+  LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(
+      cute::make_shape(M, N, K, 1));
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ElementA const*>(A_fp8), stride_A,
+          reinterpret_cast<ElementB const*>(B_fp8), stride_B,
+          act_scale, layout_SFA,
+          w_scale,   layout_SFB
+      },
+      {
+          {1.0f, 0.0f},                    // epilogue.thread (alpha, beta)
+          nullptr, stride_C,                // C unused (beta = 0)
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace(M, N, K, ws_size);
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp8_block128_gemm_cutlass_sm120_bf16out] can_implement FAIL "
+        "for M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp8_block128_gemm_cutlass_sm120_bf16out] initialize FAIL "
+        "for M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.run(stream);
+  return status;
+}
+}  // namespace
+void fp8_block128_gemm_cutlass_sm120_bf16out(
+    const void* A_fp8,
+    const void* B_fp8,
+    void*       D_bf16,
+    int M, int N, int K,
+    const float* act_block_scale,
+    const float* w_block_scale,
+    cudaStream_t stream)
+{
+  // Schedule selection: small M -> Pingpong (better latency at low
+  // arithmetic intensity); large M -> Cooperative (better throughput
+  // when there are enough M-tiles to fill the SMs).
+  cutlass::Status status;
+  if (M <= 64) {
+    status = run_gemm<PingpongGemm>(
+        A_fp8, B_fp8, D_bf16, M, N, K,
+        act_block_scale, w_block_scale, stream);
+  } else {
+    status = run_gemm<CooperativeGemm>(
+        A_fp8, B_fp8, D_bf16, M, N, K,
+        act_block_scale, w_block_scale, stream);
+  }
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp8_block128_gemm_cutlass_sm120_bf16out] run FAIL "
+        "for M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+}  // namespace gemm
+}  // namespace flash_rt

csrc/cutlass_sm120_block128_fp8_gemm.cuh ADDED Viewed

	@@ -0,0 +1,51 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS-based block-128 FP8 GEMM for SM120a (RTX 5090 / Blackwell
+// consumer).
+//
+// Native block-scaled FP8 GEMM with DeepSeek-V3 / Qwen3.6 layout:
+//   * activation: per-token (1) x per-128 K block scale
+//   * weight   : per-128 N x per-128 K block scale
+//   * output   : BF16
+//
+// Replaces the Path D dequantize-then-bf16-GEMM stop-gap with a
+// fused Tensor Core kernel from CUTLASS 4.x example 87b
+// (87b_blackwell_geforce_fp8_bf16_gemm_groupwise.cu). Same Python /
+// pybind signature shape as fp8_block128_gemm_descale_bf16out so
+// callers can swap with one-line change.
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+// Path B SM120a CUTLASS block-128 FP8 GEMM, BF16 output.
+//
+// Layout & shapes match Path D's signature:
+//   A_fp8      : (M, K)        e4m3 row-major
+//   B_fp8      : (N, K)        e4m3 row-major
+//   D_bf16     : (M, N)        bf16 row-major
+//   act_scale  : (M, K/128)    fp32 row-major
+//   w_scale    : (N/128, K/128) fp32 row-major
+//
+// Constraints: K and N must be multiples of 128. M is unrestricted.
+//
+// Internally selects a Cooperative or Pingpong CUTLASS schedule
+// based on M (Pingpong is faster when M is small, e.g. decode
+// step or short prefill). Caller does not provide scratch buffers
+// (the kernel is fused, no dequant intermediates needed).
+//
+// Stream-safe; per-shape arguments + workspace cached internally.
+void fp8_block128_gemm_cutlass_sm120_bf16out(
+    const void* A_fp8,
+    const void* B_fp8,
+    void*       D_bf16,
+    int M, int N, int K,
+    const float* act_block_scale,
+    const float* w_block_scale,
+    cudaStream_t stream);
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_block128_gemm_mma_sm89.cu ADDED Viewed

	@@ -0,0 +1,318 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Native Ada (sm_89) FP8 e4m3 -> BF16 block-128 scaled GEMM.
+// Header: fp8_block128_gemm_mma_sm89.cuh.
+//
+// Adapted from csrc/gemm/fp8_smallM_handtuned_sm120.cu (same cp.async
+// pipeline + m16n8k32 MMA tiling). Two sm_89-specific changes vs that file:
+//   1. MMA uses the plain Ada FP8 op `mma.sync.aligned.m16n8k32.row.col.
+//      f32.e4m3.e4m3.f32` (no `.kind::f8f6f4`, which is sm_120a-only).
+//   2. Per-tensor `alpha` is replaced by DeepSeek-style block-128 scaling:
+//      BLOCK_K is pinned to 128 so each K-iteration is exactly one scale
+//      block. Each k-iter accumulates into a temp, then folds
+//      act_scale[row,kb] * w_scale[n/128,kb] into the running accumulator.
+//
+// This reads the FP8 weight directly (no dequant-to-bf16 scratch), cutting
+// per-linear weight traffic ~5x vs fp8_block128_gemm_descale_bf16out while
+// keeping the per-token activation scale (no precision downgrade).
+#include "fp8_block128_gemm_mma_sm89.cuh"
+// Device-side kernel body. Shared verbatim with the standalone micro-bench
+// (benchmarks/sm89_fp8_block128_gemm), so the bench's `--mode baseline` runs
+// this exact kernel and cannot drift behind production.
+#include "fp8_bs_gemm_device.cuh"
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <stdexcept>
+namespace flash_rt {
+namespace gemm {
+namespace block128_sm89 {
+namespace {
+template <int BM, int BN, int W, int STAGES, int MIN_BLK>
+int launch_(const void* A, const void* B, void* D,
+            int M, int N, int K, const float* act_scale,
+            const float* w_scale, cudaStream_t s)
+{
+    constexpr int BK = 128;
+    constexpr int SCALE_KTILE = 8;
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    // Swizzled A/B cp.async stages (no pad) + staged scale tile.
+    int smem_bytes = STAGES * (BM + BN) * BK
+                   + (BM * SCALE_KTILE + SCALE_KTILE) * (int)sizeof(float);
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_bs_gemm_kernel<BM, BN, W, STAGES, MIN_BLK, false>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_bs_gemm_kernel<BM, BN, W, STAGES, MIN_BLK, false><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        act_scale, w_scale,
+        reinterpret_cast<__nv_bfloat16*>(D),
+        M, N, K);
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+// Residual-fold launch: D = bf16(acc) + resid, fusing the residual add into the
+// GEMM epilogue (no separate residual_add launch, no D HBM round-trip). resid
+// is [M, N] BF16 row-major, same layout as D. See fp8_bs_gemm_device.cuh.
+template <int BM, int BN, int W, int STAGES, int MIN_BLK>
+int launch_resid_(const void* A, const void* B, void* D,
+                  int M, int N, int K, const float* act_scale,
+                  const float* w_scale, const void* resid, cudaStream_t s)
+{
+    constexpr int BK = 128;
+    constexpr int SCALE_KTILE = 8;
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    int smem_bytes = STAGES * (BM + BN) * BK
+                   + (BM * SCALE_KTILE + SCALE_KTILE) * (int)sizeof(float);
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_bs_gemm_kernel<BM, BN, W, STAGES, MIN_BLK, true>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_bs_gemm_kernel<BM, BN, W, STAGES, MIN_BLK, true><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        act_scale, w_scale,
+        reinterpret_cast<__nv_bfloat16*>(D),
+        M, N, K,
+        reinterpret_cast<const __nv_bfloat16*>(resid));
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+}  // namespace
+#define DEFINE(NAME, BM, BN, W, S, MB)                                        \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,        \
+           const float* act_scale, const float* w_scale, cudaStream_t s) {    \
+    return launch_<BM, BN, W, S, MB>(A, B, D, M, N, K, act_scale, w_scale, s);\
+  }
+// Residual-fold variants (suffix _resid). D = bf16(acc) + resid. Only the
+// tiles the prefill down-proj actually selects are defined; additive — the
+// non-resid DEFINE list above is unchanged.
+#define DEFINE_RESID(NAME, BM, BN, W, S, MB)                                  \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,        \
+           const float* act_scale, const float* w_scale, const void* resid,   \
+           cudaStream_t s) {                                                  \
+    return launch_resid_<BM, BN, W, S, MB>(A, B, D, M, N, K, act_scale,       \
+                                           w_scale, resid, s);                \
+  }
+// GeGLU silu-fold launch: fuses gate+up GEMM + silu(gate)*up + per-token
+// block-128 FP8 quant into one launch (no [M,2N] BF16 transient). B is
+// gate_up_w [2*N, K] (gate rows [0,N), up rows [N,2N)); w_scale is gate_up_s
+// [2*N/128, K/128]. Outputs FP8 [M,N] + scale [M,N/128]. See device header.
+template <int BM, int BN, int W, int STAGES, int MIN_BLK>
+int launch_geglu_silu_fold_(const void* A, const void* B,
+                            int M, int N, int K, const float* act_scale,
+                            const float* w_scale, void* output, float* out_scale,
+                            cudaStream_t s)
+{
+    constexpr int BK = 128;
+    constexpr int SCALE_KTILE = 8;
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;        // over output N (== inter), NOT 2*N
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    // A/B cp.async stages + gate_smem (BM*BN bf16) + scales + amax scratch.
+    int smem_bytes = STAGES * (BM + BN) * BK
+                   + (BM * BN) * (int)sizeof(__nv_bfloat16)
+                   + (BM * SCALE_KTILE + 2 * SCALE_KTILE) * (int)sizeof(float)
+                   + (W * BM + BM) * (int)sizeof(float);
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_bs_geglu_silu_fold_kernel<BM, BN, W, STAGES, MIN_BLK>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_bs_geglu_silu_fold_kernel<BM, BN, W, STAGES, MIN_BLK><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        act_scale, w_scale,
+        reinterpret_cast<__nv_fp8_e4m3*>(output),
+        out_scale, M, N, K);
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+// A-persistent interleaved variant: stage A once, reuse ONE B smem region for
+// gate then up within each k-iter (both gate+up acc live in regs, true
+// interleaved per K-tile). Single B region -> 3 CTA/SM at s1 (vs interleaved's
+// 2, vs two-pass's 3). See fp8_bs_geglu_silu_fold_apersist_kernel.
+template <int BM, int BN, int W, int STAGES, int MIN_BLK>
+int launch_geglu_silu_fold_apersist_(const void* A, const void* B,
+                                     int M, int N, int K, const float* act_scale,
+                                     const float* w_scale, void* output,
+                                     float* out_scale, cudaStream_t s)
+{
+    constexpr int BK = 128;
+    constexpr int SCALE_KTILE = 8;
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    // Same smem layout as the two-pass variant (gate_smem region kept for layout
+    // parity though apersist doesn't use it as a handoff — gate stays in regs).
+    int smem_bytes = STAGES * (BM + BN) * BK
+                   + (BM * BN) * (int)sizeof(__nv_bfloat16)
+                   + (BM * SCALE_KTILE + 2 * SCALE_KTILE) * (int)sizeof(float)
+                   + (W * BM + BM) * (int)sizeof(float);
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_bs_geglu_silu_fold_apersist_kernel<BM, BN, W, STAGES, MIN_BLK>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_bs_geglu_silu_fold_apersist_kernel<BM, BN, W, STAGES, MIN_BLK><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        act_scale, w_scale,
+        reinterpret_cast<__nv_fp8_e4m3*>(output),
+        out_scale, M, N, K);
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+DEFINE(fp8_block128_gemm_bs_sm89_32x128x128_w4,   32, 128, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w4,   64, 128, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_64x128x128_w8,   64, 128, 8, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_128x128x128_w4, 128, 128, 4, 2, 2)
+DEFINE(fp8_block128_gemm_bs_sm89_128x128x128_w8, 128, 128, 8, 2, 2)
+DEFINE(fp8_block128_gemm_bs_sm89_32x64x128_w4,    32,  64, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_64x64x128_w4,    64,  64, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_128x64x128_w4,  128,  64, 4, 2, 2)
+DEFINE(fp8_block128_gemm_bs_sm89_16x128x128_w4,   16, 128, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_16x64x128_w4,    16,  64, 4, 2, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_32x128x128_w4_s1, 32, 128, 4, 1, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1,  64,  64, 4, 1, 4)
+DEFINE(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1, 128, 128, 8, 1, 2)
+#undef DEFINE
+// Residual-fold variants for the down-proj prefill tiles (see dispatcher
+// below): the 2B/8B down-proj selects 64x64_s1 (8B) / 64x64 (2B small-M) /
+// 32x64 (small-M) at the S ranges Phase-0 measured. Defined additively; the
+// baseline kernels above are untouched.
+DEFINE_RESID(fp8_block128_gemm_bs_sm89_32x64x128_w4_resid,    32,  64, 4, 2, 4)
+DEFINE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_resid,    64,  64, 4, 2, 4)
+DEFINE_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid, 64,  64, 4, 1, 4)
+DEFINE_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid, 128, 128, 8, 1, 2)
+#undef DEFINE_RESID
+// GeGLU silu-fold tile variants (BLOCK_N pinned to 128 = one quant block).
+#define DEFINE_GEGLU(NAME, BM, BN, W, S, MB)                                 \
+  int NAME(const void* A, const void* B, int M, int N, int K,                 \
+           const float* act_scale, const float* w_scale, void* output,        \
+           float* out_scale, cudaStream_t s) {                                \
+    return launch_geglu_silu_fold_<BM, BN, W, S, MB>(                         \
+        A, B, M, N, K, act_scale, w_scale, output, out_scale, s);             \
+  }
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2, 32, 128, 4, 2, 4)
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2, 16, 128, 4, 2, 4)
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2, 64, 128, 4, 2, 4)
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1, 128, 128, 8, 1, 2)
+// Low-smem variants (STAGES=1) to recover occupancy lost to gate_smem on sm89:
+// the s2 dual-buffer + gate_smem pushes dynamic smem >48KB → Block Limit Shared
+// Mem = 1 (8% occupancy, ncu-confirmed). s1 trades cp.async overlap for 3-4x
+// the CTA density. Primary candidates for the prefill M>=128 regime.
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1, 32, 128, 4, 1, 4)
+DEFINE_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1, 16, 128, 4, 1, 4)
+#undef DEFINE_GEGLU
+// A-persistent interleaved variant (single B smem region, gate+up acc both in
+// regs). launch wrapper shares the smem formula with the two-pass variant.
+#define DEFINE_GEGLU_AP(NAME, BM, BN, W, S, MB)                              \
+  int NAME(const void* A, const void* B, int M, int N, int K,                 \
+           const float* act_scale, const float* w_scale, void* output,        \
+           float* out_scale, cudaStream_t s) {                                \
+    return launch_geglu_silu_fold_apersist_<BM, BN, W, S, MB>(                \
+        A, B, M, N, K, act_scale, w_scale, output, out_scale, s);             \
+  }
+DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1, 32, 128, 4, 1, 2)
+DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1, 16, 128, 4, 1, 2)
+DEFINE_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2, 32, 128, 4, 2, 2)
+#undef DEFINE_GEGLU_AP
+int fp8_block128_gemm_blockscaled_sm89_bf16out(
+    const void* A, const void* B, void* D, int M, int N, int K,
+    const float* act_scale, const float* w_scale, cudaStream_t stream)
+{
+    if ((N % 128) != 0)
+        throw std::runtime_error(
+            "fp8_block128_gemm_blockscaled_sm89_bf16out requires N multiple of 128");
+    if ((K % 128) != 0)
+        throw std::runtime_error(
+            "fp8_block128_gemm_blockscaled_sm89_bf16out requires K multiple of 128");
+    // Tuned on 4090 over Qwen3-VL-8B-FP8 layer shapes (qkv 6144, o 4096,
+    // gate/up 12288, down 4096x12288) at S=79..256. BLOCK_M=32 keeps grid
+    // occupancy high at small M; BLOCK_N=64 wins until M crosses ~128, then
+    // the wider BLOCK_N=128 amortizes better. Tiny-N (<2048) prefers BLOCK_N=64.
+    //
+    // ViT prefill is a different regime: full-res FlashRT.png runs M=6256.
+    // On these large-M shapes the language-prefill heuristic is wrong for
+    // the small-N linears:
+    //   - patch_embed / proj   (N=1152, K≈1152..1536) prefer 32x128
+    //   - fc2 / merger-fc2     (N=1152, K>=4096)      prefer 64x64
+    // Keep the original small-M path intact and only branch once the grid is
+    // already abundant (M>=2048), so text prefill / decode remain unchanged.
+    if (N < 2048)
+    {
+        if (M >= 2048) {
+            if (K >= 4096)
+                return fp8_block128_gemm_bs_sm89_64x64x128_w4(
+                    A, B, D, M, N, K, act_scale, w_scale, stream);
+            return fp8_block128_gemm_bs_sm89_32x128x128_w4(
+                A, B, D, M, N, K, act_scale, w_scale, stream);
+        }
+        return fp8_block128_gemm_bs_sm89_16x64x128_w4(
+            A, B, D, M, N, K, act_scale, w_scale, stream);
+    }
+    if (M < 128)
+        return fp8_block128_gemm_bs_sm89_32x64x128_w4(
+            A, B, D, M, N, K, act_scale, w_scale, stream);
+    // Language prefill (M>=128, N>=2048) is limited by low eligible warps on
+    // Ada. A single cp.async stage reduces shared-memory pressure and wins on
+    // Qwen3-VL 2B/8B prefill shapes. Keep a short-prefill exception for the
+    // wide 8B MLP, where the 8-warp tile remains slightly faster.
+    if (N >= 8192 && K == 4096 && M < 1024)
+        return fp8_block128_gemm_bs_sm89_128x128x128_w8_s1(
+            A, B, D, M, N, K, act_scale, w_scale, stream);
+    // Small-M regime (M<256) for N<8192 linears (qkv/o/down): at M=128 the
+    // 64x64/s1 grid underfills the SMs (8B qkv 64x64_s1 = 192 blocks = 1.5/SM;
+    // 2B qkv = 128 blocks = 1/SM), so achieved occupancy is grid-limited well
+    // below the theoretical cap. The smaller 32x64 tile doubles grid_m (8B qkv
+    // -> 384 blocks = 3/SM; 2B qkv -> 256 = 2/SM) and wins despite a lower
+    // per-block warp cap — ncu shows 8B qkv M=128: 32x64 51.6us vs 64x64_s1
+    // 61.9us (-17%). Graph-captured e2e confirms: 2B S=128 -13.5%, 8B S=128
+    // -12.0%, 2B S=192 -5.5%, 8B S=192 -2.0% (gain shrinks as M approaches the
+    // 256 crossover, beyond which 64x64/s1 wins — see layer-regime micro-bench).
+    // Wide-MLP gate_up keeps its existing s1 tile (8B via 128x128_w8_s1 above;
+    // 2B via the default 64x64_s1 below) — it is best at all M.
+    if (M < 256 && N < 8192)
+        return fp8_block128_gemm_bs_sm89_32x64x128_w4(
+            A, B, D, M, N, K, act_scale, w_scale, stream);
+    if (N == 2048 && M < 1024)
+        return fp8_block128_gemm_bs_sm89_64x64x128_w4(
+            A, B, D, M, N, K, act_scale, w_scale, stream);
+    return fp8_block128_gemm_bs_sm89_64x64x128_w4_s1(
+        A, B, D, M, N, K, act_scale, w_scale, stream);
+}
+}  // namespace block128_sm89
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_block128_gemm_mma_sm89.cuh ADDED Viewed

	@@ -0,0 +1,102 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+namespace block128_sm89 {
+// Native Ada (sm_89) FP8 e4m3 -> BF16 block-128 scaled GEMM.
+//
+// Computes D_rm[M,N] = (act_fp8 @ w_fp8^T) with DeepSeek-style block-128
+// scaling applied in the mainloop:
+//   D[m,n] = sum_{kb} act_scale[m, kb] * w_scale[n/128, kb]
+//                     * sum_{k in kb} A[m,k] * B[n,k]
+//
+// Inputs (all device pointers):
+//   A         : [M, K]        FP8 e4m3 row-major   (per-token quantized act)
+//   B         : [N, K]        FP8 e4m3 row-major   (= W, ckpt weight)
+//   act_scale : [M, K/128]    fp32 row-major       (per-token block scale)
+//   w_scale   : [N/128, K/128] fp32 row-major      (weight_scale_inv)
+//   D         : [M, N]        BF16 row-major
+//
+// Drop-in replacement for fp8_block128_gemm_descale_bf16out but reads the
+// FP8 weight directly (no dequant scratch). K and N must be multiples of 128.
+// Returns 0 on success.
+#define DECL(NAME)                                                            \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,        \
+           const float* act_scale, const float* w_scale, cudaStream_t stream)
+DECL(fp8_block128_gemm_bs_sm89_32x128x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_64x128x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_64x128x128_w8);
+DECL(fp8_block128_gemm_bs_sm89_128x128x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_128x128x128_w8);
+DECL(fp8_block128_gemm_bs_sm89_32x64x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_64x64x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_128x64x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_16x128x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_16x64x128_w4);
+DECL(fp8_block128_gemm_bs_sm89_32x128x128_w4_s1);
+DECL(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1);
+DECL(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1);
+#undef DECL
+// Residual-fold tile variants (epilogue adds `resid`): D = bf16(acc) + resid.
+// resid is [M, N] BF16 row-major, same layout as D. Fuses the residual add
+// into the GEMM epilogue (no separate residual_add launch, no D HBM
+// round-trip). Additive — the non-resid kernels above are unchanged.
+#define DECL_RESID(NAME)                                                       \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,         \
+           const float* act_scale, const float* w_scale, const void* resid,    \
+           cudaStream_t stream)
+DECL_RESID(fp8_block128_gemm_bs_sm89_32x64x128_w4_resid);
+DECL_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_resid);
+DECL_RESID(fp8_block128_gemm_bs_sm89_64x64x128_w4_s1_resid);
+DECL_RESID(fp8_block128_gemm_bs_sm89_128x128x128_w8_s1_resid);
+#undef DECL_RESID
+// GeGLU silu-fold tile variants: fuse gate+up GEMM + silu(gate)*up + per-token
+// block-128 FP8 quant into one launch. B = gate_up_w [2*N, K] (gate rows
+// [0,N), up rows [N,2N)); w_scale = gate_up_s [2*N/128, K/128]. Output FP8
+// [M,N] + scale [M,N/128]. BLOCK_N pinned to 128 (one quant block per CTA).
+#define DECL_GEGLU(NAME)                                                     \
+  int NAME(const void* A, const void* B, int M, int N, int K,                 \
+           const float* act_scale, const float* w_scale, void* output,        \
+           float* out_scale, cudaStream_t stream)
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s2);
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s2);
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_64x128_w4_s2);
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_128x128_w8_s1);
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1);
+DECL_GEGLU(fp8_bs_geglu_silu_fold_sm89_16x128_w4_s1);
+#undef DECL_GEGLU
+// A-persistent interleaved variant (single B smem region, both gate+up acc in
+// registers). Same I/O contract as DECL_GEGLU.
+#define DECL_GEGLU_AP(NAME)                                                    \
+  int NAME(const void* A, const void* B, int M, int N, int K,                 \
+           const float* act_scale, const float* w_scale, void* output,        \
+           float* out_scale, cudaStream_t stream)
+DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s1);
+DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_16x128_w4_s1);
+DECL_GEGLU_AP(fp8_bs_geglu_silu_fold_apersist_sm89_32x128_w4_s2);
+#undef DECL_GEGLU_AP
+// Auto-dispatch over the tuned tile set above based on (M, N, K).
+int fp8_block128_gemm_blockscaled_sm89_bf16out(
+    const void* A, const void* B, void* D, int M, int N, int K,
+    const float* act_scale, const float* w_scale, cudaStream_t stream);
+}  // namespace block128_sm89
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_bs_gemm_device.cuh ADDED Viewed

	@@ -0,0 +1,1327 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+// Shared device-side implementation of the SM89 FP8 block-128 scaled GEMM
+// kernel. This header is the single source of truth for the kernel body: both
+// the production launcher (fp8_block128_gemm_mma_sm89.cu) and the standalone
+// micro-benchmark (benchmarks/sm89_fp8_block128_gemm) include it, so the
+// bench's `--mode baseline` runs the *exact* production kernel and cannot
+// drift behind it. When experimenting, copy this kernel into the bench's
+// candidate slot and edit there; once an experiment is accepted and folded
+// back here, the bench baseline tracks it automatically.
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+namespace flash_rt {
+namespace gemm {
+namespace block128_sm89 {
+__device__ __forceinline__ void mma_m16n8k32_e4m3(
+    float &d0, float &d1, float &d2, float &d3,
+    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
+    uint32_t b0, uint32_t b1)
+{
+    // Ada (sm_89) FP8 tensor-core op — NO .kind::f8f6f4 qualifier.
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
+        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+__device__ __forceinline__ void cp_async_16(uint32_t smem, const uint8_t* src) {
+    int b = (src == nullptr) ? 0 : 16;
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n"
+                 :: "r"(smem), "l"(src), "r"(b));
+}
+__device__ __forceinline__ uint32_t to_smem(const void* p) {
+    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
+}
+// True when the adjacent output column pair {c, c+1} is fully in bounds, so a
+// 32-bit bfloat162 store is valid. n_pair_base is even (=...+2*l) and N is a
+// multiple of 128, so &D[row*N + c] is 4-byte aligned for the vector store.
+__device__ __forceinline__ bool col_pair_ok(int c, int N) {
+    return c + 1 < N;
+}
+// ldmatrix.x4: load four 8x8 b16 fragments from smem into 4 registers/lane in
+// one instruction, replacing 4 scalar 32-bit LDS to offload the LSU pipe
+// (NCU on the scalar path: LSU 67.7%, 54.7M shared loads = 27% of all insts).
+__device__ __forceinline__ void ldmatrix_x4_b16(
+    uint32_t &d0, uint32_t &d1, uint32_t &d2, uint32_t &d3, uint32_t smem_addr)
+{
+    asm volatile(
+        "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+        : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
+        : "r"(smem_addr));
+}
+// SiLU in fp32. Matches quantize::silu_f32 (fp8_per_token_block_quant.cu:416)
+// so the GeGLU silu-fold epilogue reproduces silu_mul_merged's math exactly.
+__device__ __forceinline__ float silu_f32(float x) {
+    return x / (1.0f + expf(-x));
+}
+// BLOCK_K is pinned to 128 (one DeepSeek scale block per K-iteration).
+//  - A: [M, K] row-major FP8 e4m3, act_scale [M, K/128] fp32
+//  - B: [N, K] row-major FP8 e4m3, w_scale [N/128, K/128] fp32
+//  - D: [M, N] row-major BF16
+//  - BLOCK_N must keep each warp's 8-wide N-atoms inside one 128 scale block.
+//
+// RESID (opt-in epilogue fold): when true, the BF16 store adds a per-element
+// residual `resid[M, N]` (same BF16 layout as D): D = bf16(acc + resid).
+// This folds what would otherwise be a separate residual_add launch + an
+// extra D round-trip through HBM, mirroring #134's residual-fold epilogue.
+// When RESID=false, `resid` is unused and the `if constexpr (RESID)` branch
+// is dead-stripped at compile time, so the baseline kernel is byte-identical.
+template <int BLOCK_M, int BLOCK_N, int NUM_WARPS, int STAGES,
+          int MIN_BLOCKS_PER_SM, bool RESID = false>
+__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM)
+void fp8_bs_gemm_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,
+    const float* __restrict__ act_scale,   // [M, K/128]
+    const float* __restrict__ w_scale,     // [N/128, K/128]
+    __nv_bfloat16* __restrict__ D,
+    int M, int N, int K,
+    const __nv_bfloat16* __restrict__ resid = nullptr)  // [M, N] BF16, used iff RESID
+{
+    constexpr int BLOCK_K    = 128;
+    constexpr int THREADS    = NUM_WARPS * 32;
+    constexpr int M_ATOMS    = BLOCK_M / 16;
+    constexpr int N_ATOMS    = BLOCK_N / 8;
+    constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS;
+    constexpr int N_PAIRS_PW = N_ATOMS_PW / 2;      // ldmatrix pairs 2 N-atoms
+    constexpr int K_ATOMS    = BLOCK_K / 32;        // = 4
+    constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16;  // 8 chunks of 16 bytes
+    // 128B swizzle: chunk_sw = chunk ^ (row & SWIZZLE_MASK). Removes the old
+    // SMEM_K_PAD and the bank conflicts; applied identically on cp.async store
+    // and ldmatrix load so the round-trip is bit-exact.
+    constexpr int SWIZZLE_MASK = NUM_CHUNKS_PER_ROW - 1;  // = 7
+    static_assert(BLOCK_M % 16 == 0, "BLOCK_M multiple of 16");
+    static_assert(BLOCK_N % 8 == 0,  "BLOCK_N multiple of 8");
+    static_assert(BLOCK_N <= 128, "one CTA must fit one N scale block");
+    static_assert((BLOCK_N / 8) % NUM_WARPS == 0, "N-atoms split across warps");
+    static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0,
+                  "ldmatrix pairs 2 N-atoms: N_ATOMS_PW must be even >= 2");
+    // Stage the per-CTA activation/weight scales in shared memory with a
+    // coalesced load, so the per-k_iter scale fold reads smem instead of
+    // row-strided scalar global loads (NCU's top global-load bottleneck).
+    // Only SCALE_KTILE scale-block columns are staged at a time, re-staged on
+    // each k-tile boundary, so the smem footprint is K-independent (~2 KB) and
+    // occupancy does not regress on large-K shapes (e.g. down, K128=96).
+    constexpr int SCALE_KTILE = 8;
+    constexpr int A_TILE = BLOCK_M * BLOCK_K;       // swizzled, no pad
+    constexpr int B_TILE = BLOCK_N * BLOCK_K;
+    extern __shared__ uint8_t smem_raw[];
+    uint8_t* A_smem = smem_raw;
+    uint8_t* B_smem = A_smem + STAGES * A_TILE;
+    float* as_smem = reinterpret_cast<float*>(B_smem + STAGES * B_TILE);
+    float* ws_smem = as_smem + BLOCK_M * SCALE_KTILE;
+    const int cta_m = blockIdx.x;
+    const int cta_n = blockIdx.y;
+    const int m_base = cta_m * BLOCK_M;
+    const int n_base = cta_n * BLOCK_N;
+    const int t = threadIdx.x;
+    const int warp_id = t / 32;
+    const int lane = t % 32;
+    const int l = lane % 4;
+    const int h = lane / 4;
+    // ldmatrix.x4 lane -> fragment partition.
+    const int frag_group = lane / 8;       // 0..3 (TL,TR,BL,BR)
+    const int row_in_frag = lane % 8;      // row within an 8x8 fragment
+    const int row_block = frag_group / 2;  // top(0)/bottom(1) 8 rows
+    const int col_block = frag_group % 2;  // left(0)/right(1) 16-byte chunk
+    const int K128 = K >> 7;                        // # scale blocks along K
+    // Coalesced staging of one SCALE_KTILE-wide scale block into smem.
+    auto stage_scales = [&](int kb0) {
+        const int as_total = BLOCK_M * SCALE_KTILE;
+        for (int idx = t; idx < as_total; idx += THREADS) {
+            int r = idx / SCALE_KTILE;
+            int kc = idx - r * SCALE_KTILE;
+            int row = m_base + r;
+            int kb = kb0 + kc;
+            as_smem[idx] = (row < M && kb < K128)
+                ? act_scale[(size_t)row * K128 + kb] : 0.0f;
+        }
+        for (int kc = t; kc < SCALE_KTILE; kc += THREADS) {
+            int kb = kb0 + kc;
+            ws_smem[kc] = (kb < K128)
+                ? w_scale[(size_t)(n_base >> 7) * K128 + kb] : 0.0f;
+        }
+        __syncthreads();
+    };
+    auto issue_load = [&](int stage, int k_base) {
+        constexpr int A_CHUNKS = BLOCK_M * NUM_CHUNKS_PER_ROW;
+        constexpr int A_ITERS = (A_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < A_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= A_CHUNKS) break;
+            int row_a = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_a = idx % NUM_CHUNKS_PER_ROW;
+            int m_glob = m_base + row_a;
+            int k_glob = k_base + chunk_a * 16;
+            const uint8_t* a_src = nullptr;
+            if (m_glob < M && k_glob < K) {
+                a_src = reinterpret_cast<const uint8_t*>(&A[(size_t)m_glob * K + k_glob]);
+            }
+            int csw = chunk_a ^ (row_a & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&A_smem[stage * A_TILE + row_a * BLOCK_K + csw * 16]),
+                a_src);
+        }
+        constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW;
+        constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < B_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= B_CHUNKS) break;
+            int row_b = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_b = idx % NUM_CHUNKS_PER_ROW;
+            int n_glob = n_base + row_b;
+            int k_glob = k_base + chunk_b * 16;
+            const uint8_t* b_src = nullptr;
+            if (n_glob < N && k_glob < K) {
+                b_src = reinterpret_cast<const uint8_t*>(&B[(size_t)n_glob * K + k_glob]);
+            }
+            int csw = chunk_b ^ (row_b & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&B_smem[stage * B_TILE + row_b * BLOCK_K + csw * 16]),
+                b_src);
+        }
+    };
+    // Running (scaled) accumulators across all K-blocks.
+    float acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi)
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) acc[mi][ni][j] = 0.0f;
+    const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K;
+    #pragma unroll
+    for (int s = 0; s < STAGES - 1; ++s) {
+        int kb = s * BLOCK_K;
+        if (kb < K) issue_load(s, kb);
+        asm volatile("cp.async.commit_group;\n" ::);
+    }
+    int compute_stage = 0;
+    for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) {
+        int issue_iter = k_iter + (STAGES - 1);
+        int issue_stage = issue_iter % STAGES;
+        if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K);
+        asm volatile("cp.async.commit_group;\n" ::);
+        asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1));
+        __syncthreads();
+        // This k_iter is exactly one scale block (kb = k_iter).
+        const int kb = k_iter;
+        // Re-stage the next SCALE_KTILE-wide scale block on each tile boundary.
+        if ((kb % SCALE_KTILE) == 0) stage_scales(kb);
+        // w_scale is constant across this CTA's BLOCK_N if it fits one
+        // 128 block; index per warp's N base to stay correct for BLOCK_N>128.
+        float tacc[M_ATOMS][N_ATOMS_PW][4];
+        #pragma unroll
+        for (int mi = 0; mi < M_ATOMS; ++mi)
+            #pragma unroll
+            for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+                #pragma unroll
+                for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f;
+        uint8_t* A_stage = A_smem + compute_stage * A_TILE;
+        uint8_t* B_stage = B_smem + compute_stage * B_TILE;
+        #pragma unroll
+        for (int ka = 0; ka < K_ATOMS; ++ka) {
+            // ldmatrix.x4 loads the m16xk32 A fragment (4 regs/lane) per m-atom.
+            uint32_t A_regs[M_ATOMS][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row = mi * 16 + row_block * 8 + row_in_frag;
+                int chunk = 2 * ka + col_block;
+                int csw = chunk ^ (row & SWIZZLE_MASK);
+                ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                                to_smem(&A_stage[row * BLOCK_K + csw * 16]));
+            }
+            // ldmatrix.x4 loads two N-atoms (n16xk32) per pair.
+            uint32_t B_regs[N_PAIRS_PW][4];
+            #pragma unroll
+            for (int np = 0; np < N_PAIRS_PW; ++np) {
+                int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag;
+                int chunk = 2 * ka + col_block;
+                int csw = chunk ^ (nrow & SWIZZLE_MASK);
+                ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                                to_smem(&B_stage[nrow * BLOCK_K + csw * 16]));
+            }
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int ni0 = np * 2, ni1 = np * 2 + 1;
+                    // ldm fragment -> mma A operand: a0=d0,a1=d2,a2=d1,a3=d3.
+                    mma_m16n8k32_e4m3(
+                        tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3],
+                        A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][0], B_regs[np][1]);
+                    mma_m16n8k32_e4m3(
+                        tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3],
+                        A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][2], B_regs[np][3]);
+                }
+            }
+        }
+        // Fold block scales: D += act_scale[row,kb] * w_scale[ncol/128,kb] * tacc
+        // Scales come from the smem stage (coalesced load above), indexed by
+        // the column within the current SCALE_KTILE tile. BLOCK_N <= 128 keeps
+        // the CTA inside one 128-column weight-scale block.
+        int kbt = kb % SCALE_KTILE;
+        float ws_cta = ws_smem[kbt];
+        #pragma unroll
+        for (int mi = 0; mi < M_ATOMS; ++mi) {
+            int row0 = m_base + mi * 16 + h;
+            int row1 = row0 + 8;
+            float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt];
+            float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt];
+            #pragma unroll
+            for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+                acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta);
+                acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta);
+                acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta);
+                acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta);
+            }
+        }
+        // Do not let the next cp.async overwrite this shared-memory stage
+        // before all warps finish reading it.
+        __syncthreads();
+        compute_stage = (compute_stage + 1) % STAGES;
+    }
+    asm volatile("cp.async.wait_all;\n" ::);
+    // Epilogue: write BF16. m16n8 layout: thread (h,l) -> rows {h,h+8},
+    // cols {2*l, 2*l+1}.
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            // acc[0,1] = row0 cols {2l,2l+1}; acc[2,3] = row1 cols {2l,2l+1}.
+            // RESID epilogue fold: add the BF16 residual in-register before the
+            // bf16 store, so the residual read is fused into the GEMM epilogue
+            // and never lands as a separate D HBM round-trip + launch.
+            if constexpr (RESID) {
+                if (row0 < M && col_pair_ok(n_pair_base, N)) {
+                    __nv_bfloat162 r = *reinterpret_cast<const __nv_bfloat162*>(
+                        &resid[(size_t)row0 * N + n_pair_base]);
+                    *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row0 * N + n_pair_base]) =
+                        __floats2bfloat162_rn(acc[mi][ni][0] + __low2float(r),
+                                               acc[mi][ni][1] + __high2float(r));
+                } else if (row0 < M) {
+                    if (n_pair_base < N)
+                        D[(size_t)row0 * N + n_pair_base] = __float2bfloat16(
+                            acc[mi][ni][0] + __bfloat162float(resid[(size_t)row0 * N + n_pair_base]));
+                    if (n_pair_base + 1 < N)
+                        D[(size_t)row0 * N + n_pair_base + 1] = __float2bfloat16(
+                            acc[mi][ni][1] + __bfloat162float(resid[(size_t)row0 * N + n_pair_base + 1]));
+                }
+                if (row1 < M && col_pair_ok(n_pair_base, N)) {
+                    __nv_bfloat162 r = *reinterpret_cast<const __nv_bfloat162*>(
+                        &resid[(size_t)row1 * N + n_pair_base]);
+                    *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row1 * N + n_pair_base]) =
+                        __floats2bfloat162_rn(acc[mi][ni][2] + __low2float(r),
+                                               acc[mi][ni][3] + __high2float(r));
+                } else if (row1 < M) {
+                    if (n_pair_base < N)
+                        D[(size_t)row1 * N + n_pair_base] = __float2bfloat16(
+                            acc[mi][ni][2] + __bfloat162float(resid[(size_t)row1 * N + n_pair_base]));
+                    if (n_pair_base + 1 < N)
+                        D[(size_t)row1 * N + n_pair_base + 1] = __float2bfloat16(
+                            acc[mi][ni][3] + __bfloat162float(resid[(size_t)row1 * N + n_pair_base + 1]));
+                }
+            } else {
+                // Emit one 32-bit bfloat162 store per row instead of two scalar
+                // 16-bit stores (NCU's top store-pattern bottleneck after C1).
+                // Tail (odd last column) falls back to scalar stores.
+                if (row0 < M && col_pair_ok(n_pair_base, N)) {
+                    *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row0 * N + n_pair_base]) =
+                        __floats2bfloat162_rn(acc[mi][ni][0], acc[mi][ni][1]);
+                } else if (row0 < M) {
+                    if (n_pair_base < N)     D[(size_t)row0 * N + n_pair_base]   = __float2bfloat16(acc[mi][ni][0]);
+                    if (n_pair_base + 1 < N) D[(size_t)row0 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][1]);
+                }
+                if (row1 < M && col_pair_ok(n_pair_base, N)) {
+                    *reinterpret_cast<__nv_bfloat162*>(&D[(size_t)row1 * N + n_pair_base]) =
+                        __floats2bfloat162_rn(acc[mi][ni][2], acc[mi][ni][3]);
+                } else if (row1 < M) {
+                    if (n_pair_base < N)     D[(size_t)row1 * N + n_pair_base]   = __float2bfloat16(acc[mi][ni][2]);
+                    if (n_pair_base + 1 < N) D[(size_t)row1 * N + n_pair_base+1] = __float2bfloat16(acc[mi][ni][3]);
+                }
+            }
+        }
+    }
+}
+// ============================================================================
+// GeGLU silu-fold megakernel (Phase 2): fuses gate GEMM + up GEMM +
+// silu(gate)*up + per-token block-128 FP8 quant into ONE launch, writing FP8
+// output + scale directly — eliminating the [M, 2*N] BF16 transient that the
+// baseline gate_up GEMM would write and silu_mul_merged_to_fp8 would read back.
+//
+//   gate_up_w : [2*N, K] FP8 row-major  (gate rows [0,N); up rows [N,2N))
+//   gate_up_s : [2*N/128, K/128] fp32   (up row = gate row + N/128)
+//   A         : [M, K] FP8 (per-token quantized),  act_scale [M, K/128]
+//   output    : [M, N] FP8,  scale [M, N/128]
+//
+// Two-pass per CTA (mirrors sm100 flashrt_megakernel_geglu's "gate stays in
+// smem"): pass 1 accumulates gate over full K and stores silu(gate) as BF16
+// into a smem gate buffer; pass 2 reuses the same A/B smem staging, accumulates
+// up, then the epilogue reads gate from smem, forms v = bf16(bf16(silu(gate))*up)
+// (matching silu_mul_merged's two bf16 roundings), reduces |v| over the 128-col
+// quant block per row, and quantizes to FP8. No grid_barrier (single CTA owns
+// its full quant block: BLOCK_N == 128 == one scale block). GEMM body reuses
+// the same cp.async + ldmatrix.x4 + mma.m16n8k32 tiles as fp8_bs_gemm_kernel.
+// ============================================================================
+template <int BLOCK_M, int BLOCK_N, int NUM_WARPS, int STAGES,
+          int MIN_BLOCKS_PER_SM>
+__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM)
+void fp8_bs_geglu_silu_fold_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,        // gate_up_w [2*N, K]
+    const float* __restrict__ act_scale,        // [M, K/128]
+    const float* __restrict__ w_scale,          // gate_up_s [2*N/128, K/128]
+    __nv_fp8_e4m3* __restrict__ output,         // [M, N]
+    float* __restrict__ out_scale,              // [M, N/128]
+    int M, int N, int K)
+{
+    static_assert(BLOCK_N == 128,
+        "GeGLU silu-fold requires BLOCK_N==128 (one quant block per CTA)");
+    constexpr int BLOCK_K    = 128;
+    constexpr int THREADS    = NUM_WARPS * 32;
+    constexpr int M_ATOMS    = BLOCK_M / 16;
+    constexpr int N_ATOMS    = BLOCK_N / 8;       // 16
+    constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS;
+    constexpr int N_PAIRS_PW = N_ATOMS_PW / 2;
+    constexpr int K_ATOMS    = BLOCK_K / 32;      // 4
+    constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16;
+    constexpr int SWIZZLE_MASK = NUM_CHUNKS_PER_ROW - 1;
+    constexpr int SCALE_KTILE = 8;
+    constexpr int A_TILE = BLOCK_M * BLOCK_K;
+    constexpr int B_TILE = BLOCK_N * BLOCK_K;
+    static_assert(BLOCK_M % 16 == 0, "BLOCK_M multiple of 16");
+    static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0,
+                  "ldmatrix pairs 2 N-atoms: N_ATOMS_PW must be even >= 2");
+    extern __shared__ uint8_t smem_raw[];
+    uint8_t* A_smem = smem_raw;
+    uint8_t* B_smem = A_smem + STAGES * A_TILE;
+    // gate_smem: silu(gate) as BF16, [BLOCK_M, BLOCK_N]. One CTA-tile, written
+    // by pass 1 epilogue, read by pass 2 epilogue. The sm100 geglu's "gate
+    // stays in smem" handoff, without tcgen05/EVT.
+    __nv_bfloat16* gate_smem = reinterpret_cast<__nv_bfloat16*>(
+        B_smem + STAGES * B_TILE);
+    float* as_smem = reinterpret_cast<float*>(gate_smem + BLOCK_M * BLOCK_N);
+    float* wsg_smem = as_smem + BLOCK_M * SCALE_KTILE;   // gate w_scale row
+    float* wsu_smem = wsg_smem + SCALE_KTILE;            // up   w_scale row
+    // amax partials: 4 warps × BLOCK_M rows. Cross-warp reduce per row.
+    float* amax_smem = wsu_smem + SCALE_KTILE;
+    const int cta_m = blockIdx.x;
+    const int cta_n = blockIdx.y;
+    const int m_base = cta_m * BLOCK_M;
+    const int n_base = cta_n * BLOCK_N;          // n0, < N (output col block)
+    const int t = threadIdx.x;
+    const int warp_id = t / 32;
+    const int lane = t % 32;
+    const int l = lane % 4;
+    const int h = lane / 4;
+    const int frag_group = lane / 8;
+    const int row_in_frag = lane % 8;
+    const int row_block = frag_group / 2;
+    const int col_block = frag_group % 2;
+    const int K128 = K >> 7;
+    const int N128 = N >> 7;                       // gate w_scale blocks
+    // gate B-rows [n_base, n_base+BLOCK_N); up B-rows [n_base+N, n_base+N+BLOCK_N)
+    const int gate_b_row0 = n_base;
+    const int up_b_row0   = n_base + N;
+    const int gate_ws_row = (n_base >> 7);          // gate w_scale block row
+    const int up_ws_row   = gate_ws_row + N128;     // up w_scale block row
+    // ---- scale staging (shared by both passes; re-staged per SCALE_KTILE) ----
+    auto stage_scales = [&](int kb0) {
+        const int as_total = BLOCK_M * SCALE_KTILE;
+        for (int idx = t; idx < as_total; idx += THREADS) {
+            int r = idx / SCALE_KTILE;
+            int kc = idx - r * SCALE_KTILE;
+            int row = m_base + r;
+            int kb = kb0 + kc;
+            as_smem[idx] = (row < M && kb < K128)
+                ? act_scale[(size_t)row * K128 + kb] : 0.0f;
+        }
+        for (int kc = t; kc < SCALE_KTILE; kc += THREADS) {
+            int kb = kb0 + kc;
+            wsg_smem[kc] = (kb < K128)
+                ? w_scale[(size_t)gate_ws_row * K128 + kb] : 0.0f;
+            wsu_smem[kc] = (kb < K128)
+                ? w_scale[(size_t)up_ws_row * K128 + kb] : 0.0f;
+        }
+        __syncthreads();
+    };
+    // ---- cp.async A + (gate or up) B tile staging ----
+    // b_row0 selects which 128-row band of B [2*N, K] to stage.
+    auto issue_load = [&](int stage, int k_base, int b_row0) {
+        constexpr int A_CHUNKS = BLOCK_M * NUM_CHUNKS_PER_ROW;
+        constexpr int A_ITERS = (A_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < A_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= A_CHUNKS) break;
+            int row_a = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_a = idx % NUM_CHUNKS_PER_ROW;
+            int m_glob = m_base + row_a;
+            int k_glob = k_base + chunk_a * 16;
+            const uint8_t* a_src = nullptr;
+            if (m_glob < M && k_glob < K) {
+                a_src = reinterpret_cast<const uint8_t*>(&A[(size_t)m_glob * K + k_glob]);
+            }
+            int csw = chunk_a ^ (row_a & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&A_smem[stage * A_TILE + row_a * BLOCK_K + csw * 16]),
+                a_src);
+        }
+        constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW;
+        constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < B_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= B_CHUNKS) break;
+            int row_b = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_b = idx % NUM_CHUNKS_PER_ROW;
+            int n_glob = b_row0 + row_b;
+            int k_glob = k_base + chunk_b * 16;
+            const uint8_t* b_src = nullptr;
+            if (n_glob < 2 * N && k_glob < K) {
+                b_src = reinterpret_cast<const uint8_t*>(&B[(size_t)n_glob * K + k_glob]);
+            }
+            int csw = chunk_b ^ (row_b & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&B_smem[stage * B_TILE + row_b * BLOCK_K + csw * 16]),
+                b_src);
+        }
+    };
+    // ---- one GEMM pass over full K, accumulating into `acc` with the given
+    // w_scale smem row (gate or up). b_row0 selects the B band. ----
+    auto run_pass = [&](float (*acc)[N_ATOMS_PW][4], int b_row0,
+                        const float* ws_smem_pass) {
+        const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K;
+        #pragma unroll
+        for (int s = 0; s < STAGES - 1; ++s) {
+            int kb = s * BLOCK_K;
+            if (kb < K) issue_load(s, kb, b_row0);
+            asm volatile("cp.async.commit_group;\n" ::);
+        }
+        int compute_stage = 0;
+        for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) {
+            int issue_iter = k_iter + (STAGES - 1);
+            int issue_stage = issue_iter % STAGES;
+            if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K, b_row0);
+            asm volatile("cp.async.commit_group;\n" ::);
+            asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1));
+            __syncthreads();
+            const int kb = k_iter;
+            if ((kb % SCALE_KTILE) == 0) stage_scales(kb);
+            float tacc[M_ATOMS][N_ATOMS_PW][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi)
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+                    #pragma unroll
+                    for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f;
+            uint8_t* A_stage = A_smem + compute_stage * A_TILE;
+            uint8_t* B_stage = B_smem + compute_stage * B_TILE;
+            #pragma unroll
+            for (int ka = 0; ka < K_ATOMS; ++ka) {
+                uint32_t A_regs[M_ATOMS][4];
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    int row = mi * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (row & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                                    to_smem(&A_stage[row * BLOCK_K + csw * 16]));
+                }
+                uint32_t B_regs[N_PAIRS_PW][4];
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (nrow & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                                    to_smem(&B_stage[nrow * BLOCK_K + csw * 16]));
+                }
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    #pragma unroll
+                    for (int np = 0; np < N_PAIRS_PW; ++np) {
+                        int ni0 = np * 2, ni1 = np * 2 + 1;
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][0], B_regs[np][1]);
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][2], B_regs[np][3]);
+                    }
+                }
+            }
+            int kbt = kb % SCALE_KTILE;
+            float ws_cta = ws_smem_pass[kbt];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row0 = m_base + mi * 16 + h;
+                int row1 = row0 + 8;
+                float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt];
+                float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt];
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+                    acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta);
+                    acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta);
+                    acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta);
+                    acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta);
+                }
+            }
+            __syncthreads();
+            compute_stage = (compute_stage + 1) % STAGES;
+        }
+        asm volatile("cp.async.wait_all;\n" ::);
+    };
+    // =================== Pass 1: gate GEMM → silu(gate) in smem ===================
+    float gate_acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi)
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) gate_acc[mi][ni][j] = 0.0f;
+    run_pass(gate_acc, gate_b_row0, wsg_smem);
+    // Pass 1 epilogue: store silu(gate_acc) as BF16 into gate_smem[BM, BN].
+    // Thread (h,l) owns rows {mi*16+h, mi*16+h+8}, cols {ni*8+2l, ni*8+2l+1}
+    // within its warp's N band. Replicate silu_mul_merged's first bf16 rounding
+    // (bf16(silu(g))) so the fused path matches the split kernel's precision.
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            // gate_smem is [BLOCK_M, BLOCK_N]; local col = n_pair_base.
+            if (row0 < M) {
+                __nv_bfloat162 gs = __floats2bfloat162_rn(
+                    silu_f32(gate_acc[mi][ni][0]), silu_f32(gate_acc[mi][ni][1]));
+                *reinterpret_cast<__nv_bfloat162*>(
+                    &gate_smem[(row0 - m_base) * BLOCK_N + n_pair_base]) = gs;
+                __nv_bfloat162 gs2 = __floats2bfloat162_rn(
+                    silu_f32(gate_acc[mi][ni][2]), silu_f32(gate_acc[mi][ni][3]));
+                *reinterpret_cast<__nv_bfloat162*>(
+                    &gate_smem[(row1 - m_base) * BLOCK_N + n_pair_base]) = gs2;
+            }
+        }
+    }
+    __syncthreads();   // gate_smem visible to pass 2 epilogue in all warps
+    // gate_acc registers now free; reused for up_acc.
+    // =================== Pass 2: up GEMM → up_acc ===================
+    float up_acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi)
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) up_acc[mi][ni][j] = 0.0f;
+    run_pass(up_acc, up_b_row0, wsu_smem);
+    // =================== Pass 2 epilogue: silu(gate)*up + quant → FP8 ===================
+    // v = bf16(bf16(silu(gate)) * up), matching silu_mul_merged's two bf16
+    // roundings (silu(gate) was already bf16-rounded into gate_smem in pass 1;
+    // here we bf16-round the product). Then per-row amax over the 128-col block
+    // and quantize.
+    constexpr float kFp8Max = 448.0f;
+    // Each thread owns 8 cols (4 n-atoms × 2) for 2 rows per m-atom. Compute |v|
+    // and a per-warp partial amax per row (the warp owns 32 of the row's 128 cols).
+    // amax_smem[warp_id][row_in_cta] holds the warp's row-amax partial.
+    float v[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        int rloc0 = mi * 16 + h;          // local row in [0, BLOCK_M)
+        int rloc1 = rloc0 + 8;
+        float amax0 = 0.0f, amax1 = 0.0f;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            // gate value (already bf16(silu(gate))) from smem; up from registers.
+            if (row0 < M) {
+                __nv_bfloat162 g = *reinterpret_cast<const __nv_bfloat162*>(
+                    &gate_smem[rloc0 * BLOCK_N + n_pair_base]);
+                float gf0 = __low2float(g), gf1 = __high2float(g);
+                v[mi][ni][0] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][0]));
+                v[mi][ni][1] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][1]));
+                amax0 = fmaxf(amax0, fmaxf(fabsf(v[mi][ni][0]), fabsf(v[mi][ni][1])));
+            } else {
+                v[mi][ni][0] = 0.0f; v[mi][ni][1] = 0.0f;
+            }
+            if (row1 < M) {
+                __nv_bfloat162 g = *reinterpret_cast<const __nv_bfloat162*>(
+                    &gate_smem[rloc1 * BLOCK_N + n_pair_base]);
+                float gf0 = __low2float(g), gf1 = __high2float(g);
+                v[mi][ni][2] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][2]));
+                v[mi][ni][3] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][3]));
+                amax1 = fmaxf(amax1, fmaxf(fabsf(v[mi][ni][2]), fabsf(v[mi][ni][3])));
+            } else {
+                v[mi][ni][2] = 0.0f; v[mi][ni][3] = 0.0f;
+            }
+        }
+        // Warp-shuffle reduce the 4 lanes (l=0..3) that share row0 / row1.
+        for (int off = 2; off > 0; off >>= 1) {
+            amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffff, amax0, off));
+            amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffff, amax1, off));
+        }
+        if (l == 0) {
+            amax_smem[warp_id * BLOCK_M + rloc0] = amax0;
+            amax_smem[warp_id * BLOCK_M + rloc1] = amax1;
+        }
+    }
+    __syncthreads();
+    // Cross-warp reduce: each warp wrote its row-amax partial. Final reduce per
+    // row done by warp 0 lanes, broadcast via smem.
+    #pragma unroll
+    for (int rloc = t; rloc < BLOCK_M; rloc += THREADS) {
+        int row = m_base + rloc;
+        if (row >= M) continue;
+        float amax = 0.0f;
+        #pragma unroll
+        for (int w = 0; w < NUM_WARPS; ++w)
+            amax = fmaxf(amax, amax_smem[w * BLOCK_M + rloc]);
+        float sc = fmaxf(amax / kFp8Max, 1.0e-12f);
+        amax_smem[rloc] = sc;     // reuse slot to broadcast final scale
+        // Each active thread owns a distinct rloc in this strided loop, so each
+        // writes its own row's scale — no race. (The earlier `warp_id==0 &&
+        // lane==0` guard let only thread 0 write, leaving rows 1..BLOCK_M-1
+        // unwritten → garbage out_scale, correct-but-unscaled fp8 output.)
+        out_scale[(size_t)row * (N >> 7) + (n_base >> 7)] = sc;
+    }
+    __syncthreads();
+    // Quantize + store FP8. Thread re-reads its v[] and the row's scale.
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        int rloc0 = mi * 16 + h;
+        int rloc1 = rloc0 + 8;
+        float sc0 = (row0 < M) ? amax_smem[rloc0] : 1.0f;
+        float sc1 = (row1 < M) ? amax_smem[rloc1] : 1.0f;
+        float inv0 = 1.0f / sc0, inv1 = 1.0f / sc1;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            if (row0 < M && col_pair_ok(n_pair_base, N)) {
+                float q0 = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max);
+                float q1 = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max);
+                // pack two fp8 e4m3 into a 16-bit store
+                __nv_fp8_e4m3 p0(q0), p1(q1);
+                uint16_t pack = (uint16_t)(*reinterpret_cast<const uint8_t*>(&p1)) << 8
+                              | (uint16_t)(*reinterpret_cast<const uint8_t*>(&p0));
+                *reinterpret_cast<uint16_t*>(&output[(size_t)row0 * N + n_pair_base]) = pack;
+            } else if (row0 < M) {
+                if (n_pair_base < N) {
+                    float q = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max);
+                    output[(size_t)row0 * N + n_pair_base] = __nv_fp8_e4m3(q);
+                }
+                if (n_pair_base + 1 < N) {
+                    float q = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max);
+                    output[(size_t)row0 * N + n_pair_base + 1] = __nv_fp8_e4m3(q);
+                }
+            }
+            if (row1 < M && col_pair_ok(n_pair_base, N)) {
+                float q2 = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max);
+                float q3 = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max);
+                __nv_fp8_e4m3 p2(q2), p3(q3);
+                uint16_t pack = (uint16_t)(*reinterpret_cast<const uint8_t*>(&p3)) << 8
+                              | (uint16_t)(*reinterpret_cast<const uint8_t*>(&p2));
+                *reinterpret_cast<uint16_t*>(&output[(size_t)row1 * N + n_pair_base]) = pack;
+            } else if (row1 < M) {
+                if (n_pair_base < N) {
+                    float q = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max);
+                    output[(size_t)row1 * N + n_pair_base] = __nv_fp8_e4m3(q);
+                }
+                if (n_pair_base + 1 < N) {
+                    float q = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max);
+                    output[(size_t)row1 * N + n_pair_base + 1] = __nv_fp8_e4m3(q);
+                }
+            }
+        }
+    }
+}
+// ============================================================================
+// GeGLU silu-fold, A-persistent two-pass variant.
+//
+// Same fusion as fp8_bs_geglu_silu_fold_kernel (gate+up GEMM + silu(gate)*up +
+// per-token block-128 FP8 quant, one launch, no [M,2N] BF16 transient), but a
+// different smem/register strategy that fixes the two-pass weaknesses the ncu
+// diagnosis isolated:
+//
+//   two-pass loss = (a) A re-loaded twice (2*M*K HBM) + (b) 2x pipeline drain.
+//   interleaved loss = 2x B smem (both gate+up staged) -> 1 CTA/SM occupancy.
+//
+// A-persistent: stage A into smem ONCE (reused by both the gate pass and the up
+// pass), but keep only ONE B smem region that is filled with B_gate for the
+// gate pass and then RE-FILLED with B_up for the up pass (sequential, not
+// simultaneous). So:
+//   - A loaded once from HBM (the interleaved HBM win), held in smem across
+//     both passes -> no A re-load. act_scale also staged once.
+//   - B smem = a single STAGES*BN*BK region (NOT doubled) -> fits 4 CTA/SM.
+//   - only ONE accumulator live at a time (gate_acc -> store silu(gate) to a
+//     small smem gate buffer -> reuse regs for up_acc) -> ~70 regs (two-pass
+//     register profile), not the ~140 of true interleaved.
+//
+// The catch: A must fit in smem for the whole K-walk (A is [BM, K], not
+// [BM, BK]), so this only works when K is small enough that BM*K fp8 + the rest
+// stays under the smem budget — i.e. the Qwen3-VL gate_up shapes where K=hidden
+// (2B K=2048, 8B K=4096). For BM=32: A_persist = 32*K = 64KB (8B) / 32KB (2B).
+// 8B 64KB alone already exceeds a CTA's smem, so A-persistent is only viable
+// for the 2B shape (K=2048) at BM<=32, OR by staging A in K-chunks and walking
+// gate+up together within each K-chunk (chunked-interleaved). The chunked form
+// is implemented here: A is staged per BLOCK_K tile like the baseline, but BOTH
+// the gate MMA and the up MMA for that K-tile run before the tile is evicted —
+// i.e. gate and up advance K-tile-by-K-tile together (true interleaved per
+// K-tile, the sm100 geglu pattern), yet B is staged in ONE region reused for
+// gate-then-up WITHIN the k-iter (load B_gate, gate-MMA, load B_up into the
+// SAME region, up-MMA). That keeps B smem single (no 2x) AND loads A once AND
+// holds both gate_acc+up_acc in regs (interleaved) — but pays by serializing
+// the two B loads within a k-iter (no overlap between B_gate and B_up loads).
+//
+// Net vs two-pass: A loaded once (saves M*K HBM), one pipeline drain (K_ITERS
+// stalls not 2*K_ITERS), but B_gate/B_up loads are serial within each k-iter.
+// Net vs interleaved: B smem halved (2 CTA/SM recoverable to 3-4), but loses
+// B_gate||B_up load overlap. On sm89 (HBM-bound, no TMA) the smem/occupancy
+// recovery usually dominates, so this is the predicted winner.
+// ============================================================================
+template <int BLOCK_M, int BLOCK_N, int NUM_WARPS, int STAGES,
+          int MIN_BLOCKS_PER_SM>
+__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM)
+void fp8_bs_geglu_silu_fold_apersist_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,        // gate_up_w [2*N, K]
+    const float* __restrict__ act_scale,        // [M, K/128]
+    const float* __restrict__ w_scale,          // gate_up_s [2*N/128, K/128]
+    __nv_fp8_e4m3* __restrict__ output,         // [M, N]
+    float* __restrict__ out_scale,              // [M, N/128]
+    int M, int N, int K)
+{
+    static_assert(BLOCK_N == 128,
+        "GeGLU silu-fold requires BLOCK_N==128 (one quant block per CTA)");
+    constexpr int BLOCK_K    = 128;
+    constexpr int THREADS    = NUM_WARPS * 32;
+    constexpr int M_ATOMS    = BLOCK_M / 16;
+    constexpr int N_ATOMS    = BLOCK_N / 8;       // 16
+    constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS;
+    constexpr int N_PAIRS_PW = N_ATOMS_PW / 2;
+    constexpr int K_ATOMS    = BLOCK_K / 32;      // 4
+    constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16;
+    constexpr int SWIZZLE_MASK = NUM_CHUNKS_PER_ROW - 1;
+    constexpr int SCALE_KTILE = 8;
+    constexpr int A_TILE = BLOCK_M * BLOCK_K;
+    constexpr int B_TILE = BLOCK_N * BLOCK_K;
+    static_assert(BLOCK_M % 16 == 0, "BLOCK_M multiple of 16");
+    static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0,
+                  "ldmatrix pairs 2 N-atoms: N_ATOMS_PW must be even >= 2");
+    extern __shared__ uint8_t smem_raw[];
+    uint8_t* A_smem = smem_raw;                   // STAGES * A_TILE
+    uint8_t* B_smem = A_smem + STAGES * A_TILE;   // STAGES * B_TILE (reused gate/up)
+    // gate_smem: silu(gate) BF16, [BM, BN]. Written by gate epilogue, read by
+    // the final silu(gate)*up epilogue (NOT per k-iter — only once at the end).
+    __nv_bfloat16* gate_smem = reinterpret_cast<__nv_bfloat16*>(
+        B_smem + STAGES * B_TILE);
+    float* as_smem = reinterpret_cast<float*>(gate_smem + BLOCK_M * BLOCK_N);
+    float* wsg_smem = as_smem + BLOCK_M * SCALE_KTILE;   // gate w_scale row
+    float* wsu_smem = wsg_smem + SCALE_KTILE;            // up   w_scale row
+    float* amax_smem = wsu_smem + SCALE_KTILE;
+    const int cta_m = blockIdx.x;
+    const int cta_n = blockIdx.y;
+    const int m_base = cta_m * BLOCK_M;
+    const int n_base = cta_n * BLOCK_N;
+    const int gate_b_row0 = n_base;
+    const int up_b_row0   = n_base + N;
+    const int t = threadIdx.x;
+    const int warp_id = t / 32;
+    const int lane = t % 32;
+    const int l = lane % 4;
+    const int h = lane / 4;
+    const int frag_group = lane / 8;
+    const int row_in_frag = lane % 8;
+    const int row_block = frag_group / 2;
+    const int col_block = frag_group % 2;
+    const int K128 = K >> 7;
+    const int N128 = N >> 7;
+    const int gate_ws_row = (n_base >> 7);
+    const int up_ws_row   = gate_ws_row + N128;
+    auto stage_scales = [&](int kb0) {
+        const int as_total = BLOCK_M * SCALE_KTILE;
+        for (int idx = t; idx < as_total; idx += THREADS) {
+            int r = idx / SCALE_KTILE;
+            int kc = idx - r * SCALE_KTILE;
+            int row = m_base + r;
+            int kb = kb0 + kc;
+            as_smem[idx] = (row < M && kb < K128)
+                ? act_scale[(size_t)row * K128 + kb] : 0.0f;
+        }
+        for (int kc = t; kc < SCALE_KTILE; kc += THREADS) {
+            int kb = kb0 + kc;
+            wsg_smem[kc] = (kb < K128)
+                ? w_scale[(size_t)gate_ws_row * K128 + kb] : 0.0f;
+            wsu_smem[kc] = (kb < K128)
+                ? w_scale[(size_t)up_ws_row * K128 + kb] : 0.0f;
+        }
+        __syncthreads();
+    };
+    // Stage A + one B band (gate or up) into smem. b_row0 picks the band.
+    auto issue_load = [&](int stage, int k_base, int b_row0) {
+        constexpr int A_CHUNKS = BLOCK_M * NUM_CHUNKS_PER_ROW;
+        constexpr int A_ITERS = (A_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < A_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= A_CHUNKS) break;
+            int row_a = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_a = idx % NUM_CHUNKS_PER_ROW;
+            int m_glob = m_base + row_a;
+            int k_glob = k_base + chunk_a * 16;
+            const uint8_t* a_src = nullptr;
+            if (m_glob < M && k_glob < K) {
+                a_src = reinterpret_cast<const uint8_t*>(&A[(size_t)m_glob * K + k_glob]);
+            }
+            int csw = chunk_a ^ (row_a & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&A_smem[stage * A_TILE + row_a * BLOCK_K + csw * 16]),
+                a_src);
+        }
+        constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW;
+        constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < B_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= B_CHUNKS) break;
+            int row_b = idx / NUM_CHUNKS_PER_ROW;
+            int chunk_b = idx % NUM_CHUNKS_PER_ROW;
+            int n_glob = b_row0 + row_b;
+            int k_glob = k_base + chunk_b * 16;
+            const uint8_t* b_src = nullptr;
+            if (n_glob < 2 * N && k_glob < K) {
+                b_src = reinterpret_cast<const uint8_t*>(&B[(size_t)n_glob * K + k_glob]);
+            }
+            int csw = chunk_b ^ (row_b & SWIZZLE_MASK);
+            cp_async_16(
+                to_smem(&B_smem[stage * B_TILE + row_b * BLOCK_K + csw * 16]),
+                b_src);
+        }
+    };
+    // MMA pass over the staged A/B tiles for one k-iter, accumulating into `acc`
+    // with the given w_scale smem row.
+    auto mma_tile = [&](float (*acc)[N_ATOMS_PW][4], int compute_stage,
+                        const float* ws_smem_pass) {
+        const int kb = (compute_stage);  // caller passes k_iter; recompute below
+        (void)kb;
+        float tacc[M_ATOMS][N_ATOMS_PW][4];
+        #pragma unroll
+        for (int mi = 0; mi < M_ATOMS; ++mi)
+            #pragma unroll
+            for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+                #pragma unroll
+                for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f;
+        uint8_t* A_stage = A_smem + compute_stage * A_TILE;
+        uint8_t* B_stage = B_smem + compute_stage * B_TILE;
+        #pragma unroll
+        for (int ka = 0; ka < K_ATOMS; ++ka) {
+            uint32_t A_regs[M_ATOMS][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row = mi * 16 + row_block * 8 + row_in_frag;
+                int chunk = 2 * ka + col_block;
+                int csw = chunk ^ (row & SWIZZLE_MASK);
+                ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                                to_smem(&A_stage[row * BLOCK_K + csw * 16]));
+            }
+            uint32_t B_regs[N_PAIRS_PW][4];
+            #pragma unroll
+            for (int np = 0; np < N_PAIRS_PW; ++np) {
+                int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag;
+                int chunk = 2 * ka + col_block;
+                int csw = chunk ^ (nrow & SWIZZLE_MASK);
+                ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                                to_smem(&B_stage[nrow * BLOCK_K + csw * 16]));
+            }
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int ni0 = np * 2, ni1 = np * 2 + 1;
+                    mma_m16n8k32_e4m3(
+                        tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3],
+                        A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][0], B_regs[np][1]);
+                    mma_m16n8k32_e4m3(
+                        tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3],
+                        A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][2], B_regs[np][3]);
+                }
+            }
+        }
+        return tacc;  // caller folds scales into acc
+    };
+    // Running accumulators for gate and up, both live across the whole K-loop
+    // (true interleaved: both gate and up advance K-tile-by-K-tile together).
+    float gate_acc[M_ATOMS][N_ATOMS_PW][4];
+    float up_acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi)
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                gate_acc[mi][ni][j] = 0.0f;
+                up_acc[mi][ni][j] = 0.0f;
+            }
+    const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K;
+    // Prefetch STAGES-1 A tiles (A is shared by both passes — issued once).
+    // B is NOT prefetched here; within each k-iter we issue B_gate then B_up
+    // into the SAME smem region after the previous iter's B is consumed.
+    #pragma unroll
+    for (int s = 0; s < STAGES - 1; ++s) {
+        int kb = s * BLOCK_K;
+        if (kb < K) issue_load(s, kb, gate_b_row0);   // first prefetch = gate B
+        asm volatile("cp.async.commit_group;\n" ::);
+    }
+    int compute_stage = 0;
+    for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) {
+        int issue_iter = k_iter + (STAGES - 1);
+        int issue_stage = issue_iter % STAGES;
+        // Issue the NEXT A tile + the NEXT gate-B tile (the up-B for this k_iter
+        // is loaded inside the gate-MMA sync below, reusing B_smem after gate
+        // MMA reads finish).
+        if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K, gate_b_row0);
+        asm volatile("cp.async.commit_group;\n" ::);
+        asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1));
+        __syncthreads();
+        const int kb = k_iter;
+        if ((kb % SCALE_KTILE) == 0) stage_scales(kb);
+        // ---- gate MMA on the staged A + B_gate ----
+        {
+            float tacc[M_ATOMS][N_ATOMS_PW][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi)
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+                    #pragma unroll
+                    for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f;
+            uint8_t* A_stage = A_smem + compute_stage * A_TILE;
+            uint8_t* B_stage = B_smem + compute_stage * B_TILE;
+            #pragma unroll
+            for (int ka = 0; ka < K_ATOMS; ++ka) {
+                uint32_t A_regs[M_ATOMS][4];
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    int row = mi * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (row & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                                    to_smem(&A_stage[row * BLOCK_K + csw * 16]));
+                }
+                uint32_t B_regs[N_PAIRS_PW][4];
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (nrow & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                                    to_smem(&B_stage[nrow * BLOCK_K + csw * 16]));
+                }
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    #pragma unroll
+                    for (int np = 0; np < N_PAIRS_PW; ++np) {
+                        int ni0 = np * 2, ni1 = np * 2 + 1;
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][0], B_regs[np][1]);
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][2], B_regs[np][3]);
+                    }
+                }
+            }
+            int kbt = kb % SCALE_KTILE;
+            float ws_cta = wsg_smem[kbt];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row0 = m_base + mi * 16 + h;
+                int row1 = row0 + 8;
+                float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt];
+                float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt];
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+                    gate_acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta);
+                    gate_acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta);
+                    gate_acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta);
+                    gate_acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta);
+                }
+            }
+        }
+        __syncthreads();   // B_smem safe to overwrite with B_up
+        // ---- load B_up into the SAME B_smem region, then up MMA ----
+        // (A is still resident in A_smem[compute_stage]; not reloaded from HBM.)
+        {
+            // issue B_up into B_smem[compute_stage] (A_smem left untouched)
+            constexpr int B_CHUNKS = BLOCK_N * NUM_CHUNKS_PER_ROW;
+            constexpr int B_ITERS = (B_CHUNKS + THREADS - 1) / THREADS;
+            int k_base = k_iter * BLOCK_K;
+            #pragma unroll
+            for (int it = 0; it < B_ITERS; ++it) {
+                int idx = it * THREADS + t;
+                if (idx >= B_CHUNKS) break;
+                int row_b = idx / NUM_CHUNKS_PER_ROW;
+                int chunk_b = idx % NUM_CHUNKS_PER_ROW;
+                int n_glob = up_b_row0 + row_b;
+                int k_glob = k_base + chunk_b * 16;
+                const uint8_t* b_src = nullptr;
+                if (n_glob < 2 * N && k_glob < K) {
+                    b_src = reinterpret_cast<const uint8_t*>(&B[(size_t)n_glob * K + k_glob]);
+                }
+                int csw = chunk_b ^ (row_b & SWIZZLE_MASK);
+                cp_async_16(
+                    to_smem(&B_smem[compute_stage * B_TILE + row_b * BLOCK_K + csw * 16]),
+                    b_src);
+            }
+            asm volatile("cp.async.commit_group;\n" ::);
+            asm volatile("cp.async.wait_group %0;\n" :: "n"(0));
+            __syncthreads();
+            float tacc[M_ATOMS][N_ATOMS_PW][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi)
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+                    #pragma unroll
+                    for (int j = 0; j < 4; ++j) tacc[mi][ni][j] = 0.0f;
+            uint8_t* A_stage = A_smem + compute_stage * A_TILE;
+            uint8_t* B_stage = B_smem + compute_stage * B_TILE;
+            #pragma unroll
+            for (int ka = 0; ka < K_ATOMS; ++ka) {
+                uint32_t A_regs[M_ATOMS][4];
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    int row = mi * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (row & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                                    to_smem(&A_stage[row * BLOCK_K + csw * 16]));
+                }
+                uint32_t B_regs[N_PAIRS_PW][4];
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int nrow = warp_id * N_ATOMS_PW * 8 + np * 16 + row_block * 8 + row_in_frag;
+                    int chunk = 2 * ka + col_block;
+                    int csw = chunk ^ (nrow & SWIZZLE_MASK);
+                    ldmatrix_x4_b16(B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                                    to_smem(&B_stage[nrow * BLOCK_K + csw * 16]));
+                }
+                #pragma unroll
+                for (int mi = 0; mi < M_ATOMS; ++mi) {
+                    #pragma unroll
+                    for (int np = 0; np < N_PAIRS_PW; ++np) {
+                        int ni0 = np * 2, ni1 = np * 2 + 1;
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni0][0], tacc[mi][ni0][1], tacc[mi][ni0][2], tacc[mi][ni0][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][0], B_regs[np][1]);
+                        mma_m16n8k32_e4m3(
+                            tacc[mi][ni1][0], tacc[mi][ni1][1], tacc[mi][ni1][2], tacc[mi][ni1][3],
+                            A_regs[mi][0], A_regs[mi][2], A_regs[mi][1], A_regs[mi][3],
+                            B_regs[np][2], B_regs[np][3]);
+                    }
+                }
+            }
+            int kbt = kb % SCALE_KTILE;
+            float ws_cta = wsu_smem[kbt];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row0 = m_base + mi * 16 + h;
+                int row1 = row0 + 8;
+                float as0 = as_smem[(mi * 16 + h) * SCALE_KTILE + kbt];
+                float as1 = as_smem[(mi * 16 + h + 8) * SCALE_KTILE + kbt];
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+                    up_acc[mi][ni][0] += tacc[mi][ni][0] * (as0 * ws_cta);
+                    up_acc[mi][ni][1] += tacc[mi][ni][1] * (as0 * ws_cta);
+                    up_acc[mi][ni][2] += tacc[mi][ni][2] * (as1 * ws_cta);
+                    up_acc[mi][ni][3] += tacc[mi][ni][3] * (as1 * ws_cta);
+                }
+            }
+        }
+        __syncthreads();
+        compute_stage = (compute_stage + 1) % STAGES;
+    }
+    asm volatile("cp.async.wait_all;\n" ::);
+    // ============ Epilogue: silu(gate)*up + per-row amax + quant ============
+    // gate_acc and up_acc both live in registers. Replicate silu_mul_merged's
+    // two bf16 roundings: bf16(silu(gate)) then bf16(silu_bf * up).
+    constexpr float kFp8Max = 448.0f;
+    float v[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        int rloc0 = mi * 16 + h;
+        int rloc1 = rloc0 + 8;
+        float amax0 = 0.0f, amax1 = 0.0f;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            if (row0 < M) {
+                float gf0 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][0])));
+                float gf1 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][1])));
+                v[mi][ni][0] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][0]));
+                v[mi][ni][1] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][1]));
+                amax0 = fmaxf(amax0, fmaxf(fabsf(v[mi][ni][0]), fabsf(v[mi][ni][1])));
+            } else { v[mi][ni][0] = 0.0f; v[mi][ni][1] = 0.0f; }
+            if (row1 < M) {
+                float gf0 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][2])));
+                float gf1 = __bfloat162float(__float2bfloat16(silu_f32(gate_acc[mi][ni][3])));
+                v[mi][ni][2] = __bfloat162float(__float2bfloat16(gf0 * up_acc[mi][ni][2]));
+                v[mi][ni][3] = __bfloat162float(__float2bfloat16(gf1 * up_acc[mi][ni][3]));
+                amax1 = fmaxf(amax1, fmaxf(fabsf(v[mi][ni][2]), fabsf(v[mi][ni][3])));
+            } else { v[mi][ni][2] = 0.0f; v[mi][ni][3] = 0.0f; }
+        }
+        for (int off = 2; off > 0; off >>= 1) {
+            amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffff, amax0, off));
+            amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffff, amax1, off));
+        }
+        if (l == 0) {
+            amax_smem[warp_id * BLOCK_M + rloc0] = amax0;
+            amax_smem[warp_id * BLOCK_M + rloc1] = amax1;
+        }
+    }
+    __syncthreads();
+    #pragma unroll
+    for (int rloc = t; rloc < BLOCK_M; rloc += THREADS) {
+        int row = m_base + rloc;
+        if (row >= M) continue;
+        float amax = 0.0f;
+        #pragma unroll
+        for (int w = 0; w < NUM_WARPS; ++w)
+            amax = fmaxf(amax, amax_smem[w * BLOCK_M + rloc]);
+        float sc = fmaxf(amax / kFp8Max, 1.0e-12f);
+        amax_smem[rloc] = sc;
+        // Each active thread owns a distinct rloc — write its row's scale
+        // directly (no single-thread guard; see two-pass variant for the bug
+        // this fixes).
+        out_scale[(size_t)row * (N >> 7) + (n_base >> 7)] = sc;
+    }
+    __syncthreads();
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        int rloc0 = mi * 16 + h;
+        int rloc1 = rloc0 + 8;
+        float sc0 = (row0 < M) ? amax_smem[rloc0] : 1.0f;
+        float sc1 = (row1 < M) ? amax_smem[rloc1] : 1.0f;
+        float inv0 = 1.0f / sc0, inv1 = 1.0f / sc1;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            if (row0 < M && col_pair_ok(n_pair_base, N)) {
+                float q0 = fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max);
+                float q1 = fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max);
+                __nv_fp8_e4m3 p0(q0), p1(q1);
+                uint16_t pack = (uint16_t)(*reinterpret_cast<const uint8_t*>(&p1)) << 8
+                              | (uint16_t)(*reinterpret_cast<const uint8_t*>(&p0));
+                *reinterpret_cast<uint16_t*>(&output[(size_t)row0 * N + n_pair_base]) = pack;
+            } else if (row0 < M) {
+                if (n_pair_base < N) output[(size_t)row0 * N + n_pair_base] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][0] * inv0, -kFp8Max), kFp8Max));
+                if (n_pair_base + 1 < N) output[(size_t)row0 * N + n_pair_base + 1] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][1] * inv0, -kFp8Max), kFp8Max));
+            }
+            if (row1 < M && col_pair_ok(n_pair_base, N)) {
+                float q2 = fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max);
+                float q3 = fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max);
+                __nv_fp8_e4m3 p2(q2), p3(q3);
+                uint16_t pack = (uint16_t)(*reinterpret_cast<const uint8_t*>(&p3)) << 8
+                              | (uint16_t)(*reinterpret_cast<const uint8_t*>(&p2));
+                *reinterpret_cast<uint16_t*>(&output[(size_t)row1 * N + n_pair_base]) = pack;
+            } else if (row1 < M) {
+                if (n_pair_base < N) output[(size_t)row1 * N + n_pair_base] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][2] * inv1, -kFp8Max), kFp8Max));
+                if (n_pair_base + 1 < N) output[(size_t)row1 * N + n_pair_base + 1] = __nv_fp8_e4m3(fminf(fmaxf(v[mi][ni][3] * inv1, -kFp8Max), kFp8Max));
+            }
+        }
+    }
+    (void)gate_smem;  // apersist keeps gate in regs; gate_smem unused (kept for layout parity)
+    (void)mma_tile;   // helper retained for future chunked variant; unused in this path
+}
+}  // namespace block128_sm89
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_gemv_m1_sm120.cu ADDED Viewed

	@@ -0,0 +1,159 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Dedicated M=1 FP8 e4m3 -> BF16 GEMV for sm_120a decode (batch=1 token).
+//
+// The hand-tuned MMA GEMMs pad M=1 to BLOCK_M=16 (m16n8k32), computing 16
+// rows to use 1 — fine for compute (memory-bound) but the N=2560 shapes only
+// spawn N/BLOCK_N blocks and starve the SMs. This GEMV assigns one warp per
+// output row: A[1,K] is staged once into smem (hot in L2 across blocks), each
+// warp streams its B row in 16-byte coalesced chunks and warp-reduces the dot
+// product. BLOCK_N effectively 1-per-warp => N/WARPS_PER_BLOCK blocks (e.g.
+// N=2560, W=8 -> 320 blocks) saturates occupancy without a split-K reduction.
+#include "fp8_gemv_m1_sm120.cuh"
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+namespace flash_rt {
+namespace gemm {
+namespace gemv_m1 {
+namespace {
+// One warp per output row n. A staged in smem as raw fp8 (K bytes). B row read
+// in uint4 (16 fp8) coalesced chunks, stride 32 across the warp. K assumed a
+// multiple of 16 (all Higgs/Qwen3 GEMM K: 2560/4096/9728).
+template <int WARPS_PER_BLOCK>
+__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 8)
+void gemv_fp8_m1_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,   // [K]
+    const __nv_fp8_e4m3* __restrict__ B,   // [N, K]
+    __nv_bfloat16* __restrict__ D,         // [N]
+    int N, int K, float alpha)
+{
+    extern __shared__ __nv_fp8_e4m3 sA[];   // [K]
+    const int tid     = threadIdx.x;
+    const int lane    = tid & 31;
+    const int warp    = tid >> 5;
+    const int threads = WARPS_PER_BLOCK * 32;
+    const int K16     = K >> 4;             // # of 16-byte (uint4) groups
+    // Cooperatively stage A into smem, 16 bytes per thread.
+    uint4* sA16 = reinterpret_cast<uint4*>(sA);
+    const uint4* A16 = reinterpret_cast<const uint4*>(A);
+    for (int i = tid; i < K16; i += threads) sA16[i] = A16[i];
+    __syncthreads();
+    const int n = blockIdx.x * WARPS_PER_BLOCK + warp;
+    if (n >= N) return;
+    const uint4* Brow = reinterpret_cast<const uint4*>(B) + (size_t)n * K16;
+    const __nv_fp8_e4m3* sAf = sA;
+    float acc = 0.0f;
+    for (int i = lane; i < K16; i += 32) {
+        uint4 bpack = Brow[i];
+        const __nv_fp8_e4m3* bp = reinterpret_cast<const __nv_fp8_e4m3*>(&bpack);
+        const __nv_fp8_e4m3* ap = sAf + (i << 4);
+        #pragma unroll
+        for (int j = 0; j < 16; ++j) {
+            acc += float(ap[j]) * float(bp[j]);
+        }
+    }
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    }
+    if (lane == 0) D[n] = __float2bfloat16(acc * alpha);
+}
+// GEMV with fused residual accumulate: D[n] += acc * alpha (in-place into the
+// residual stream). The residual is per-element local — no cross-block
+// dependency like the norm — so it folds into the epilogue for free, removing
+// the separate residual_add launch. Each output n is written by one warp lane.
+template <int WARPS_PER_BLOCK>
+__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 8)
+void gemv_fp8_m1_resadd_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,
+    __nv_bfloat16* __restrict__ D,   // residual stream, accumulated in place
+    int N, int K, float alpha)
+{
+    extern __shared__ __nv_fp8_e4m3 sA[];
+    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
+    const int threads = WARPS_PER_BLOCK * 32, K16 = K >> 4;
+    uint4* sA16 = reinterpret_cast<uint4*>(sA);
+    const uint4* A16 = reinterpret_cast<const uint4*>(A);
+    for (int i = tid; i < K16; i += threads) sA16[i] = A16[i];
+    __syncthreads();
+    const int n = blockIdx.x * WARPS_PER_BLOCK + warp;
+    if (n >= N) return;
+    const uint4* Brow = reinterpret_cast<const uint4*>(B) + (size_t)n * K16;
+    float acc = 0.0f;
+    for (int i = lane; i < K16; i += 32) {
+        uint4 bpack = Brow[i];
+        const __nv_fp8_e4m3* bp = reinterpret_cast<const __nv_fp8_e4m3*>(&bpack);
+        const __nv_fp8_e4m3* ap = sA + (i << 4);
+        #pragma unroll
+        for (int j = 0; j < 16; ++j) acc += float(ap[j]) * float(bp[j]);
+    }
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
+    if (lane == 0) D[n] = __float2bfloat16(__bfloat162float(D[n]) + acc * alpha);
+}
+template <int W>
+int launch_resadd_(const void* A, const void* B, void* D,
+                   int N, int K, float alpha, cudaStream_t stream) {
+    dim3 grid((N + W - 1) / W);
+    size_t smem = (size_t)K * sizeof(__nv_fp8_e4m3);
+    gemv_fp8_m1_resadd_kernel<W><<<grid, W * 32, smem, stream>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D), N, K, alpha);
+    return 0;
+}
+template <int W>
+int launch_(const void* A, const void* B, void* D,
+            int /*M*/, int N, int K, float alpha, cudaStream_t stream) {
+    dim3 grid((N + W - 1) / W);
+    dim3 block(W * 32);
+    size_t smem = (size_t)K * sizeof(__nv_fp8_e4m3);
+    gemv_fp8_m1_kernel<W><<<grid, block, smem, stream>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D),
+        N, K, alpha);
+    return 0;
+}
+}  // namespace
+#define DEFINE(NAME, W)                                                         \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,          \
+           float alpha, cudaStream_t stream) {                                  \
+    return launch_<W>(A, B, D, M, N, K, alpha, stream);                         \
+  }
+DEFINE(gemv_fp8_m1_w4,  4)
+DEFINE(gemv_fp8_m1_w8,  8)
+DEFINE(gemv_fp8_m1_w16, 16)
+#define DEFINE_RA(NAME, W)                                                      \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,          \
+           float alpha, cudaStream_t stream) {                                  \
+    return launch_resadd_<W>(A, B, D, N, K, alpha, stream);                     \
+  }
+DEFINE_RA(gemv_fp8_m1_resadd_w4, 4)
+DEFINE_RA(gemv_fp8_m1_resadd_w8, 8)
+#undef DEFINE
+#undef DEFINE_RA
+}  // namespace gemv_m1
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_gemv_m1_sm120.cuh ADDED Viewed

	@@ -0,0 +1,31 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+namespace gemv_m1 {
+// Dedicated M=1 FP8 e4m3 -> BF16 GEMV for sm_120a decode shapes.
+// Inputs: FP8 A [1,K] row-major, FP8 B [N,K] row-major (= W.T), BF16 D [1,N].
+// alpha = a_scale * w_scale (per-tensor). M is ignored (M=1 assumed).
+// Warp-per-output-row: each warp reduces one B row against A (held in smem),
+// 16-byte vectorized coalesced B loads. No MMA / no BLOCK_M padding tax.
+// Returns 0 on success.
+#define DECL(NAME) \
+  int NAME(const void* A, const void* B, void* D, \
+           int M, int N, int K, float alpha, cudaStream_t stream)
+DECL(gemv_fp8_m1_w4);
+DECL(gemv_fp8_m1_w8);
+DECL(gemv_fp8_m1_w16);
+DECL(gemv_fp8_m1_resadd_w4);  // D[n] += acc*alpha (fused residual)
+DECL(gemv_fp8_m1_resadd_w8);
+#undef DECL
+}  // namespace gemv_m1
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_gemv_m1_sm89.cu ADDED Viewed

	@@ -0,0 +1,188 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// M=1 FP8 e4m3 -> BF16 block-128 scaled GEMV for SM89 Qwen3-VL decode.
+// Header: fp8_gemv_m1_sm89.cuh.
+//
+// Split out of the SM120 per-tensor GEMV (fp8_gemv_m1_sm120) so the SM89
+// block-128 decode path owns its own file: per-token activation scale
+// [K/128] and DeepSeek-style weight block scale [N/128, K/128], applied in
+// the warp reduction. Warp-per-output-row, A held in smem, 16-byte coalesced
+// B loads. No MMA / no padding tax (M=1).
+#include "fp8_gemv_m1_sm89.cuh"
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+namespace flash_rt {
+namespace gemm {
+namespace gemv_m1_sm89 {
+namespace {
+template <int WARPS_PER_BLOCK>
+__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 8)
+void gemv_fp8_block128_m1_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,   // [K]
+    const __nv_fp8_e4m3* __restrict__ B,   // [N, K]
+    __nv_bfloat16* __restrict__ D,         // [N]
+    int N, int K,
+    const float* __restrict__ act_scale,   // [K/128]
+    const float* __restrict__ w_scale,     // [N/128, K/128]
+    float alpha)
+{
+    extern __shared__ __nv_fp8_e4m3 sA[];
+    const int tid     = threadIdx.x;
+    const int lane    = tid & 31;
+    const int warp    = tid >> 5;
+    const int threads = WARPS_PER_BLOCK * 32;
+    const int K16     = K >> 4;
+    const int K128    = K >> 7;
+    uint4* sA16 = reinterpret_cast<uint4*>(sA);
+    const uint4* A16 = reinterpret_cast<const uint4*>(A);
+    for (int i = tid; i < K16; i += threads) sA16[i] = A16[i];
+    __syncthreads();
+    const int n = blockIdx.x * WARPS_PER_BLOCK + warp;
+    if (n >= N) return;
+    const uint4* Brow = reinterpret_cast<const uint4*>(B) + (size_t)n * K16;
+    const __nv_fp8_e4m3* sAf = sA;
+    const float* w_scale_row = w_scale + (size_t)(n >> 7) * K128;
+    float acc = 0.0f;
+    for (int i = lane; i < K16; i += 32) {
+        const int kb = i >> 3;
+        const float s = act_scale[kb] * w_scale_row[kb] * alpha;
+        uint4 bpack = Brow[i];
+        const __nv_fp8_e4m3* bp =
+            reinterpret_cast<const __nv_fp8_e4m3*>(&bpack);
+        const __nv_fp8_e4m3* ap = sAf + (i << 4);
+        #pragma unroll
+        for (int j = 0; j < 16; ++j) {
+            acc += float(ap[j]) * float(bp[j]) * s;
+        }
+    }
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    }
+    if (lane == 0) D[n] = __float2bfloat16(acc);
+}
+template <int W>
+int launch_block128_(const void* A, const void* B, void* D,
+                     int /*M*/, int N, int K,
+                     const float* act_scale, const float* w_scale,
+                     float alpha, cudaStream_t stream) {
+    dim3 grid((N + W - 1) / W);
+    dim3 block(W * 32);
+    size_t smem = (size_t)K * sizeof(__nv_fp8_e4m3);
+    gemv_fp8_block128_m1_kernel<W><<<grid, block, smem, stream>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D),
+        N, K, act_scale, w_scale, alpha);
+    return 0;
+}
+// BF16-input variant: skips activation FP8 quantization.
+// A is BF16, B is FP8 with block-128 weight scale. No act_scale needed.
+template <int WARPS_PER_BLOCK>
+__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 8)
+void gemv_fp8_block128_m1_bf16in_kernel(
+    const __nv_bfloat16* __restrict__ A,   // [K] BF16
+    const __nv_fp8_e4m3* __restrict__ B,   // [N, K] FP8
+    __nv_bfloat16* __restrict__ D,         // [N]
+    int N, int K,
+    const float* __restrict__ w_scale)     // [N/128, K/128]
+{
+    extern __shared__ __nv_bfloat16 sA_bf16[];
+    const int tid     = threadIdx.x;
+    const int lane    = tid & 31;
+    const int warp    = tid >> 5;
+    const int threads = WARPS_PER_BLOCK * 32;
+    const int K16     = K >> 4;
+    const int K128    = K >> 7;
+    // Load BF16 activation (2 bytes each) via uint32 pairs.
+    uint* sU = reinterpret_cast<uint*>(sA_bf16);
+    const uint* AU = reinterpret_cast<const uint*>(A);
+    const int Khalf = K >> 1;
+    for (int i = tid; i < Khalf; i += threads) sU[i] = AU[i];
+    __syncthreads();
+    const int n = blockIdx.x * WARPS_PER_BLOCK + warp;
+    if (n >= N) return;
+    const uint4* Brow = reinterpret_cast<const uint4*>(B) + (size_t)n * K16;
+    const float* w_scale_row = w_scale + (size_t)(n >> 7) * K128;
+    float acc = 0.0f;
+    for (int i = lane; i < K16; i += 32) {
+        const int kb = i >> 3;
+        const float ws = w_scale_row[kb];
+        uint4 bpack = Brow[i];
+        const __nv_fp8_e4m3* bp =
+            reinterpret_cast<const __nv_fp8_e4m3*>(&bpack);
+        const __nv_bfloat16* ap = sA_bf16 + (i << 4);
+        float dot = 0.0f;
+        #pragma unroll
+        for (int j = 0; j < 16; ++j) {
+            dot += __bfloat162float(ap[j]) * float(bp[j]);
+        }
+        acc += dot * ws;
+    }
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    }
+    if (lane == 0) D[n] = __float2bfloat16(acc);
+}
+template <int W>
+int launch_block128_bf16in_(const void* A, const void* B, void* D,
+                            int /*M*/, int N, int K,
+                            const float* w_scale, cudaStream_t stream) {
+    dim3 grid((N + W - 1) / W);
+    dim3 block(W * 32);
+    size_t smem = (size_t)K * sizeof(__nv_bfloat16);
+    gemv_fp8_block128_m1_bf16in_kernel<W><<<grid, block, smem, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D),
+        N, K, w_scale);
+    return 0;
+}
+}  // namespace
+#define DEFINE_BLOCK128(NAME, W)                                               \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,         \
+           const float* act_scale, const float* w_scale, float alpha,          \
+           cudaStream_t stream) {                                              \
+    return launch_block128_<W>(A, B, D, M, N, K, act_scale, w_scale, alpha,    \
+                               stream);                                        \
+  }
+DEFINE_BLOCK128(gemv_fp8_block128_m1_w4, 4)
+DEFINE_BLOCK128(gemv_fp8_block128_m1_w8, 8)
+DEFINE_BLOCK128(gemv_fp8_block128_m1_w16, 16)
+#undef DEFINE_BLOCK128
+#define DEFINE_BLOCK128_BF16IN(NAME, W)                                        \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,         \
+           const float* w_scale, cudaStream_t stream) {                        \
+    return launch_block128_bf16in_<W>(A, B, D, M, N, K, w_scale, stream);      \
+  }
+DEFINE_BLOCK128_BF16IN(gemv_fp8_block128_m1_bf16in_w8, 8)
+DEFINE_BLOCK128_BF16IN(gemv_fp8_block128_m1_bf16in_w16, 16)
+#undef DEFINE_BLOCK128_BF16IN
+}  // namespace gemv_m1_sm89
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_gemv_m1_sm89.cuh ADDED Viewed

	@@ -0,0 +1,38 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+namespace gemv_m1_sm89 {
+// M=1 FP8 e4m3 -> BF16 GEMV with per-token activation block scale [K/128] and
+// per-weight 128x128 block scale [N/128, K/128]. Matches official Qwen3-VL FP8
+// checkpoints that store `.weight` + `.weight_scale_inv`. Decode-shape sibling
+// of the M>1 fp8_block128_gemm_mma_sm89 kernel. Warp-per-output-row, A staged
+// in smem, 16-byte coalesced B loads. Returns 0 on success.
+#define DECL_BLOCK128(NAME) \
+  int NAME(const void* A, const void* B, void* D, \
+           int M, int N, int K, const float* act_scale, \
+           const float* w_scale, float alpha, cudaStream_t stream)
+DECL_BLOCK128(gemv_fp8_block128_m1_w4);
+DECL_BLOCK128(gemv_fp8_block128_m1_w8);
+DECL_BLOCK128(gemv_fp8_block128_m1_w16);
+#undef DECL_BLOCK128
+// BF16-input variants: A is BF16, B is FP8. No act_scale, only w_scale.
+#define DECL_BLOCK128_BF16IN(NAME) \
+  int NAME(const void* A, const void* B, void* D, \
+           int M, int N, int K, const float* w_scale, cudaStream_t stream)
+DECL_BLOCK128_BF16IN(gemv_fp8_block128_m1_bf16in_w8);
+DECL_BLOCK128_BF16IN(gemv_fp8_block128_m1_bf16in_w16);
+#undef DECL_BLOCK128_BF16IN
+}  // namespace gemv_m1_sm89
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_smallM_handtuned_ldmatrix_sm120.cu ADDED Viewed

	@@ -0,0 +1,377 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Hand-tuned FP8 e4m3 -> BF16 GEMM v2 for sm_120a small-M motus shapes.
+// Adds 128B swizzle smem layout + ldmatrix.x4.m8n8.b16 reads to clear the
+// 4-way smem bank conflict that bottlenecks v1 (`fp8_smallM_handtuned`).
+//
+// Restrictions for this version:
+//   - BLOCK_K = 128 (natural 128B swizzle row stride)
+//   - N_ATOMS_PER_WARP must be even (paired into one ldmatrix.x4 each)
+//
+// MMA path identical to v1 (inline PTX m16n8k32 e4m3 e4m3 f32).
+#include "fp8_smallM_handtuned_ldmatrix_sm120.cuh"
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+namespace flash_rt {
+namespace gemm {
+namespace smallM_ld {
+namespace {
+__device__ __forceinline__ void mma_m16n8k32_e4m3(
+    float &d0, float &d1, float &d2, float &d3,
+    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
+    uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
+        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+__device__ __forceinline__ void cp_async_16(uint32_t smem, const uint8_t* src) {
+    int b = (src == nullptr) ? 0 : 16;
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n"
+                 :: "r"(smem), "l"(src), "r"(b));
+}
+__device__ __forceinline__ uint32_t to_smem(const void* p) {
+    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
+}
+__device__ __forceinline__ void ldmatrix_x4_b16(
+    uint32_t &d0, uint32_t &d1, uint32_t &d2, uint32_t &d3,
+    uint32_t smem_addr)
+{
+    asm volatile(
+        "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+        : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
+        : "r"(smem_addr));
+}
+// 128B swizzle: byte_addr_swizzled = row*128 + (chunk16 XOR (row & 7))*16
+// where chunk16 = byte_col / 16 in [0, 7] (one chunk = 16 bytes).
+// Applied identically on cp.async store and ldmatrix load to round-trip cleanly.
+template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_WARPS,
+          int STAGES = 2, int MIN_BLOCKS_PER_SM = 4>
+__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM)
+void fp8_gemm_ld_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,
+    __nv_bfloat16* __restrict__ D,
+    int M, int N, int K, float alpha)
+{
+    static_assert(BLOCK_K == 64 || BLOCK_K == 128 || BLOCK_K == 256
+                  || BLOCK_K == 512,
+                  "BLOCK_K must be 64/128/256/512 (BK=512 uses repeated 128B swizzle period)");
+    constexpr int NUM_CHUNKS_PER_ROW = BLOCK_K / 16;       // 4, 8, 16, or 32
+    // Swizzle mask: 128B period (8 chunks). For BK=256/512 the pattern
+    // repeats every 8 chunks; ldmatrix.x4 reads within one period at a
+    // time so bank-conflict-free remains.
+    constexpr int SWIZZLE_MASK = (NUM_CHUNKS_PER_ROW <= 8)
+                                 ? (NUM_CHUNKS_PER_ROW - 1) : 7;
+    constexpr int THREADS = NUM_WARPS * 32;
+    constexpr int M_ATOMS = BLOCK_M / 16;
+    constexpr int N_ATOMS = BLOCK_N / 8;
+    constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS;
+    static_assert(BLOCK_M % 16 == 0, "BLOCK_M must be multiple of 16");
+    static_assert(BLOCK_N % 8 == 0,  "BLOCK_N must be multiple of 8");
+    static_assert(N_ATOMS_PW >= 2 && N_ATOMS_PW % 2 == 0,
+                  "N atoms per warp must be even >=2 (paired for ldmatrix.x4)");
+    constexpr int N_PAIRS_PW = N_ATOMS_PW / 2;
+    constexpr int K_ATOMS = BLOCK_K / 32;  // = 4 for BLOCK_K=128
+    constexpr int A_TILE_BYTES = BLOCK_M * BLOCK_K;
+    constexpr int B_TILE_BYTES = BLOCK_N * BLOCK_K;
+    extern __shared__ __align__(128) uint8_t smem_raw[];
+    uint8_t* A_smem = smem_raw;
+    uint8_t* B_smem = A_smem + STAGES * A_TILE_BYTES;
+    const int cta_m  = blockIdx.x;
+    const int cta_n  = blockIdx.y;
+    const int m_base = cta_m * BLOCK_M;
+    const int n_base = cta_n * BLOCK_N;
+    const int t = threadIdx.x;
+    const int warp_id = t / 32;
+    const int lane = t % 32;
+    // Lane partition for ldmatrix.x4 addressing (lane -> fragment).
+    const int frag_group = lane / 8;      // 0..3 (TL,TR,BL,BR per ldmatrix)
+    const int row_in_frag = lane % 8;     // row within fragment 0..7
+    const int row_block = frag_group / 2; // top(0) / bot(1)
+    const int col_block = frag_group % 2; // left(0) / right(1)
+    // Lane partition for mma epilogue write.
+    const int h = lane / 4;  // 0..7
+    const int l = lane % 4;  // 0..3
+    auto issue_load = [&](int stage, int k_base) {
+        // A tile: BLOCK_M rows x BLOCK_K bytes, each thread issues 16-byte chunks.
+        constexpr int A_CHUNKS = BLOCK_M * (BLOCK_K / 16);
+        constexpr int A_ITERS  = (A_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < A_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= A_CHUNKS) break;
+            int row_a   = idx / (BLOCK_K / 16);
+            int chunk_a = idx % (BLOCK_K / 16);
+            int m_g = m_base + row_a;
+            int k_g = k_base + chunk_a * 16;
+            const uint8_t* src = nullptr;
+            if (m_g < M && k_g < K) {
+                src = reinterpret_cast<const uint8_t*>(&A[m_g * K + k_g]);
+            }
+            int chunk_sw = chunk_a ^ (row_a & SWIZZLE_MASK);
+            uint32_t dst = to_smem(
+                &A_smem[stage * A_TILE_BYTES + row_a * BLOCK_K + chunk_sw * 16]);
+            cp_async_16(dst, src);
+        }
+        // B tile: BLOCK_N rows x BLOCK_K bytes.
+        constexpr int B_CHUNKS = BLOCK_N * (BLOCK_K / 16);
+        constexpr int B_ITERS  = (B_CHUNKS + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < B_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= B_CHUNKS) break;
+            int row_b   = idx / (BLOCK_K / 16);
+            int chunk_b = idx % (BLOCK_K / 16);
+            int n_g = n_base + row_b;
+            int k_g = k_base + chunk_b * 16;
+            const uint8_t* src = nullptr;
+            if (n_g < N && k_g < K) {
+                src = reinterpret_cast<const uint8_t*>(&B[n_g * K + k_g]);
+            }
+            int chunk_sw = chunk_b ^ (row_b & SWIZZLE_MASK);
+            uint32_t dst = to_smem(
+                &B_smem[stage * B_TILE_BYTES + row_b * BLOCK_K + chunk_sw * 16]);
+            cp_async_16(dst, src);
+        }
+    };
+    // Per-warp accumulators: M_ATOMS rows of mma x N_ATOMS_PW cols x 4 fp32.
+    float acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi)
+    #pragma unroll
+    for (int ni = 0; ni < N_ATOMS_PW; ++ni)
+    #pragma unroll
+    for (int j = 0; j < 4; ++j) acc[mi][ni][j] = 0.0f;
+    // Prefetch STAGES-1 chunks.
+    const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K;
+    #pragma unroll
+    for (int s = 0; s < STAGES - 1; ++s) {
+        if (s * BLOCK_K < K) issue_load(s, s * BLOCK_K);
+        asm volatile("cp.async.commit_group;\n" ::);
+    }
+    int compute_stage = 0;
+    for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) {
+        int issue_iter = k_iter + (STAGES - 1);
+        int issue_stage = issue_iter % STAGES;
+        if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K);
+        asm volatile("cp.async.commit_group;\n" ::);
+        asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1));
+        __syncthreads();
+        uint8_t* A_stage = A_smem + compute_stage * A_TILE_BYTES;
+        uint8_t* B_stage = B_smem + compute_stage * B_TILE_BYTES;
+        // K-atom inner loop. Per k_a: ldmatrix A (per m-atom) and B (per N-pair).
+        #pragma unroll
+        for (int k_a = 0; k_a < K_ATOMS; ++k_a) {
+            uint32_t A_regs[M_ATOMS][4];
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int row_in_tile = mi * 16 + row_block * 8 + row_in_frag;
+                int chunk = 2 * k_a + col_block;
+                int chunk_sw = chunk ^ (row_in_tile & SWIZZLE_MASK);
+                uint32_t addr = to_smem(
+                    &A_stage[row_in_tile * BLOCK_K + chunk_sw * 16]);
+                ldmatrix_x4_b16(
+                    A_regs[mi][0], A_regs[mi][1], A_regs[mi][2], A_regs[mi][3],
+                    addr);
+                // ldmatrix output mapping vs mma m16n8k32 A operand:
+                //   ldm d0=TL → mma a0
+                //   ldm d1=TR → mma a2
+                //   ldm d2=BL → mma a1
+                //   ldm d3=BR → mma a3
+            }
+            uint32_t B_regs[N_PAIRS_PW][4];
+            #pragma unroll
+            for (int np = 0; np < N_PAIRS_PW; ++np) {
+                int n_base_pair = warp_id * N_ATOMS_PW * 8 + np * 16;
+                int n_row_in_tile = n_base_pair + row_block * 8 + row_in_frag;
+                int chunk = 2 * k_a + col_block;
+                int chunk_sw = chunk ^ (n_row_in_tile & SWIZZLE_MASK);
+                uint32_t addr = to_smem(
+                    &B_stage[n_row_in_tile * BLOCK_K + chunk_sw * 16]);
+                ldmatrix_x4_b16(
+                    B_regs[np][0], B_regs[np][1], B_regs[np][2], B_regs[np][3],
+                    addr);
+                // ldm output mapping for paired N-atoms:
+                //   d0 = TL = N-atom0's b0 (rows 0-7, K-cols 0-15)
+                //   d1 = TR = N-atom0's b1 (rows 0-7, K-cols 16-31)
+                //   d2 = BL = N-atom1's b0 (rows 8-15, K-cols 0-15)
+                //   d3 = BR = N-atom1's b1 (rows 8-15, K-cols 16-31)
+            }
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                #pragma unroll
+                for (int np = 0; np < N_PAIRS_PW; ++np) {
+                    int ni0 = np * 2;
+                    int ni1 = np * 2 + 1;
+                    // N-atom 0: B = (b0=B_regs[np][0], b1=B_regs[np][1])
+                    mma_m16n8k32_e4m3(
+                        acc[mi][ni0][0], acc[mi][ni0][1],
+                        acc[mi][ni0][2], acc[mi][ni0][3],
+                        A_regs[mi][0], A_regs[mi][2],
+                        A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][0], B_regs[np][1]);
+                    // N-atom 1: B = (b0=B_regs[np][2], b1=B_regs[np][3])
+                    mma_m16n8k32_e4m3(
+                        acc[mi][ni1][0], acc[mi][ni1][1],
+                        acc[mi][ni1][2], acc[mi][ni1][3],
+                        A_regs[mi][0], A_regs[mi][2],
+                        A_regs[mi][1], A_regs[mi][3],
+                        B_regs[np][2], B_regs[np][3]);
+                }
+            }
+        }
+        compute_stage = (compute_stage + 1) % STAGES;
+    }
+    asm volatile("cp.async.wait_all;\n" ::);
+    // Epilogue: write 4 fp32 acc per lane to D[BF16].
+    // mma m16n8 output per lane (h=lane/4, l=lane%4):
+    //   d0,d1 -> row h,         cols 2l, 2l+1
+    //   d2,d3 -> row h+8,       cols 2l, 2l+1
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int col_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            if (row0 < M) {
+                if (col_base     < N)
+                    D[row0 * N + col_base    ] = __float2bfloat16(acc[mi][ni][0] * alpha);
+                if (col_base + 1 < N)
+                    D[row0 * N + col_base + 1] = __float2bfloat16(acc[mi][ni][1] * alpha);
+            }
+            if (row1 < M) {
+                if (col_base     < N)
+                    D[row1 * N + col_base    ] = __float2bfloat16(acc[mi][ni][2] * alpha);
+                if (col_base + 1 < N)
+                    D[row1 * N + col_base + 1] = __float2bfloat16(acc[mi][ni][3] * alpha);
+            }
+        }
+    }
+}
+template <int BM, int BN, int BK, int W, int STAGES = 2, int MIN_BLK = 4>
+int launch_(const void* A, const void* B, void* D,
+            int M, int N, int K, float alpha, cudaStream_t s)
+{
+    if (K % BK != 0) return 2;
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    int smem_bytes = STAGES * (BM + BN) * BK;
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_gemm_ld_kernel<BM, BN, BK, W, STAGES, MIN_BLK>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_gemm_ld_kernel<BM, BN, BK, W, STAGES, MIN_BLK><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D),
+        M, N, K, alpha);
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+}  // namespace
+#define DEFINE(NAME, BM, BN, BK, W, S)                                           \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,           \
+           float alpha, cudaStream_t stream) {                                   \
+    return launch_<BM, BN, BK, W, S, 4>(A, B, D, M, N, K, alpha, stream);        \
+  }
+#define DEFINE_BIG(NAME, BM, BN, BK, W, S)                                       \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,           \
+           float alpha, cudaStream_t stream) {                                   \
+    return launch_<BM, BN, BK, W, S, 1>(A, B, D, M, N, K, alpha, stream);        \
+  }
+DEFINE(ld_fp8_gemm_16x64x128_w4,    16,  64, 128, 4, 2)
+DEFINE(ld_fp8_gemm_16x128x128_w4,   16, 128, 128, 4, 2)
+DEFINE(ld_fp8_gemm_16x256x128_w8,   16, 256, 128, 8, 2)
+DEFINE(ld_fp8_gemm_32x64x128_w4,    32,  64, 128, 4, 2)
+DEFINE(ld_fp8_gemm_32x128x128_w4,   32, 128, 128, 4, 2)
+DEFINE(ld_fp8_gemm_32x128x128_w8,   32, 128, 128, 8, 2)
+DEFINE(ld_fp8_gemm_16x64x128_w4_s3, 16,  64, 128, 4, 3)
+DEFINE(ld_fp8_gemm_16x128x128_w4_s3,16, 128, 128, 4, 3)
+DEFINE(ld_fp8_gemm_32x64x128_w4_s3, 32,  64, 128, 4, 3)
+DEFINE(ld_fp8_gemm_32x128x128_w4_s3,32, 128, 128, 4, 3)
+DEFINE(ld_fp8_gemm_16x192x128_w4,   16, 192, 128, 4, 2)
+DEFINE(ld_fp8_gemm_32x192x128_w4,   32, 192, 128, 4, 2)
+DEFINE(ld_fp8_gemm_16x64x128_w4_s4, 16,  64, 128, 4, 4)
+DEFINE(ld_fp8_gemm_16x64x128_w4_s5, 16,  64, 128, 4, 5)
+DEFINE(ld_fp8_gemm_32x64x128_w4_s4, 32,  64, 128, 4, 4)
+DEFINE(ld_fp8_gemm_32x64x128_w4_s5, 32,  64, 128, 4, 5)
+DEFINE(ld_fp8_gemm_16x128x128_w4_s4,16, 128, 128, 4, 4)
+DEFINE(ld_fp8_gemm_32x128x128_w4_s4,32, 128, 128, 4, 4)
+// BK=256 variants — large K-tile, fewer K-iters, more compute per CTA.
+DEFINE(ld_fp8_gemm_16x64x256_w4,    16,  64, 256, 4, 2)
+DEFINE(ld_fp8_gemm_16x128x256_w4,   16, 128, 256, 4, 2)
+DEFINE(ld_fp8_gemm_32x64x256_w4,    32,  64, 256, 4, 2)
+DEFINE(ld_fp8_gemm_32x128x256_w4,   32, 128, 256, 4, 2)
+DEFINE(ld_fp8_gemm_16x64x256_w4_s3, 16,  64, 256, 4, 3)
+// BK=64 variants — small K-tile, more K-iters, finer pipeline grain.
+DEFINE(ld_fp8_gemm_16x64x64_w4,     16,  64, 64,  4, 2)
+DEFINE(ld_fp8_gemm_16x128x64_w4,    16, 128, 64,  4, 2)
+DEFINE(ld_fp8_gemm_32x64x64_w4,     32,  64, 64,  4, 2)
+DEFINE(ld_fp8_gemm_16x64x64_w4_s3,  16,  64, 64,  4, 3)
+DEFINE(ld_fp8_gemm_16x64x64_w4_s4,  16,  64, 64,  4, 4)
+// und_qkv (M=188, N=9216, K=512) untried variants: bigger BM reduces
+// CTA count (188/64=3 m_tiles vs 188/32=6); bigger BK reduces K_iter
+// overhead; BK=512 single-iter eliminates pipeline overhead at K=512.
+DEFINE(ld_fp8_gemm_64x64x128_w4,    64,  64, 128, 4, 2)
+DEFINE(ld_fp8_gemm_64x128x128_w4,   64, 128, 128, 4, 2)
+DEFINE(ld_fp8_gemm_64x64x256_w4,    64,  64, 256, 4, 2)
+DEFINE(ld_fp8_gemm_64x128x256_w4,   64, 128, 256, 4, 2)
+DEFINE(ld_fp8_gemm_64x64x256_w4_s3, 64,  64, 256, 4, 3)
+DEFINE(ld_fp8_gemm_32x64x256_w4_s3, 32,  64, 256, 4, 3)
+DEFINE(ld_fp8_gemm_32x128x256_w4_s3,32, 128, 256, 4, 3)
+DEFINE(ld_fp8_gemm_128x64x128_w4,   128, 64, 128, 4, 2)
+DEFINE(ld_fp8_gemm_128x128x128_w4,  128, 128, 128, 4, 2)
+// BK=512 single-iter benched 14.35us, worse than BK=256 pipelined 12.35us.
+// Kept the relaxed static_assert (BK=512 now allowed) but no variants
+// instantiated — they lose to existing BK=256 cp.async pipeline.
+#undef DEFINE
+#undef DEFINE_BIG
+}  // namespace smallM_ld
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_smallM_handtuned_ldmatrix_sm120.cuh ADDED Viewed

	@@ -0,0 +1,68 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Hand-tuned FP8 e4m3 GEMM v2 — adds 128B swizzle smem layout + ldmatrix.x4
+// loads to clear bank conflicts that bottleneck v1 (`fp8_smallM_handtuned`).
+//
+// All variants restricted to BLOCK_K = 128 (natural 128B swizzle stride).
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+namespace smallM_ld {
+#define DECL(NAME) \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K, \
+           float alpha, cudaStream_t stream)
+DECL(ld_fp8_gemm_16x64x128_w4);
+DECL(ld_fp8_gemm_16x128x128_w4);
+DECL(ld_fp8_gemm_16x256x128_w8);
+DECL(ld_fp8_gemm_32x64x128_w4);
+DECL(ld_fp8_gemm_32x128x128_w4);
+DECL(ld_fp8_gemm_32x128x128_w8);
+DECL(ld_fp8_gemm_16x64x128_w4_s3);
+DECL(ld_fp8_gemm_16x128x128_w4_s3);
+DECL(ld_fp8_gemm_32x64x128_w4_s3);
+DECL(ld_fp8_gemm_32x128x128_w4_s3);
+DECL(ld_fp8_gemm_16x192x128_w4);
+DECL(ld_fp8_gemm_32x192x128_w4);
+DECL(ld_fp8_gemm_16x64x128_w4_s4);
+DECL(ld_fp8_gemm_16x64x128_w4_s5);
+DECL(ld_fp8_gemm_32x64x128_w4_s4);
+DECL(ld_fp8_gemm_32x64x128_w4_s5);
+DECL(ld_fp8_gemm_16x128x128_w4_s4);
+DECL(ld_fp8_gemm_32x128x128_w4_s4);
+DECL(ld_fp8_gemm_16x64x256_w4);
+DECL(ld_fp8_gemm_16x128x256_w4);
+DECL(ld_fp8_gemm_32x64x256_w4);
+DECL(ld_fp8_gemm_32x128x256_w4);
+DECL(ld_fp8_gemm_16x64x256_w4_s3);
+DECL(ld_fp8_gemm_16x64x64_w4);
+DECL(ld_fp8_gemm_16x128x64_w4);
+DECL(ld_fp8_gemm_32x64x64_w4);
+DECL(ld_fp8_gemm_16x64x64_w4_s3);
+DECL(ld_fp8_gemm_16x64x64_w4_s4);
+// und_qkv attack variants (M=188, K=512)
+DECL(ld_fp8_gemm_64x64x128_w4);
+DECL(ld_fp8_gemm_64x128x128_w4);
+DECL(ld_fp8_gemm_64x64x256_w4);
+DECL(ld_fp8_gemm_64x128x256_w4);
+DECL(ld_fp8_gemm_64x64x256_w4_s3);
+DECL(ld_fp8_gemm_32x64x256_w4_s3);
+DECL(ld_fp8_gemm_32x128x256_w4_s3);
+DECL(ld_fp8_gemm_128x64x128_w4);
+DECL(ld_fp8_gemm_128x128x128_w4);
+#undef DECL
+}  // namespace smallM_ld
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_smallM_handtuned_sm120.cu ADDED Viewed

	@@ -0,0 +1,338 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// Hand-tuned FP8 e4m3 -> BF16 GEMM for sm_120a small-M motus shapes.
+// Inline-PTX m16n8k32 mma + 2-stage cp.async pipeline, no cutlass collective
+// builder overhead. Modeled after V5split kernel_A pattern; epilogue is just
+// alpha * acc -> BF16 (no bias / GELU / quant).
+//
+// Motivation: cutlass scaffold has ~8 us launch/setup floor on sm_120 for
+// small-M kernels, even with smallest tiles. cuBLASLt nvjet ~5 us. To break
+// below cuBLASLt, must avoid the scaffold entirely.
+//
+// Per-tensor scale (A_scale, W_scale as float scalars folded into alpha).
+// Returns 0 on success.
+#include "fp8_smallM_handtuned_sm120.cuh"
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+namespace flash_rt {
+namespace gemm {
+namespace smallM_hand {
+namespace {
+__device__ __forceinline__ void mma_m16n8k32_e4m3(
+    float &d0, float &d1, float &d2, float &d3,
+    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
+    uint32_t b0, uint32_t b1)
+{
+    asm volatile(
+        "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
+        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
+        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
+}
+__device__ __forceinline__ void cp_async_16(uint32_t smem, const uint8_t* src) {
+    int b = (src == nullptr) ? 0 : 16;
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n"
+                 :: "r"(smem), "l"(src), "r"(b));
+}
+__device__ __forceinline__ uint32_t to_smem(const void* p) {
+    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
+}
+// Generic FP8 GEMM kernel parameterized on tile shape + pipeline stages.
+//  - A: [M, K] row-major FP8 e4m3
+//  - B: [N, K] row-major FP8 e4m3  (= W.T col-major layout)
+//  - D: [M, N] row-major BF16
+//  - alpha = a_scale * w_scale (per-tensor)
+//  - STAGES = pipeline depth (2 or 3)
+//  - MIN_BLOCKS_PER_SM = launch_bounds hint (1 for big-smem variants)
+template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_WARPS,
+          int STAGES = 2, int MIN_BLOCKS_PER_SM = 4>
+__global__ __launch_bounds__(NUM_WARPS * 32, MIN_BLOCKS_PER_SM)
+void fp8_gemm_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,
+    const __nv_fp8_e4m3* __restrict__ B,
+    __nv_bfloat16* __restrict__ D,
+    int M, int N, int K,
+    float alpha)
+{
+    static_assert(BLOCK_K % 32 == 0, "BLOCK_K must be multiple of 32");
+    static_assert(BLOCK_N % 8 == 0,  "BLOCK_N must be multiple of 8");
+    static_assert(BLOCK_M % 16 == 0, "BLOCK_M must be multiple of 16 (mma m=16)");
+    static_assert((BLOCK_N / 8) % NUM_WARPS == 0, "N-atoms must split evenly across warps");
+    constexpr int THREADS    = NUM_WARPS * 32;
+    constexpr int M_ATOMS    = BLOCK_M / 16;        // m-atom rows per CTA
+    constexpr int N_ATOMS    = BLOCK_N / 8;
+    constexpr int N_ATOMS_PW = N_ATOMS / NUM_WARPS;
+    constexpr int K_ATOMS    = BLOCK_K / 32;
+    constexpr int SMEM_K_PAD = BLOCK_K + 16;        // +16 byte padding to avoid bank conflict
+    // smem layout: [stage][row][col_padded]; STAGES stages.
+    extern __shared__ uint8_t smem_raw[];
+    uint8_t* A_smem = smem_raw;
+    uint8_t* B_smem = A_smem + STAGES * BLOCK_M * SMEM_K_PAD;
+    const int cta_m = blockIdx.x;
+    const int cta_n = blockIdx.y;
+    const int m_base = cta_m * BLOCK_M;
+    const int n_base = cta_n * BLOCK_N;
+    const int t = threadIdx.x;
+    const int warp_id = t / 32;
+    const int lane = t % 32;
+    const int l = lane % 4;
+    const int h = lane / 4;
+    auto issue_load = [&](int stage, int k_base) {
+        // Load A [BLOCK_M, BLOCK_K] FP8 = BLOCK_M * BLOCK_K bytes.
+        constexpr int A_TOTAL_16B = BLOCK_M * BLOCK_K / 16;
+        constexpr int A_ITERS = (A_TOTAL_16B + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < A_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= A_TOTAL_16B) break;
+            int row_a = idx / (BLOCK_K / 16);
+            int koff_a = (idx % (BLOCK_K / 16)) * 16;
+            int m_glob = m_base + row_a;
+            int k_glob = k_base + koff_a;
+            const uint8_t* a_src = nullptr;
+            if (m_glob < M && k_glob < K) {
+                a_src = reinterpret_cast<const uint8_t*>(&A[m_glob * K + k_glob]);
+            }
+            cp_async_16(
+                to_smem(&A_smem[stage * BLOCK_M * SMEM_K_PAD
+                                + row_a * SMEM_K_PAD + koff_a]),
+                a_src);
+        }
+        // Load B [BLOCK_N, BLOCK_K] FP8 (B is [N, K] row-major).
+        constexpr int B_TOTAL_16B = BLOCK_N * BLOCK_K / 16;
+        constexpr int B_ITERS = (B_TOTAL_16B + THREADS - 1) / THREADS;
+        #pragma unroll
+        for (int it = 0; it < B_ITERS; ++it) {
+            int idx = it * THREADS + t;
+            if (idx >= B_TOTAL_16B) break;
+            int row_b = idx / (BLOCK_K / 16);
+            int koff_b = (idx % (BLOCK_K / 16)) * 16;
+            int n_glob = n_base + row_b;
+            int k_glob = k_base + koff_b;
+            const uint8_t* b_src = nullptr;
+            if (n_glob < N && k_glob < K) {
+                b_src = reinterpret_cast<const uint8_t*>(&B[n_glob * K + k_glob]);
+            }
+            cp_async_16(
+                to_smem(&B_smem[stage * BLOCK_N * SMEM_K_PAD
+                                + row_b * SMEM_K_PAD + koff_b]),
+                b_src);
+        }
+    };
+    // Per-warp accumulators: M_ATOMS * N_ATOMS_PW * 4 fp32
+    float acc[M_ATOMS][N_ATOMS_PW][4];
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) acc[mi][ni][j] = 0.0f;
+        }
+    }
+    // Prefetch STAGES-1 chunks before main loop (deep pipeline).
+    const int K_ITERS = (K + BLOCK_K - 1) / BLOCK_K;
+    #pragma unroll
+    for (int s = 0; s < STAGES - 1; ++s) {
+        int kb = s * BLOCK_K;
+        if (kb < K) issue_load(s, kb);
+        asm volatile("cp.async.commit_group;\n" ::);
+    }
+    int compute_stage = 0;
+    for (int k_iter = 0; k_iter < K_ITERS; ++k_iter) {
+        int k_base = k_iter * BLOCK_K;
+        // Issue next load STAGES-1 ahead.
+        int issue_iter = k_iter + (STAGES - 1);
+        int issue_stage = issue_iter % STAGES;
+        if (issue_iter < K_ITERS) issue_load(issue_stage, issue_iter * BLOCK_K);
+        asm volatile("cp.async.commit_group;\n" ::);
+        // Wait until STAGES-1 prior loads are still in flight, current ready.
+        asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 1));
+        __syncthreads();
+        #pragma unroll
+        for (int k_iter = 0; k_iter < K_ATOMS; ++k_iter) {
+            int kA0 = k_iter * 32 + 4 * l;
+            int kA2 = k_iter * 32 + 4 * l + 16;
+            #pragma unroll
+            for (int mi = 0; mi < M_ATOMS; ++mi) {
+                int rA0 = mi * 16 + h;
+                int rA1 = mi * 16 + h + 8;
+                uint32_t A0 = *reinterpret_cast<const uint32_t*>(
+                    &A_smem[compute_stage * BLOCK_M * SMEM_K_PAD + rA0 * SMEM_K_PAD + kA0]);
+                uint32_t A1 = *reinterpret_cast<const uint32_t*>(
+                    &A_smem[compute_stage * BLOCK_M * SMEM_K_PAD + rA1 * SMEM_K_PAD + kA0]);
+                uint32_t A2 = *reinterpret_cast<const uint32_t*>(
+                    &A_smem[compute_stage * BLOCK_M * SMEM_K_PAD + rA0 * SMEM_K_PAD + kA2]);
+                uint32_t A3 = *reinterpret_cast<const uint32_t*>(
+                    &A_smem[compute_stage * BLOCK_M * SMEM_K_PAD + rA1 * SMEM_K_PAD + kA2]);
+                #pragma unroll
+                for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+                    int co_n = warp_id * N_ATOMS_PW * 8 + ni * 8 + h;
+                    uint32_t B0 = *reinterpret_cast<const uint32_t*>(
+                        &B_smem[compute_stage * BLOCK_N * SMEM_K_PAD + co_n * SMEM_K_PAD + kA0]);
+                    uint32_t B1 = *reinterpret_cast<const uint32_t*>(
+                        &B_smem[compute_stage * BLOCK_N * SMEM_K_PAD + co_n * SMEM_K_PAD + kA2]);
+                    mma_m16n8k32_e4m3(
+                        acc[mi][ni][0], acc[mi][ni][1], acc[mi][ni][2], acc[mi][ni][3],
+                        A0, A1, A2, A3, B0, B1);
+                }
+            }
+        }
+        compute_stage = (compute_stage + 1) % STAGES;
+    }
+    asm volatile("cp.async.wait_all;\n" ::);
+    // Epilogue: alpha * acc -> BF16 -> HBM.
+    // m16n8 layout: thread (h, l): rows {h, h+8}, cols {2*l, 2*l+1}.
+    #pragma unroll
+    for (int mi = 0; mi < M_ATOMS; ++mi) {
+        int row0 = m_base + mi * 16 + h;
+        int row1 = row0 + 8;
+        #pragma unroll
+        for (int ni = 0; ni < N_ATOMS_PW; ++ni) {
+            int n_pair_base = n_base + warp_id * N_ATOMS_PW * 8 + ni * 8 + 2 * l;
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                int row = (j < 2) ? row0 : row1;
+                int col = n_pair_base + (j & 1);
+                if (row < M && col < N) {
+                    float v = acc[mi][ni][j] * alpha;
+                    D[row * N + col] = __float2bfloat16(v);
+                }
+            }
+        }
+    }
+}
+template <int BM, int BN, int BK, int W, int STAGES = 2, int MIN_BLK = 4>
+int launch_(const void* A, const void* B, void* D,
+            int M, int N, int K, float alpha, cudaStream_t s)
+{
+    int grid_m = (M + BM - 1) / BM;
+    int grid_n = (N + BN - 1) / BN;
+    dim3 grid(grid_m, grid_n, 1);
+    dim3 block(W * 32, 1, 1);
+    int smem_bytes = STAGES * (BM + BN) * (BK + 16);
+    // sm_120 default dynamic smem is 48 KB; opt-in to higher (up to ~228 KB).
+    if (smem_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(
+            (const void*)&fp8_gemm_kernel<BM, BN, BK, W, STAGES, MIN_BLK>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    }
+    fp8_gemm_kernel<BM, BN, BK, W, STAGES, MIN_BLK><<<grid, block, smem_bytes, s>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(A),
+        reinterpret_cast<const __nv_fp8_e4m3*>(B),
+        reinterpret_cast<__nv_bfloat16*>(D),
+        M, N, K, alpha);
+    cudaError_t err = cudaGetLastError();
+    return (err == cudaSuccess) ? 0 : 1;
+}
+}  // namespace
+// Variant instantiations.
+#define DEFINE(NAME, BM, BN, BK, W, S)                                           \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,           \
+           float alpha, cudaStream_t stream) {                                   \
+    return launch_<BM, BN, BK, W, S, 4>(A, B, D, M, N, K, alpha, stream);        \
+  }
+// Big-smem variant — uses MIN_BLOCKS_PER_SM=1 to relax register pressure.
+#define DEFINE_BIG(NAME, BM, BN, BK, W, S)                                       \
+  int NAME(const void* A, const void* B, void* D, int M, int N, int K,           \
+           float alpha, cudaStream_t stream) {                                   \
+    return launch_<BM, BN, BK, W, S, 1>(A, B, D, M, N, K, alpha, stream);        \
+  }
+// 2-stage pipeline baseline.
+DEFINE(fp8_gemm_16x64x128_w4,   16,  64, 128, 4, 2)
+DEFINE(fp8_gemm_16x128x128_w4,  16, 128, 128, 4, 2)
+DEFINE(fp8_gemm_16x256x128_w8,  16, 256, 128, 8, 2)
+DEFINE(fp8_gemm_32x64x128_w4,   32,  64, 128, 4, 2)
+DEFINE(fp8_gemm_32x128x128_w4,  32, 128, 128, 4, 2)
+DEFINE(fp8_gemm_32x128x128_w8,  32, 128, 128, 8, 2)
+// 3-stage pipeline (better cp.async overlap).
+DEFINE(fp8_gemm_16x64x128_w4_s3,   16,  64, 128, 4, 3)
+DEFINE(fp8_gemm_16x128x128_w4_s3,  16, 128, 128, 4, 3)
+DEFINE(fp8_gemm_32x64x128_w4_s3,   32,  64, 128, 4, 3)
+DEFINE(fp8_gemm_32x128x128_w4_s3,  32, 128, 128, 4, 3)
+// BLOCK_K=256 (fewer K-iters, bigger cp.async chunks per iter).
+DEFINE(fp8_gemm_16x64x256_w4,   16,  64, 256, 4, 2)
+DEFINE(fp8_gemm_16x128x256_w4,  16, 128, 256, 4, 2)
+DEFINE(fp8_gemm_32x64x256_w4,   32,  64, 256, 4, 2)
+DEFINE(fp8_gemm_32x128x256_w4,  32, 128, 256, 4, 2)
+// Wider BLOCK_N for big-N shapes (action_qkv, und_qkv: N=9216).
+DEFINE(fp8_gemm_16x192x128_w4,  16, 192, 128, 4, 2)
+DEFINE(fp8_gemm_16x192x128_w8,  16, 192, 128, 8, 2)
+DEFINE(fp8_gemm_32x192x128_w4,  32, 192, 128, 4, 2)
+// 4-stage pipeline.
+DEFINE(fp8_gemm_16x64x128_w4_s4,   16,  64, 128, 4, 4)
+DEFINE(fp8_gemm_32x64x128_w4_s4,   32,  64, 128, 4, 4)
+// Wider BLOCK_N=384 (needs N % 384, 8-warp config).
+DEFINE(fp8_gemm_16x384x128_w8,     16, 384, 128, 8, 2)
+DEFINE(fp8_gemm_32x384x128_w8,     32, 384, 128, 8, 2)
+// More warps, smaller BLOCK_N (more N-tiles parallelism per CTA).
+DEFINE(fp8_gemm_16x64x128_w8,      16,  64, 128, 8, 2)
+DEFINE(fp8_gemm_32x64x128_w8,      32,  64, 128, 8, 2)
+// 32x64x128 with 8-stage pipeline (deep cp.async overlap for K-bound shapes).
+DEFINE(fp8_gemm_32x64x128_w4_s5,   32,  64, 128, 4, 5)
+// BLOCK_K=64 variants — better pipeline overlap for K-small shapes (K=512).
+DEFINE(fp8_gemm_16x64x64_w4,       16,  64,  64, 4, 2)
+DEFINE(fp8_gemm_16x128x64_w4,      16, 128,  64, 4, 2)
+DEFINE(fp8_gemm_32x64x64_w4,       32,  64,  64, 4, 2)
+DEFINE(fp8_gemm_32x128x64_w4,      32, 128,  64, 4, 2)
+DEFINE(fp8_gemm_16x64x64_w4_s3,    16,  64,  64, 4, 3)
+DEFINE(fp8_gemm_16x64x64_w4_s4,    16,  64,  64, 4, 4)
+// Big-smem (BLOCK_N=384/512) — MIN_BLOCKS_PER_SM=1 relaxes register pressure.
+// Targets multi-wave shapes (und_qkv: 9216 N) to reduce wave count.
+DEFINE_BIG(fp8_gemm_16x384x128_w4_big,  16, 384, 128, 4, 2)
+DEFINE_BIG(fp8_gemm_32x384x128_w4_big,  32, 384, 128, 4, 2)
+DEFINE_BIG(fp8_gemm_16x512x128_w8_big,  16, 512, 128, 8, 2)
+DEFINE_BIG(fp8_gemm_16x256x128_w4_big,  16, 256, 128, 4, 2)
+DEFINE_BIG(fp8_gemm_32x256x128_w4_big,  32, 256, 128, 4, 2)
+// BLOCK_M=64 / 128 variants — for M=138 shapes to drop wave count to 1.
+// und_qkv (M=138 N=9216): BLOCK_M=64 -> 3 M-tiles, BLOCK_M=128 -> 2 M-tiles.
+// Combined with BLOCK_N=128: 216 / 144 total CTAs => 1-1.3 waves on 170 SMs.
+DEFINE(fp8_gemm_64x64x128_w4,    64,  64, 128, 4, 2)
+DEFINE(fp8_gemm_64x128x128_w4,   64, 128, 128, 4, 2)
+DEFINE(fp8_gemm_64x128x128_w8,   64, 128, 128, 8, 2)
+DEFINE(fp8_gemm_128x64x128_w4,  128,  64, 128, 4, 2)
+DEFINE(fp8_gemm_128x128x128_w4, 128, 128, 128, 4, 2)
+DEFINE(fp8_gemm_128x128x128_w8, 128, 128, 128, 8, 2)
+DEFINE_BIG(fp8_gemm_64x256x128_w4_big,  64, 256, 128, 4, 2)
+DEFINE_BIG(fp8_gemm_64x256x128_w8_big,  64, 256, 128, 8, 2)
+DEFINE_BIG(fp8_gemm_128x256x128_w8_big, 128, 256, 128, 8, 2)
+#undef DEFINE
+#undef DEFINE_BIG
+}  // namespace smallM_hand
+}  // namespace gemm
+}  // namespace flash_rt

csrc/fp8_smallM_handtuned_sm120.cuh ADDED Viewed

	@@ -0,0 +1,89 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+namespace smallM_hand {
+// Hand-tuned FP8 e4m3 -> BF16 GEMM for sm_120a small-M motus shapes.
+// Inputs: FP8 A [M,K] row-major, FP8 B [N,K] row-major (= W.T), BF16 D [M,N].
+// alpha = a_scale * w_scale (per-tensor).
+// Returns 0 on success.
+#define DECL(NAME) \
+  int NAME(const void* A, const void* B, void* D, \
+           int M, int N, int K, float alpha, cudaStream_t stream)
+// 2-stage pipeline.
+DECL(fp8_gemm_16x64x128_w4);
+DECL(fp8_gemm_16x128x128_w4);
+DECL(fp8_gemm_16x256x128_w8);
+DECL(fp8_gemm_32x64x128_w4);
+DECL(fp8_gemm_32x128x128_w4);
+DECL(fp8_gemm_32x128x128_w8);
+// 3-stage pipeline.
+DECL(fp8_gemm_16x64x128_w4_s3);
+DECL(fp8_gemm_16x128x128_w4_s3);
+DECL(fp8_gemm_32x64x128_w4_s3);
+DECL(fp8_gemm_32x128x128_w4_s3);
+// BLOCK_K=256.
+DECL(fp8_gemm_16x64x256_w4);
+DECL(fp8_gemm_16x128x256_w4);
+DECL(fp8_gemm_32x64x256_w4);
+DECL(fp8_gemm_32x128x256_w4);
+// BLOCK_N=192 (for N=9216 shapes).
+DECL(fp8_gemm_16x192x128_w4);
+DECL(fp8_gemm_16x192x128_w8);
+DECL(fp8_gemm_32x192x128_w4);
+// 4-stage pipeline.
+DECL(fp8_gemm_16x64x128_w4_s4);
+DECL(fp8_gemm_32x64x128_w4_s4);
+// Wider BLOCK_N=384 (needs N % 384).
+DECL(fp8_gemm_16x384x128_w8);
+DECL(fp8_gemm_32x384x128_w8);
+// 8-warp variants of 16x64 / 32x64.
+DECL(fp8_gemm_16x64x128_w8);
+DECL(fp8_gemm_32x64x128_w8);
+// 5-stage pipeline.
+DECL(fp8_gemm_32x64x128_w4_s5);
+// BLOCK_K=64.
+DECL(fp8_gemm_16x64x64_w4);
+DECL(fp8_gemm_16x128x64_w4);
+DECL(fp8_gemm_32x64x64_w4);
+DECL(fp8_gemm_32x128x64_w4);
+DECL(fp8_gemm_16x64x64_w4_s3);
+DECL(fp8_gemm_16x64x64_w4_s4);
+// Big-smem BLOCK_N variants.
+DECL(fp8_gemm_16x384x128_w4_big);
+DECL(fp8_gemm_32x384x128_w4_big);
+DECL(fp8_gemm_16x512x128_w8_big);
+DECL(fp8_gemm_16x256x128_w4_big);
+DECL(fp8_gemm_32x256x128_w4_big);
+// BLOCK_M=64 / 128 — wave reduction for M=138 shapes (und_qkv main target).
+DECL(fp8_gemm_64x64x128_w4);
+DECL(fp8_gemm_64x128x128_w4);
+DECL(fp8_gemm_64x128x128_w8);
+DECL(fp8_gemm_128x64x128_w4);
+DECL(fp8_gemm_128x128x128_w4);
+DECL(fp8_gemm_128x128x128_w8);
+DECL(fp8_gemm_64x256x128_w4_big);
+DECL(fp8_gemm_64x256x128_w8_big);
+DECL(fp8_gemm_128x256x128_w8_big);
+#undef DECL
+}  // namespace smallM_hand
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm_types_sm110.h ADDED Viewed

	@@ -0,0 +1,362 @@

+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS SM100-family FP8 GEMM templates used by the SM110a Thor build.
+// The public package path consumes the BF16-output Sq/T1/Wide variants below;
+// the FP16 and FP32 variants remain available to the native implementation.
+// T1/T2 use explicit TmaWarpSpecialized2Sm scheduling for tactic control.
+#pragma once
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/fusion/operations.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/util/packed_stride.hpp"
+using namespace cute;
+// Type aliases
+using cutlass_fp8 = cutlass::float_e4m3_t;
+using cutlass_fp16 = cutlass::half_t;
+// ============================================================
+//  PlainGemm: 256×128×64, Cluster 2×2×1
+//  Standard FP8→FP16 GEMM (Identity epilogue)
+// ============================================================
+namespace sm100_plain {
+using Tile = Shape<_256, _128, _64>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_plain
+// ============================================================
+//  GeluGemm: 256×128×64 + GELU epilogue
+// ============================================================
+namespace sm100_gelu {
+using Tile = Shape<_256, _128, _64>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::GELU, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_gelu
+// ============================================================
+//  SqGemm: 256×256×128 — deeper K pipeline for large GEMMs
+// ============================================================
+namespace sm100_sq {
+using Tile = Shape<_256, _256, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_sq
+// ============================================================
+//  WideGemm: 256×128×128 — deeper K for FFN down projection
+// ============================================================
+namespace sm100_wide {
+using Tile = Shape<_256, _128, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_wide
+// ============================================================
+//  T1Gemm: 128×256×128, Cluster 2×1×1, TmaWarpSpecialized2Sm
+//  EXACT match for Myelin's s128x256 best tactic (2SM)
+// ============================================================
+namespace sm100_t1 {
+using Tile = Shape<_128, _256, _128>;
+using Cluster = Shape<_2, _1, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::TmaWarpSpecialized2Sm, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_t1
+// ============================================================
+//  T2Gemm: 256×256×128, Cluster 2×1×1, TmaWarpSpecialized2Sm
+// ============================================================
+namespace sm100_t2 {
+using Tile = Shape<_256, _256, _128>;
+using Cluster = Shape<_2, _1, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass_fp16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::TmaWarpSpecialized2Sm, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_t2
+// ============================================================
+//  FP32 Output Variants — for models with activations > FP16 range
+//  (e.g., Pi0-FAST Gemma 2B deep layers where residual > 65504)
+//  Same tile configs, only output dtype changed: cutlass_fp16 → float
+// ============================================================
+using cutlass_fp32 = float;
+using cutlass_bf16 = cutlass::bfloat16_t;
+// ============================================================
+//  BF16 Output Variants — for models trained in BF16 with large activations
+//  Same FP8 inputs/accumulation, BF16 output (range ±3.4e38)
+// ============================================================
+namespace sm100_sq_bf16out {
+using Tile = Shape<_256, _256, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_bf16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_sq_bf16out
+namespace sm100_wide_bf16out {
+using Tile = Shape<_256, _128, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_bf16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_wide_bf16out
+// Wide projection with a per-column BF16 bias. This keeps the public
+// row-major [N,K] weight contract and removes the layout-dependent cuBLASLt
+// penalty on PI0.5/SigLIP down projections.
+namespace sm100_wide_bias_bf16out {
+using Tile = Shape<_256, _128, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombPerColBias<
+    cutlass_bf16, float, cutlass_bf16, cutlass_bf16>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_wide_bias_bf16out
+namespace sm100_wide_bias_gelu_bf16out {
+using Tile = Shape<_256, _128, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombPerColBiasEltAct<
+    cutlass::epilogue::thread::GELU_taylor,
+    cutlass_bf16, float, cutlass_bf16, cutlass_bf16>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_wide_bias_gelu_bf16out
+namespace sm100_t1_bf16out {
+using Tile = Shape<_128, _256, _128>;
+using Cluster = Shape<_2, _1, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_bf16, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass_bf16, cutlass::layout::RowMajor, 8,
+    cutlass::epilogue::TmaWarpSpecialized2Sm, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_t1_bf16out
+namespace sm100_sq_f32out {
+using Tile = Shape<_256, _256, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp32, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp32, cutlass::layout::RowMajor, 4,
+    cutlass_fp32, cutlass::layout::RowMajor, 4,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_sq_f32out
+namespace sm100_wide_f32out {
+using Tile = Shape<_256, _128, _128>;
+using Cluster = Shape<_2, _2, _1>;
+using Fusion = cutlass::epilogue::fusion::LinCombEltAct<
+    cutlass::epilogue::thread::Identity, cutlass_fp32, float>;
+using Epi = typename cutlass::epilogue::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+    float, float, cutlass_fp32, cutlass::layout::RowMajor, 4,
+    cutlass_fp32, cutlass::layout::RowMajor, 4,
+    cutlass::epilogue::collective::EpilogueScheduleAuto, Fusion>::CollectiveOp;
+using Main = typename cutlass::gemm::collective::CollectiveBuilder<
+    cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
+    cutlass_fp8, cutlass::layout::RowMajor, 16,
+    cutlass_fp8, cutlass::layout::ColumnMajor, 16,
+    float, Tile, Cluster,
+    cutlass::gemm::collective::StageCountAutoCarveout<
+        static_cast<int>(sizeof(typename Epi::SharedStorage))>,
+    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
+    cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>, Main, Epi>>;
+}  // namespace sm100_wide_f32out

examples/README.md ADDED Viewed

	@@ -0,0 +1,3 @@


1	+ # Examples
2	+
3	+ See the package README for minimal Hub usage.

flake.lock ADDED Viewed

	@@ -0,0 +1,118 @@

+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1785676244,
+        "narHash": "sha256-os1+/tdJLsC/iYG6t0F49Q3FOCaMB5fFp5cyDVF1d2o=",
+        "owner": "LiangSu8899",
+        "repo": "kernels",
+        "rev": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+        "type": "github"
+      },
+      "original": {
+        "owner": "LiangSu8899",
+        "repo": "kernels",
+        "rev": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix ADDED Viewed

	@@ -0,0 +1,20 @@

+{
+  description = "Flake for FlashRT FP8 GEMM kernels";
+  inputs = {
+    # Based on huggingface/kernels@e9152aa with only the stale CUTLASS 4.5.2
+    # fixed-output hash updated. Return to upstream after the hash fix lands.
+    kernel-builder.url =
+      "github:LiangSu8899/kernels/d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8";
+  };
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

scripts/README.md ADDED Viewed

	@@ -0,0 +1,3 @@


1	+ # Scripts
2	+
3	+ Package-specific helper scripts are not required yet.

tests/README.md ADDED Viewed

	@@ -0,0 +1,9 @@

+# Tests
+```bash
+python fp8-gemm/tests/test_fp8_gemm.py --backend source --mode full
+```
+SM110 full mode adds PI0.5, GROOT, Cosmos Edge, and LingBot VLA projection
+shapes plus forced Sq/T1/Wide correctness rows. Use `--backend installed` with
+the exact artifact directory for the release gate.

tests/test_fp8_gemm.py ADDED Viewed

	@@ -0,0 +1,698 @@

+#!/usr/bin/env python3
+"""Correctness tests for fp8-gemm."""
+from __future__ import annotations
+import argparse
+import importlib
+import json
+import math
+import os
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "fp8-gemm"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+SHAPES = {
+    "decode_m1_k512_n512": (1, 512, 512),
+    "decode_m1_k4096_n2048": (1, 4096, 2048),
+    "decode_m1_k4096_n8192": (1, 4096, 8192),
+    "small_m8_k1024_n2048": (8, 1024, 2048),
+    "small_m16_k4096_n4096": (16, 4096, 4096),
+    "small_m32_k4096_n8192": (32, 4096, 8192),
+    "small_m64_k512_n1024": (64, 512, 1024),
+}
+SM110_SHAPES = {
+    "large_m_boundary_65": (65, 2048, 2048),
+    # PI0.5 / PI0 decoder and encoder projection families.
+    "pi05_action_qkv": (51, 2048, 2560),
+    "pi05_action_o": (51, 2048, 2048),
+    "pi05_action_gate_up": (51, 2048, 16384),
+    "pi05_action_down": (51, 8192, 2048),
+    # GROOT N1.6/N1.7 DiT, backbone, and vision rows.
+    "groot_dit_qkv": (51, 1536, 4608),
+    "groot_n17_llm_o": (277, 2048, 2048),
+    "groot_n17_llm_gate_up": (277, 2048, 16384),
+    "groot_n17_llm_down": (277, 8192, 2048),
+    "groot_n17_vit_o": (1024, 1024, 1024),
+    # Cosmos Edge and LingBot projection families.
+    "cosmos_edge_action": (64, 2048, 9216),
+    "lingbot_vision_o": (1024, 1280, 1280),
+    "lingbot_action_gate_up": (105, 2048, 16384),
+    # PI0.5 Thor prefill tower, full real row envelope.
+    "pi05_prefill_qkv": (712, 2048, 2560),
+    "pi05_prefill_o": (970, 2048, 2048),
+    "pi05_prefill_gate_up": (768, 2048, 32768),
+    "pi05_prefill_down": (768, 16384, 2048),
+}
+MODES = {
+    "smoke": ["decode_m1_k512_n512", "small_m8_k1024_n2048"],
+    "headline": [
+        "decode_m1_k4096_n2048",
+        "decode_m1_k4096_n8192",
+        "small_m16_k4096_n4096",
+        "small_m32_k4096_n8192",
+    ],
+    "full": list(SHAPES.keys()),
+}
+@dataclass
+class Metrics:
+    shape: str
+    M: int
+    K: int
+    N: int
+    variant: int
+    tile: str
+    max_abs: float
+    mean_abs: float
+    p99_abs: float
+    cosine: float
+    dtype: str
+    tolerance: str
+    passed: bool
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+    @staticmethod
+    def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
+        return select_tile(m, n, k, variant)
+    def fp8_linear_bf16(self, x, w, alpha=1.0, out=None, variant=0):
+        if out is None:
+            out = torch.empty((x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16)
+        self._ops.fp8_linear_bf16(x, w, float(alpha), int(variant), out)
+        return out
+    def fp8_linear_residual_bf16(self, x, w, residual, alpha=1.0, variant=0):
+        self._ops.fp8_linear_residual_bf16(x, w, float(alpha), int(variant), residual)
+        return residual
+    def fp8_linear_bias_bf16(self, x, w, bias, alpha=1.0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16
+            )
+        self._ops.fp8_linear_bias_bf16(x, w, bias, float(alpha), out)
+        return out
+    def fp8_linear_bias_residual_bf16(
+        self, x, w, bias, residual, alpha=1.0
+    ):
+        self._ops.fp8_linear_bias_residual_bf16(
+            x, w, bias, float(alpha), residual
+        )
+        return residual
+    def fp8_linear_bias_gelu_bf16(self, x, w, bias, alpha=1.0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16
+            )
+        self._ops.fp8_linear_bias_gelu_bf16(x, w, bias, float(alpha), out)
+        return out
+    def fp8_blockwise_linear_bf16(
+        self, x, w, input_scale, weight_scale, out=None
+    ):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], w.shape[0]),
+                device=x.device,
+                dtype=torch.bfloat16,
+            )
+        self._ops.fp8_blockwise_linear_bf16(
+            x, w, input_scale, weight_scale, out
+        )
+        return out
+    def fp8_blockwise_swiglu_quantize_fp8(
+        self, x, gate_up_weight, input_scale, gate_up_weight_scale,
+        output=None, output_scale=None,
+    ):
+        n = gate_up_weight.shape[0] // 2
+        if output is None:
+            output = torch.empty(
+                (x.shape[0], n), device=x.device, dtype=torch.float8_e4m3fn
+            )
+        if output_scale is None:
+            output_scale = torch.empty(
+                (x.shape[0], n // 128), device=x.device, dtype=torch.float32
+            )
+        self._ops.fp8_blockwise_swiglu_quantize_fp8(
+            x, gate_up_weight, input_scale, gate_up_weight_scale,
+            output, output_scale,
+        )
+        return output, output_scale
+def _current_arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return "12.0a" if (major, minor) == (12, 0) else f"{major}.{minor}"
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    namespace = "fp8_gemm_source_test"
+    cutlass_include = Path(
+        os.environ.get(
+            "CUTLASS_INCLUDE",
+            str(
+                ROOT.parent
+                / "flashrt_pr31_review"
+                / "third_party"
+                / "cutlass"
+                / "include"
+            ),
+        )
+    )
+    if not (cutlass_include / "cutlass" / "cutlass.h").is_file():
+        raise RuntimeError(
+            "CUTLASS 4 include path is required; set CUTLASS_INCLUDE"
+        )
+    capability = torch.cuda.get_device_capability(0)
+    if capability == (8, 9):
+        cuda_sources = [
+            str(PACKAGE / "csrc" / "fp8_block128_gemm_mma_sm89.cu"),
+            str(PACKAGE / "csrc" / "fp8_gemv_m1_sm89.cu"),
+        ]
+        source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM89_ONLY"
+    elif capability == (11, 0):
+        cuda_sources = [
+            str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu"),
+            str(PACKAGE / "csrc" / "cublaslt_fp8_bias_sm110.cu"),
+        ]
+        source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY"
+    else:
+        cuda_sources = [
+            str(PACKAGE / "csrc" / "fp8_gemv_m1_sm120.cu"),
+            str(PACKAGE / "csrc" / "fp8_smallM_handtuned_sm120.cu"),
+            str(PACKAGE / "csrc" / "fp8_smallM_handtuned_ldmatrix_sm120.cu"),
+            str(PACKAGE / "csrc" / "cutlass_sm120_block128_fp8_gemm.cu"),
+        ]
+        source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM120_ONLY"
+    load(
+        name=namespace,
+        sources=[str(PACKAGE / "torch-ext" / "torch_binding.cpp"), *cuda_sources],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"),
+            str(REGISTRATION_INCLUDE),
+            str(cutlass_include),
+            str(cutlass_include.parent / "tools" / "util" / "include"),
+        ],
+        extra_cflags=["-O3", "-DNDEBUG", "-DCUDA_KERNEL", source_define],
+        extra_cuda_cflags=[
+            "-O3", "-DNDEBUG", "--expt-relaxed-constexpr", "--use_fast_math",
+            "-U__CUDA_NO_HALF_OPERATORS__",
+            "-U__CUDA_NO_HALF_CONVERSIONS__",
+            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+            "-U__CUDA_NO_HALF2_OPERATORS__",
+            "-DCUDA_KERNEL", source_define
+        ],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("fp8_gemm")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def select_tile(m: int, n: int, k: int, variant: int = 0) -> str:
+    if torch.cuda.get_device_capability(0) == (11, 0):
+        forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"}
+        if variant not in {0, *forced}:
+            raise RuntimeError("SM110 variant must be in [0, 3]")
+        if variant:
+            return forced[variant]
+        if m >= 512 and k == 2048 and 2048 <= n <= 2560:
+            return "sm110_sq_bf16"
+        if m >= 512 and n >= 16 * k:
+            return "sm110_t1_bf16"
+        if m >= 512 and k >= 4 * n:
+            return "sm110_wide_bf16"
+        if n >= 8 * k:
+            return "sm110_wide_bf16"
+        if m >= 128 and k >= 4 * n:
+            return "sm110_sq_bf16"
+        if n == k and m >= 512:
+            return "sm110_sq_bf16" if k <= 1024 else "sm110_wide_bf16"
+        if n == k and m >= 128:
+            return "sm110_wide_bf16"
+        return "sm110_t1_bf16"
+    if m == 1:
+        if variant == 4:
+            return "gemv_fp8_m1_w4"
+        if variant == 8:
+            return "gemv_fp8_m1_w8"
+        if variant == 16:
+            return "gemv_fp8_m1_w16"
+        if n <= 2048:
+            return "gemv_fp8_m1_w4"
+        if n <= 8192:
+            return "gemv_fp8_m1_w8"
+        return "gemv_fp8_m1_w16"
+    if m <= 16:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
+        if n % 256 == 0:
+            return "ld_fp8_gemm_16x256x128_w8"
+        if n % 192 == 0:
+            return "ld_fp8_gemm_16x192x128_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_16x128x128_w4"
+        return "ld_fp8_gemm_16x64x128_w4"
+    if m <= 32:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_32x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_32x64x256_w4"
+        if n % 192 == 0:
+            return "ld_fp8_gemm_32x192x128_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_32x128x128_w4"
+        return "ld_fp8_gemm_32x64x128_w4"
+    if m <= 64:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_64x128x128_w4"
+        return "ld_fp8_gemm_64x64x128_w4"
+    if m <= 64:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_64x128x128_w4"
+        return "ld_fp8_gemm_64x64x128_w4"
+    raise RuntimeError("unsupported M")
+def make_inputs(m: int, k: int, n: int, seed: int):
+    gen = torch.Generator(device="cuda")
+    gen.manual_seed(seed)
+    x_bf16 = (torch.randn((m, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16)
+    w_bf16 = (torch.randn((n, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16)
+    x = x_bf16.to(torch.float8_e4m3fn)
+    w = w_bf16.to(torch.float8_e4m3fn)
+    return x, w
+def reference(x: torch.Tensor, w: torch.Tensor, alpha: float) -> torch.Tensor:
+    return ((x.float() @ w.float().T) * float(alpha)).to(torch.bfloat16)
+def compare(got: torch.Tensor, expected: torch.Tensor) -> tuple[float, float, float, float]:
+    diff = (got.float() - expected.float()).abs().flatten()
+    max_abs = float(diff.max().item())
+    mean_abs = float(diff.mean().item())
+    p99_rank = max(1, min(diff.numel(), math.ceil(0.99 * diff.numel())))
+    p99_abs = float(diff.kthvalue(p99_rank).values.item())
+    cos = float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item())
+    return max_abs, mean_abs, p99_abs, cos
+def check_threshold(max_abs: float, mean_abs: float, p99_abs: float, cos: float) -> bool:
+    return max_abs <= 0.5 and mean_abs <= 0.02 and p99_abs <= 0.25 and cos >= 0.999
+def run_case(ops, name: str, shape: tuple[int, int, int], variant: int = 0) -> Metrics:
+    m, k, n = shape
+    x, w = make_inputs(m, k, n, seed=1000 + m + k + n + variant)
+    alpha = 1.0
+    expected = reference(x, w, alpha)
+    got = ops.fp8_linear_bf16(x, w, alpha=alpha, variant=variant)
+    torch.cuda.synchronize()
+    max_abs, mean_abs, p99_abs, cos = compare(got, expected)
+    tile = ops.select_fp8_linear_tile(m, n, k, variant)
+    passed = check_threshold(max_abs, mean_abs, p99_abs, cos)
+    return Metrics(
+        shape=name,
+        M=m,
+        K=k,
+        N=n,
+        variant=variant,
+        tile=tile,
+        max_abs=max_abs,
+        mean_abs=mean_abs,
+        p99_abs=p99_abs,
+        cosine=cos,
+        dtype=str(got.dtype),
+        tolerance="max_abs<=0.5 mean_abs<=0.02 p99_abs<=0.25 cosine>=0.999",
+        passed=passed,
+    )
+def run_residual_case(ops) -> Metrics:
+    m, k, n = (1, 4096, 2048)
+    x, w = make_inputs(m, k, n, seed=2026)
+    residual = torch.randn((1, n), device="cuda", dtype=torch.bfloat16) * 0.1
+    expected = (residual.float() + reference(x, w, 1.0).float()).to(torch.bfloat16)
+    got = residual.clone()
+    variant = 0 if torch.cuda.get_device_capability(0) == (11, 0) else 8
+    ops.fp8_linear_residual_bf16(x, w, got, alpha=1.0, variant=variant)
+    torch.cuda.synchronize()
+    max_abs, mean_abs, p99_abs, cos = compare(got, expected)
+    passed = check_threshold(max_abs, mean_abs, p99_abs, cos)
+    return Metrics(
+        shape="decode_residual_m1_k4096_n2048",
+        M=m,
+        K=k,
+        N=n,
+        variant=variant,
+        tile=(
+            "sm110_t1_bf16_residual"
+            if torch.cuda.get_device_capability(0) == (11, 0)
+            else "gemv_fp8_m1_resadd_w8"
+        ),
+        max_abs=max_abs,
+        mean_abs=mean_abs,
+        p99_abs=p99_abs,
+        cosine=cos,
+        dtype=str(got.dtype),
+        tolerance="max_abs<=0.5 mean_abs<=0.02 p99_abs<=0.25 cosine>=0.999",
+        passed=passed,
+    )
+def run_bias_cases(ops) -> int:
+    count = 0
+    shapes = [
+        (512, 1152, 4304),
+        (768, 4304, 1152),
+        (768, 1152, 3456),
+    ]
+    for m, k, n in shapes:
+        x, w = make_inputs(m, k, n, seed=7000 + m + k + n)
+        bias = (torch.randn((n,), device="cuda") * 0.1).to(torch.bfloat16)
+        alpha = 0.75
+        base = (x.float() @ w.float().T) * alpha
+        got = ops.fp8_linear_bias_bf16(x, w, bias, alpha=alpha)
+        expected = (base + bias.float()).to(torch.bfloat16)
+        maximum, mean, p99, cosine = compare(got, expected)
+        assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, (
+            "bias", m, k, n, maximum, mean, p99, cosine
+        )
+        residual = (torch.randn((m, n), device="cuda") * 0.1).to(
+            torch.bfloat16
+        )
+        residual_before = residual.clone()
+        got_residual = ops.fp8_linear_bias_residual_bf16(
+            x, w, bias, residual, alpha=alpha
+        )
+        expected_residual = (
+            residual_before.float() + base + bias.float()
+        ).to(torch.bfloat16)
+        maximum, mean, p99, cosine = compare(got_residual, expected_residual)
+        assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, (
+            "bias_residual", m, k, n, maximum, mean, p99, cosine
+        )
+        got_gelu = ops.fp8_linear_bias_gelu_bf16(x, w, bias, alpha=alpha)
+        expected_gelu = torch.nn.functional.gelu(
+            base + bias.float(), approximate="tanh"
+        ).to(torch.bfloat16)
+        maximum, mean, p99, cosine = compare(got_gelu, expected_gelu)
+        assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, (
+            "bias_gelu", m, k, n, maximum, mean, p99, cosine
+        )
+        count += 3
+    m, k, n = (512, 1152, 4304)
+    x, w = make_inputs(m, k, n, seed=8801)
+    bias = torch.randn((n,), device="cuda", dtype=torch.bfloat16)
+    def invoke(input, weight, bias):
+        return ops.fp8_linear_bias_bf16(input, weight, bias)
+    eager = invoke(x, w, bias)
+    compiled = torch.compile(invoke, fullgraph=True)(x, w, bias)
+    torch.testing.assert_close(compiled, eager, rtol=0.0, atol=0.0)
+    graph_out = torch.empty_like(eager)
+    ops.fp8_linear_bias_bf16(x, w, bias, out=graph_out)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        ops.fp8_linear_bias_bf16(x, w, bias, out=graph_out)
+    graph.replay()
+    torch.testing.assert_close(graph_out, eager, rtol=0.0, atol=0.0)
+    return count + 2
+def run_blockwise_case(
+    ops, name: str, shape: tuple[int, int, int]
+) -> Metrics:
+    m, k, n = shape
+    gen = torch.Generator(device="cuda").manual_seed(5000 + m + k + n)
+    x = (torch.randn((m, k), device="cuda", generator=gen) * 0.4).to(
+        torch.float8_e4m3fn
+    )
+    w = (torch.randn((n, k), device="cuda", generator=gen) * 0.4).to(
+        torch.float8_e4m3fn
+    )
+    input_scale = (
+        0.005
+        + 0.02
+        * torch.rand((m, k // 128), device="cuda", generator=gen)
+    ).float().contiguous()
+    weight_scale = (
+        0.005
+        + 0.02
+        * torch.rand((n // 128, k // 128), device="cuda", generator=gen)
+    ).float().contiguous()
+    expanded_input_scale = input_scale.repeat_interleave(128, dim=1)
+    expanded_weight_scale = weight_scale.repeat_interleave(
+        128, dim=0
+    ).repeat_interleave(128, dim=1)
+    expected = (
+        (x.float() * expanded_input_scale)
+        @ (w.float() * expanded_weight_scale).T
+    ).to(torch.bfloat16)
+    got = ops.fp8_blockwise_linear_bf16(
+        x, w, input_scale, weight_scale
+    )
+    torch.cuda.synchronize()
+    max_abs, mean_abs, p99_abs, cos = compare(got, expected)
+    passed = (
+        max_abs <= 0.0625
+        and mean_abs <= 0.003
+        and p99_abs <= 0.015625
+        and cos >= 0.9999
+    )
+    return Metrics(
+        shape=name,
+        M=m,
+        K=k,
+        N=n,
+        variant=0,
+        tile=(
+            "mma_sm89_block128"
+            if torch.cuda.get_device_capability(0) == (8, 9)
+            else "cutlass_sm120_block128"
+        ),
+        max_abs=max_abs,
+        mean_abs=mean_abs,
+        p99_abs=p99_abs,
+        cosine=cos,
+        dtype=str(got.dtype),
+        tolerance=(
+            "max_abs<=0.0625 mean_abs<=0.003 "
+            "p99_abs<=0.015625 cosine>=0.9999"
+        ),
+        passed=passed,
+    )
+def run_blockwise_compile_case(ops) -> None:
+    m, k, n = (51, 1536, 1536)
+    gen = torch.Generator(device="cuda").manual_seed(9153)
+    x = (torch.randn((m, k), device="cuda", generator=gen) * 0.4).to(
+        torch.float8_e4m3fn
+    )
+    w = (torch.randn((n, k), device="cuda", generator=gen) * 0.4).to(
+        torch.float8_e4m3fn
+    )
+    input_scale = torch.rand(
+        (m, k // 128), device="cuda", generator=gen, dtype=torch.float32
+    ).mul_(0.02).add_(0.005)
+    weight_scale = torch.rand(
+        (n // 128, k // 128),
+        device="cuda",
+        generator=gen,
+        dtype=torch.float32,
+    ).mul_(0.02).add_(0.005)
+    def invoke(input, weight, input_scale, weight_scale):
+        return ops.fp8_blockwise_linear_bf16(
+            input, weight, input_scale, weight_scale
+        )
+    eager = invoke(x, w, input_scale, weight_scale)
+    compiled = torch.compile(invoke, fullgraph=True)(
+        x, w, input_scale, weight_scale
+    )
+    torch.testing.assert_close(compiled, eager, rtol=0.0, atol=0.0)
+def run_sm89_swiglu_case(ops, m: int, n: int, k: int) -> None:
+    gen = torch.Generator(device="cuda").manual_seed(8900 + m + n + k)
+    x = (torch.randn((m, k), device="cuda", generator=gen) * 0.3).to(
+        torch.float8_e4m3fn
+    )
+    weight = (
+        torch.randn((2 * n, k), device="cuda", generator=gen) * 0.3
+    ).to(torch.float8_e4m3fn)
+    input_scale = torch.rand(
+        (m, k // 128), device="cuda", generator=gen
+    ).mul_(0.02).add_(0.005)
+    weight_scale = torch.rand(
+        (2 * n // 128, k // 128), device="cuda", generator=gen
+    ).mul_(0.02).add_(0.005)
+    output, output_scale = ops.fp8_blockwise_swiglu_quantize_fp8(
+        x, weight, input_scale, weight_scale
+    )
+    expanded_x_scale = input_scale.repeat_interleave(128, dim=1)
+    expanded_w_scale = weight_scale.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1)
+    x_f32 = x.float() * expanded_x_scale
+    weight_f32 = weight.float() * expanded_w_scale
+    gate, up = (x_f32 @ weight_f32.t()).split(n, dim=1)
+    expected = (
+        torch.nn.functional.silu(gate).bfloat16() * up.bfloat16()
+    ).bfloat16()
+    actual = (
+        output.float() * output_scale.repeat_interleave(128, dim=1)
+    ).bfloat16()
+    maximum, mean, p99, cosine = compare(actual, expected)
+    assert output.dtype == torch.float8_e4m3fn
+    assert output_scale.dtype == torch.float32
+    assert torch.isfinite(output_scale).all() and (output_scale > 0).all()
+    assert cosine >= 0.999 and mean <= 0.01 and p99 <= 0.05, (
+        m, n, k, maximum, mean, p99, cosine
+    )
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=sorted(MODES), default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    capability = torch.cuda.get_device_capability(0)
+    if capability not in {(8, 9), (11, 0), (12, 0)}:
+        raise SystemExit(
+            "fp8-gemm source tests require SM89, SM110, or SM120; "
+            f"got SM{capability[0]}{capability[1]}"
+        )
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    rows = []
+    if capability in {(11, 0), (12, 0)}:
+        rows.extend(run_case(ops, name, SHAPES[name]) for name in MODES[args.mode])
+        rows.append(run_residual_case(ops))
+    if capability == (11, 0) and args.mode == "full":
+        rows.extend(
+            run_case(ops, name, shape) for name, shape in SM110_SHAPES.items()
+        )
+        rows.extend(
+            run_case(
+                ops,
+                f"sm110_forced_variant_{variant}",
+                SM110_SHAPES["pi05_action_gate_up"],
+                variant,
+            )
+            for variant in (1, 2, 3)
+        )
+        bias_count = run_bias_cases(ops)
+    else:
+        bias_count = 0
+    if capability in {(8, 9), (12, 0)}:
+        blockwise_shapes = [
+            ("blockwise_decode", (1, 1024, 1024)),
+            ("blockwise_action", (51, 1536, 1536)),
+        ]
+        if args.mode == "full":
+            blockwise_shapes += [
+                ("blockwise_groot", (277, 2048, 2048)),
+                ("blockwise_vision", (1024, 1152, 1152)),
+                ("blockwise_video", (2520, 3072, 3072)),
+                ("blockwise_qwen_mlp", (128, 4096, 12288)),
+            ]
+        rows.extend(
+            run_blockwise_case(ops, name, shape)
+            for name, shape in blockwise_shapes
+        )
+        run_blockwise_compile_case(ops)
+    if capability == (8, 9):
+        for m, n, k in [
+            (1, 128, 128), (16, 512, 1024), (31, 1536, 1536),
+            (32, 2048, 4096), (51, 4096, 4096), (128, 4096, 4096),
+            (256, 4096, 4096),
+        ]:
+            run_sm89_swiglu_case(ops, m, n, k)
+        try:
+            x = torch.zeros((257, 128), device="cuda", dtype=torch.float8_e4m3fn)
+            w = torch.zeros((256, 128), device="cuda", dtype=torch.float8_e4m3fn)
+            xs = torch.ones((257, 1), device="cuda", dtype=torch.float32)
+            ws = torch.ones((2, 1), device="cuda", dtype=torch.float32)
+            ops.fp8_blockwise_swiglu_quantize_fp8(x, w, xs, ws)
+        except RuntimeError as error:
+            assert "M <= 256" in str(error)
+        else:
+            raise AssertionError("M=257 must be rejected")
+    failed = [row for row in rows if not row.passed]
+    payload = {
+        "passed": len(rows) - len(failed) + bias_count,
+        "failed": len(failed),
+        "rows": [asdict(row) for row in rows],
+        "bias_checks": bias_count,
+    }
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    if args.json_out:
+        output_path = Path(args.json_out)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_text(
+            json.dumps(payload, indent=2, sort_keys=True) + "\n"
+        )
+    if failed:
+        raise SystemExit(1)
+if __name__ == "__main__":
+    main()

torch-ext/README.md ADDED Viewed

	@@ -0,0 +1,6 @@

+# torch-ext
+Python package: `fp8_gemm`
+The package registers Torch custom ops and fake implementations for
+`torch.compile` tracing.

torch-ext/fp8_gemm/__init__.py ADDED Viewed

	@@ -0,0 +1,345 @@

+"""FlashRT FP8 GEMM kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bf16"))
+def _fp8_linear_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out must have shape (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_residual_bf16"))
+def _fp8_linear_residual_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    alpha: float,
+    variant: int,
+    residual: torch.Tensor,
+) -> None:
+    if input.shape[0] != 1:
+        raise RuntimeError("residual path supports only M=1")
+    if residual.shape != (1, weight.shape[0]):
+        raise RuntimeError("residual must have shape (1, weight.shape[0])")
+    return None
+def _check_bias_linear_shapes(input, weight, bias, out) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input and weight K dimensions must match")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias must have shape (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out must have shape (input.shape[0], weight.shape[0])")
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(input, weight, bias, alpha: float, out) -> None:
+    _check_bias_linear_shapes(input, weight, bias, out)
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_residual_bf16"))
+def _fp8_linear_bias_residual_bf16_fake(
+    input, weight, bias, alpha: float, residual
+) -> None:
+    _check_bias_linear_shapes(input, weight, bias, residual)
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_bf16"))
+def _fp8_linear_bias_gelu_bf16_fake(input, weight, bias, alpha: float, out) -> None:
+    _check_bias_linear_shapes(input, weight, bias, out)
+@torch.library.register_fake(add_op_namespace_prefix("fp8_blockwise_linear_bf16"))
+def _fp8_blockwise_linear_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    m, k = input.shape
+    n = weight.shape[0]
+    if weight.shape[1] != k or n % 128 or k % 128:
+        raise RuntimeError("weight shape is invalid or N/K are not divisible by 128")
+    if input_scale.shape != (m, k // 128):
+        raise RuntimeError("input_scale must have shape (M, K / 128)")
+    if weight_scale.shape != (n // 128, k // 128):
+        raise RuntimeError("weight_scale must have shape (N / 128, K / 128)")
+    if out.shape != (m, n):
+        raise RuntimeError("out must have shape (M, N)")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_blockwise_swiglu_quantize_fp8")
+)
+def _fp8_blockwise_swiglu_quantize_fp8_fake(
+    input: torch.Tensor,
+    gate_up_weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    gate_up_weight_scale: torch.Tensor,
+    output: torch.Tensor,
+    output_scale: torch.Tensor,
+) -> None:
+    m, k = input.shape
+    if gate_up_weight.dim() != 2 or gate_up_weight.shape[0] % 2:
+        raise RuntimeError("gate_up_weight must have shape (2*N, K)")
+    n = gate_up_weight.shape[0] // 2
+    if gate_up_weight.shape[1] != k or n % 128 or k % 128:
+        raise RuntimeError("gate_up_weight shape is invalid or N/K are not divisible by 128")
+    if input_scale.shape != (m, k // 128):
+        raise RuntimeError("input_scale must have shape (M, K/128)")
+    if gate_up_weight_scale.shape != (2 * n // 128, k // 128):
+        raise RuntimeError("gate_up_weight_scale must have shape (2*N/128, K/128)")
+    if output.shape != (m, n) or output_scale.shape != (m, n // 128):
+        raise RuntimeError("output buffers have invalid shapes")
+    return None
+def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
+    """Return the FlashRT tile selected by the public dispatcher."""
+    m = int(m)
+    n = int(n)
+    k = int(k)
+    variant = int(variant)
+    if m <= 0 or n <= 0 or k <= 0:
+        raise RuntimeError("m, n, and k must be positive")
+    if k % 16 != 0:
+        raise RuntimeError("k must be divisible by 16")
+    capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
+    if capability == (11, 0):
+        forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"}
+        if variant not in {0, *forced}:
+            raise RuntimeError("SM110 variant must be 0 (auto), 1 (Sq), 2 (T1), or 3 (Wide)")
+        if n % 16 or k % 16:
+            raise RuntimeError("SM110 requires n and k divisible by 16")
+        if variant:
+            return forced[variant]
+        if m >= 512 and k == 2048 and 2048 <= n <= 2560:
+            return "sm110_sq_bf16"
+        if m >= 512 and n >= 16 * k:
+            return "sm110_t1_bf16"
+        if m >= 512 and k >= 4 * n:
+            return "sm110_wide_bf16"
+        if n >= 8 * k:
+            return "sm110_wide_bf16"
+        if m >= 128 and k >= 4 * n:
+            return "sm110_sq_bf16"
+        if n == k and m >= 512:
+            return "sm110_sq_bf16" if k <= 1024 else "sm110_wide_bf16"
+        if n == k and m >= 128:
+            return "sm110_wide_bf16"
+        return "sm110_t1_bf16"
+    if m == 1:
+        if k % 32:
+            raise RuntimeError("SM120 requires k divisible by 32")
+        if variant == 4:
+            return "gemv_fp8_m1_w4"
+        if variant == 8:
+            return "gemv_fp8_m1_w8"
+        if variant == 16:
+            return "gemv_fp8_m1_w16"
+        if variant != 0:
+            raise RuntimeError("M=1 variant must be 0, 4, 8, or 16")
+        if n <= 2048:
+            return "gemv_fp8_m1_w4"
+        if n <= 8192:
+            return "gemv_fp8_m1_w8"
+        return "gemv_fp8_m1_w16"
+    if variant != 0:
+        raise RuntimeError("small-M dispatcher currently supports variant=0 only")
+    if k % 32:
+        raise RuntimeError("SM120 requires k divisible by 32")
+    if m <= 16:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
+        if n % 256 == 0:
+            return "ld_fp8_gemm_16x256x128_w8"
+        if n % 192 == 0:
+            return "ld_fp8_gemm_16x192x128_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_16x128x128_w4"
+        return "ld_fp8_gemm_16x64x128_w4"
+    if m <= 32:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_32x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_32x64x256_w4"
+        if n % 192 == 0:
+            return "ld_fp8_gemm_32x192x128_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_32x128x128_w4"
+        return "ld_fp8_gemm_32x64x128_w4"
+    if m <= 64:
+        if k % 256 == 0:
+            return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4"
+        if n % 128 == 0:
+            return "ld_fp8_gemm_64x128x128_w4"
+        return "ld_fp8_gemm_64x64x128_w4"
+    raise RuntimeError("only M=1 decode or 2 <= M <= 64 small-M rows are supported")
+def fp8_linear_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 0,
+) -> torch.Tensor:
+    """Compute ``(input @ weight.T) * alpha`` with BF16 output.
+    ``input`` and ``weight`` must be FP8 E4M3 CUDA tensors with shapes
+    ``(M, K)`` and ``(N, K)``. ``alpha`` is a host float, normally the product
+    of static per-tensor input and weight scales. SM110 uses the production
+    CUTLASS Sq/T1/Wide dispatcher over full model row counts; SM120 uses the
+    hand-tuned M<=64 path.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bf16(input, weight, float(alpha), int(variant), out)
+    return out
+def fp8_linear_residual_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = 0,
+) -> torch.Tensor:
+    """In-place ``residual += (input @ weight.T) * alpha`` for M=1 decode."""
+    ops.fp8_linear_residual_bf16(input, weight, float(alpha), int(variant), residual)
+    return residual
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 FP8 linear with fused BF16 bias and BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(input, weight, bias, float(alpha), out)
+    return out
+def fp8_linear_bias_residual_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+) -> torch.Tensor:
+    """SM110 fused ``residual += alpha * input @ weight.T + bias``."""
+    ops.fp8_linear_bias_residual_bf16(
+        input, weight, bias, float(alpha), residual
+    )
+    return residual
+def fp8_linear_bias_gelu_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 FP8 linear with fused BF16 bias and GELU epilogue."""
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_gelu_bf16(input, weight, bias, float(alpha), out)
+    return out
+def fp8_blockwise_linear_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Block-128 scaled FP8 linear with BF16 output on SM89/SM120."""
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_blockwise_linear_bf16(
+        input, weight, input_scale, weight_scale, out
+    )
+    return out
+def fp8_blockwise_swiglu_quantize_fp8(
+    input: torch.Tensor,
+    gate_up_weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    gate_up_weight_scale: torch.Tensor,
+    *,
+    output: torch.Tensor | None = None,
+    output_scale: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM89 block-128 FP8 gate/up GEMM + SiLU + FP8 requant producer."""
+    n = gate_up_weight.shape[0] // 2
+    if output is None:
+        output = torch.empty(
+            (input.shape[0], n), device=input.device, dtype=torch.float8_e4m3fn
+        )
+    if output_scale is None:
+        output_scale = torch.empty(
+            (input.shape[0], n // 128), device=input.device, dtype=torch.float32
+        )
+    ops.fp8_blockwise_swiglu_quantize_fp8(
+        input, gate_up_weight, input_scale, gate_up_weight_scale,
+        output, output_scale
+    )
+    return output, output_scale
+__all__ = [
+    "fp8_linear_bf16",
+    "fp8_linear_residual_bf16",
+    "fp8_linear_bias_bf16",
+    "fp8_linear_bias_residual_bf16",
+    "fp8_linear_bias_gelu_bf16",
+    "fp8_blockwise_linear_bf16",
+    "fp8_blockwise_swiglu_quantize_fp8",
+    "select_fp8_linear_tile",
+]

torch-ext/torch_binding.cpp ADDED Viewed

	@@ -0,0 +1,577 @@

+// SPDX-License-Identifier: Apache-2.0
+#include <torch/all.h>
+#include <torch/library.h>
+#include <limits>
+#include <sstream>
+#include <string>
+#if defined(CUDA_KERNEL)
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#endif
+#if !defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) && \
+    !defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+#include "fp8_gemv_m1_sm120.cuh"
+#include "fp8_smallM_handtuned_ldmatrix_sm120.cuh"
+#include "fp8_smallM_handtuned_sm120.cuh"
+#include "cutlass_sm120_block128_fp8_gemm.cuh"
+#endif
+#if !defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY) && \
+    !defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+#include "fp8_block128_gemm_mma_sm89.cuh"
+#include "fp8_gemv_m1_sm89.cuh"
+#endif
+#if !defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) && \
+    !defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY)
+#include "cutlass_sm110_fp8_gemm.cuh"
+#endif
+#include "cublaslt_fp8_bias_sm110.cuh"
+#include "registration.h"
+#include "torch_binding.h"
+namespace {
+using KernelFn = int (*)(const void*, const void*, void*, int, int, int, float, cudaStream_t);
+using Sm110KernelFn = int (*)(void*, void*, void*, int, int, int, float, float,
+                              cudaStream_t);
+void check_cuda_contiguous(torch::Tensor const& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+void check_fp8_matrix(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == c10::ScalarType::Float8_e4m3fn,
+              name, " must have dtype torch.float8_e4m3fn");
+  TORCH_CHECK(tensor.dim() == 2, name, " must have shape (rows, cols)");
+  TORCH_CHECK(tensor.size(0) > 0 && tensor.size(1) > 0,
+              name, " dimensions must be positive");
+}
+void check_bf16_matrix(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16,
+              name, " must have dtype torch.bfloat16");
+  TORCH_CHECK(tensor.dim() == 2, name, " must have shape (rows, cols)");
+}
+void check_fp32_matrix(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat32,
+              name, " must have dtype torch.float32");
+  TORCH_CHECK(tensor.dim() == 2, name, " must be rank 2");
+}
+void check_bf16_vector(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16,
+              name, " must have dtype torch.bfloat16");
+  TORCH_CHECK(tensor.dim() == 1, name, " must be rank 1");
+}
+int checked_positive_int(int64_t value, const char* name) {
+  TORCH_CHECK(value > 0 && value <= std::numeric_limits<int>::max(),
+              name, " must fit in positive int");
+  return static_cast<int>(value);
+}
+void check_common(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& out) {
+  check_fp8_matrix(input, "input");
+  check_fp8_matrix(weight, "weight");
+  check_bf16_matrix(out, "out");
+  TORCH_CHECK(input.get_device() == weight.get_device(),
+              "input and weight must be on the same CUDA device");
+  TORCH_CHECK(input.get_device() == out.get_device(),
+              "input and out must be on the same CUDA device");
+  TORCH_CHECK(input.size(1) == weight.size(1),
+              "input.shape[1] must equal weight.shape[1]");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({input.size(0), weight.size(0)}),
+              "out must have shape (input.shape[0], weight.shape[0])");
+  TORCH_CHECK(input.size(1) % 16 == 0,
+              "K must be divisible by 16 for FP8 tensor-core kernels");
+}
+std::string tile_name_for_shape(int M, int N, int K, int variant) {
+  if (M == 1) {
+    if (variant == 4) return "gemv_fp8_m1_w4";
+    if (variant == 8) return "gemv_fp8_m1_w8";
+    if (variant == 16) return "gemv_fp8_m1_w16";
+    TORCH_CHECK(variant == 0, "M=1 variant must be 0, 4, 8, or 16");
+    if (N <= 2048) return "gemv_fp8_m1_w4";
+    if (N <= 8192) return "gemv_fp8_m1_w8";
+    return "gemv_fp8_m1_w16";
+  }
+  TORCH_CHECK(variant == 0,
+              "small-M public dispatcher currently supports variant=0 only; "
+              "use benchmark scripts for tile sweeps before promoting a forced variant");
+  if (M <= 16) {
+    if (K % 256 == 0) {
+      if (N % 128 == 0) return "ld_fp8_gemm_16x128x256_w4";
+      return "ld_fp8_gemm_16x64x256_w4";
+    }
+    if (N % 256 == 0) return "ld_fp8_gemm_16x256x128_w8";
+    if (N % 192 == 0) return "ld_fp8_gemm_16x192x128_w4";
+    if (N % 128 == 0) return "ld_fp8_gemm_16x128x128_w4";
+    return "ld_fp8_gemm_16x64x128_w4";
+  }
+  if (M <= 32) {
+    if (K % 256 == 0) {
+      if (N % 128 == 0) return "ld_fp8_gemm_32x128x256_w4";
+      return "ld_fp8_gemm_32x64x256_w4";
+    }
+    if (N % 192 == 0) return "ld_fp8_gemm_32x192x128_w4";
+    if (N % 128 == 0) return "ld_fp8_gemm_32x128x128_w4";
+    return "ld_fp8_gemm_32x64x128_w4";
+  }
+  if (M <= 64) {
+    if (K % 256 == 0) {
+      if (N % 128 == 0) return "ld_fp8_gemm_64x128x256_w4";
+      return "ld_fp8_gemm_64x64x256_w4";
+    }
+    if (N % 128 == 0) return "ld_fp8_gemm_64x128x128_w4";
+    return "ld_fp8_gemm_64x64x128_w4";
+  }
+  TORCH_CHECK(false, "M > 64 is not exposed in fp8-gemm v1; pending tile tuning");
+  TORCH_CHECK(false, "unsupported M");
+}
+KernelFn kernel_for_tile(std::string const& tile, bool residual) {
+#if defined(CUDA_KERNEL) && !defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) && \
+    !defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+  namespace gemv = flash_rt::gemm::gemv_m1;
+  namespace hand = flash_rt::gemm::smallM_hand;
+  namespace ld = flash_rt::gemm::smallM_ld;
+  if (tile == "gemv_fp8_m1_w4") return residual ? gemv::gemv_fp8_m1_resadd_w4 : gemv::gemv_fp8_m1_w4;
+  if (tile == "gemv_fp8_m1_w8") return residual ? gemv::gemv_fp8_m1_resadd_w8 : gemv::gemv_fp8_m1_w8;
+  if (tile == "gemv_fp8_m1_w16") {
+    TORCH_CHECK(!residual, "residual path supports only GEMV w4/w8 variants");
+    return gemv::gemv_fp8_m1_w16;
+  }
+  TORCH_CHECK(!residual, "residual path supports M=1 GEMV only");
+  if (tile == "ld_fp8_gemm_16x64x128_w4") return ld::ld_fp8_gemm_16x64x128_w4;
+  if (tile == "ld_fp8_gemm_16x128x128_w4") return ld::ld_fp8_gemm_16x128x128_w4;
+  if (tile == "ld_fp8_gemm_16x256x128_w8") return ld::ld_fp8_gemm_16x256x128_w8;
+  if (tile == "ld_fp8_gemm_16x192x128_w4") return ld::ld_fp8_gemm_16x192x128_w4;
+  if (tile == "ld_fp8_gemm_16x64x256_w4") return ld::ld_fp8_gemm_16x64x256_w4;
+  if (tile == "ld_fp8_gemm_16x128x256_w4") return ld::ld_fp8_gemm_16x128x256_w4;
+  if (tile == "ld_fp8_gemm_32x64x128_w4") return ld::ld_fp8_gemm_32x64x128_w4;
+  if (tile == "ld_fp8_gemm_32x128x128_w4") return ld::ld_fp8_gemm_32x128x128_w4;
+  if (tile == "ld_fp8_gemm_32x192x128_w4") return ld::ld_fp8_gemm_32x192x128_w4;
+  if (tile == "ld_fp8_gemm_32x64x256_w4") return ld::ld_fp8_gemm_32x64x256_w4;
+  if (tile == "ld_fp8_gemm_32x128x256_w4") return ld::ld_fp8_gemm_32x128x256_w4;
+  if (tile == "ld_fp8_gemm_64x64x128_w4") return ld::ld_fp8_gemm_64x64x128_w4;
+  if (tile == "ld_fp8_gemm_64x128x128_w4") return ld::ld_fp8_gemm_64x128x128_w4;
+  if (tile == "ld_fp8_gemm_64x64x256_w4") return ld::ld_fp8_gemm_64x64x256_w4;
+  if (tile == "ld_fp8_gemm_64x128x256_w4") return ld::ld_fp8_gemm_64x128x256_w4;
+#else
+  (void)tile;
+  (void)residual;
+#endif
+  TORCH_CHECK(false, "unsupported FP8 GEMM tile: ", tile);
+}
+const char* sm110_tile_name_for_shape(int M, int N, int K, int variant) {
+  if (variant == 1) return "sm110_sq_bf16";
+  if (variant == 2) return "sm110_t1_bf16";
+  if (variant == 3) return "sm110_wide_bf16";
+  // Thor sweep envelope (PI0.5/GROOT/Cosmos Edge/LingBot): Wide wins
+  // N>=8K expansions and larger square projections. Sq wins smaller square
+  // vision projections and larger-row K>=4N contractions; T1 wins the
+  // remaining projection/down paths.
+  // The forced variants remain available for diagnostic tile sweeps.
+  if (M >= 512 && K == 2048 && N >= 2048 && N <= 2560) {
+    return "sm110_sq_bf16";
+  }
+  if (M >= 512 && N >= 16 * K) {
+    return "sm110_t1_bf16";
+  }
+  if (M >= 512 && K >= 4 * N) {
+    return "sm110_wide_bf16";
+  }
+  if (N >= 8 * K) return "sm110_wide_bf16";
+  if (M >= 128 && K >= 4 * N) return "sm110_sq_bf16";
+  if (N == K && M >= 512) {
+    return K <= 1024 ? "sm110_sq_bf16" : "sm110_wide_bf16";
+  }
+  if (N == K && M >= 128) {
+    return "sm110_wide_bf16";
+  }
+  return "sm110_t1_bf16";
+}
+Sm110KernelFn sm110_kernel_for_shape(int M, int N, int K, int variant) {
+#if defined(CUDA_KERNEL) && !defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) && \
+    !defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY)
+  const char* tile = sm110_tile_name_for_shape(M, N, K, variant);
+  if (std::string(tile) == "sm110_wide_bf16") return &cutlass_fp8_wide_bf16out;
+  if (std::string(tile) == "sm110_t1_bf16") return &cutlass_fp8_t1_bf16out;
+  return &cutlass_fp8_sq_bf16out;
+#else
+  (void)M;
+  (void)N;
+  (void)K;
+  (void)variant;
+  TORCH_CHECK(false, "SM110 FP8 GEMM source is not present in this build");
+#endif
+}
+void launch(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    double alpha,
+    int64_t variant64,
+    torch::Tensor& out,
+    bool residual) {
+  check_common(input, weight, out);
+  const int M = checked_positive_int(input.size(0), "M");
+  const int K = checked_positive_int(input.size(1), "K");
+  const int N = checked_positive_int(weight.size(0), "N");
+  const int variant = static_cast<int>(variant64);
+  if (residual) {
+    TORCH_CHECK(M == 1, "fp8_linear_residual_bf16 supports only M=1");
+  }
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(input.device());
+  auto* props = at::cuda::getDeviceProperties(input.get_device());
+  TORCH_CHECK((props->major == 11 && props->minor == 0) ||
+                  (props->major == 12 && props->minor == 0),
+              "fp8_linear_bf16 requires SM110 or SM120; got SM",
+              props->major, props->minor);
+  auto stream = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
+  if (props->major == 11) {
+    TORCH_CHECK(variant >= 0 && variant <= 3,
+                "SM110 variant must be 0 (auto), 1 (Sq), 2 (T1), or 3 (Wide)");
+    TORCH_CHECK(N % 16 == 0 && K % 16 == 0,
+                "SM110 CUTLASS FP8 GEMM requires N and K divisible by 16");
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY)
+    TORCH_CHECK(false, "SM110 FP8 GEMM source is not present in this build");
+#else
+    Sm110KernelFn fn = sm110_kernel_for_shape(M, N, K, variant);
+    const int rc = fn(input.data_ptr(), weight.data_ptr(), out.data_ptr(),
+                      M, N, K, static_cast<float>(alpha),
+                      residual ? 1.0f : 0.0f, stream);
+    TORCH_CHECK(rc == 0, sm110_tile_name_for_shape(M, N, K, variant),
+                " failed with rc=", rc);
+#endif
+  } else {
+    TORCH_CHECK(K % 32 == 0,
+                "SM120 FP8 GEMM requires K divisible by 32");
+    TORCH_CHECK(M <= 64,
+                "SM120 per-tensor FP8 path supports only M <= 64; got M=", M);
+    if (residual) {
+      TORCH_CHECK(M == 1, "SM120 residual path supports only M=1");
+    }
+    const std::string tile = tile_name_for_shape(M, N, K, variant);
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+    TORCH_CHECK(false, "SM120 per-tensor FP8 source is not present in this build");
+#else
+    KernelFn fn = kernel_for_tile(tile, residual);
+    const int rc = fn(input.data_ptr(), weight.data_ptr(), out.data_ptr(),
+                      M, N, K, static_cast<float>(alpha), stream);
+    TORCH_CHECK(rc == 0, tile, " failed with rc=", rc);
+#endif
+  }
+#else
+  TORCH_CHECK(false, "fp8-gemm was not built with CUDA support");
+#endif
+}
+void launch_bias(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& out,
+    double beta,
+    FlashRtFp8BiasEpilogue epilogue,
+    const char* op_name) {
+  check_common(input, weight, out);
+  check_bf16_vector(bias, "bias");
+  TORCH_CHECK(bias.size(0) == weight.size(0),
+              "bias must have shape (weight.shape[0],)");
+  TORCH_CHECK(input.get_device() == bias.get_device(),
+              "input and bias must be on the same CUDA device");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(input.device());
+  auto* props = at::cuda::getDeviceProperties(input.get_device());
+  TORCH_CHECK(props->major == 11 && props->minor == 0,
+              op_name, " requires SM110; got SM", props->major, props->minor);
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY)
+  TORCH_CHECK(false, "SM110 FP8 bias GEMM source is not present in this build");
+#else
+  auto stream = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
+  const int M = checked_positive_int(input.size(0), "M");
+  const int N = checked_positive_int(weight.size(0), "N");
+  const int K = checked_positive_int(input.size(1), "K");
+  int rc;
+  if (M >= 512 && K >= 3 * N) {
+    rc = epilogue == FlashRtFp8BiasEpilogue::kBiasGelu
+             ? cutlass_fp8_wide_bias_gelu_bf16out(
+                   input.data_ptr(), weight.data_ptr(), bias.data_ptr(),
+                   out.data_ptr(), M, N, K, static_cast<float>(alpha), stream)
+             : cutlass_fp8_wide_bias_bf16out(
+                   input.data_ptr(), weight.data_ptr(), bias.data_ptr(),
+                   out.data_ptr(), M, N, K, static_cast<float>(alpha),
+                   static_cast<float>(beta), stream);
+  } else {
+    rc = fp8_linear_bias_sm110_bf16(
+        input.data_ptr(), weight.data_ptr(), bias.data_ptr(), out.data_ptr(),
+        M, N, K, static_cast<float>(alpha), static_cast<float>(beta),
+        epilogue, stream);
+  }
+  TORCH_CHECK(rc == 0, op_name, " failed with rc=", rc);
+#endif
+#else
+  TORCH_CHECK(false, "fp8-gemm was not built with CUDA support");
+#endif
+}
+}  // namespace
+void fp8_linear_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    double alpha,
+    int64_t variant,
+    torch::Tensor& out) {
+  launch(input, weight, alpha, variant, out, false);
+}
+void fp8_linear_residual_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    double alpha,
+    int64_t variant,
+    torch::Tensor& residual) {
+  launch(input, weight, alpha, variant, residual, true);
+}
+void fp8_linear_bias_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& out) {
+  launch_bias(input, weight, bias, alpha, out, 0.0,
+              FlashRtFp8BiasEpilogue::kBias, "fp8_linear_bias_bf16");
+}
+void fp8_linear_bias_residual_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& residual) {
+  launch_bias(input, weight, bias, alpha, residual, 1.0,
+              FlashRtFp8BiasEpilogue::kBias,
+              "fp8_linear_bias_residual_bf16");
+}
+void fp8_linear_bias_gelu_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& out) {
+  launch_bias(input, weight, bias, alpha, out, 0.0,
+              FlashRtFp8BiasEpilogue::kBiasGelu,
+              "fp8_linear_bias_gelu_bf16");
+}
+void fp8_blockwise_linear_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& input_scale,
+    torch::Tensor const& weight_scale,
+    torch::Tensor& out) {
+  check_fp8_matrix(input, "input");
+  check_fp8_matrix(weight, "weight");
+  check_fp32_matrix(input_scale, "input_scale");
+  check_fp32_matrix(weight_scale, "weight_scale");
+  check_bf16_matrix(out, "out");
+  const int64_t M = input.size(0);
+  const int64_t K = input.size(1);
+  const int64_t N = weight.size(0);
+  TORCH_CHECK(weight.size(1) == K,
+              "weight must have shape (N, input.shape[1])");
+  TORCH_CHECK(K % 128 == 0 && N % 128 == 0,
+              "N and K must be divisible by 128");
+  TORCH_CHECK(input_scale.sizes() ==
+                  torch::IntArrayRef({M, K / 128}),
+              "input_scale must have shape (M, K / 128)");
+  TORCH_CHECK(weight_scale.sizes() ==
+                  torch::IntArrayRef({N / 128, K / 128}),
+              "weight_scale must have shape (N / 128, K / 128)");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({M, N}),
+              "out must have shape (M, N)");
+  TORCH_CHECK(input.get_device() == weight.get_device() &&
+                  input.get_device() == input_scale.get_device() &&
+                  input.get_device() == weight_scale.get_device() &&
+                  input.get_device() == out.get_device(),
+              "all tensors must be on the same CUDA device");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(input.device());
+  auto* props = at::cuda::getDeviceProperties(input.get_device());
+  TORCH_CHECK((props->major == 8 && props->minor == 9) ||
+                  (props->major == 12 && props->minor == 0),
+              "fp8_blockwise_linear_bf16 requires SM89 or SM120; got SM",
+              props->major, props->minor);
+  auto stream = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
+  if (props->major == 8) {
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+    TORCH_CHECK(false, "SM89 blockwise kernels are not present in this source-test build");
+#else
+    int rc;
+    if (M == 1) {
+      namespace gemv89 = flash_rt::gemm::gemv_m1_sm89;
+      if (N <= 2048) {
+        rc = gemv89::gemv_fp8_block128_m1_w4(
+            input.data_ptr(), weight.data_ptr(), out.data_ptr(), 1,
+            checked_positive_int(N, "N"), checked_positive_int(K, "K"),
+            input_scale.data_ptr<float>(), weight_scale.data_ptr<float>(),
+            1.0f, stream);
+      } else if (N <= 8192) {
+        rc = gemv89::gemv_fp8_block128_m1_w8(
+            input.data_ptr(), weight.data_ptr(), out.data_ptr(), 1,
+            checked_positive_int(N, "N"), checked_positive_int(K, "K"),
+            input_scale.data_ptr<float>(), weight_scale.data_ptr<float>(),
+            1.0f, stream);
+      } else {
+        rc = gemv89::gemv_fp8_block128_m1_w16(
+            input.data_ptr(), weight.data_ptr(), out.data_ptr(), 1,
+            checked_positive_int(N, "N"), checked_positive_int(K, "K"),
+            input_scale.data_ptr<float>(), weight_scale.data_ptr<float>(),
+            1.0f, stream);
+      }
+    } else {
+      rc = flash_rt::gemm::block128_sm89::
+          fp8_block128_gemm_blockscaled_sm89_bf16out(
+              input.data_ptr(), weight.data_ptr(), out.data_ptr(),
+              checked_positive_int(M, "M"), checked_positive_int(N, "N"),
+              checked_positive_int(K, "K"), input_scale.data_ptr<float>(),
+              weight_scale.data_ptr<float>(), stream);
+    }
+    TORCH_CHECK(rc == 0, "SM89 blockwise FP8 linear failed with rc=", rc);
+#endif
+  } else {
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM89_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+    TORCH_CHECK(false, "SM120 blockwise kernel is not present in this source-test build");
+#else
+    flash_rt::gemm::fp8_block128_gemm_cutlass_sm120_bf16out(
+        input.data_ptr(), weight.data_ptr(), out.data_ptr(),
+        checked_positive_int(M, "M"), checked_positive_int(N, "N"),
+        checked_positive_int(K, "K"), input_scale.data_ptr<float>(),
+        weight_scale.data_ptr<float>(), stream);
+#endif
+  }
+#else
+  TORCH_CHECK(false, "fp8-gemm was not built with CUDA support");
+#endif
+}
+void fp8_blockwise_swiglu_quantize_fp8(
+    torch::Tensor const& input,
+    torch::Tensor const& gate_up_weight,
+    torch::Tensor const& input_scale,
+    torch::Tensor const& gate_up_weight_scale,
+    torch::Tensor& output,
+    torch::Tensor& output_scale) {
+  check_fp8_matrix(input, "input");
+  check_fp8_matrix(gate_up_weight, "gate_up_weight");
+  check_fp32_matrix(input_scale, "input_scale");
+  check_fp32_matrix(gate_up_weight_scale, "gate_up_weight_scale");
+  check_fp8_matrix(output, "output");
+  check_fp32_matrix(output_scale, "output_scale");
+  const int64_t M = input.size(0);
+  const int64_t K = input.size(1);
+  TORCH_CHECK(gate_up_weight.size(0) % 2 == 0 &&
+                  gate_up_weight.size(1) == K,
+              "gate_up_weight must have shape (2*N, K)");
+  const int64_t N = gate_up_weight.size(0) / 2;
+  TORCH_CHECK(M > 0 && M <= 256,
+              "SM89 fused SwiGLU producer supports 1 <= M <= 256");
+  TORCH_CHECK(N % 128 == 0 && K % 128 == 0,
+              "N and K must be divisible by 128");
+  TORCH_CHECK(input_scale.sizes() == torch::IntArrayRef({M, K / 128}),
+              "input_scale must have shape (M, K / 128)");
+  TORCH_CHECK(gate_up_weight_scale.sizes() ==
+                  torch::IntArrayRef({2 * N / 128, K / 128}),
+              "gate_up_weight_scale must have shape (2*N/128, K/128)");
+  TORCH_CHECK(output.sizes() == torch::IntArrayRef({M, N}),
+              "output must have shape (M, N)");
+  TORCH_CHECK(output_scale.sizes() == torch::IntArrayRef({M, N / 128}),
+              "output_scale must have shape (M, N/128)");
+  TORCH_CHECK(input.get_device() == gate_up_weight.get_device() &&
+                  input.get_device() == input_scale.get_device() &&
+                  input.get_device() == gate_up_weight_scale.get_device() &&
+                  input.get_device() == output.get_device() &&
+                  input.get_device() == output_scale.get_device(),
+              "all tensors must be on the same CUDA device");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(input.device());
+  auto* props = at::cuda::getDeviceProperties(input.get_device());
+  TORCH_CHECK(props->major == 8 && props->minor == 9,
+              "fp8_blockwise_swiglu_quantize_fp8 requires SM89; got SM",
+              props->major, props->minor);
+#if defined(FLASHRT_FP8_GEMM_SOURCE_SM120_ONLY) || \
+    defined(FLASHRT_FP8_GEMM_SOURCE_SM110_ONLY)
+  TORCH_CHECK(false, "SM89 fused producer is not present in this source-test build");
+#else
+  auto stream = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
+  int rc = flash_rt::gemm::block128_sm89::
+      fp8_bs_geglu_silu_fold_sm89_32x128_w4_s1(
+          input.data_ptr(), gate_up_weight.data_ptr(),
+          checked_positive_int(M, "M"), checked_positive_int(N, "N"),
+          checked_positive_int(K, "K"), input_scale.data_ptr<float>(),
+          gate_up_weight_scale.data_ptr<float>(), output.data_ptr(),
+          output_scale.data_ptr<float>(), stream);
+  TORCH_CHECK(rc == 0, "SM89 fused SwiGLU FP8 producer failed with rc=", rc);
+#endif
+#else
+  TORCH_CHECK(false, "fp8-gemm was not built with CUDA support");
+#endif
+}
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("fp8_linear_bf16(Tensor input, Tensor weight, float alpha, int variant, Tensor! out) -> ()");
+  ops.def("fp8_linear_residual_bf16(Tensor input, Tensor weight, float alpha, int variant, Tensor! residual) -> ()");
+  ops.def("fp8_linear_bias_bf16(Tensor input, Tensor weight, Tensor bias, float alpha, Tensor! out) -> ()");
+  ops.def("fp8_linear_bias_residual_bf16(Tensor input, Tensor weight, Tensor bias, float alpha, Tensor! residual) -> ()");
+  ops.def("fp8_linear_bias_gelu_bf16(Tensor input, Tensor weight, Tensor bias, float alpha, Tensor! out) -> ()");
+  ops.def("fp8_blockwise_linear_bf16("
+          "Tensor input, Tensor weight, Tensor input_scale, "
+          "Tensor weight_scale, Tensor! out) -> ()");
+  ops.def("fp8_blockwise_swiglu_quantize_fp8("
+          "Tensor input, Tensor gate_up_weight, Tensor input_scale, "
+          "Tensor gate_up_weight_scale, Tensor! output, Tensor! output_scale) -> ()");
+#if defined(CUDA_KERNEL)
+  ops.impl("fp8_linear_bf16", torch::kCUDA, &fp8_linear_bf16);
+  ops.impl("fp8_linear_residual_bf16", torch::kCUDA, &fp8_linear_residual_bf16);
+  ops.impl("fp8_linear_bias_bf16", torch::kCUDA, &fp8_linear_bias_bf16);
+  ops.impl("fp8_linear_bias_residual_bf16", torch::kCUDA,
+           &fp8_linear_bias_residual_bf16);
+  ops.impl("fp8_linear_bias_gelu_bf16", torch::kCUDA,
+           &fp8_linear_bias_gelu_bf16);
+  ops.impl("fp8_blockwise_linear_bf16",
+           torch::kCUDA,
+           &fp8_blockwise_linear_bf16);
+  ops.impl("fp8_blockwise_swiglu_quantize_fp8",
+           torch::kCUDA,
+           &fp8_blockwise_swiglu_quantize_fp8);
+#endif
+}
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h ADDED Viewed

	@@ -0,0 +1,54 @@

+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <torch/all.h>
+void fp8_linear_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    double alpha,
+    int64_t variant,
+    torch::Tensor& out);
+void fp8_linear_residual_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    double alpha,
+    int64_t variant,
+    torch::Tensor& residual);
+void fp8_linear_bias_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& out);
+void fp8_linear_bias_residual_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& residual);
+void fp8_linear_bias_gelu_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& bias,
+    double alpha,
+    torch::Tensor& out);
+void fp8_blockwise_linear_bf16(
+    torch::Tensor const& input,
+    torch::Tensor const& weight,
+    torch::Tensor const& input_scale,
+    torch::Tensor const& weight_scale,
+    torch::Tensor& out);
+void fp8_blockwise_swiglu_quantize_fp8(
+    torch::Tensor const& input,
+    torch::Tensor const& gate_up_weight,
+    torch::Tensor const& input_scale,
+    torch::Tensor const& gate_up_weight_scale,
+    torch::Tensor& output,
+    torch::Tensor& output_scale);