[jinchen62]

JIRA ID: AIHPBLAS-4726

Reopen #12250 because GemmFromAnywhere was merged to develop and the branch was deleted.

Summary

hipBLASLt loses on decode GEMMs with very few tokens. This ports the rocBLAS-internal wvSpltK wave-split-K kernels into hipBLASLt as GFA CustomKernels for BF16 TN with n ≤ 4 tokens on gfx950. It also routes those problems to them through library logic, so hipBLASLt picks them out of the box.

TN with a skinny n is how PyTorch-based inference calls hipBLASLt. A linear layer computes y[T, H] = x[T, K] · W[H, K]ᵀ in row-major terms. Column-major hipBLASLt computes the transpose, yᵀ = W · xᵀ, that is (H, K) · (K, T) = (H, T): TN with m = H (output features), n = T (tokens) and k = K. torch.mm(x, W.t()) on a GLM-5.2 decode shape logs:

hipblaslt-bench --api_method c -m 2624 -n 1 -k 6144 --lda 6144 --ldb 6144 --ldc 2624 --ldd 2624 ... --transA T --transB N --a_type bf16_r --b_type bf16_r --c_type bf16_r --d_type bf16_r ... --compute_type f32_r

F.linear on a sliced activation (qkv[:, :6144] of a 6200-wide tensor) reaches it with --ldb 6200, so the kernels take every leading dimension as a kernarg.

Earlier revisions also carried FP16 NN kernels with m = tokens. Those followed hipblaslt-bench runs that fed the GEMM list's row-major M straight into hipBLASLt's m, a call no framework makes, so they are dropped.

Kernels

wvSpltK_bf16_tn_m1, wvSpltK_bf16_tn_m2 and wvSpltK_bf16_tn_m4, ported from rocBLAS's TRANSA=false instantiation. The suffix is rocBLAS's M, the token count, which is hipBLASLt's n here.

• 64-lane K-split, 16 waves/WG and persistent WGs (grid = CU count). Tokens are staged in 128 KB of LDS, and any past that are read from global memory, so the kernels handle any K.
• v_dot2c_f32_bf16, no MFMA. gfx950 only: gfx942 has no BF16 dot instruction.
• Only the .s files ship. They are emitted at code object v4, the version hipBLASLt links its library code objects at.

Kernel	n	YTILE	UNRL	LDS	VGPRs
wvSpltK_bf16_tn_m1	1	2	2	128 KB	44
wvSpltK_bf16_tn_m2	2	2	2	128 KB	58
wvSpltK_bf16_tn_m4	0 < n ≤ 4	1	4	128 KB	112

• CuCount is a kernarg, not the hardcoded 80 (MI308). It and the launch grid are one value: the device's CU count, capped like Tensile's persistent grids by the caller's SM-count target (HIPBLASLT_MATMUL_DESC_SM_COUNT_TARGET) and TENSILE_PERSISTENT_MAX_CUS. So the persistent stride can never disagree with the grid.
• C and D are separate pointers. rocBLAS called these kernels in place and read the beta term through the output pointer. hipBLASLt also dispatches out-of-place, where that would read uninitialized D.
• The m4 tile (YTILE=1, UNRL=4) is tuned for gfx950's 256 CUs. On the GLM-5.2 decode shapes a wider tile leaves too few waves to hide memory latency; with weights streamed from HBM, this tile is 1.2× faster than YTILE=3, UNRL=2.
• The kernels take 72 bytes of kernargs. They order everything the staging and K loop read first and preload those 14 dwords into SGPRs, so only the epilogue's arguments are fetched with s_load.

Predicates

Enforced in custom.config:

• batch == 1. There are no batch strides in the kernarg list, so a strided-batched problem would leave every batch past the first unwritten.
• m > 8 (below one wave tile the tail fixup underflows), K % 8 == 0, and unit strides on all four tensors.
• m1 / m2: n == 1 / n == 2. m4: 0 < n ≤ 4. hipBLASLt quick-returns m == 0 but not n == 0, because grouped GEMM relies on it.

Library logic

gfx950/gfx950/Range/gfx950_Cijk_Alik_Bljk_BBS_BH_UserArgs.yaml routes n = 1, 2 and 3–4 to m1, m2 and m4, for m ≥ 9, 8 ≤ K ≤ 4096 and batch 1.

