[Issue]: rmsnorm kernel issue in SGLang with aiter v0.1.10.post1

Problem Description

Dear AITER developers, I encountered a problem when using AITER v0.1.10.post1 in SGLang v0.5.8.

The new add_rmsnorm kernel introduced in AITER v0.1.10.post1 produces different outputs when running with chunked prefill enabled vs disabled with a vlm model lmms-lab/llava-onevision-qwen2-7b-ov. The outputs should be identical regardless of the chunked prefill setting.
The older rmsnorm2d_fwd_with_add_ck kernel produces consistent outputs in both modes.

add_rmsnorm

[2026-02-03 09:19:17] 
calling add_rmsnorm(out                          = torch.Size([7311, 3584]) torch.float16 cuda:0 0x7b1126200000,
                    input                        = torch.Size([7311, 3584]) torch.float16 cuda:0 0x7b112a000000,
                    residual_in                  = torch.Size([7311, 3584]) torch.float16 cuda:0 0x7b112de00000,
                    residual_out                 = torch.Size([7311, 3584]) torch.float16 cuda:0 0x7b111b400000,
                    weight                       = torch.Size([3584]) torch.float16 cuda:0 0x7b4372ca5000,
                    epsilon                      = 1e-06)
[2026-02-03 09:19:17] Prefill batch, #new-seq: 1, #new-token: 7311, #cached-token: 14, token usage: 0.00, #running-req: 0, #queue-req: 0, input throughput (token/s): 7618.93, cuda graph: False
[2026-02-03 09:19:17] Decode batch, #running-req: 1, #token: 7349, token usage: 0.00, cuda graph: True, gen throughput (token/s): 49.60, #queue-req: 0
[2026-02-03 09:19:17] INFO:     127.0.0.1:54368 - "POST /generate HTTP/1.1" 200 OK
Traceback (most recent call last):
  File "/sglang-checkout/python/sglang/srt/utils/common.py", line 2776, in retry
    return fn()
  File "/sglang-checkout/python/sglang/test/test_utils.py", line 2009, in <lambda>
    lambda: super(CustomTestCase, self)._callTestMethod(method),
  File "/usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
    method()
  File "/sglang-checkout/test/registered/vlm/test_vision_chunked_prefill.py", line 249, in test_chunked_prefill
    self._test_chunked_prefill(batches=[False, True], num_frames=[1, [2, 6, 8, 10]])
  File "/sglang-checkout/test/registered/vlm/test_vision_chunked_prefill.py", line 235, in _test_chunked_prefill
    self.assertEqual(len(output_chunked), len(output_no_chunked))
  File "/usr/lib/python3.10/unittest/case.py", line 845, in assertEqual
    assertion_func(first, second, msg=msg)
  File "/usr/lib/python3.10/unittest/case.py", line 838, in _baseAssertEqual
    raise self.failureException(msg)
AssertionError: 164 != 155
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size 1024 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size -1 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
output with chunked prefill:
The image is a still from a video, featuring a person standing on a stage with a dark background. The individual appears to be a male, dressed in a
output without chunked prefill:
The image depicts a scene from a presentation or event. In the foreground, there is a person standing on a stage. The individual is dressed in a black shirt
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size 1024 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size -1 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
output with chunked prefill:
The image depicts a person standing on a stage with a dark background. The individual is dressed in a black outfit and appears to be speaking or presenting. A large
output without chunked prefill:
The image is a still from a video, as indicated by the low resolution and the presence of a play button overlay in the center, suggesting that it is paused
FAILED

rmsnorm2d_fwd_with_add_ck

[2026-02-03 09:06:25] 
calling rmsnorm2d_fwd_with_add_ck(out                          = torch.Size([7311, 3584]) torch.float16 cuda:0 0x733a7e200000,
                                  input                        = torch.Size([7311, 3584]) torch.float16 cuda:0 0x733a82000000,
                                  residual_in                  = torch.Size([7311, 3584]) torch.float16 cuda:0 0x733a85e00000,
                                  residual_out                 = torch.Size([7311, 3584]) torch.float16 cuda:0 0x733a73400000,
                                  weight                       = torch.Size([3584]) torch.float16 cuda:0 0x736ccf2a5000,
                                  epsilon                      = 1e-06,
                                  use_model_sensitive_rmsnorm  = 0)
[2026-02-03 09:06:25] Prefill batch, #new-seq: 1, #new-token: 7311, #cached-token: 14, token usage: 0.00, #running-req: 0, #queue-req: 0, input throughput (token/s): 9663.87, cuda graph: False
[2026-02-03 09:06:25] Decode batch, #running-req: 1, #token: 7349, token usage: 0.00, cuda graph: True, gen throughput (token/s): 62.05, #queue-req: 0
[2026-02-03 09:06:25] INFO:     127.0.0.1:37080 - "POST /generate HTTP/1.1" 200 OK
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size 1024 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
CI_OFFLINE: Per-run marker found and current validation passed (requires_hf_quant_config=False), enabling offline mode - lmms-lab/llava-onevision-qwen2-7b-ov
command=python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chunked-prefill-size -1 --device cuda --host 127.0.0.1 --port 11000
CI_OFFLINE: Launching server HF_HUB_OFFLINE=1 model=lmms-lab/llava-onevision-qwen2-7b-ov
output with chunked prefill:
The video features a person standing on a stage with a dark background. The individual is dressed in a black outfit and appears to be speaking or presenting. The stage
output without chunked prefill:
The video features a person standing on a stage with a dark background. The individual is dressed in a black outfit and appears to be speaking or presenting. The stage
ROUGE-L scores: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
Average ROUGE-L score: 1.0000
output with chunked prefill:
['The video features a close-up shot of a person holding a small, rectangular electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up shot of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up shot of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt and']
output without chunked prefill:
['The video features a close-up shot of a person holding a small, rectangular electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up shot of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up shot of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt', 'The video features a close-up of a person holding a small, white electronic device, which appears to be an iPod. The individual is wearing a black shirt and']
ROUGE-L scores: [1.0, 1.0, 1.0, 1.0]
Average ROUGE-L score: 1.0000
PASSED

Reproduction

LOGFILE=sglang/log/test_vision_chunked_prefill_with_aiter_0.1.10.post1_debug_$(date -u +%Y%m%d_%H%M%S).log
{
  set -x
    export SGLANG_AMD_CI=1 
    export SGLANG_IS_IN_CI=1
    export SGLANG_IS_IN_CI_AMD=1
    export SGLANG_USE_AITER=1
    python3 -m pytest \
      test/registered/vlm/test_vision_chunked_prefill.py \
      -v -s
} 2>&1 | tee "$LOGFILE"

Operating System

Ubuntu 24.04 LTS

CPU

AMD EPYC 9655 96-Core Processor x 2

GPU

MI325X x 8

ROCm Version

ROCm700

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

Start by running test/registered/vlm/test_vision_chunked_prefill.py with AITER enabled and compare add_rmsnorm against rmsnorm2d_fwd_with_add_ck on the reported SGLang and model setup. Done means chunked and non-chunked prefill produce identical outputs and the test passes.