Fuse shared-length query scaling during inference

[autumnust]

Problem identified

LogScale and GatedLogScale allocate a device length scalar, prepare its logarithm and head coefficients, and materialize separate query-scaling intermediates on every attention call. Gated scaling adds separate tanh, addition and multiplication passes. This costs launches and GPU memory traffic during ordinary CUDA inference.

Current status

This draft is stacked on PR #1152, branch codex/autoresearch-cache-pinning at ffcf34ee1. That parent restores fitted-cache pinning; query scaling still uses the separate eager operations. All improvements below are incremental against that parent, with unchanged model, precision, workload and benchmark settings.

Solution

Fuse shared integer-length preparation, per-head conversion, query multiplication and optional gating into one optional Triton kernel. Preserve float32 length arithmetic and each eager dtype-rounding boundary; leave the gate network unchanged. Native PyTorch remains the fallback for gradients, tensor-valued lengths, noncontiguous inputs, other dtypes/backends, or unavailable Triton. No public API or model-state format changes.

The mechanism applies to ordinary shared-length attention calls, with no benchmark identifiers or exact-case shape dispatch. All 151 public API checks pass, including dtype/layout/length fallbacks, gradient equality and graph capture; all pre-commit hooks pass.

Quantifiable improvement

4.30% lower aggregate prediction latency across six cases, including 4.63% on the initial four-case suite. Full prediction arrays and official metric errors match exactly in all 22 case pairs. Three complete alternating pairs were run per case; MIC and Coffee have two additional pairs to examine timing variability. Every observation is retained.

Kumo-Tabular Large v1.0.1, 16 estimators, FP16 autocast, one RTX PRO 6000 Blackwell Server Edition GPU and two fixed CPUs. Measurements use the official validation → cached refit → formal prediction lifecycle and timer boundaries, without profiling. Times are separate medians; relative changes use median paired ratios. Ranges show observed spread, not confidence intervals. Negative reductions indicate slowdowns.

Case	Pairs	Prediction ms, parent → candidate	Reduction, median [range]	Fit s, parent → candidate	Peak GPU MiB, P / C	Peak host GiB, P / C
airfoil_self_noise	3	59.04 → 58.28	1.42% [-0.48–3.86%]	1.96 → 1.90	3829 / 3757	3.08 / 3.07
MIC	5	239.91 → 231.24	4.00% [2.18–16.50%]	8.38 → 8.09	4155 / 4089	16.31 / 16.31
coffee_rating_prediction	5	256.83 → 237.67	6.37% [-1.30–12.79%]	20.64 → 19.20	7661 / 7445	28.47 / 28.46
Bioresponse	3	2049.37 → 1913.32	6.64% [6.09–6.70%]	57.01 → 52.35	10209 / 9953	52.65 / 52.65
credit-g	3	68.42 → 67.14	2.18% [1.87–3.38%]	2.58 → 2.46	5587 / 5459	4.30 / 4.30
polish_companies_bankruptcy	3	464.93 → 442.23	5.05% [4.88–5.32%]	13.23 → 12.50	5719 / 5585	11.08 / 11.07

Bioresponse saves 136.1 ms between median prediction times. Airfoil's small change is mixed and close to noise; one Coffee candidate run is 1.30% slower. Neither case shows a reproducible regression above 5%. MIC/Coffee reference timings have wider spread, which is why their additional pairs are included. Credit-g and Polish Companies Bankruptcy were selected independently and withheld from development.

The GPU allowance is fixed at 81,569,985,331 bytes (75.97 GiB), with a 64 GiB host limit and no swap. All runs satisfy the resource policy and cleanup checks. GPU figures are sampled total device usage including driver-reserved memory; framework peaks are also retained. Bioresponse's maximum framework allocation falls from 4.982 to 4.743 GiB. These results do not establish a universal capacity limit or speedup on other hardware.

Reproduction

Final head: ba793705be3cbeb25a6600f4e4bacf4ea38d2c96. Initial-suite runs use 834cfc5ad; the final commit only guards CUDA-specific imports. CUDA execution bodies are identical, and final API/generalization checks run on the final head. PyTorch 2.14.1+cu130; TabArena a334f8e33d337a3c51df18ec3ab35c37d5e37256; checkpoint 3c3e10bbdb590ace29e7026847f92db3c603096d.

