[PERF][kernels] Accumulate portable linear_logp backward dW with FP32 addmm_

[hongleng]

Scope request only. I will not start implementation until a maintainer confirms
that this optimization is worth pursuing, agrees with the boundary below, and
clarifies coordination with #208.

Current production path

The shared portable, non-tensor-parallel linear_logp backward is reachable from
the DeepSpeed training path:

DeepSpeedTrainingWorker.train
  -> _extract_logps
  -> kernel_registry.get_op("linear_logp")
  -> TritonLinearLogpOp / CUDA fallback
  -> _LinearLogpFunction.backward
  -> chunked_linear_logp_backward

Relevant code:

• rl_engine/executors/deepspeed_trainer.py
• rl_engine/kernels/ops/triton/loss/linear_logp.py
• rl_engine/kernels/ops/pytorch/loss/linear_logp.py

This helper is used by the portable CUDA/ROCm path and by the non-TP SM90
fallback. The normal SM90 fused backward extension has its own implementation
and does not use this helper.

Current accumulation

chunked_linear_logp_backward initializes an FP32 gradient table and, for every
chunk, materializes a GEMM result before adding it to that table:

grad_w = torch.zeros((vocab, hidden_dim), dtype=torch.float32, device=weight.device)

for chunk in chunks:
    grad_w += torch.matmul(left, right).float()

The proposed change preserves the FP32 zero-initialization but lets GEMM
accumulate directly into grad_w when both operands are already FP32 and
ambient autocast is disabled:

if left.dtype == right.dtype == torch.float32 and not device_autocast_enabled:
    grad_w.addmm_(left, right)
else:
    grad_w += torch.matmul(left, right).float()

The existing path remains the fallback for autocast-enabled or non-FP32
operands. dW operands, including hidden.float(), will only be constructed when
the weight gradient is requested.

No public API, forward computation, registry, TP metadata, or collective changes
are proposed.

Static traffic model

Let:

C = floor(2^24 / V)
K = ceil(N / C)

Considering only gradient-table traffic:

current   = 4*V*D + 16*K*V*D bytes
candidate = 4*V*D +  8*K*V*D bytes
saving    = 8*K*V*D bytes

For N=4096, D=2048, V=32768, default chunking gives K=8. The proposed
accumulation therefore removes eight framework-visible aten::add_ device
operations and reduces logical gradient-table traffic by 4.0 GiB.

These are static logical quantities. They are not claims of eight fewer GPU
kernels, 4 GiB lower measured peak VRAM, or 4 GiB lower physical HBM traffic.

Local feasibility evidence

A throwaway prototype compared the current and proposed accumulation on an RTX
3050 Ti / SM86 with PyTorch 2.12.1+cu126.

For forced four-chunk cases across BF16, FP16, and FP32:

• complete backward helper: 1.076–1.175x
• isolated dW: 1.139–1.235x
• no stable peak-memory improvement

The prototype also covers:

• bias and no-bias;
• N=0;
• hidden-only, weight-only, and bias-only gradients;
• zero upstream gradients and signed-zero comparison;
• same-input repeat determinism;
• ambient autocast fallback;
• higher-order autograd smoke.

The isolated dW result is diagnostic only. It is not proposed as the performance
acceptance result.

Correctness boundary

The proposal deliberately retains grad_w=zeros(...). An earlier experiment
that assigned the first chunk directly changed some zero gradients from +0 to
-0, so that version is excluded.

The intended contract is:

• current-vs-candidate agreement under the repository's existing tolerance;
• same-input repeat determinism;
• preserved N=0 behavior;
• zero-upstream numerical zeros with matching signbit;
• autocast-enabled execution retains the current accumulation path;
• no extra dW casts or allocations when the weight gradient is not requested.

Universal bitwise equality for arbitrary random inputs is not required: FP32
GEMM accumulation may differ by approximately one ULP.

The device-specific autocast check must remain compatible with the repository's
minimum supported PyTorch version (2.4.1).

Proposed target validation

Before implementation is promoted to a PR, I will compare base and candidate on
the maintainer-selected authoritative portable target.

The headline measurement will be portable operator forward+backward latency,
not isolated dW or helper-only latency. Validation will include:

• representative default-chunking RL shapes;
• paired base-vs-candidate measurements;
• bias and no-bias;
• BF16, plus supported FP16/FP32 checks;
• a 1/2/4/8-chunk diagnostic sweep;
• incremental peak allocation and profiler/device-op evidence;
• the normal SM90 fused path as an unchanged control when testing H100.

I propose proceeding only if the portable operator shows a stable, repeatable
positive gain without a material regression on the other measured production
shapes. I am not proposing a single fixed percentage requirement for every
shape before target data is available.

Expected implementation scope

Only:

1. rl_engine/kernels/ops/pytorch/loss/linear_logp.py
2. focused coverage in tests/test_linear_logp.py

The existing benchmarks/benchmark_linear_logp.py will be reused for base/head
measurements. No new operator, configuration flag, benchmark file, or user-facing
documentation is planned unless requested during review.

Non-goals and overlap

• No tensor-parallel backward or collective changes; [Feat][Kernel] Add TP Fused linear logp Triton #208 owns that surface.
• No implementation work will start until [Feat][Kernel] Add TP Fused linear logp Triton #208 is merged or a maintainer
  explicitly confirms that parallel work is appropriate.
• No SM90 fused-extension changes.
• No forward, API, registry, trainer, or dispatch changes.

A current issue/PR search found no equivalent non-TP addmm_ accumulation work.

Maintainer decision requested

1. Is this portable/non-TP backward optimization worth pursuing within the
   two-file scope above?
2. Which portable hardware target should be authoritative for the base-vs-head
   performance decision: A100, H100 forced fallback, a ROCm target, or another
   GPU-CI target?
3. Should implementation wait for [Feat][Kernel] Add TP Fused linear logp Triton #208 to merge, or is parallel work acceptable?