• The ranges stop at K = 4096, short of the kernels' own bound. Past that, hipBLASLt's own kernels match or beat wvSpltK on MI355X (see Perf).
• Tuning tools still see the kernels at any K: hipBLASLt's all-solutions list (getAllAlgos, which aiter's hipb_findallsols enumerates) includes them past the cap.
• The file carries the plain-GEMM problem type. A custom kernel takes its logic file's type, and these kernels support no bias, activation or scale vector, so those problems keep their current kernels.
• It reuses the header of the shipped plain-GEMM logic (gfx950/gfx950/Equality/gfx950_Cijk_Alik_Bljk_BBS_BH_UserArgs.yaml), so both share a placeholder library and tuned Equality sizes keep priority.
• A matched range whose kernel fails its own predicates (for example K % 8 ≠ 0) falls through to the next library.
• General-batched (pointer-array) problems never use these kernels: hipBLASLt skips custom kernels for that batch mode in every selection path (see below).
• The file uses the MI350 (0x75a0) header. MI355X searches its own tuned logic (gfx950_id75a3) first, so its tuned sizes are unaffected. On MI350, the ranges come before 19 skinny sizes tuned in its Bias library (see Perf).
• The prediction-library and dynamic stream-K modes skip Equality and Range rows, so those modes don't use these kernels.

Tensile and host-side changes

• CustomGridSize::ComputeUnits and a ComputeUnits arg semantic, so a persistent kernel can take the CU count as both grid and kernarg. The count honors the same CU budget as Tensile's persistent grids.
• Wired the classic Tensile AssertSizeEqual, AssertSizeGreaterThan, AssertSizeLessThan and AssertStride{A,B,C,D}Equal keys through extractDimPredicate onto the existing SizeEqual, SizeGreaterThan, SizeLessThan and Stride*Equal runtime predicates. The C++ side and its serialization already existed; only the Python mapping was missing. One registry drives validation and predicate emission, with a test that fails if a key is validated but not wired.
• Custom kernels in pointer-array mode: getAllSolutions already skipped handwritten custom kernels for general-batched problems, which they cannot run. The heuristic (getSolutions, behind getBestSolutions / getBestRawSolutions) and isSolutionSupported now do too. The heuristic widens its search by the number it drops, so callers still get their requested count. Before this, a pointer-array GEMM with batch_count = 1 in a skinny range hit an illegal memory access.

Which frameworks reach these kernels

• PyTorch: BF16 torch.mm(x, W.t()) and F.linear without bias, as above.
• SGLang on ROCm (SGLANG_USE_AITER=1): linear layers go through aiter's tgemm.mm, which reaches hipBLASLt only for shapes missing from aiter's tuned configs (its torch fallback) or through its hipb_mm backend.
  • In the current SGLang ROCm image (v0.5.20-rocm724-mi35x), aiter has tuned flydsl/triton kernels for 6 of the 7 GLM-5.2 decode shapes at ≤ 4 tokens. On those, its kernels are about 10% faster than wvSpltK with weights streamed from HBM.
  • The seventh, 6144×16384, falls back to hipBLASLt. It is past the K cap and keeps hipBLASLt's existing kernel, so GLM-5.2 decode under SGLang+aiter is unchanged.
• vLLM on ROCm: decode with ≤ 5 tokens runs vLLM's own wvSplitK (VLLM_ROCM_USE_SKINNY_GEMM, on by default), so it doesn't reach hipBLASLt.

Test plan

All on MI355X gfx950 with the shipped kernels:

• Tensile full-element validation of custom_rocblas_gemv_bf16_tn_m{1,2,4}.yaml (8, 8 and 11 sizes) at β=0, and at α=2/β=2 with random out-of-place C: all PASSED. The sizes include padded leading dimensions and K past the LDS stage (73728, 36864 and 24576).
• hipBLASLt out of the box through hipblaslt-bench: 21 cases pick the expected kernel and pass CPU verification.
  • Picks wvSpltK: n = 1–4 up to the K = 4096 range edge, m up to 32768, sliced activations (ldb > K) and β = 0.5.
  • Falls back: K past the cap (4104, 6144 and 16384), n = 5, bias, GELU, batch 2, K % 8 ≠ 0, m ≤ 8, FP16 NN, and general-batched (pointer-array) problems with batch_count = 1.
