[Feat] Support input_rmsnorm_quant & qknorm_quant fused for Atom dpsk v2 model

[qichu-yun]

Motivation

Enable DeepSeek V2 FP4 input_rmsnorm_quant and qknorm_quant fusion in ATOM across both Triton and non-Triton GEMM paths. This reduces separate RMSNorm + quant kernel launches while keeping FP4 activation scale layouts and cached KV projection paths correct.

Technical Details

• Enable qknorm quant fusion for FP8/FP4 without requiring Triton GEMM.
• Route FP4 qknorm and input RMSNorm quant through fused RMSNorm+MXFP4 quant kernels.
• Select the correct MXFP4 activation/scale layout for preshuffle Triton, non-shuffle Triton, and ASM GEMM paths.
• Add FP4-safe cached KV projection gather fallback for MLA attention.
• Align the SGLang MLA prepare path with the native ATOM qknorm quant fusion behavior.

Test Plan

server:

#!/bin/bash
export ATOM_USE_TRITON_MOE=1
export ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1
export ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION=1

rm -rf /root/.cache/ /tmp/torchinductor_root
model_path=/shared/data/amd_int/models/deepseek-ai/DeepSeek-R1-0528-MXFP4-v2
python -m atom.entrypoints.openai_server \
    --model $model_path \
    --host localhost \
    --server-port 8000 \
    --tensor-parallel-size 8 \
    --kv_cache_dtype fp8 \
    --gpu-memory-utilization 0.8 \
    --no-enable_prefix_caching

client:

#!/bin/bash
set -euo pipefail
 
addr=localhost
port=8000
url=http://${addr}:${port}/v1/completions
 
model_path="/shared/data/amd_int/models/deepseek-ai/DeepSeek-R1-0528-MXFP4-v2"
num_concurrent="${LM_EVAL_CONCURRENT:-16}"
max_gen_toks="${LM_EVAL_MAX_GEN_TOKS:-512}"


lm_eval --model local-completions \
    --model_args "{\"base_url\": \"${url}\", \"model\": \"${model_path}\", \"num_concurrent\": ${num_concurrent}, \"max_retries\": 1, \"max_gen_toks\": ${max_gen_toks}, \"tokenized_requests\": false}" \
    --tasks gsm8k \
    --batch_size auto \
    --num_fewshot 5 \
    --trust_remote_code \

Test Result

For preshuffle FP4 gemm kernel:

For non-preshuffle FP4 gemm kernel:
export ATOM_USE_FP4_NON_SHUFFLE_TRITON_GEMM=1

Accuracy:

Backend	Model	NON_SHUFFLE=0	NON_SHUFFLE=1	Delta
atom	DeepSeek-R1-0528-MXFP4-v2	0.9439	0.9469	+0.0030
atom	DeepSeek-R1-0528-MXFP4-MTP-MoEFP4	0.9507	0.9553	+0.0046
atom	DeepSeek-R1-0528	0.9629	0.9629	+0.0000

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[valarLip]

…

[qichu-yun]

…

have deleted it....