enable minimax m3 fp8 index cache and fuse index_score_kernel with partial_topk kernel by ganyi1996ppo · Pull Request #1473 · ROCm/ATOM

ganyi1996ppo · 2026-07-06T07:30:36Z

Motivation

deps on aiter PR ROCm/aiter#4098
launch script:

export HIP_VISIBLE_DEVICES=0,1,2,3
export ATOM_FORCE_ATTN_TRITON=1

python -m atom.entrypoints.openai_server --model $main_model \
  -tp 4 --server-port 8014 --trust-remote-code --gpu-memory-utilization 0.85 --block-size 128 --no-enable_prefix_caching \
  --max-num-batched-tokens 16384 --max-model-len 16384\
  --kv_cache_dtype fp8 --index-cache-dtype fp8 \
  --torch-profiler-dir ./trace

accuracy:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9424|±  |0.0064|
|     |       |strict-match    |     5|exact_match|↑  |0.9431|±  |0.0064|

Perf
This PR

============ Serving Benchmark Result ============
Successful requests:                     256       
Failed requests:                         0         
Maximum request concurrency:             32        
Benchmark duration (s):                  138.58    
Total input tokens:                      2097152   
Total generated tokens:                  262144    
Request throughput (req/s):              1.85      
Output token throughput (tok/s):         1891.65   
Peak output token throughput (tok/s):    3392.00   
Peak concurrent requests:                64.00     
Total token throughput (tok/s):          17024.81  
---------------Time to First Token----------------
Mean TTFT (ms):                          3490.40   
Median TTFT (ms):                        3491.54   
P99 TTFT (ms):                           6514.43   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          13.52     
Median TPOT (ms):                        13.55     
P99 TPOT (ms):                           16.92     
---------------Inter-token Latency----------------
Mean ITL (ms):                           13.50     
Median ITL (ms):                         10.47     
P99 ITL (ms):                            12.96     
==================================================

before

============ Serving Benchmark Result ============
Successful requests:                     256       
Failed requests:                         0         
Maximum request concurrency:             32        
Benchmark duration (s):                  142.15    
Total input tokens:                      2097152   
Total generated tokens:                  262144    
Request throughput (req/s):              1.80      
Output token throughput (tok/s):         1844.15   
Peak output token throughput (tok/s):    5088.00   
Peak concurrent requests:                64.00     
Total token throughput (tok/s):          16597.33  
---------------Time to First Token----------------
Mean TTFT (ms):                          3500.17   
Median TTFT (ms):                        3498.67   
P99 TTFT (ms):                           6520.38   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          13.94     
Median TPOT (ms):                        13.91     
P99 TPOT (ms):                           17.34     
---------------Inter-token Latency----------------
Mean ITL (ms):                           13.93     
Median ITL (ms):                         10.98     
P99 ITL (ms):                            12.66     
==================================================

Technical Details

Test Plan

Test Result

Submission Checklist

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

Copilot

Pull request overview

This PR updates the MiniMax-M3 sparse attention/index-topk path to support FP8 for the index-key cache and to reduce decode-time overhead by fusing index-score computation with partial top-k selection in Triton.

Changes:

Enable FP8 dtype for the MiniMax-M3 index cache allocation/binding when KV cache is FP8.
Add FP8-aware dot-product handling in the index block-score kernel.
Fuse decode index-score + partial top-k into a new Triton kernel and adjust decode chunking parameters.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
atom/plugin/vllm/attention/minimax_m3_attnetion.py	Passes FP8 cache dtype through to the vLLM indexer-cache layer when KV cache dtype indicates FP8.
atom/model_ops/minimax_m3/index_topk.py	Adds FP8 handling in scoring and fuses decode scoring with partial top-k to remove the intermediate score tensor.
atom/model_ops/attentions/aiter_attention.py	Allocates the sparse attention index cache in FP8 when the configured KV cache dtype indicates FP8.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

        self.index_cache_layer = MiniMaxM3SparseIndexerCache(
            layer_name=f"{self.layer_name}.index_cache",
            head_dim=index_head_dim,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=cache_dtype if str(cache_dtype).startswith("fp8") else "auto",
        )


+            index_cache_dtype = (
+                dtypes.d_dtypes[config.kv_cache_dtype]
+                if str(config.kv_cache_dtype).startswith("fp8")
+                else config.torch_dtype
+            )


Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

        self.index_cache_layer = MiniMaxM3SparseIndexerCache(
            layer_name=f"{self.layer_name}.index_cache",
            head_dim=index_head_dim,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=(
+                cache_dtype if str(cache_dtype).startswith("fp8") else "auto"
+            ),
        )


Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

            kv_cache_dtype=kv_cache_dtype,