• hipblaslt-test: new pre_checkin case matmul_skinny_bf16_tn, 16 tests, all PASSED. They cover n = 1–4, strided and general-batched, padded leading dimensions, and β 0 and 2.
• CU budget: with --sm_count_target 64 / 128 and TENSILE_PERSISTENT_MAX_CUS=48, wvSpltK launches on 64 / 128 / 48 workgroups (checked with rocprofv3) and passes CPU verification. CustomKernel_test.cpp covers the grid and kernarg.
• Direct HIP vs CPU FP32 reference, on an earlier build of the same kernels: 156/156, plus 79/79 with α ≠ 1. These cover padded and misaligned leading dimensions, β ≠ 0 and the global-memory fallback. D is poisoned, so a stray store fails loudly.
• Unit suites: custom-kernel, library-logic corpus and ValidParameters (1220 passed); StaggerU census (417 passed). The logic-file test skips in TheRock's installed-artifact layout, where the logic tree is absent.

Perf (MI355X gfx950, hipBLASLt out of the box)

Setup:

• Compares hipBLASLt's pick with vs. without the new Range rows (hipblaslt-bench --requested_solution 2 -i 100 -j 30 --use_gpu_timer), both timed back to back in one process.
• Numbers are wvSpltK's speedup over hipBLASLt's previous pick, as weights hot / weights streamed from HBM (--rotating 512). Shapes outside a range run hipBLASLt's own kernel and count as 1.00×.

GLM-5.2 decode shapes (N × K: 2624×6144, 2048×2048, 6144×2048, 4096×2048, 6144×4096, 16384×2048, 6144×16384), geomean over the seven. 2624×6144 and 6144×16384 are past the cap.

tokens	1	2	3	4
BF16 TN	1.24× / 1.16×	1.17× / 1.11×	1.01× / 1.04×	1.01× / 1.03×

In range, n = 3–4 still lose on 6144×2048, 6144×4096 and 16384×2048 with weights hot (0.86–0.99×), and on 16384×2048 streamed (0.89–0.90×).

N × K sweep (N 1024–32768, K 1024–16384; 30 shapes per row, 29 for n = 4 because 1024×1024 is tuned):

n	shapes in range	in range	whole sweep	worst in range
1	18/30	1.33× / 1.25×	1.19× / 1.14×	1.03× / 0.99×
2	18/30	1.21× / 1.16×	1.12× / 1.09×	0.97× / 0.95×
3	18/30	1.06× / 1.10×	1.04× / 1.06×	0.62× / 0.73×
4	17/29	1.03× / 1.08×	1.02× / 1.04×	0.61× / 0.73×

• Why the cap: measured without it, wvSpltK at K ≥ 8192 was 1.03× / 1.00× at n = 1 and 0.89–0.94× at n ≥ 2. The GLM-5.2 shape 6144×16384 fell to 0.65× / 0.74× at n = 4.
• Inside the ranges, tn_m4 loses at very large N with short K: 0.61–0.62× hot and 0.73× streamed at 32768×1024.

MI350-tuned sizes:

• On MI350, the ranges take priority over 19 BF16 TN sizes tuned in the MI350 Bias library.
• Timed on MI355X as a stand-in, against the kernel MI355X's own tuning picks for each size, wvSpltK is 1.19× / 1.23× geomean.
• 2 (hot) and 3 (streamed) of the 19 are more than 5% slower. The worst are 0.83× hot (n = 4, 2880×4096) and 0.74× streamed (n = 1, 18432×3072).

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[jinchen62]

@nakajee Switched to AssertSizeEqual: [{0: 1}] as suggested.

[codecov]

Codecov Report

❌ Patch coverage is 87.30159% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rary/src/amd_detail/rocblaslt/src/tensile_host.cpp	70.37%	4 Missing and 4 partials ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop   #12329      +/-   ##
===========================================
+ Coverage    71.93%   72.13%   +0.20%     
===========================================
  Files         2863     2863              
  Lines       472708   472459     -249     
  Branches     69829    69885      +56     
===========================================
+ Hits        340007   340782     +775     
+ Misses      109118   107910    -1208     
- Partials     23583    23767     +184     

