Skip to content

Geometric policy blending on cuda backend - #2440

Open
Kovax007 wants to merge 2 commits into
LeelaChessZero:masterfrom
Kovax007:gpu-policy-blend
Open

Geometric policy blending on cuda backend#2440
Kovax007 wants to merge 2 commits into
LeelaChessZero:masterfrom
Kovax007:gpu-policy-blend

Conversation

@Kovax007

@Kovax007 Kovax007 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Geometric blending means a weighted sum of the heads' logits, taken before the softmax. It returns a single policy vector that nothing downstream can distinguish from an unblended one. PolicyTemperature is applied identically to every head, so it factors out of the mixture entirely and keeps its usual meaning.

Weights are set as backend options, and the heads that carry weight are decided once when the backend is built:

--backend-opts=policy_blend_optimistic=0.55
--backend-opts=policy_blend_optimistic=0.55,policy_blend_soft=0.1

The two also exist as pro-only UCI options (BackendPolicyBlendOptimistic, BackendPolicyBlendSoft) so a tuner can move them as ordinary parameters. They are hidden to pro options as the changes only can apply once the backend constructed. Weights are renormalised onto the simplex if they sum past one, so the blend stays a weighted mean rather than becoming a temperature. A network that lacks the requested head refuses to load, rather than silently not blending.

With no weight set, nothing changes: no extra head is built, no extra buffer is allocated, and the priors are bit-identical to master. That was checked against a build of this branch point over 100 positions, at defaults and under a tuned configuration, with identical bestmoves and no output of any kind. Cost ~0.5% fewer nodes per second with a 256x10 net when enabled. It iss the extra head's own projections plus the mixing kernel.

Tested at: https://bench-direct.lczero.org/test/1153/ Whether blending helps, and at what weights, depends on the net this is a capability rather than a general strength gain. On the net tested it was worth 20.2 +/- 3.1 Elo over 10000 games at 10+0.1 against the same net unblended.

Cuda backends only for now. The same graph would be three Mul and two Add nodes in the onnx converter, which would cover the onnx and xla backends and their execution providers without a kernel, if this turns out to be worth having.

Developed with the help of Claude agents for research, implementation and verification; all design decisions, testing and review are my own.

Geometric blending means a weighted sum of the heads' logits, taken before
the softmax. It  returns a single policy vector that nothing downstream can distinguish
from an unblended one. PolicyTemperature is applied identically to every head, so it
factors out of the mixture entirely and keeps its usual meaning.

Weights are set as backend options, and the heads that carry weight are decided
once when the backend is built:

    --backend-opts=policy_blend_optimistic=0.55
    --backend-opts=policy_blend_optimistic=0.55,policy_blend_soft=0.1

The two also exist as pro-only UCI options (`BackendPolicyBlendOptimistic`,
`BackendPolicyBlendSoft`) so a tuner can move them as ordinary parameters.
They are hidden to pro options as the changes only can apply once the backend
constructed. Weights are renormalised onto the simplex if they sum past one,
so the blend stays a weighted mean rather than becoming a temperature.
A network that lacks the requested head refuses to load, rather than silently not blending.

With no weight set, nothing changes: no extra head is built, no extra buffer is
allocated, and the priors are bit-identical to master. That was checked against
a build of this branch point over 100 positions, at defaults and under a tuned
configuration, with identical bestmoves and no output of any kind.
Cost ~0.5% fewer nodes per second with a 256x10 net when enabled.
It iss the extra head's own projections plus the mixing kernel.

Tested at: https://bench-direct.lczero.org/test/1153/
Whether blending helps, and at what weights, depends on the net this is a capability rather than a general strength gain. On the net tested it was worth 20.2 +/- 3.1 Elo over 10000 games at 10+0.1 against the same net unblended.

Cuda backends only for now. The same graph would be three Mul and two Add nodes
in the onnx converter, which would cover the onnx and xla backends and their
execution providers without a kernel, if this turns out to be worth having.
@Kovax007
Kovax007 requested review from Menkib64 and borg323 and a lite review from Copilot and removed request for Menkib64 and borg323 August 9, 2026 18:23

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds CUDA-backend support for geometric (logit-space, pre-softmax) blending of multiple policy heads into a single policy output, configurable via backend options and exposed as pro-only UCI parameters.

