feat(qsa): Support pyhip qsa indexer for prefill/decode, sparse gqa for prefill

[apinge]

Motivation

This PR is based on tingqli/pyhip#57 . This PR introduces qsa indexer optimization for prefill/decode and sparse gqa optimization for prefill.

MTP draft extend and eager verify(no cuda graph verify) still use original sglang indexer. A follow-up PR will fix it.

Co-authored-by: Cheng Luo Cheng.Luo@amd.com

Modifications

1. introduce pyhip.ops.qsa.flydsl.attention see src/pyhip/ops/qsa/flydsl/readme.md in pyhip
2. introduce pyhip.ops.qsa.flydsl.indexer.prefill_indexer and pyhip.ops.qsa.flydsl.indexer.decode_indexer as qsa indexer.
3. for mtp, fix error overwriting for ring buffer by doubling the pending ring size
4. In prefix chunk-prefill, remove the .tolist() fence and move slot mapping into metadata.

Accuracy Tests

TP2

python3 -m sglang.test.run_eval --eval-name gsm8k --num-examples 1319 --max-tokens 16384 --port 9080
[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
[aiter] import [module_aiter_core] under /opt/aiter/aiter/jit/module_aiter_core.so
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None self.stop=None self.record_meta_info=True
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [12:12<00:00,  1.79it/s]
Total latency: 732.915 s
Score: 0.981
Output throughput: 872.229 token/s
[METRIC] gsm8k_score=0.9809741248097412 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
[METRIC] gsm8k_latency=732.9151995830471 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
====================
Speculative decoding: no per-request spec_accept_length in responses (non-speculative server, or --api completion which lacks return_meta_info).
====================
Writing report to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.html
{'score:std': np.float64(0.13661585289966702), 'score': np.float64(0.9809741248097412), 'latency': 732.9151995830471, 'output_throughput': 872.2291478791523}
Writing results to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.json

TP2+MTP

 python3 -m sglang.test.run_eval --eval-name gsm8k --num-examples 1319 --max-tokens 16384 --port 9080
[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
[aiter] import [module_aiter_core] under /opt/aiter/aiter/jit/module_aiter_core.so
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None self.stop=None self.record_meta_info=True
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [09:48<00:00,  2.23it/s]
Total latency: 588.718 s
Score: 0.979
Output throughput: 1055.547 token/s
[METRIC] gsm8k_score=0.9794520547945206 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
[METRIC] gsm8k_latency=588.7182018849999 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
====================
Speculative accept length (per-request, from meta_info): n=1314 mean=3.1367 min=2.5217 max=3.5776
====================
Writing report to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.html
{'score:std': np.float64(0.14186517244663), 'score': np.float64(0.9794520547945206), 'latency': 588.7182018849999, 'output_throughput': 1055.547455489389}
Writing results to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.json

TP4

python3 -m sglang.test.run_eval --eval-name gsm8k --num-examples 1319 --max-tokens 16384 --port 9080
[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
[aiter] import [module_aiter_core] under /opt/aiter/aiter/jit/module_aiter_core.so
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None self.stop=None self.record_meta_info=True
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [09:54<00:00,  2.21it/s]
Total latency: 594.771 s
Score: 0.974
Output throughput: 1102.404 token/s
[METRIC] gsm8k_score=0.974124809741248 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
[METRIC] gsm8k_latency=594.7711844240548 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
====================
Speculative decoding: no per-request spec_accept_length in responses (non-speculative server, or --api completion which lacks return_meta_info).
====================
Writing report to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.html
{'score:std': np.float64(0.1587629200658179), 'score': np.float64(0.974124809741248), 'latency': 594.7711844240548, 'output_throughput': 1102.4037767312552}
Writing results to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.json

TP4 MTP

python3 -m sglang.test.run_eval --eval-name gsm8k --num-examples 1319 --max-tokens 16384 --port 9080
[aiter] WARNING: NUMA balancing is enabled, which may cause errors. It is recommended to disable NUMA balancing by running "sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'" for more details: https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html#disable-numa-auto-balancing
[aiter] import [module_aiter_core] under /opt/aiter/aiter/jit/module_aiter_core.so
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None self.stop=None self.record_meta_info=True
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1314/1314 [09:08<00:00,  2.40it/s]
Total latency: 548.394 s
Score: 0.979
Output throughput: 1178.383 token/s
[METRIC] gsm8k_score=0.9794520547945206 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
[METRIC] gsm8k_latency=548.3940234169131 labels={"model": "Qwen/Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16", "eval": "gsm8k"}
====================
Speculative accept length (per-request, from meta_info): n=1314 mean=3.1386 min=2.4643 max=3.8705
====================
Writing report to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.html
{'score:std': np.float64(0.14186517244662997), 'score': np.float64(0.9794520547945206), 'latency': 548.3940234169131, 'output_throughput': 1178.3826453351348}
Writing results to /tmp/gsm8k_Qwen_Qwen3.8-Flash-Next-PTPC-FP8-PLE-BF16.json

Speed Tests and Profiling

All following profile and testing were run on TP2 without MTP.

E2E

Before

#Input tokens: 384000
#Output tokens: 11200
Starting warmup with 10 sequences...
Warmup completed with 10 sequences. Starting main benchmark run...
100%|██████████| 32/32 [02:41<00:00,  5.05s/it]

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 1         
Successful requests:                     32        
Benchmark duration (s):                  161.60    
Total input tokens:                      384000    
Total input text tokens:                 384000    
Total generated tokens:                  11200     
Total generated tokens (retokenized):    11179     
Request throughput (req/s):              0.20      
Input token throughput (tok/s):          2376.17   
Output token throughput (tok/s):         69.30     
Peak output token throughput (tok/s):    86.00     
Peak concurrent requests:                2         
Total token throughput (tok/s):          2445.48   
Concurrency:                             1.00      
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   5048.98   
Median E2E Latency (ms):                 5045.76   
P90 E2E Latency (ms):                    5073.75   
P95 E2E Latency (ms):                    5079.48   
P99 E2E Latency (ms):                    5087.47   
---------------Time to First Token----------------
Mean TTFT (ms):                          942.94    
Median TTFT (ms):                        942.97    
P90 TTFT (ms):                           954.86    
P95 TTFT (ms):                           957.88    
P99 TTFT (ms):                           959.24    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          11.77     
Median TPOT (ms):                        11.76     
P90 TPOT (ms):                           11.82     
P95 TPOT (ms):                           11.83     
P99 TPOT (ms):                           11.84     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           11.77     
Median ITL (ms):                         11.69     
P90 ITL (ms):                            11.75     
P95 ITL (ms):                            11.83     
P99 ITL (ms):                            13.94     
Max ITL (ms):                            30.25     
==================================================

After

#Input tokens: 384000
#Output tokens: 11200
Starting warmup with 10 sequences...
Warmup completed with 10 sequences. Starting main benchmark run...
100%|██████████| 32/32 [02:05<00:00,  3.93s/it]

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 1         
Successful requests:                     32        
Benchmark duration (s):                  125.84    
Total input tokens:                      384000    
Total input text tokens:                 384000    
Total generated tokens:                  11200     
Total generated tokens (retokenized):    11170     
Request throughput (req/s):              0.25      
Input token throughput (tok/s):          3051.45   
Output token throughput (tok/s):         89.00     
Peak output token throughput (tok/s):    112.00    
Peak concurrent requests:                2         
Total token throughput (tok/s):          3140.45   
Concurrency:                             1.00      
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3931.45   
Median E2E Latency (ms):                 3931.53   
P90 E2E Latency (ms):                    3953.20   
P95 E2E Latency (ms):                    3956.95   
P99 E2E Latency (ms):                    3960.41   
---------------Time to First Token----------------
Mean TTFT (ms):                          785.17    
Median TTFT (ms):                        784.95    
P90 TTFT (ms):                           792.47    
P95 TTFT (ms):                           793.73    
P99 TTFT (ms):                           807.14    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.02      
Median TPOT (ms):                        9.01      
P90 TPOT (ms):                           9.08      
P95 TPOT (ms):                           9.08      
P99 TPOT (ms):                           9.10      
---------------Inter-Token Latency----------------
Mean ITL (ms):                           9.02      
Median ITL (ms):                         8.96      
P90 ITL (ms):                            9.00      
P95 ITL (ms):                            9.05      
P99 ITL (ms):                            10.32     
Max ITL (ms):                            30.00     
==================================================

prefill indexer + sparse gqa

qsa prefill 12K in conc1
before	indexer	7ms 202µs
	others	618µs
	gqa	9ms 23µs
sum		16.841ms
after	indexer	297us
	data prepare+other	172us
	gqa(union+sparse)	2ms594us
sum		3.63ms

before

indexer

others

gqa

after

indexer

data prepare+other

gqa(union+sparse)

decode indexer

gqa decode 1conc 12K
before	indexer	318.865us
after	indexer	20.5 us

before

after

Checklist

• Format your code according to the Format code with pre-commit.
• 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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[Copilot]

🟡 Changes recommended

The pinned PyHIP revision lacks the QSA modules imported by the runtime and AMD test.

1 open finding

• Pin PyHIP to a revision containing the QSA API · New

What changed in this PR

Adds optional PyHIP QSA acceleration for gfx942 prefill/decode and improves sparse prefill handling.

Changes:

• Adds PyHIP QSA attention/indexer integration and validation.
• Expands QSA pending rings and optimizes packed prefill K/V gathering.
• Fixes fast-top-k overflow handling with AMD coverage.

File	Description
test/​registered/​kernels/​test_qsa.py	Tests verify-window ring compression.
test/​registered/​kernels/​test_pyhip_qsa.py	Tests PyHIP attention, indexer, and graphs.
test/​registered/​kernels/​ops/​elementwise/​test_fast_topk.py	Tests top-k overflow on AMD.
scripts/​launch_qwen38_flash_next_fp8_mi308x_pure_tp1.sh	Updates model and PyHIP settings.
scripts/​launch_qwen38_flash_next_fp8_mi308x_pure_tp_4_or_8_or_2.sh	Enables PyHIP QSA for multi-TP launches.
python/​sglang/​test/​qsa_validation.py	Adds numerical validation utilities.
python/​sglang/​srt/​mem_cache/​qsa_kv_pool.py	Enlarges per-request QSA rings.
python/​sglang/​srt/​layers/​attention/​qwen_sparse_attn_backend.py	Integrates PyHIP sparse prefill attention.
python/​sglang/​srt/​layers/​attention/​qsa/​qsa_indexer.py	Dispatches supported frames to PyHIP.
python/​sglang/​srt/​layers/​attention/​qsa/​pyhip_backend.py	Implements PyHIP adapters.
python/​sglang/​srt/​layers/​attention/​qsa/​metadata.py	Adds expanded ring addressing.
python/​sglang/​srt/​layers/​attention/​qsa/​graph_metadata.py	Updates graph ring metadata.
python/​sglang/​srt/​environ.py	Defines PyHIP QSA environment flags.
python/​sglang/​kernels/​jit/​csrc/​elementwise/​fast_topk.cuh	Adds full-row overflow fallback.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[yicyang]

Shall we add device check here ?

[apinge]

There is unnecessary to double check the device in PyHIPIndexer's __init__ . Because there is a check _supports_frame in sglang/python/sglang/srt/layers/attention/qsa/pyhip_backend.py.

[apinge]

Rethinking this, and you make a lot of sense. I used a macro instead of a device check here 336e193 and 0a9558e

[sammysun0711]

what happened if not self._supports_frame is True, forward return None, but pyhip QSA backend does not do anything?

[sammysun0711]

Should we update this PyHIPIndexer static method _supports_frame to can_use_for_forward? Move condition check in python/sglang/srt/layers/attention/qsa/qsa_indexer.py line 616:

if hidden_states.is_cuda and self._pyhip_indexer is not None:
        output = self._pyhip_indexer.forward(
            indexer=self,
            hidden=hidden_states,
            positions=positions,
            logical=logical_positions,
            batch=forward_batch,
            metadata=indexer_metadata,
            state_slots=state_slots,
        )
        if output is not None:
            return output

So that we could have a proper fallback if pyhip index not supported.

[apinge]

That makes some sense, but function-wise, it works as a fallback. can_use_for_forward is the first function in self._pyhip_indexer.forward if it returns None, the indexer can run the original sglang functions.

QSAIndexer.forward_cuda()
  └─ self._pyhip_indexer.forward(...)
       └─ PyHIPIndexer.can_use_for_forward(...)
            ├─ False → return None
            ├─ EXTEND → _prefill() → runtime.prefill_indexer(...)
            ├─ graph DECODE / TARGET_VERIFY
            │          → _decode() → runtime.decode_indexer(...)
            └─ other mode → return None

QSAIndexer.forward_cuda()
  └─ output is None
       → origin sglang project_qk / compress / MQA / top-k

[apinge]

I've added qsa indexer support for decode eager mode(c9d2ef4), and left MTP draft verify (not in cuda graph) and MTP draft extend for follow-up PRs.