Flag	Coverage Δ		*Carryforward flag
TensileLite-CPP	49.33% <ø> (ø)
TensileLite-Unit	83.29% <100.00%> (+0.01%)	⬆️
hipBLAS	90.62% <ø> (ø)		Carriedforward from 9e24e6d
hipBLASLt	38.96% <70.37%> (+0.04%)	⬆️
hipCUB	82.68% <ø> (ø)		Carriedforward from 9e24e6d
hipDNN	86.64% <ø> (+0.01%)	⬆️	Carriedforward from 9e24e6d
hipFFT	54.32% <ø> (+10.89%)	⬆️	Carriedforward from 9e24e6d
hipRAND	76.12% <ø> (ø)		Carriedforward from 9e24e6d
hipSOLVER	68.96% <ø> (ø)		Carriedforward from 9e24e6d
hipSPARSE	86.99% <ø> (ø)		Carriedforward from 9e24e6d
rocBLAS	48.25% <ø> (ø)		Carriedforward from 9e24e6d
rocFFT	48.78% <ø> (+1.16%)	⬆️	Carriedforward from 9e24e6d
rocRAND	57.45% <ø> (ø)		Carriedforward from 9e24e6d
rocSOLVER	76.83% <ø> (ø)		Carriedforward from 9e24e6d
rocSPARSE	74.75% <ø> (+0.12%)	⬆️	Carriedforward from 9e24e6d
rocThrust	93.15% <ø> (ø)		Carriedforward from 9e24e6d

*This pull request uses carry forward flags. Click here to find out more.

Files with missing lines	Coverage Δ
...s/hipblaslt/tensilelite/Tensile/AddCustomConfig.py	97.13% <100.00%> (+0.01%)	⬆️
...aslt/tensilelite/Tensile/Common/ValidParameters.py	100.00% <100.00%> (ø)
...ects/hipblaslt/tensilelite/Tensile/Contractions.py	88.03% <100.00%> (+0.29%)	⬆️
...cts/hipblaslt/tensilelite/Tensile/CustomKernels.py	97.42% <ø> (ø)
...rary/src/amd_detail/rocblaslt/src/tensile_host.cpp	43.65% <70.37%> (+0.17%)	⬆️

... and 81 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jinchen62]

m vs. n: which dimension is the token count

The skinny shapes in the GLM-5.2 decode GEMM list we were given are written in PyTorch's row-major naming, and the first revision of this PR used them with the wrong mapping. The two naming conventions map like this:

	GEMM list / PyTorch (row-major)	hipBLASLt (column-major)
tokens	M	n
output features (H)	N	m
input features	K	k

The list writes each GEMM as C[M, N] = A[M, K] @ B[K, N] and times it as torch.mm(A, B.t()), where B is the [N, K] weight. hipBLASLt is column-major, so it computes the transpose: Cᵀ = W · Aᵀ, i.e. (H, K) · (K, T) = (H, T). That is a TN GEMM with m = H, n = tokens and k = K. The (H, T) column-major output has the same memory layout as the (T, H) row-major output PyTorch expects. hipBLASLt's call log (HIPBLASLT_LOG_MASK=32) for torch.mm on the 2624×6144 shape with one token shows exactly that:

hipblaslt-bench --api_method c -m 2624 -n 1 -k 6144 --lda 6144 --ldb 6144 --ldc 2624 --ldd 2624 ... --transA T --transB N --a_type bf16_r --b_type bf16_r --c_type bf16_r --d_type bf16_r ... --compute_type f32_r

So a true decode is n = 1 in hipBLASLt terms. F.linear on a sliced activation (qkv[:, :6144] of a 6200-wide tensor) arrives the same way, with --ldb 6200.

Where "M = 1, 2, 4" came from. The hipblaslt-bench runs that came with the list fed its M straight into hipBLASLt's m, in an NN problem; their kernels are Cijk_Ailk_Bljk_…_MT1x4x256, an NN tile one row tall. With one token the two layouts happen to touch the same memory. With two or more tokens, NN with m = tokens expects transposed activations and output. Either way, hipBLASLt treats NN and TN as different problem types, and torch.mm never issues NN with m = tokens. So those benchmark numbers, and any kernels tuned from them, don't apply to the PyTorch call.

What changed in the PR. The original FP16 NN kernels (m = tokens) followed that NN mapping, and the list is BF16 anyway, so they are dropped. The PR now ships BF16 TN kernels for n ≤ 4 on gfx950, in the TN BF16 logic. That is the exact problem torch.mm(x, W.t()), F.linear and aiter's hipb_mm produce, and it's verified end to end through hipBLASLt on that call, including sliced activations. The kernel names keep rocBLAS's suffix, wvSpltK_bf16_tn_m1/m2/m4, where rocBLAS's M is the token count, i.e. hipBLASLt's n. The description's "Which frameworks reach these kernels" section covers when SGLang (with aiter) and vLLM actually send these GEMMs to hipBLASLt.

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-10-09 — Included in TheRock via ROCm/TheRock#8881.