Evidence is retained on 4u8g-tur-0037 under /raid/leisun/repos/structured_dm/execution-notes/inference_efficiency_auto_research/experiments/r3-query-scaling-preparation/: verified-review.json, cuda-execution-equivalence.json, reproduction.md, verified source bundles, and runs 6/7/9. All headline comparisons use GPU 7 and its frozen policy; earlier GPU 6 screening is excluded.

From the recorded paired-run harness (40ff0ea), using the official protected harness (2b2051b) and qualified environment:

python -m measurement.kumo_tabular_arena.paired_runs \
  --baseline "$SDM_PARENT" --candidate "$SDM_CANDIDATE" \
  --harness "$OFFICIAL_HARNESS" --policy "$GPU7_POLICY" \
  --datasets airfoil_self_noise MIC coffee_rating_prediction Bioresponse \
  --pairs 3 --output "$FRESH_OUTPUT"

Repeat with --datasets MIC coffee_rating_prediction --pairs 2 and --datasets credit-g polish_companies_bankruptcy --pairs 3. Requalify hardware-specific limits on another host.

[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.

[autumnust]

I went through this PR and find it in general capturing the right problem, implementation seems clean enogh.
The solution itself wasn't that complicated and there's established pattern for triton-lang based kernel fusion in this code base, the performance bump, on top of the cache-pining one, is just marginal (I use 3% as the threshold to determine if a finding should survive)

@rusty1s @akihironitta what do you guys think here?

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/structured-data-models/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 8b196d05-3cee-4a0c-abac-fb783a29e6fa

📥 Commits

Reviewing files that changed from the base of the PR and between ffcf34e and ba79370.

📒 Files selected for processing: 4

• sdm/_kernels/log_scale.py
• sdm/_kernels/triton/log_scale.py
• sdm/nn/scaling.py
• test/nn/test_scaling.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

📝 Summary

Summary by CodeRabbit

• Improvements
  • Log-based query scaling now handles both standard and gated scaling through a unified path. Results continue to account for key length, per-head scale, and, when provided, the gate.
  • Added coverage for scaling outputs across supported dtypes, layouts, key-length inputs, and gradient-enabled use.

The change adds a shared log-scaling operation with eager PyTorch and Triton paths. LogScale and GatedLogScale use this operation. Tests cover outputs, gradients, dtypes, layouts, key-length inputs, and compilation.

Changes: Log scaling

Layer / File(s)	Summary
Kernel and dispatch sdm/_kernels/log_scale.py, sdm/_kernels/triton/log_scale.py	Adds eager and Triton implementations. The dispatcher uses Triton when its input, device, and gradient checks pass; otherwise, it uses eager PyTorch operations.
Scaling module integration and validation sdm/nn/scaling.py, test/nn/test_scaling.py	Routes LogScale and GatedLogScale through the shared operation. Adds reference comparisons for outputs and gradients, plus coverage for dtypes, layouts, key-length inputs, CUDA fallbacks, and compilation.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ScalingModule as LogScale or GatedLogScale
  participant Dispatcher as log_scale
  participant Triton as Triton log_scale
  participant Eager as Eager PyTorch
  ScalingModule->>Dispatcher: Pass query, key_len, head_scale, and optional gate
  alt Triton conditions pass
    Dispatcher->>Triton: Launch scaling kernel
    Triton-->>Dispatcher: Return scaled query
  else Triton conditions do not pass
    Dispatcher->>Eager: Apply clamped log scaling and optional gate
    Eager-->>Dispatcher: Return scaled query
  end
  Dispatcher-->>ScalingModule: Return scaled query

Loading

Merge Risk: ⚪ Minimal · up to ba793

No actionable merge-blocking risk was identified in the scaling change. Complete normal test checks before merging.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 21.43% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 14 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description check		The description clearly explains the fused Triton query-scaling implementation, fallback behavior, correctness requirements, and reported performance results.
Title check		The title directly and concisely identifies the main change: fusing shared-length query scaling for inference.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}