… MI35x test

Previous commit defaulted the AITER MoE gate/up tile layout to SEPARATED
because that is what `Mxfp4MoEMethod` and the gptoss_fp4 tuned FlyDSL
kernels need. However other AITER MXFP4 paths still rely on the
INTERLEAVE layout that the pre-fix code produced, so flipping the global
default would silently regress them.

Restore INTERLEAVE as the default and have the gpt-oss path explicitly
opt out via `SGLANG_USE_AITER_MOE_GU_ITLV=0`:

- `environ.py`: change `SGLANG_USE_AITER_MOE_GU_ITLV` default from False
  to True; update the docstring to describe the new semantics.
- `aiter.py`: update the inline comment near the gate_mode selection to
  describe the new default and the opt-out semantics; logic unchanged.
- `test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py`: add
  `SGLANG_USE_AITER_MOE_GU_ITLV=0` to the env_vars of both MXFP4
  ModelConfigs (`openai/gpt-oss-20b`, `openai/gpt-oss-120b`) so the
  AITER MXFP4 + gpt-oss accuracy path keeps using the SEPARATED layout
  that matches the new Mxfp4MoEMethod shuffle. Other ROCm gpt-oss tests
  (`mi30x` BF16, NVIDIA CUDA paths) do not touch AITER MXFP4 and are
  left unchanged.

The MXFP4 MoE method pads hidden_size up to a multiple of 256
(2880 -> 3072 for gpt-oss-120b) via torch.nn.functional.pad inside
Mxfp4MoEMethod.apply, which fires a separate pad kernel per layer per
forward pass on top of the post-attention RMSNorm + residual add.

Fold the pad into the preceding RMSNorm by:

* RMSNorm gains an optional `x_pad_to_multiple` ctor arg that wires the
  aiter Triton `fused_add_rmsnorm_pad` kernel in `forward_aiter`. The
  fused kernel produces a (M, padded_N) output and an (M, N) residual_out
  in one launch.