+            index_cache_dtype=self.index_cache_dtype,
            k_scale=fused_k_scale,
            v_scale=fused_v_scale,
            asm_layout=True,


Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Signed-off-by: ganyi <ygan@amd.com>

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

        self.index_cache_layer = MiniMaxM3SparseIndexerCache(
            layer_name=f"{self.layer_name}.index_cache",
            head_dim=index_head_dim,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=(
+                cache_dtype if str(cache_dtype).startswith("fp8") else "auto"


valarLip · 2026-07-15T09:01:20Z

    return bool(getattr(impl, "is_indexed_sparse_attention", False))


+def _resolve_minimax_m3_index_cache_dtype(config) -> torch.dtype:


change to _resolve_index_cache_dtype

Signed-off-by: ganyi <ygan@amd.com>

Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

+            kv_cache_dtype=(
+                cache_dtype if str(cache_dtype).startswith("fp8") else "auto"
+            ),


Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

    for blk in tl.range(chunk_start_block, chunk_end_block):
        page = tl.load(bt_row + blk).to(tl.int64)
-        pos = blk * BLOCK_SIZE_K + off_k
-        pos_mask = pos < causal_len
        k = tl.load(


valarLip · 2026-07-16T15:23:10Z

+            type=str,
+            default="auto",
+            help=(
+                "MiniMax-M3 index cache type. 'auto' uses the model dtype; "


"MiniMax-M3 index cache type"?
it's better default =kv_cache_dtype?

Signed-off-by: ganyi <ygan@amd.com>

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

        factors.append(vllm_factors)
        factors.append(self.tensor_parallel_size)
        factors.append(self.enable_dp_attention)
+        factors.append(self.index_cache_dtype)


        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        cache_config: str = "bf16",
+        index_cache_config: str = "auto",


        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        cache_config: str = "bf16",
+        index_cache_config: str = "auto",


        config: PretrainedConfig,
        prefix: str,
        cache_config: str = "bf16",
+        index_cache_config: str = "auto",


Signed-off-by: ganyi <ygan@amd.com>

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

        self.index_cache_layer = MiniMaxM3SparseIndexerCache(
            layer_name=f"{self.layer_name}.index_cache",
            head_dim=index_head_dim,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=(
+                cache_dtype if str(cache_dtype).startswith("fp8") else "auto"
+            ),
        )


        cache_config = atom_config.kv_cache_dtype
+        index_cache_config = atom_config.index_cache_dtype
        quant_config = atom_config.quant_config


+    _decode_index_score_topk_partial_kernel[(total_q, num_idx_heads, num_topk_chunks)](
+        idx_q,
+        index_kv_cache,
        topk_score_partial,
        topk_idx_partial,


Signed-off-by: ganyi <ygan@amd.com>

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.

        factors.append(vllm_factors)
        factors.append(self.tensor_parallel_size)
        factors.append(self.enable_dp_attention)
+        factors.append(self.index_cache_dtype)


        self.index_cache_layer = MiniMaxM3SparseIndexerCache(
            layer_name=f"{self.layer_name}.index_cache",
            head_dim=index_head_dim,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=(
+                cache_dtype if str(cache_dtype).startswith("fp8") else "auto"
+            ),
        )


Copilot AI review requested due to automatic review settings July 6, 2026 07:30

ganyi1996ppo marked this pull request as draft July 6, 2026 07:30

Copilot started reviewing on behalf of ganyi1996ppo July 6, 2026 07:31 View session

Copilot AI reviewed Jul 6, 2026

View reviewed changes

ganyi1996ppo marked this pull request as ready for review July 7, 2026 02:15

zufayu requested a review from amd-ruitang3 July 7, 2026 02:59

Copilot AI review requested due to automatic review settings July 7, 2026 04:53

Copilot started reviewing on behalf of ganyi1996ppo July 7, 2026 04:53 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings July 7, 2026 04:58

valarLip previously approved these changes Jul 7, 2026

View reviewed changes

Copilot started reviewing on behalf of ganyi1996ppo July 7, 2026 04:59 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Comment thread atom/model_ops/attention_mha.py

Comment on lines 1168 to 1172

kv_cache_dtype=kv_cache_dtype,

index_cache_dtype=self.index_cache_dtype,

k_scale=fused_k_scale,

v_scale=fused_v_scale,

asm_layout=True,

Copilot AI review requested due to automatic review settings July 13, 2026 11:29

ganyi1996ppo force-pushed the ganyi/fuse_score_topk branch from 97326c9 to 2d2bef9 Compare July 13, 2026 11:29

Copilot started reviewing on behalf of ganyi1996ppo July 13, 2026 11:30 View session

Copilot AI reviewed Jul 13, 2026

View reviewed changes

zejunchen-zejun mentioned this pull request Jul 15, 2026

[atom][atom-vllm] optimize MTP for GLM 5.2 #1531

Merged

ganyi1996ppo added 6 commits July 15, 2026 14:42

fuse score topk

9049bdb

Signed-off-by: ganyi <ygan@amd.com>

fp8 index cache

44b930c

Signed-off-by: ganyi <ygan@amd.com>

fp8 index score calculate in fp8

291ed54

Signed-off-by: ganyi <ygan@amd.com>

black

95512d9

Signed-off-by: ganyi <ygan@amd.com>

add index_cache_dtype as parameter

e432f54

Signed-off-by: ganyi <ygan@amd.com>

add fp8 kv cache and fp8 index cache to recipe

fab5a7a

Signed-off-by: ganyi <ygan@amd.com>

Copilot AI review requested due to automatic review settings July 15, 2026 06:42

ganyi1996ppo force-pushed the ganyi/fuse_score_topk branch from 2d2bef9 to fab5a7a Compare July 15, 2026 06:42

Copilot started reviewing on behalf of ganyi1996ppo July 15, 2026 06:42 View session

Copilot AI reviewed Jul 15, 2026

View reviewed changes

valarLip reviewed Jul 15, 2026

View reviewed changes

Copilot started work on behalf of ganyi1996ppo July 15, 2026 09:42 View session

resolve comments

ba2030a

Signed-off-by: ganyi <ygan@amd.com>

Copilot AI review requested due to automatic review settings July 15, 2026 09:43

ganyi1996ppo dismissed valarLip’s stale review via ba2030a July 15, 2026 09:43

Copilot started reviewing on behalf of ganyi1996ppo July 15, 2026 09:43 View session

Fix SGLang readiness check to validate expected model

e57b297

Copilot finished work on behalf of ganyi1996ppo July 15, 2026 09:49

Copilot AI reviewed Jul 15, 2026

View reviewed changes

Comment thread atom/plugin/vllm/attention/minimax_m3_attnetion.py

Comment on lines +239 to +241

kv_cache_dtype=(

cache_dtype if str(cache_dtype).startswith("fp8") else "auto"

),

Copilot AI review requested due to automatic review settings July 15, 2026 09:50

Copilot started reviewing on behalf of ganyi1996ppo July 15, 2026 09:51 View session

Copilot AI reviewed Jul 15, 2026

View reviewed changes

Comment thread atom/model_ops/minimax_m3/index_topk.py

Comment on lines 424 to 426

for blk in tl.range(chunk_start_block, chunk_end_block):

page = tl.load(bt_row + blk).to(tl.int64)

pos = blk * BLOCK_SIZE_K + off_k

pos_mask = pos < causal_len

k = tl.load(

valarLip reviewed Jul 16, 2026

View reviewed changes

make index-cache-dtype default as same as kv-cache-dtype

d745f45

Signed-off-by: ganyi <ygan@amd.com>

Copilot AI review requested due to automatic review settings July 17, 2026 02:01

Copilot started reviewing on behalf of ganyi1996ppo July 17, 2026 02:02 View session

Copilot AI reviewed Jul 17, 2026

View reviewed changes

convert auto for bf16 type

6bc1923

Signed-off-by: ganyi <ygan@amd.com>

Copilot AI review requested due to automatic review settings July 17, 2026 02:08

Copilot started reviewing on behalf of ganyi1996ppo July 17, 2026 02:08 View session

Copilot AI reviewed Jul 17, 2026

View reviewed changes

black

4bccc90

Signed-off-by: ganyi <ygan@amd.com>

Copilot AI review requested due to automatic review settings July 17, 2026 02:18

Copilot started reviewing on behalf of ganyi1996ppo July 17, 2026 02:19 View session

Copilot AI reviewed Jul 17, 2026

View reviewed changes

valarLip approved these changes Jul 17, 2026

View reviewed changes

valarLip merged commit c101aa0 into main Jul 17, 2026
51 of 58 checks passed

valarLip deleted the ganyi/fuse_score_topk branch July 17, 2026 06:49

		return bool(getattr(impl, "is_indexed_sparse_attention", False))


		def _resolve_minimax_m3_index_cache_dtype(config) -> torch.dtype:

Conversation

ganyi1996ppo commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

valarLip Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

ganyi1996ppo Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

valarLip Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

ganyi1996ppo Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ganyi1996ppo commented Jul 6, 2026 •

edited

Loading