Keep the sparse Jensen-Shannon and Gram kernel epilogues in the input precision

[maxwbuckley]

Two float kernels that evaluate their epilogue in fp64 because of a double
literal or an int exponent. Same class of issue as #2531. This PR combines
what were #2535 and #2536, as requested in
#2535 (comment), so that the two
land in a single CI run and a single merge.

---

1. Sparse Jensen-Shannon post-processing (lp_distance.cuh)

Problem

The final step of the sparse Jensen-Shannon distance is a raft::linalg::map over the whole a_nrows x b_nrows output:

[=] __device__(value_t input) { return raft::sqrt(0.5 * input); }

0.5 is a double, so for value_t == float the multiply promotes the expression and raft::sqrt resolves to the fp64 overload.

Unlike a stray double in an add or a multiply, an fp64 sqrt is not one instruction — it is a Newton-Raphson refinement sequence. The SASS emitted for the float instantiation of the map kernel was:

DMUL        R14, R4, 0.5      <-- the literal
MUFU.RSQ64H R13, R15
DMUL        R16, R12, R12
DFMA        R16, R14, -R16, 1
DFMA        R18, R16, R18, 0.5
DMUL        R16, R12, R16
DFMA        R16, R18, R16, R12
DMUL        R18, R14, R16
DFMA        R22, R18, -R18, R14
DFMA        R4,  R22, R20, R18

Ten fp64 instructions per output element, each running at 1/64 the fp32 rate on consumer parts.

Fix

Write the constant as value_t(0.5) so overload resolution picks the float path. The same region becomes:

FMUL     R0, R4, 0.5
MUFU.RSQ R7, R0
FMUL.FTZ R5, R0, R7
FMUL.FTZ R7, R7, 0.5
FFMA     R0, -R5, R5, R0
FFMA     R5, R0, R7, R5

Numerics

0.5 is exactly representable in binary, so the double instantiation is bit-for-bit unchanged. The float instantiation now rounds once instead of twice, so individual results move by at most an ulp.

Measurements

RTX 5090 (sm_120a), CUDA 13.2, random CSR inputs, median of 5 runs, libcuvs.so swapped between the two builds and interleaved:

rows x cols, nnz/row	map kernel	speedup	full pairwise_distance	speedup
4096 x 4096, 32	263 -> 25 us	10.4x	6.40 -> 6.16 ms	1.04x
8192 x 4096, 32	1046 -> 319 us	3.28x	25.25 -> 24.54 ms	1.03x
16384 x 4096, 8	4168 -> 1396 us	2.99x	29.16 -> 26.38 ms	1.11x
16384 x 16384, 32	4168 -> 1394 us	2.99x	130.08 -> 127.32 ms	1.02x

Output checksums (mean over the full distance matrix) are identical to six decimal places in all four shapes.

At 16384 rows the map kernel now moves 2.1 GB in 1.39 ms, i.e. it sits at the memory roofline rather than being fp64-throughput-bound — which is why it stops at 3x. The 4096 case reaches 10.4x because its output fits in L2.

End-to-end gains are smaller because the balanced COO SpMV that produces the map's input is 85-95% of the call.

---

2. Polynomial and RBF Gram epilogues (kernel_matrices.cu)

polynomial_kernel, polynomial_kernel_nopad and rbf_kernel_expanded all
evaluate their epilogue in fp64 even when instantiated on float:

• pow(gain * x + offset, exponent) has exp_t == int, so overload
  resolution picks double pow(double, double) and promotes the whole
  expression.
• exp(-1.0 * gain * (...)) promotes on the -1.0 literal, and then
  resolves to the fp64 exp.

pow and exp are not single instructions; both expand to a polynomial
evaluation, so the promotion multiplies out. Before this change the float
instantiations contained 92 and 19 fp64 instructions respectively; after, they
contain none. The double instantiations are unchanged (98 and 22).

Casting the exponent to math_t and dropping the -1.0 selects the float
overloads without changing the double path.

Cost of the epilogue