* GptOssDecoderLayer constructs `post_attention_layernorm` with the
  MoE-input alignment (256 for MXFP4) only when env
  `SGLANG_FUSE_RMSNORM_PAD=1`, the aiter backend is active, the quant
  method is MXFP4, and TP=1 (the only attn-scatter mode where the
  pre-padded layernorm output flows through `_simple` without touching
  AllReduce/scatter helpers that haven't been taught the padded width).
* `Mxfp4MoEMethod.apply` skips the explicit pad when the input already
  arrives at `self.hidden_size`.
* `GptOssSparseMoeBlock.forward_normal` slices the (pre-padded) input
  back to the unpadded width for the router/topk GEMMs, and trims the
  experts output to the unpadded width before postprocess_layer (which
  pairs hidden_states with the unpadded residual).

Gated behind `SGLANG_FUSE_RMSNORM_PAD` (default off) so other models /
TP>1 / non-aiter setups keep the existing code path bit-for-bit.

Numerical: per-call RMSNorm output matches the baseline bit-exact at
decode batch sizes and within one bf16 ULP at 16k-token prefill
(reduction-order difference between the two Triton kernels). End-to-end
GSM8K 3-shot on gpt-oss-120b is within ±1 sigma of baseline:

  flexible-extract  baseline 0.8855 ± 0.0088  fused 0.8863 ± 0.0087
  strict-match      baseline 0.3654 ± 0.0133  fused 0.3525 ± 0.0132

Co-authored-by: Cursor <[email protected]>

End-to-end integration of Atom's pa_decode_gluon decode kernel into SGLang
via a SHUFFLE 5D-vectorized KV cache layout, opt-in behind the new env var
SGLANG_KV_CACHE_LAYOUT=vectorized_5d.

Pool layout (memory_pool.py, swa_memory_pool.py):
  K: (num_blocks, H_kv, D/X, page, X)
  V: (num_blocks, H_kv, page/X, D, X)        X = 16 / dtype_bytes (=8 bf16)
Both the full-attn and SWA sub-pools follow the env-driven layout; aiter
mha_batch_prefill_func and pa_decode_gluon consume this shape natively.

New writer (layers/attention/utils.py reshape_and_cache_shuffle_5d +
launch_reshape_and_cache_shuffle_5d): scatters per-token (num_tokens, H_kv, D)
inputs into the 5D pool slots.

aiter_backend wiring:
 - save_kv_cache routes through token_to_kv_pool.set_kv_buffer when the
   pool is 5D so the SHUFFLE writer is invoked (SWAKVPool dispatches per
   layer to the right inner pool with the proper full-to-swa loc translation).
 - mha_batch_prefill_func gains an extend_no_prefix shortcut: with
   --disable-radix-cache (no prefix reuse) we feed the kernel the raw
   incoming K/V in 3D linear mode, bypassing the 5D pool on the prefill
   read path.
 - forward_decode adds a 5D branch that calls pa_decode_gluon with Atom
   ps=True / get_recommended_splits convention. Full layers pass
   forward_metadata.kv_indices + seq_lens; SWA layers pass swa_page_table
   + seq_lens + sliding_window=sliding_window_size (ctx_part=128,
   max_part_num=1, mirroring Atom).

models/utils.py: disable enable_fused_set_kv_buffer when the pool is 5D.
The fused QKV+RoPE+set_kv writer .view()s the pool as 4D NHD, which
silently passes numel checks against a 5D buffer but corrupts the data.

GSM8K (lm_eval, 3-shot, chat template, gpt-oss-120b on MI355x TP=1):
  Plan A v2:    flexible-extract 0.8908 / strict-match 0.3381
  Baseline NHD: flexible-extract 0.8855 / strict-match 0.3647

Perf (8K/1K, c=128, 256 prompts):
  Plan A v2:    median TPOT 6.39 ms, output throughput 2210 tok/s
  Baseline NHD: median TPOT 7.06 ms, output throughput 2006 tok/s
  ~10% TPOT improvement; pa_decode_gluon kernel itself is ~46% faster
  than unified_attention on this workload.

Co-authored-by: Cursor <[email protected]>

Profile of Plan A v2 decode (8K/1K c=128) showed ~38ms of elementwise
overhead per 5 decode steps coming from two per-layer copies:

  1. MHATokenToKVPool.set_kv_buffer in vectorized_5d mode called
     .contiguous() on cache_k / cache_v. The writer kernel uses
     key.stride(0) directly for the source token stride and assumes only
     that head/dim are contiguous within each token (stride(1)=head_size,
     stride(2)=1). Both hold for K/V from a QKV split + RoPE in the
     upstream attention path even when the outer per-token stride is
     non-canonical, so the protective copies were no-ops semantically
     while still launching a large elementwise kernel per layer.

  2. forward_decode in the vectorized_5d branch allocated a private
     o_out = empty_like(q_in) for pa_decode_gluon and then did
     o.copy_(o_out.view_as(o)) per layer. Replaced by a zero-copy
     o.view(-1, num_q_heads, v_head_dim) passed directly as the kernel
     output.

Profile diff (8K/1K c=128, 5 decode steps):
  vectorized_elementwise:  25341 us -> 542 us   (-98%)
  manual_unroll:           13068 us -> 868 us   (-93%)

Benchmark (8K/1K c=128, 256 prompts):
  v2:           median TPOT 6.39 ms, 2210 tok/s
  v3 (this):    median TPOT 6.14 ms, 2300 tok/s
  NHD baseline: median TPOT 7.06 ms, 2007 tok/s
  -> v3 is 13% faster TPOT and 14.6% higher output throughput than
     baseline NHD.

GSM8K (lm_eval, 3-shot, chat template) holds:
  v3:           flexible 0.8878 / strict 0.3525
  v2:           flexible 0.8908 / strict 0.3381
  NHD baseline: flexible 0.8855 / strict 0.3647

Co-authored-by: Cursor <[email protected]>

… pool

The HIP fused rotary-embedding + KV-cache-write kernel
(fused_qk_rope_reshape_and_cache) already supports the SHUFFLE 5D cache
layout natively:

  flash_layout=True:  key_cache shape (T_cache, page, KH, D)        [NHD 4D]
  flash_layout=False: key_cache shape (T_cache, KH, D//x, page, x)  [5D]
                      value_cache supports the matching transposed
                      5D shape (T_cache, KH, page//x, D, x).

We had disabled enable_fused_set_kv_buffer for 5D pools because the
HIP-path arg builder reshaped k_buffer / v_buffer with the 4D NHD
.view(), which silently passes numel checks against a 5D buffer but
corrupts the data. With this commit:

 * enable_fused_set_kv_buffer no longer special-cases 5D — both HIP and
   CUDA paths can use the fused kernel for any pool layout.
 * create_fused_set_kv_buffer_arg detects k_buffer.ndim == 5 and passes
   the raw 5D buffer (no view) for the SHUFFLE layout; otherwise it
   keeps the legacy NHD 4D paged view.
 * forward_cuda in rotary_embedding/base.py detects key_cache.ndim==5
   and flips flash_layout accordingly.

Effect: per layer we now run ONE fused kernel instead of separate
rotary_embedding_kernel + reshape_and_cache_shuffle_5d launches.

Benchmark (8K/1K c=128, 256 prompts):
  v3 (no fused): median TPOT 6.139 ms, 2300 tok/s
  v4 (fused):    median TPOT 6.058 ms, 2307 tok/s
  NHD baseline:  median TPOT 7.055 ms, 2007 tok/s
  -> v4 is 14.1% faster TPOT and 15% higher output throughput than NHD.

GSM8K (lm_eval, 3-shot, chat template):
  v4:           flexible 0.8840 / strict 0.3541
  v3:           flexible 0.8878 / strict 0.3525
  NHD baseline: flexible 0.8855 / strict 0.3647

Co-authored-by: Cursor <[email protected]>

Two fixes to unblock fp8_e4m3 KV cache on the 5D vectorized layout, plus
a small UX default.

1) pa_decode_gluon was missing key_scale / value_scale on fp8 KV.
   The kernel reads fp8 bytes and, without per-tensor dequant scales,
   feeds the raw fp8 values straight into the dot product. The write
   side (fused_qk_rope_reshape_and_cache with apply_scale=True) was
   already correct, so this is purely a decode-side regression.

   Fix: when self.kv_cache_dtype == fp8_dtype, pass layer.k_scale /
   layer.v_scale (falling back to the backend's per-tensor scales) into
   pa_decode_gluon. Shape [1] per-tensor scales are supported natively.

2) MHATokenToKVPool 5D path asserted store_dtype == dtype, which blocked
   fp8 storage with bf16 compute. The X vectorization width is a
   property of the on-pool *storage* dtype (16 / store_dtype.itemsize),
   so compute the width from store_dtype and drop the assert. For bf16
   storage X stays 8, for fp8 storage X becomes 16.

3) server_args _handle_page_size: default page_size to 64 when
   SGLANG_KV_CACHE_LAYOUT=vectorized_5d (the layout the kernel is tuned
   for), so users don't have to pass --page-size 64 explicitly.

GSM8K (lm_eval, 3-shot, chat template, gpt-oss-120b, 1xMI300):
  NHD       + fp8  : flexible 0.8893 / strict 0.3381  (control)
  Plan A 5D + fp8  : flexible 0.8802 / strict 0.3306  (after fix)
  Plan A 5D + fp8  : ~0.01 / ~0.00                    (before fix)

Co-authored-by: Cursor <[email protected]>

The pre-Plan-A note said SWA layers stayed on the NHD unified_attention
path because the SWA sub-pool was forced to NHD. This hasn't been true
since SWAKVPool started honoring SGLANG_KV_CACHE_LAYOUT for both
sub-pools — SWA layers now go through pa_decode_gluon with
swa_page_table + sliding_window=layer.sliding_window_size, same as
full-attn but with a non-zero sliding_window arg. Update the comment so
it reflects what the code actually does.

Co-authored-by: Cursor <[email protected]>

forward_extend on the SHUFFLE 5D KV pool was crashing whenever a batch
arrived with non-zero extend_prefix_lens (chunked-prefill 2nd+ segment,
prefix-cache hit, multi-turn, …). The fallback aiter
mha_batch_prefill_func paged path it dropped into doesn't ship a
compiled kernel for (page_size=64, bf16/fp8, paged 5D) and aborts with
"no matching kernel found", which surfaces as a HSA_STATUS_ERROR_EXCEPTION
hardware exception. The repro is just:

  SGLANG_KV_CACHE_LAYOUT=vectorized_5d SGLANG_USE_AITER_MOE_GU_ITLV=False \
  python3 -m sglang.launch_server \
    --model-path /path/to/gpt-oss-120b/ --tp 1 --trust-remote-code \
    --mem-fraction-static 0.85 \
    --prefill-attention-backend aiter --decode-attention-backend aiter \
    --page-size 64 --disable-radix-cache --port 8000

  lm_eval --model local-chat-completions ... --tasks gsm8k --num_fewshot 3

(--disable-radix-cache is irrelevant — the scheduler still produces
non-first chunks during chunked prefill.)

Fix: for the 5D pool route the non-no-prefix case through a new
gather-and-linearize path that mirrors the SHUFFLE writer in reverse:

  * Add gather_shuffle_5d_to_linear / launch_gather_shuffle_5d_to_linear
    in layers/attention/utils.py. Same source addressing as the writer
    so it is bit-exact symmetric. Output is (T, num_heads, head_size)
    in the cache's store_dtype (uint8 bytes for fp8 pools, bf16 for
    bf16 pools).
  * In aiter_backend.forward_extend, when self.kv_cache_is_vectorized_5d
    take the per-token slot ids straight from
    forward_metadata.kv_indices[:seq_lens_sum] (or .swa_page_table for
    SWA layers — both are already per-token slot lists populated by
    create_flashinfer_kv_indices_triton, no page-table indirection),
    triton-gather K/V from the right sub-pool, then run
    mha_batch_prefill_func in 3D LINEAR mode (page_size=1) which does
    have a working kernel.
  * For fp8 store layers we reinterpret the gathered uint8 tensor as
    fp8 (zero-copy view), cast q to fp8, and pass q/k/v_descale to
    aiter — its LINEAR mha_batch_prefill_func supports fp8 K/V/Q
    natively so no host-side dequant is needed.

GSM8K (lm_eval, 3-shot, chat template, gpt-oss-120b, TP=1, MI300, the
exact repro command above, but with --kv-cache-dtype auto vs fp8_e4m3):
  bf16 KV: flexible 0.8863 / strict 0.3442
  fp8  KV: flexible 0.8848 / strict 0.3374
Matches the pre-A4 single-chunk baseline (0.8802 / 0.3306) within
seed noise, and the server no longer crashes on the 2nd chunk.

Co-authored-by: Cursor <[email protected]>

…e clear the flag purpose