[AMD] Drop the redundant scale zero-fill before AITER per-tensor FP8 quant

[siliangchen-amd]

[AMD] Drop the redundant scale zero-fill before AITER per-tensor FP8 quant

Motivation

dynamic_per_tensor_quant_fp8 zero-fills its one-element scale buffer before
handing it to the backend:

scale = torch.zeros(1, device=input.device, dtype=torch.float32)
if _use_aiter:
    dynamic_per_tensor_quant(output, input, scale)
elif _has_vllm:
    torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
else:
    _native_dynamic_per_tensor_quant_fp8(output, input, scale)

Only one of those three backends reads the incoming value.

backend	how scale is written	needs the pre-fill
vLLM	segmented_max_reduction atomicMax-es into it	yes
AITER	launches initializeScale(scale, 1, 0.0f) first	no
native	ends in scale.view(-1).copy_(...), a full overwrite	no

For AITER, csrc/kernels/quant_kernels.cu
dynamic_per_tensor_quant opens with its own fill in both the fp8 and int8
branches:

aiter::initializeScale<<<dim3(1), dim3(64), 0, stream>>>(
    reinterpret_cast<float*>(scale.data_ptr()), 1, 0.0f);
aiter::data_to_scale_kernel<input_dtype, opus::fp8_t><<<grid, block, 0, stream>>>(
    reinterpret_cast<float*>(scale.data_ptr()), ..., cols);

So on the AITER path the buffer is zeroed twice per call, once by torch and
once by the kernel that was going to do it anyway. The torch one is a real
vectorized_elementwise_kernel<FillFunctor<float>> launch to write a single
float, and it happens on every quantized GEMM input — four per decoder layer,
so 128 per decode step on a 32-layer model, replayed out of the captured decode
CUDA graph for the life of the server.

Change

Move the allocation inside each branch and zero only where a backend depends on
it. The vLLM branch is untouched in behaviour and keeps torch.zeros; AITER and
the native fallback take torch.empty.

Correctness

Reading the source is not enough, since the shipped AITER is a binary wheel, so
the installed artifact was tested adversarially: pre-fill scale with values
that would visibly corrupt a max-reduction, then compare against the
zero-filled baseline.

shape 4096x4096  baseline scale = 2.666667e-01
   pre-fill zeros (baseline)   -> identical
   pre-fill +inf               -> identical
   pre-fill 1e30               -> identical
   pre-fill NaN                -> identical
   pre-fill -inf               -> identical
   pre-fill -1e30              -> identical

Bit-identical quantized output and identical scale for every poison value, on
all three shapes tested (4096x4096, 1024x5120, 17x4096). If the kernel ever
read the incoming buffer, +inf and NaN would not survive a max-reduction
unchanged.

The profiler shows the launch that disappears, and the one that stays:

torch.zeros (current)                        torch.empty (patched)
  10  aten::fill_                              10  aiter::dynamic_per_tensor_quant
  10  vectorized_elementwise_kernel<           10  aiter::initializeScale
        FillFunctor<float>>                    10  aiter::data_to_scale_kernel
  10  aiter::dynamic_per_tensor_quant          10  aiter::scaled_quant_kernel
  10  aiter::initializeScale
  10  aiter::data_to_scale_kernel
  10  aiter::scaled_quant_kernel

aiter::initializeScale is present on both sides — the zeroing still happens,
just once.

Results

8x MI325X (gfx942), sglang main, SGLANG_USE_AITER=1, --quantization fp8,
sglang.bench_serving random 1024/1024, concurrency 64, 320 prompts, TP 8
(TP 4 where the head count does not divide by 8). Three reps per variant
against one server boot; the table reports the median of reps 2-3, with rep 1
dropped as server warm-up on both sides.

model	arch	size	zeros (tok/s)	empty (tok/s)	delta	mean TPOT
Llama-3.1-8B-Instruct	Llama-3.1	8B	8772	9593	+9.36%	6.80 -> 6.21 ms
Qwen3-8B	Qwen3	8B	8106	8427	+3.96%	7.41 -> 7.11 ms
Qwen3-32B	Qwen3	32B	4345	4469	+2.85%	13.55 -> 13.09 ms
Mistral-7B-Instruct-v0.3	Mistral	7B	9522	9734	+2.22%	6.29 -> 6.14 ms
Qwen3-30B-A3B	Qwen3 MoE	30B-A3B	5925	6007	+1.39%	10.14 -> 9.98 ms
NeuralPizza-7B-V0.1	Mistral	7B	9363	9939	+6.14%	6.38 -> 6.00 ms
DeepSeek-R1-Distill-Llama-8B	Llama-3.1	8B	8851	9238	+4.37%	6.75 -> 6.42 ms
WeirdCompound-v1.5-24b	Mistral	24B	6846	7062	+3.15%	8.54 -> 8.32 ms
Human-Like-Qwen2.5-7B-Instruct	Qwen2	7B	10208	10446	+2.33%	5.75 -> 5.59 ms
Olmo-3-7B-Instruct	Olmo3	7B	4191	4269	+1.86%	14.54 -> 14.29 ms

Ten of ten improve, mean +3.76%, median +3.00%, spread +1.39% to +9.36%. The
top five rows are stock upstream checkpoints across four architecture families
and two sparsity regimes; the rest are community fine-tunes that widen the
architecture coverage.

Mean TPOT falls on every one of the ten, which is the signal the mechanism
predicts: the removed work is per decode step, not per request. The dense
8B models gain most, consistent with a fixed per-step launch cost amortised
against a smaller step.

Scope and limits

• Only the AITER and native branches change. The vLLM path keeps its zero-fill
  because segmented_max_reduction accumulates into the buffer, and it is
  unexercised by the numbers above (all runs are SGLANG_USE_AITER=1).
• Measured on gfx942 only. The change is not arch-specific, but the numbers are.
• The per-token path (dynamic_per_token_quant) is untouched.
• A microbenchmark of the call site in isolation prices the removed kernel at
  roughly 3.7 us of device time against a ~30 us call, but its totals were too
  noisy to quote as the headline; the serving numbers above are the claim.

---

CI States

Latest PR Test (Base): ⏳ Run #35695900578
Latest PR Test (Extra): ⏳ Run #35695900183
Latest PR Test (AMD ROCm 10): ⏳ Run #35695900526

[BBuf]

Can you clean up the comments or delete the comments directly?

[siliangchen-amd]

Thanks for the review @BBuf — trimmed. The block is down to three lines covering the one thing the code cannot show: why the vLLM branch has to keep its pre-zero (its segmented_max_reduction accumulates with atomicMax, so dropping the fill there would be a silent correctness bug, not just slower). Everything else — the AITER initializeScale reference, the adversarial +inf/NaN check against the shipped wheel, and the per-model numbers — now lives only in the PR description.

The diff is +6/-1 on one file.

Could you also add the run-ci label when you get a chance?

[BBuf]

Thanks for the review @BBuf — trimmed. The block is down to three lines covering the one thing the code cannot show: why the vLLM branch has to keep its pre-zero (its segmented_max_reduction accumulates with atomicMax, so dropping the fill there would be a silent correctness bug, not just slower). Everything else — the AITER initializeScale reference, the adversarial +inf/NaN check against the shipped wheel, and the per-model numbers — now lives only in the PR description.

The diff is +6/-1 on one file.

Could you also add the run-ci label when you get a chance?

Done