Measured with KernelFactory::create(...)->evaluate() on random column-major
inputs, RTX 5090 (sm_120a), CUDA 13.2, median of 5. linear is the same GEMM
with no epilogue; tanh has the same epilogue structure but no double
literal, so both act as controls.

float:

  shape             kernel   before     after    speedup
  8192^2,  d=128    linear    0.316 ms   0.315 ms   1.00x   (control)
  8192^2,  d=128    poly      8.669 ms   0.652 ms  13.3x
  8192^2,  d=128    rbf       2.189 ms   0.746 ms   2.93x
  8192^2,  d=128    tanh      0.662 ms   0.664 ms   1.00x   (control)
  16384^2, d=32     linear    0.696 ms   0.695 ms   1.00x   (control)
  16384^2, d=32     poly     33.952 ms   2.112 ms  16.1x
  16384^2, d=32     rbf       7.831 ms   2.198 ms   3.56x
  16384^2, d=32     tanh      2.113 ms   2.112 ms   1.00x   (control)
  4096^2,  d=1024   poly      2.725 ms   0.703 ms   3.88x
  4096^2,  d=1024   rbf       1.160 ms   0.752 ms   1.54x

double is unchanged to within 0.05% on every shape.

Before this change the polynomial epilogue cost 24-43x the GEMM it decorates,
and the float path (33.95 ms) was within 1.4x of the double path
(47.56 ms). Afterwards poly, tanh and linear+epilogue all land within a few
percent of each other, i.e. the epilogue is memory-bound as it should be.

Effect on cuML's SVM

cuML's SVC calls these through its kernel cache. Built cuML 26.10 against this
branch (CPM_cuvs_SOURCE), blobs data as in cuML's own bench/sg/svc.cu,
median of 3, run to convergence. Solver iteration counts are reported so that a
changed convergence path is not mistaken for a speedup.

  float, converged        before      after     speedup   n_iter
  50000x1000  poly         9.43 ms    6.26 ms    1.51x    500 -> 500
  50000x1000  rbf         12.04 ms   10.44 ms    1.15x    1319 -> 1319
  50000x2     poly       808.60 ms  686.71 ms    1.18x    338727 -> 323526
  50000x2     rbf         89.42 ms   62.36 ms    1.43x    10436 -> 9219
  50000x1000  linear       4.50 ms    4.38 ms    1.03x    192 -> 192  (control)
  50000x1000  tanh        12.89 ms   12.59 ms    1.02x    771 -> 771  (control)

The two 50000x1000 rows converge in an identical number of iterations, so
1.51x and 1.15x are like-for-like. The 50000x2 rows changed iteration count;
per iteration they are 1.12x and 1.27x. The 2048x100000 shape is ~1.01x, since
at d=100000 the GEMM dominates. All double SVM configurations are unchanged.

Numerics

The float path now rounds once instead of going through fp64. Measured
against the previous fp64-then-round result over the value range these kernels
actually see (4M samples):

  polynomial degree=2   max relative error 1.19e-07   (1 ulp)
  polynomial degree=3   max relative error 1.19e-07   (1 ulp)
  polynomial degree=4   max relative error 1.19e-07   (1 ulp)
  rbf                   max relative error 1.67e-07   (2 ulp)

double results are bit-identical.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxwbuckley]

@lowener thank you :)

[lowener]

Can you merge PR #2536 with this one? It's the same kind of changes and that way we only have one run of CI and one merge.

[maxwbuckley]

@lowener done — #2536 is now folded into this PR and I've closed it.

The branch is three commits: the Jensen-Shannon raft::sqrt fix (lp_distance.cuh), the polynomial/RBF epilogue fix (kernel_matrices.cu), and one commit for the copyright headers on both files. The two source files are byte-identical to what was on the separate branches, so nothing changed beyond the combination. Description above now covers both.

Ready for a fresh /ok to test ba762fd363d19a8ac68fb824ce70fbc7f415bacc whenever you are.

[lowener]

LGTM

[lowener]

/ok to test ba762fd