[AMD] Skip x_scale.transpose+contiguous before bpreshuffle GEMM via upstream pre-transposed scale

[Jacob0226]

Depends on aiter PR ROCm/aiter#3041 — must be merged first.

Motivation

#23319 ([AMD] Use bpreshuffle FP8 blockscale GEMM to replace ABScale GEMM, merged 2026-04-23) introduced a per-call x_scale.transpose(-1, -2).contiguous().view(*x_scale.shape) on the bpreshuffle path of aiter_w8a8_block_fp8_linear (fp8_utils.py:784). On GLM-5.1-FP8 / MI355X / TP=8, this dispatches an at::native::elementwise_kernel_manual_unroll<direct_copy_kernel_cuda> kernel before EVERY bpreshuffle CK GEMM, costing ~4.165 µs per call.

Trace from DeepseekV2AttentionMLA:

ReplicatedLinear (qkv_a_proj)         direct_copy 4.25 + CK GEMM 9.85   →  2 kernels
ColumnParallelLinear (q_b_proj)       direct_copy 4.25 + CK GEMM 9.85   →  2 kernels
Indexer > ReplicatedLinear (wk)       direct_copy 4.25 + CK GEMM 9.85   →  2 kernels
RowParallelLinear (o_proj)            direct_copy 4.25 + CK GEMM 9.85   →  2 kernels
                                                                           =====
                                                                           4 redundant copies / layer

The aiter side already supports a transpose_scale=True flag on fused_rms_fp8_group_quant that writes the scale tensor directly in (num_bs_cols, M) column-major layout — exactly what the bpreshuffle GEMM consumes. SGLang upstream callers were just not using it. This PR opts them in and tags the returned scale tensor with a _aiter_bpreshuffle_layout=True attribute marker so the GEMM dispatch in fp8_utils.py can skip the redundant transpose+copy.

Modifications

The optimization is a producer/consumer pair connected by a tensor attribute marker _aiter_bpreshuffle_layout.

1. Consumer — fp8_utils.py::aiter_w8a8_block_fp8_linear gates the transpose+contiguous on the marker. If the marker matches the GEMM's expected layout (bpreshuffle ↔ transposed), skip the copy. Callers that do not set the marker get the pre-patch behavior (backward compatible).

2. Producer — updates seven call sites of fused_rms_fp8_group_quant (2 in communicator.py, 2 in forward_mla.py, 3 in forward_mha.py) to pass transpose_scale=True and set result_tuple[1]._aiter_bpreshuffle_layout = True on the returned scale tensor.

3. Producer (o_proj path) — forward_mla.py::forward_absorb_core calls fused_flatten_fp8_group_quant(transpose_scale=True) and tags the returned scale the same way. This requires the companion aiter PR [Triton] fused_flatten_fp8_group_quant: add transpose_scale param ROCm/aiter#3041.

Validation

Trace verification (MI355X TP=8 GLM-5.1-FP8 NSA TileLang decode, prof_in8192_out1024_conc4_p8 / TP-0-DECODE)

DeepseekV2AttentionMLA section — number of direct_copy kernels per layer:

Baseline:  4 / 4  GEMMs hit direct_copy
This PR:   0 / 4  (all four direct_copy kernels eliminated)

DeepseekV2AttentionMLA total time per layer:

Baseline:  ~177 µs
This PR:   157.10 µs   (–19.9 µs)

Accuracy Tests

GLM-5.1-FP8 launch cmd:

export SAFETENSORS_FAST_GPU=1
export SGLANG_ROCM_FUSED_DECODE_MLA=0
export ROCM_QUICK_REDUCE_QUANTIZATION=INT4
python3 -m sglang.launch_server \
  --model GLM-5.1-FP8 \
  --tp 8 --port 8552 \
  --tool-call-parser glm47 --reasoning-parser glm45 \
  --watchdog-timeout 1200 \
  --mem-fraction-static 0.85 \
  --kv-cache-dtype fp8_e4m3 --disable-radix-cache \
  --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 8}' \
  --nsa-prefill-backend tilelang --nsa-decode-backend tilelang

MI355X GSM8K (TP=8, 1200 questions, parallel=1200), measured against the same baseline as the Speed Tests section below:

Baseline:  0.948
This PR:   0.943

Within run-to-run variance (~±1%); consistent with the optimization being byte-equivalent at the GEMM level.

Speed Tests and Profiling

Baseline: sgl-project/sglang main + #23562 (preshuffled paged MQA + page_size=64, by @1am9trash, still open) + #24125 (cat-skip, by @Jacob0226, has 2 approvals). Both are NSA TileLang prerequisites for GLM-5.1-FP8 on MI355X — the model + GPU combo this PR targets. Docker image: rocm/sgl-dev:v0.5.10.post1-rocm720-mi35x-20260503. The optimization itself is logically independent of both #23562 and #24125; the numbers would shift slightly without them but the direction holds.

Bench cmd: sglang.bench_serving --dataset-name random --random-range-ratio 0.8 --random-input <8192|1024> --random-output 1024 --max-concurrency <C> --num-prompt 10*C --output-file /dev/null

vs baseline (10 ISL/OSL/concurrency points)

Workload	Cnc	Δ TPUT	Δ TPOT
in1024_out1024	4	+3.83%	-3.84%
in1024_out1024	8	+3.57%	-3.55%
in1024_out1024	16	+3.91%	-3.83%
in1024_out1024	32	+3.51%	-3.41%
in1024_out1024	64	+3.19%	-3.07%
in8192_out1024	4	+3.57%	-3.57%
in8192_out1024	8	+2.65%	-2.63%
in8192_out1024	16	+2.73%	-2.99%
in8192_out1024	32	+2.49%	-2.42%
in8192_out1024	64	+1.66%	-1.58%

• Headline: throughput +3.1% / TPOT -3.1% (faster) averaged across all 10 points.
• Best gain at small input + high concurrency (decode-bound).
• Smallest gain at large input + high concurrency (prefill-heavy).

Dependencies

• Aiter PR [Triton] fused_flatten_fp8_group_quant: add transpose_scale param ROCm/aiter#3041: required for the o_proj fused_flatten_fp8_group_quant opt-in. Without it, that call site will TypeError at runtime.

Checklist

• Format your code according to the Format code with pre-commit (pre-commit run --files <changed files>: all hooks pass).
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!