From f9b7e46f91fdcd50d14122c992ad0114043b6725 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 03:01:29 +0000 Subject: [PATCH 1/9] perf(gemm): tile short-token FFN weight gradients --- csrc/cuda/gemm/det_gemm_kernel.cu | 104 +++++++++++++++++++++++++++++- tests/test_det_gemm.py | 2 + 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index e622f563..9a46fdd7 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -116,6 +116,88 @@ __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) { + extern __shared__ __align__(1024) nv_bf16 smem[]; + nv_bf16* sX = 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) { @@ -478,11 +560,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 From 8fbd3637c071e59527d71f6a02e361ced87b6f88 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 04:08:39 +0000 Subject: [PATCH 2/9] fix(gemm): avoid CUDA shared-memory symbol collision --- csrc/cuda/gemm/det_gemm_kernel.cu | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 9a46fdd7..9697bfd6 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -137,8 +137,12 @@ __global__ void det_gemm_db_small_k(const nv_bf16* __restrict__ X, int tokens, int in_features, int out_features) { - extern __shared__ __align__(1024) nv_bf16 smem[]; - nv_bf16* sX = smem; + // 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; @@ -252,8 +256,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); From 949b088b61a3ab42bc1c0f5f8fa00544131a3558 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 05:36:49 +0000 Subject: [PATCH 3/9] docs(perf): add H100 FFN kernel profile charts --- .../after_kernel_flamegraph.svg | 34 +++++++++++++++++++ .../before_kernel_flamegraph.svg | 33 ++++++++++++++++++ .../qwen_ffn_h100_trace/kernel_summary.md | 28 +++++++++++++++ 3 files changed, 95 insertions(+) create mode 100644 benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md diff --git a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg new file mode 100644 index 00000000..fb755584 --- /dev/null +++ b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg @@ -0,0 +1,34 @@ + + + + After — H100 CUDA kernel-time flamegraph + Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process + Total selected det_gemm kernel time: 519.80 ms + + + + + + det_gemm_sm90_kernel + 391.45 ms · 643 launches + det_gemm_naive + 116.34 ms · 84 launches + + small_k + 12.01 ms + + + 0 ms + 519.80 ms + + + det_gemm_db_small_k replaces the red naive dW region with a 128×128 tiled outer-product kernel. + Changed kernel: 157.66 ms → 12.01 ms (13.13× kernel-time speedup; 92.38% reduction). + Source: nsys stats --report cuda_gpu_kern_sum --format csv + diff --git a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg new file mode 100644 index 00000000..72a32624 --- /dev/null +++ b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg @@ -0,0 +1,33 @@ + + + + Before — H100 CUDA kernel-time flamegraph + Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process + Total selected det_gemm kernel time: 666.24 ms + + + + + + det_gemm_sm90_kernel + 392.26 ms · 643 launches + det_gemm_naive<bf16,true> + 157.66 ms · 90 launches + det_gemm_naive + 116.32 ms · 84 launches + + + 0 ms + 666.24 ms + + + The short-token dW work is dominated by the 16×16 scalar naive kernel (red). + 90 launches repeatedly reload the same token rows, which is the target of the tiled outer-product kernel. + Source: nsys stats --report cuda_gpu_kern_sum --format csv + 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. From adb575db4b57b64f7a11818776269784cb8ed159 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 05:41:04 +0000 Subject: [PATCH 4/9] docs(perf): label kernel breakdown charts accurately --- ...er_kernel_flamegraph.svg => after_kernel_time_breakdown.svg} | 2 +- ...e_kernel_flamegraph.svg => before_kernel_time_breakdown.svg} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename benchmarks/results/qwen_ffn_h100_trace/{after_kernel_flamegraph.svg => after_kernel_time_breakdown.svg} (98%) rename benchmarks/results/qwen_ffn_h100_trace/{before_kernel_flamegraph.svg => before_kernel_time_breakdown.svg} (98%) diff --git a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg similarity index 98% rename from benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg rename to benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg index fb755584..f6fffcae 100644 --- a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg +++ b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg @@ -7,7 +7,7 @@ .small { font: 13px sans-serif; fill: #334155; } .axis { font: 12px sans-serif; fill: #64748b; } - After — H100 CUDA kernel-time flamegraph + After — H100 CUDA kernel-time breakdown Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process Total selected det_gemm kernel time: 519.80 ms diff --git a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg similarity index 98% rename from benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg rename to benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg index 72a32624..2abe99e8 100644 --- a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg +++ b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg @@ -7,7 +7,7 @@ .small { font: 13px sans-serif; fill: #334155; } .axis { font: 12px sans-serif; fill: #64748b; } - Before — H100 CUDA kernel-time flamegraph + Before — H100 CUDA kernel-time breakdown Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process Total selected det_gemm kernel time: 666.24 ms From 84c50caf3012e93e6c912e69a18daa336223172b Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 03:01:29 +0000 Subject: [PATCH 5/9] perf(gemm): tile short-token FFN weight gradients --- csrc/cuda/gemm/det_gemm_kernel.cu | 104 +++++++++++++++++++++++++++++- tests/test_det_gemm.py | 2 + 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index e622f563..9a46fdd7 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -116,6 +116,88 @@ __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) { + extern __shared__ __align__(1024) nv_bf16 smem[]; + nv_bf16* sX = 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) { @@ -478,11 +560,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 From 5e21101d7692d1b5ac8d791f00e0941bb90c6e09 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 04:08:39 +0000 Subject: [PATCH 6/9] fix(gemm): avoid CUDA shared-memory symbol collision --- csrc/cuda/gemm/det_gemm_kernel.cu | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 9a46fdd7..9697bfd6 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -137,8 +137,12 @@ __global__ void det_gemm_db_small_k(const nv_bf16* __restrict__ X, int tokens, int in_features, int out_features) { - extern __shared__ __align__(1024) nv_bf16 smem[]; - nv_bf16* sX = smem; + // 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; @@ -252,8 +256,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); From 41d89be2ee3ceb60baea170f607a74d0acdb42e2 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 05:36:49 +0000 Subject: [PATCH 7/9] docs(perf): add H100 FFN kernel profile charts --- .../after_kernel_flamegraph.svg | 34 +++++++++++++++++++ .../before_kernel_flamegraph.svg | 33 ++++++++++++++++++ .../qwen_ffn_h100_trace/kernel_summary.md | 28 +++++++++++++++ 3 files changed, 95 insertions(+) create mode 100644 benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md diff --git a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg new file mode 100644 index 00000000..fb755584 --- /dev/null +++ b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg @@ -0,0 +1,34 @@ + + + + After — H100 CUDA kernel-time flamegraph + Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process + Total selected det_gemm kernel time: 519.80 ms + + + + + + det_gemm_sm90_kernel + 391.45 ms · 643 launches + det_gemm_naive + 116.34 ms · 84 launches + + small_k + 12.01 ms + + + 0 ms + 519.80 ms + + + det_gemm_db_small_k replaces the red naive dW region with a 128×128 tiled outer-product kernel. + Changed kernel: 157.66 ms → 12.01 ms (13.13× kernel-time speedup; 92.38% reduction). + Source: nsys stats --report cuda_gpu_kern_sum --format csv + diff --git a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg new file mode 100644 index 00000000..72a32624 --- /dev/null +++ b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg @@ -0,0 +1,33 @@ + + + + Before — H100 CUDA kernel-time flamegraph + Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process + Total selected det_gemm kernel time: 666.24 ms + + + + + + det_gemm_sm90_kernel + 392.26 ms · 643 launches + det_gemm_naive<bf16,true> + 157.66 ms · 90 launches + det_gemm_naive + 116.32 ms · 84 launches + + + 0 ms + 666.24 ms + + + The short-token dW work is dominated by the 16×16 scalar naive kernel (red). + 90 launches repeatedly reload the same token rows, which is the target of the tiled outer-product kernel. + Source: nsys stats --report cuda_gpu_kern_sum --format csv + 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. From 868d12fde21b1f8cd444748d7eec55b8779dbe8f Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 05:41:04 +0000 Subject: [PATCH 8/9] docs(perf): label kernel breakdown charts accurately --- ...er_kernel_flamegraph.svg => after_kernel_time_breakdown.svg} | 2 +- ...e_kernel_flamegraph.svg => before_kernel_time_breakdown.svg} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename benchmarks/results/qwen_ffn_h100_trace/{after_kernel_flamegraph.svg => after_kernel_time_breakdown.svg} (98%) rename benchmarks/results/qwen_ffn_h100_trace/{before_kernel_flamegraph.svg => before_kernel_time_breakdown.svg} (98%) diff --git a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg similarity index 98% rename from benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg rename to benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg index fb755584..f6fffcae 100644 --- a/benchmarks/results/qwen_ffn_h100_trace/after_kernel_flamegraph.svg +++ b/benchmarks/results/qwen_ffn_h100_trace/after_kernel_time_breakdown.svg @@ -7,7 +7,7 @@ .small { font: 13px sans-serif; fill: #334155; } .axis { font: 12px sans-serif; fill: #64748b; } - After — H100 CUDA kernel-time flamegraph + After — H100 CUDA kernel-time breakdown Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process Total selected det_gemm kernel time: 519.80 ms diff --git a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg similarity index 98% rename from benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg rename to benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg index 72a32624..2abe99e8 100644 --- a/benchmarks/results/qwen_ffn_h100_trace/before_kernel_flamegraph.svg +++ b/benchmarks/results/qwen_ffn_h100_trace/before_kernel_time_breakdown.svg @@ -7,7 +7,7 @@ .small { font: 13px sans-serif; fill: #334155; } .axis { font: 12px sans-serif; fill: #64748b; } - Before — H100 CUDA kernel-time flamegraph + Before — H100 CUDA kernel-time breakdown Nsight Systems cuda_gpu_kern_sum; tokens=1,8; aggregated over the profiled benchmark process Total selected det_gemm kernel time: 666.24 ms From 9495a460a398aad096ce6fafa3462ed2cfc34d72 Mon Sep 17 00:00:00 2001 From: frank-2077 Date: Fri, 28 Aug 2026 12:44:00 +0000 Subject: [PATCH 9/9] perf(gemm): optimize long-token FFN deterministic path --- .../bench_long_token_std.py | 82 +++++++++++++++++++ .../long_token_kernel_breakdown.svg | 42 ++++++++++ .../long_token_latency.svg | 56 +++++++++++++ .../qwen_ffn_h100_trace/long_token_summary.md | 68 +++++++++++++++ csrc/cuda/gemm/det_gemm_kernel.cu | 65 +++++++-------- 5 files changed, 279 insertions(+), 34 deletions(-) create mode 100644 benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py create mode 100644 benchmarks/results/qwen_ffn_h100_trace/long_token_kernel_breakdown.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/long_token_latency.svg create mode 100644 benchmarks/results/qwen_ffn_h100_trace/long_token_summary.md 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/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 @@ + + + + Qwen3 FFN long-token kernel time — H100 + Nsight Systems cuda_gpu_kern_sum · deterministic GEMM aggregate (81 × bool0 + 21 × bool1 launches) + BF16 · hidden=4096 · intermediate=12288 · baseline: origin/pr-351 @ adb575d · lower is better + + Deterministic GEMM total time per profiled process (ms) + + + + + + 180013509004500 + + + + + 212.642155.6721.37× + 431.163329.5991.31× + 868.660674.8381.29× + 1713.3461378.4691.24× + 1024204840968192 + + baselineoptimized + Token countBaseline totalOptimized totalReductionBaseline shareOptimized share + + 1024212.642 ms155.672 ms26.8%88.5%85.0% + 2048431.163 ms329.599 ms23.6%90.2%87.8% + 4096868.660 ms674.838 ms22.3%90.9%89.0% + 81921713.346 ms1378.469 ms19.5%91.0%89.4% + 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 @@ + + + + Qwen3 FFN long-token latency — H100 + BF16 · hidden=4096 · intermediate=12288 · CUDA-event median across 3 trials + Each trial: 5 warmups, 20 forward iterations, 10 forward+backward iterations · baseline: origin/pr-351 @ adb575d + + Deterministic forward (ms) + + + + + + + 55 + 44 + 33 + 22 + 11 + 0 + + + + + 6.5134.8531.34× + 12.86610.0311.28× + 25.49320.0341.27× + 50.05240.4211.24× + 1024204840968192 + + Deterministic forward + backward (ms) + + + + + + 170127.58542.50 + + + + + 19.97915.1711.31× + 39.81831.3641.27× + 78.73162.3421.26× + 155.348126.4411.23× + 1024204840968192 + baselineoptimizedLower is better · speedup = baseline / optimized + 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 | + +![Long-token latency](long_token_latency.svg) + +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% | + +![Long-token kernel breakdown](long_token_kernel_breakdown.svg) + +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 9697bfd6..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) { @@ -231,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]) @@ -290,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) @@ -303,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(); @@ -352,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; } } @@ -387,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])); } } }