Changes:

  • Introduces pro-only UCI options for policy-head blend weights and threads them into backend construction/config hashing.
  • Extends CUDA network construction/execution to optionally build/evaluate extra policy heads and blend their logits on-GPU before downloading the policy output.
  • Adds a CUDA kernel for blending policy logits and allocates per-head device buffers only when blending is enabled.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
src/neural/wrapper.cc Forwards pro-only UCI blend weights into backend options used to build the network.
src/neural/shared_params.h Declares new shared backend OptionIds for blend weights.
src/neural/shared_params.cc Defines/Registers new pro-only UCI options for blend weights.
src/neural/backend.cc Includes blend weights in backend configuration hashing.
src/neural/backends/cuda/network_cuda.cc Implements extra-head construction/eval and logit-space blending inside CUDA backend.
src/neural/backends/cuda/layers.h Extends AttentionPolicyHead API to support shared policy embedding.
src/neural/backends/cuda/layers.cc Implements shared-embedding behavior in AttentionPolicyHead construction/eval.
src/neural/backends/cuda/kernels.h Declares blendPolicyLogits kernel API.
src/neural/backends/cuda/inputs_outputs.h Adds extra policy-head enums/flags and allocates device buffers for extra heads when needed.
src/neural/backends/cuda/common_kernels.cu Implements blendPolicyLogits CUDA kernel and instantiations.
Suppressed comments (1)

src/neural/backends/cuda/network_cuda.cc:648

  • These comment lines are incomplete sentences (missing an object after “recomputes its own …”), which makes the description hard to follow. Please complete the wording so the option’s behavior is unambiguous.
      // Reusing the selected head's policy embedding is what makes an extra
      // head cheap. Turned off, every extra head recomputes its own
      // Exposed so the two can be compared, and as a way out if a
      // future net trips the detection below.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/neural/wrapper.cc
Comment on lines +189 to +198
const float blend_optimistic =
options.Get<float>(SharedBackendParams::kPolicyBlendOptimistic);
if (blend_optimistic > 0.0f) {
network_options.Set<float>("policy_blend_optimistic", blend_optimistic);
}
const float blend_soft =
options.Get<float>(SharedBackendParams::kPolicyBlendSoft);
if (blend_soft > 0.0f) {
network_options.Set<float>("policy_blend_soft", blend_soft);
}
Comment thread src/neural/backends/cuda/network_cuda.cc
Comment on lines +1439 to +1443
if (shared_embedding_) {
// Another head owns the embedding weights and computes the embedding.
assert(weights.pol_encoder.empty());
ip_pol_w_ = nullptr;
ip_pol_b_ = nullptr;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cudaFree(nullprt) on the nullptr assignment is what prevents the double free

Comment thread src/neural/backends/cuda/network_cuda.cc
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (3)

src/neural/wrapper.cc:197

  • Checking > 0 discards an explicitly configured zero. If backend-opts supplied a nonzero blend, setoption ... value 0 therefore cannot disable it because the original backend option remains in network_options. Use IsDefault to distinguish an untouched UCI default from an explicit zero.
  if (blend_optimistic > 0.0f) {
    network_options.Set<float>("policy_blend_optimistic", blend_optimistic);
  }
  const float blend_soft =
      options.Get<float>(SharedBackendParams::kPolicyBlendSoft);
  if (blend_soft > 0.0f) {
    network_options.Set<float>("policy_blend_soft", blend_soft);

src/neural/wrapper.cc:190

  • These values are consumed only while Create builds the CUDA network, but NetworkAsBackend::UpdateConfiguration still returns UPDATE_OK when either UCI blend option changes. Consequently, setoption updates after construction leave the existing heads, weights, and captured graphs unchanged, so tuner moves do not actually affect evaluations. Persist the constructed blend values and return NEED_RESTART when either changes.
  const float blend_optimistic =
      options.Get<float>(SharedBackendParams::kPolicyBlendOptimistic);

src/neural/backends/cuda/network_cuda.cc:478

  • This sizes scratch memory from every extra head present in the file, even when both blend weights are zero and no extra head is built. A larger unused head can therefore increase scratch_size_ (and the three tensor allocations derived from it) on the default path. Restrict this calculation to heads with nonzero blend weight so disabled blending has no memory overhead.
    for (const char* extra_head : kExtraPolicyHeadNames) {
      if (!weights.policy_heads.contains(extra_head)) continue;
      attentionPolicySize = std::max(
          attentionPolicySize,
          getMaxAttentionHeadSize(weights.policy_heads.at(extra_head),
                                  max_batch_size_) *
              sizeof(DataType));

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants