diff --git a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg
new file mode 100644
index 00000000..f6fffcae
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg
@@ -0,0 +1,34 @@
+
diff --git a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg
new file mode 100644
index 00000000..2abe99e8
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg
@@ -0,0 +1,33 @@
+
diff --git a/benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py b/benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py
new file mode 100644
index 00000000..54c642c8
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py
@@ -0,0 +1,82 @@
+import json
+import sys
+
+import torch
+
+from rl_engine.kernels.ops.pytorch.ffn.ffn import qwen3_ffn
+
+
+H, I = 4096, 12288
+WARMUP = 5
+FW_ITERS = 20
+FB_ITERS = 10
+
+
+def timed(fn, warmup, iters):
+ for _ in range(warmup):
+ fn()
+ torch.cuda.synchronize()
+ samples = []
+ for _ in range(iters):
+ start = torch.cuda.Event(enable_timing=True)
+ end = torch.cuda.Event(enable_timing=True)
+ start.record()
+ fn()
+ end.record()
+ end.synchronize()
+ samples.append(start.elapsed_time(end))
+ values = torch.tensor(samples, dtype=torch.float64)
+ return {
+ "median_ms": float(values.median()),
+ "mean_ms": float(values.mean()),
+ "min_ms": float(values.min()),
+ "max_ms": float(values.max()),
+ "samples_ms": samples,
+ }
+
+
+def main():
+ torch.cuda.set_device(0)
+ torch.manual_seed(2026)
+ result = {
+ "device": torch.cuda.get_device_name(0),
+ "torch": torch.__version__,
+ "cuda": torch.version.cuda,
+ "shape": {"hidden": H, "intermediate": I, "dtype": "bfloat16"},
+ "warmup": WARMUP,
+ "forward_iters": FW_ITERS,
+ "forward_backward_iters": FB_ITERS,
+ "rows": [],
+ }
+ for tokens in map(int, sys.argv[1:]):
+ x = torch.randn(tokens, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+ gate = torch.randn(I, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+ up = torch.randn(I, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+ down = torch.randn(H, I, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+ dout = torch.randn(tokens, H, device="cuda", dtype=torch.bfloat16)
+ for deterministic in (True, False):
+ def forward():
+ return qwen3_ffn(x, gate, up, down, deterministic=deterministic)
+
+ def forward_backward():
+ y = forward()
+ torch.autograd.grad(y, (x, gate, up, down), dout)
+
+ row = {
+ "tokens": tokens,
+ "mode": "det" if deterministic else "prod",
+ "forward": timed(forward, WARMUP, FW_ITERS),
+ "forward_backward": timed(forward_backward, WARMUP, FB_ITERS),
+ }
+ result["rows"].append(row)
+ print(
+ f"tokens={tokens} mode={row['mode']} "
+ f"forward_median={row['forward']['median_ms']:.4f} "
+ f"fwd_bwd_median={row['forward_backward']['median_ms']:.4f}",
+ flush=True,
+ )
+ print(json.dumps(result))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md b/benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md
new file mode 100644
index 00000000..d111e4af
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md
@@ -0,0 +1,28 @@
+# Qwen FFN H100 kernel profile
+
+These charts are aggregated from the Nsight Systems `cuda_gpu_kern_sum` report
+for the same `tokens=1,8`, `hidden=4096`, `intermediate=12288` profiling run.
+The benchmark correctness checks passed in both worktrees.
+
+| Profile | Kernel | Total time | Launches | Average |
+|---|---|---:|---:|---:|
+| Before | `det_gemm_sm90_kernel` | 392.262584 ms | 643 | 610.051 us |
+| Before | `det_gemm_naive` | 157.657338 ms | 90 | 1,751.748 us |
+| Before | `det_gemm_naive` | 116.315558 ms | 84 | 1,384.709 us |
+| After | `det_gemm_sm90_kernel` | 391.452104 ms | 643 | 608.790 us |
+| After | `det_gemm_naive` | 116.338858 ms | 84 | 1,384.986 us |
+| After | `det_gemm_db_small_k` | 12.006408 ms | 90 | 133.405 us |
+
+The optimized kernel replaces the 90 `det_gemm_naive` launches:
+
+- 157.657338 ms → 12.006408 ms;
+- 13.13× kernel-time speedup;
+- 92.38% reduction for the replaced kernel;
+- 145.651 ms saved in the selected deterministic GEMM kernel aggregate.
+
+The unchanged SM90 and `det_gemm_naive` rows provide a useful
+control: the observed reduction is localized to the intended short-token dW
+path rather than a general profiling artifact.
+
+Raw `.nsys-rep` files remain on the H100 profiling node because they are large;
+the SVGs and this reproducible kernel summary are the review artifacts.
diff --git a/benchmarks/results/qwen_ffn_h100_trace/long_token_kernel_breakdown.svg b/benchmarks/results/qwen_ffn_h100_trace/long_token_kernel_breakdown.svg
new file mode 100644
index 00000000..4c8a1a54
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/long_token_kernel_breakdown.svg
@@ -0,0 +1,42 @@
+
diff --git a/benchmarks/results/qwen_ffn_h100_trace/long_token_latency.svg b/benchmarks/results/qwen_ffn_h100_trace/long_token_latency.svg
new file mode 100644
index 00000000..5e863428
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/long_token_latency.svg
@@ -0,0 +1,56 @@
+
diff --git a/benchmarks/results/qwen_ffn_h100_trace/long_token_summary.md b/benchmarks/results/qwen_ffn_h100_trace/long_token_summary.md
new file mode 100644
index 00000000..e6a02426
--- /dev/null
+++ b/benchmarks/results/qwen_ffn_h100_trace/long_token_summary.md
@@ -0,0 +1,68 @@
+# Qwen FFN long-token performance
+
+This report compares the current long-token optimization with the PR #351
+baseline (`origin/pr-351`, `adb575d`). Both versions use the same H100 node,
+CUDA/PyTorch environment, BF16 inputs, and Qwen3-8B FFN shape
+`hidden=4096, intermediate=12288`.
+
+## CUDA-event latency
+
+The headline numbers use CUDA events, not Nsight-instrumented timings. Each
+trial uses 5 warmups, 20 forward iterations, and 10 forward+backward
+iterations. Values below are the median of three trial medians.
+
+| Tokens | Det forward baseline (ms) | Det forward optimized (ms) | Speedup | Det fwd+bwd baseline (ms) | Det fwd+bwd optimized (ms) | Speedup |
+|---:|---:|---:|---:|---:|---:|---:|
+| 1024 | 6.5133 | 4.8531 | 1.34x | 19.9793 | 15.1710 | 1.32x |
+| 2048 | 12.8660 | 10.0312 | 1.28x | 39.8176 | 31.3641 | 1.27x |
+| 4096 | 25.4928 | 20.0340 | 1.27x | 78.7306 | 62.3420 | 1.26x |
+| 8192 | 50.0523 | 40.4205 | 1.24x | 155.3476 | 126.4405 | 1.23x |
+
+
+
+The production path is unchanged and was measured as a control. Its three-trial
+medians were within normal run-to-run variation (forward: -0.6%, -0.5%, +0.0%,
+-2.5%; forward+backward: -1.2%, -6.7%, -3.7%, +1.5% for 1024/2048/4096/8192).
+These control measurements are not used as an optimization claim.
+
+## Nsight Systems kernel attribution
+
+The following aggregates come from `cuda_gpu_kern_sum`. Each row sums the two
+deterministic GEMM templates (`(bool)0`: 81 launches and `(bool)1`: 21 launches)
+in the profiled process. `Time (%)` is the share of all GPU kernel time in that
+profile, not end-to-end application latency.
+
+| Tokens | Baseline det GEMM total (ms) | Optimized det GEMM total (ms) | Reduction | Baseline GPU share | Optimized GPU share |
+|---:|---:|---:|---:|---:|---:|
+| 1024 | 212.642 | 155.672 | 26.8% | 88.7% | 85.0% |
+| 2048 | 431.163 | 329.599 | 23.6% | 90.4% | 87.8% |
+| 4096 | 868.660 | 674.838 | 22.3% | 91.1% | 89.0% |
+| 8192 | 1713.346 | 1378.469 | 19.5% | 91.2% | 89.4% |
+
+
+
+The deterministic GEMM remains the dominant GPU workload. The reduction stays
+positive as token count grows, while the percentage benefit decreases because
+the long-token GEMM work increasingly dominates the fixed launch and reduction
+overheads.
+
+## Reproduction
+
+The CUDA-event benchmark is in
+`bench_long_token_std.py`:
+
+```bash
+PYTHONPATH=$PWD python benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py \
+ 1024 2048 4096 8192
+```
+
+Run it once with `origin/pr-351` built and once with the candidate kernel. The
+Nsight CSVs were generated with:
+
+```bash
+nsys stats --report cuda_gpu_kern_sum --format csv \
+ --force-export=true .nsys-rep
+```
+
+Raw `.nsys-rep` and JSON timing logs are kept outside the repository; the SVGs
+and this summary are the review artifacts.
diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu
index e622f563..34b08030 100644
--- a/csrc/cuda/gemm/det_gemm_kernel.cu
+++ b/csrc/cuda/gemm/det_gemm_kernel.cu
@@ -71,22 +71,6 @@ __device__ __forceinline__ nv_bf16 bf16_add(nv_bf16 a, nv_bf16 b) {
return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b));
}
-// True iff [lo, hi) is a node of the mid-split tree over [0, n).
-__device__ __forceinline__ bool is_mid_split_node(int lo, int hi, int n) {
- int a = 0, b = n;
- while (b - a > 1) {
- if (a == lo && b == hi) return true;
- const int m = a + (b - a) / 2;
- if (hi <= m)
- b = m;
- else if (lo >= m)
- a = m;
- else
- return false;
- }
- return a == lo && b == hi;
-}
-
__device__ nv_bf16 k_tree_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B,
int row, int col, int N, int K, int lo, int hi) {
if (hi - lo <= K_TREE_LEAF) {
@@ -116,6 +100,92 @@ __global__ void det_gemm_naive(const nv_bf16* __restrict__ A,
__bfloat162float(k_tree_naive(A, B, row, col, N, K, 0, K)));
}
+// Weight gradients in the FFN are an outer-product GEMM when the token batch
+// is small: dW = dY^T @ X. The regular scalar fallback assigns one 16x16
+// block to each output tile. For Qwen3's [out,in] projections that creates
+// roughly 200k blocks at M=1/8 and repeatedly reloads the same token rows.
+// This path keeps the exact scalar accumulation order (tokens are visited in
+// ascending order and the result is rounded once to BF16), but stages one
+// token tile and computes a 128x128 output tile with 256 threads. It also
+// consumes X in its native [tokens,in] layout, so no X.T.contiguous() helper
+// allocation is needed.
+constexpr int SMALL_K_TILE = K_TREE_LEAF;
+constexpr int SMALL_M_TILE = 128;
+constexpr int SMALL_N_TILE = 128;
+constexpr int SMALL_THREADS = 256;
+
+template
+__global__ void det_gemm_db_small_k(const nv_bf16* __restrict__ X,
+ const nv_bf16* __restrict__ dY,
+ output_t* __restrict__ dW,
+ int tokens,
+ int in_features,
+ int out_features) {
+ // Give this dynamic shared-memory symbol a kernel-specific name. CUDA
+ // 12.4 diagnoses same-TU extern __shared__ declarations with different
+ // element types as incompatible, even though they belong to different
+ // kernels.
+ extern __shared__ __align__(1024) nv_bf16 small_k_smem[];
+ nv_bf16* sX = small_k_smem;
+ nv_bf16* sY = sX + SMALL_K_TILE * SMALL_N_TILE;
+
+ const int tid = threadIdx.x;
+ const int in_base = blockIdx.x * SMALL_N_TILE;
+ const int out_base = blockIdx.y * SMALL_M_TILE;
+
+ // The shared tile is padded to 32 tokens. Zero padding makes the launch
+ // shape independent of the token count while the loop below still visits
+ // exactly the original [0,tokens) reduction range.
+ for (int index = tid; index < SMALL_K_TILE * SMALL_N_TILE; index += blockDim.x) {
+ const int token = index / SMALL_N_TILE;
+ const int feature = index % SMALL_N_TILE;
+ const int global_feature = in_base + feature;
+ sX[index] = (token < tokens && global_feature < in_features)
+ ? X[token * in_features + global_feature]
+ : __float2bfloat16(0.0f);
+ }
+ for (int index = tid; index < SMALL_K_TILE * SMALL_M_TILE; index += blockDim.x) {
+ const int token = index / SMALL_M_TILE;
+ const int output = index % SMALL_M_TILE;
+ const int global_output = out_base + output;
+ sY[index] = (token < tokens && global_output < out_features)
+ ? dY[token * out_features + global_output]
+ : __float2bfloat16(0.0f);
+ }
+ __syncthreads();
+
+ // dW is physically [out_features,in_features]. Each thread computes eight
+ // elements; neighboring threads therefore issue contiguous stores.
+ for (int index = tid; index < SMALL_M_TILE * SMALL_N_TILE; index += blockDim.x) {
+ const int output = index / SMALL_N_TILE;
+ const int feature = index % SMALL_N_TILE;
+ const int global_output = out_base + output;
+ const int global_feature = in_base + feature;
+ if (global_output >= out_features || global_feature >= in_features) continue;
+
+ float acc = 0.0f;
+ for (int token = 0; token < tokens; ++token)
+ acc += __bfloat162float(sX[token * SMALL_N_TILE + feature]) *
+ __bfloat162float(sY[token * SMALL_M_TILE + output]);
+ dW[global_output * in_features + global_feature] = cast_output(acc);
+ }
+}
+
+template
+void launch_db_small_k(const nv_bf16* X,
+ const nv_bf16* dY,
+ output_t* dW,
+ int tokens,
+ int in_features,
+ int out_features,
+ cudaStream_t stream) {
+ dim3 block(SMALL_THREADS);
+ dim3 grid(cdiv(in_features, SMALL_N_TILE), cdiv(out_features, SMALL_M_TILE));
+ constexpr int smem_elements = SMALL_K_TILE * (SMALL_N_TILE + SMALL_M_TILE);
+ det_gemm_db_small_k<<>>(
+ X, dY, dW, tokens, in_features, out_features);
+}
+
template
void launch_naive(const nv_bf16* A, const nv_bf16* B, output_t* C,
int M, int N, int K, cudaStream_t stream) {
@@ -145,6 +215,21 @@ constexpr int K_TILES = BK / MMA_K; // 2
constexpr int KK_GROUPS = BK / 32; // 1
constexpr int TREE_DEPTH = 16;
+__device__ __forceinline__ int mid_tree_merge_count(int leaf, int n) {
+ int lo = 0, hi = n, count = 0;
+ while (hi - lo > 1) {
+ const int mid = lo + (hi - lo) / 2;
+ if (leaf < mid) {
+ hi = mid;
+ count = 0;
+ } else {
+ lo = mid;
+ ++count;
+ }
+ }
+ return count;
+}
+
__device__ __forceinline__ void ldmatrix_x4(uint32_t regs[4], uint32_t addr) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
: "=r"(regs[0]), "=r"(regs[1]), "=r"(regs[2]), "=r"(regs[3])
@@ -170,8 +255,8 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
const int col_base = blockIdx.x * BN;
const int kd = K / BK;
- extern __shared__ __align__(1024) char smem[];
- nv_bf16* sA = reinterpret_cast(smem);
+ extern __shared__ __align__(1024) char sm90_smem[];
+ nv_bf16* sA = reinterpret_cast(sm90_smem);
nv_bf16* sB = reinterpret_cast(sA + STAGES * BM * BK);
int* mbar_base = reinterpret_cast(sB + STAGES * BN * BK);
@@ -204,9 +289,8 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
for (int s = 0; s < STAGES; ++s) phase[s] = 0;
float tile_acc[M_TILES][N_TILES][4];
- nv_bf16 tree_v[M_TILES][N_TILES][4];
- nv_bf16 tree_stk[TREE_DEPTH][M_TILES][N_TILES][4];
- int tree_lo[TREE_DEPTH], tree_hi[TREE_DEPTH];
+ __nv_bfloat162 tree_v[M_TILES][N_TILES][2];
+ __nv_bfloat162 tree_stk[TREE_DEPTH][M_TILES][N_TILES][2];
int sp = 0;
if (tid == 0)
@@ -217,7 +301,7 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
for (int k = 0; k < kd; ++k) { // fixed ascending tile order, NO split-K
const int buf = k % STAGES;
if (tid == 0 && k + (STAGES - 1) < kd) issue_load(k + (STAGES - 1));
- det_gemm::mbar_wait(mbar[buf], phase[buf]);
+ if (tid == 0) det_gemm::mbar_wait(mbar[buf], phase[buf]);
phase[buf] ^= 1;
__syncthreads();
@@ -266,29 +350,28 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
#pragma unroll
for (int n = 0; n < N_TILES; ++n)
#pragma unroll
- for (int i = 0; i < 4; ++i) tree_v[mi][n][i] = __float2bfloat16(tile_acc[mi][n][i]);
+ for (int i = 0; i < 2; ++i)
+ tree_v[mi][n][i] = __floats2bfloat162_rn(tile_acc[mi][n][2 * i + 0],
+ tile_acc[mi][n][2 * i + 1]);
- int lo = k, hi = k + 1;
- while (sp > 0 && tree_hi[sp - 1] == lo && is_mid_split_node(tree_lo[sp - 1], hi, kd)) {
+ const int merge_count = mid_tree_merge_count(k, kd);
+ for (int merge = 0; merge < merge_count; ++merge) {
#pragma unroll
for (int mi = 0; mi < M_TILES; ++mi)
#pragma unroll
for (int n = 0; n < N_TILES; ++n)
#pragma unroll
- for (int i = 0; i < 4; ++i)
- tree_v[mi][n][i] = bf16_add(tree_stk[sp - 1][mi][n][i], tree_v[mi][n][i]);
- lo = tree_lo[sp - 1];
+ for (int i = 0; i < 2; ++i)
+ tree_v[mi][n][i] = __hadd2(tree_stk[sp - 1][mi][n][i], tree_v[mi][n][i]);
--sp;
}
- if (hi < kd) {
+ if (k + 1 < kd) {
#pragma unroll
for (int mi = 0; mi < M_TILES; ++mi)
#pragma unroll
for (int n = 0; n < N_TILES; ++n)
#pragma unroll
- for (int i = 0; i < 4; ++i) tree_stk[sp][mi][n][i] = tree_v[mi][n][i];
- tree_lo[sp] = lo;
- tree_hi[sp] = hi;
+ for (int i = 0; i < 2; ++i) tree_stk[sp][mi][n][i] = tree_v[mi][n][i];
++sp;
}
}
@@ -301,15 +384,15 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap,
const int col = col_base + n * MMA_N + (lane % 4) * 2;
if (row < M && col + 1 < N) {
C[output_offset(row, col + 0, M, N)] =
- cast_output(__bfloat162float(tree_v[mi][n][0]));
+ cast_output(__low2float(tree_v[mi][n][0]));
C[output_offset(row, col + 1, M, N)] =
- cast_output(__bfloat162float(tree_v[mi][n][1]));
+ cast_output(__high2float(tree_v[mi][n][0]));
}
if (row + 8 < M && col + 1 < N) {
C[output_offset(row + 8, col + 0, M, N)] =
- cast_output(__bfloat162float(tree_v[mi][n][2]));
+ cast_output(__low2float(tree_v[mi][n][1]));
C[output_offset(row + 8, col + 1, M, N)] =
- cast_output(__bfloat162float(tree_v[mi][n][3]));
+ cast_output(__high2float(tree_v[mi][n][1]));
}
}
}
@@ -478,11 +561,29 @@ torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc) {
torch::Tensor det_gemm_db_transposed(torch::Tensor a, torch::Tensor dc) {
check_in(a, "A"); check_in(dc, "dC");
- dc = dc.contiguous();
TORCH_CHECK(a.dim() == 2 && dc.dim() == 2,
"det_gemm_db_transposed: expect A[M,K] and dC[M,N]");
+ const int tokens = a.size(0);
+ const int in_features = a.size(1);
+ const int out_features = dc.size(1);
+ TORCH_CHECK(dc.size(0) == tokens, "det_gemm_db_transposed: M mismatch");
+
+ // The normal SM90 path expects A^T to be physically contiguous. For short
+ // token batches, materializing that transpose and launching the 16x16
+ // scalar fallback dominates the actual outer-product work. The tiled path
+ // reads A directly and preserves the same ascending-token reduction order.
+ if (tokens > 0 && tokens < SMALL_K_TILE) {
+ a = a.contiguous();
+ dc = dc.contiguous();
+ auto output = torch::empty({out_features, in_features}, a.options());
+ auto stream = at::cuda::getCurrentCUDAStream();
+ launch_db_small_k(
+ bf16(a), bf16(dc), bf16o(output), tokens, in_features, out_features, stream);
+ return output;
+ }
+
+ dc = dc.contiguous();
auto at = a.t().contiguous();
- TORCH_CHECK(dc.size(0) == at.size(1), "det_gemm_db_transposed: M mismatch");
// Preserve the exact A^T @ dC MMA/tree evaluation and change only the final
// address mapping so the canonical [N,K] weight-gradient is born contiguous.
return gemm_dispatch(at, dc, RhsLayout::kKN, OutputLayout::kNM);
diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py
index b5555413..633009ef 100644
--- a/tests/test_det_gemm.py
+++ b/tests/test_det_gemm.py
@@ -97,6 +97,8 @@ def test_rhs_transposed_materializes_unaligned_contiguous_view_for_tma():
@pytest.mark.parametrize(
"shape",
[
+ (1, 128, 128), # short-K tiled backward path
+ (8, 128, 128), # short-K tiled backward path
(128, 128, 128), # aligned SM90 path
(128, 96, 64), # SM90 logical-M padding and dim-1 crop
(31, 70, 65), # scalar fallback