Sync all amd_202602 problems with AMD-AIM upstream

Mark Saroufim · Mark Saroufim · commit 55af3b9bc963 · 2026-03-04T21:56:11.000-08:00
- moe-mxfp4: updated test cases and benchmark configs (PRs #11-#14) - mixed-mla: added tp parameter, updated reference/submission/task (PR #15) - mxfp4-mm: updated benchmark configs - Removed known non-determinism disclaimer (fixed by quantization changes)
diff --git a/problems/amd_202602/mixed-mla/reference.py b/problems/amd_202602/mixed-mla/reference.py
@@ -6,7 +6,7 @@
 output v_head_dim = kv_lora_rank = 512.
 
 The input provides:
-  q:       (total_q, 16, 576) bfloat16 — absorbed query
+  q:       (total_q, num_heads, 576) bfloat16 — absorbed query (num_heads = 128 // tp)
   kv_data: dict with KV cache in three formats:
     "bf16":  Tensor  (total_kv, 1, 576)  bfloat16          — highest precision
     "fp8":   (Tensor, Tensor)  kv_buffer fp8 + scalar scale — per-tensor quantized
@@ -37,7 +37,7 @@
 # DeepSeek R1 latent MQA constants (forward_absorb path)
 # https://huggingface.co/deepseek-ai/DeepSeek-R1-0528/blob/main/config.json
 # ---------------------------------------------------------------------------
-NUM_HEADS = 16
+TOTAL_NUM_HEADS = 128
 NUM_KV_HEADS = 1
 KV_LORA_RANK = 512
 QK_ROPE_HEAD_DIM = 64
@@ -285,17 +285,23 @@ def _aiter_mla_decode(
 # generate_input / ref_kernel / check_implementation
 # ---------------------------------------------------------------------------
 
-def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, seed: int) -> input_t:
+def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, tp: int, seed: int) -> input_t:
     """
     Generate absorbed q and compressed kv_buffer for MLA decode.
 
+    Args:
+        tp: tensor parallelism degree (4 or 8). num_heads = TOTAL_NUM_HEADS // tp.
+
     Returns all three KV cache formats in kv_data dict:
       kv_data = {
         "bf16":  Tensor               — (total_kv, 1, 576) bfloat16
         "fp8":   (Tensor, Tensor)     — kv_buffer fp8 + scalar scale
         "mxfp4": (Tensor, Tensor)     — kv_buffer fp4x2 + fp8_e8m0 scale
       }
     """
+    assert TOTAL_NUM_HEADS % tp == 0, f"TOTAL_NUM_HEADS ({TOTAL_NUM_HEADS}) must be divisible by tp ({tp})"
+    num_heads = TOTAL_NUM_HEADS // tp
+
     gen = torch.Generator(device="cuda")
     gen.manual_seed(seed)
 
@@ -304,7 +310,7 @@ def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, seed: int) -> in
 
     # Absorbed query: (total_q, num_heads, 576) bf16
     q = torch.randn(
-        (total_q, NUM_HEADS, QK_HEAD_DIM),
+        (total_q, num_heads, QK_HEAD_DIM),
         dtype=torch.bfloat16, device="cuda", generator=gen,
     ) * 0.02
 
@@ -332,7 +338,7 @@ def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, seed: int) -> in
 
     config = {
         "batch_size": batchsize,
-        "num_heads": NUM_HEADS,
+        "num_heads": num_heads,
         "num_kv_heads": NUM_KV_HEADS,
         "qk_head_dim": QK_HEAD_DIM,
         "kv_lora_rank": KV_LORA_RANK,
diff --git a/problems/amd_202602/mixed-mla/submission.py b/problems/amd_202602/mixed-mla/submission.py
@@ -4,7 +4,8 @@
 Implement custom_kernel() to beat the aiter a8w8 reference (fp8 Q + fp8 KV).
 
 DeepSeek R1 forward_absorb MLA config:
-  num_heads        = 16     (query heads, after TP split)
+  total_num_heads  = 128    (query heads before TP split)
+  num_heads        = 128 // tp  (query heads per device, tp=4 → 32, tp=8 → 16)
   num_kv_heads     = 1      (shared latent KV head)
   kv_lora_rank     = 512    (latent dim)
   qk_rope_head_dim = 64     (RoPE dim)
@@ -17,24 +18,24 @@
   - First 512 dims (kv_lora_rank) used as values (for output computation)
 
 Input tuple:
-  q:          (total_q, 16, 576)       bfloat16 — absorbed query
+  q:          (total_q, num_heads, 576) bfloat16 — absorbed query
   kv_data:    dict with three KV cache formats:
     kv_data["bf16"]  — Tensor (total_kv, 1, 576) bfloat16
     kv_data["fp8"]   — (Tensor, Tensor): kv_buffer fp8 (total_kv,1,576) + scalar scale
     kv_data["mxfp4"] — (Tensor, Tensor): kv_buffer fp4x2 (total_kv,1,288) + fp8_e8m0 scale
   qo_indptr:  (batch_size + 1,)        int32    — query segment pointers
   kv_indptr:  (batch_size + 1,)        int32    — KV segment pointers
-  config:     dict with MLA parameters
+  config:     dict with MLA parameters (includes num_heads computed from tp)
 
 Output:
-  attention output: (total_q, 16, 512) bfloat16
+  attention output: (total_q, num_heads, 512) bfloat16
 
 The reference uses aiter's a8w8 persistent MLA kernel (fp8 Q + fp8 KV),
 which is ~2-3x faster than bf16. To beat it, consider:
   1. Use mxfp4 KV cache for even lower memory bandwidth
      - Fuse dequantization with attention to avoid bf16 materialization
   2. Custom kernel with tighter memory access patterns
-  3. MQA: 1 KV head shared across 16 query heads — minimize redundant memory loads
+  3. MQA: 1 KV head shared across num_heads query heads — minimize redundant memory loads
   4. Variable-length batching: indptr-based segmented attention
   5. Split K/V from buffer: full 576 dims for keys, first 512 dims for values
 """
diff --git a/problems/amd_202602/mixed-mla/task.py b/problems/amd_202602/mixed-mla/task.py
@@ -5,13 +5,14 @@
 #
 # Input: (q, kv_data, qo_indptr, kv_indptr, config)
 #   q:          (total_q, num_heads, qk_head_dim)   bfloat16
+#               num_heads = 128 // tp (tp=4 → 32, tp=8 → 16)
 #   kv_data:    dict with three KV cache formats:
 #     "bf16":   Tensor  (total_kv, 1, 576)           bfloat16
 #     "fp8":    (Tensor, Tensor)  kv_buffer fp8 (total_kv, 1, 576) + scalar scale
 #     "mxfp4":  (Tensor, Tensor)  kv_buffer fp4x2 (total_kv, 1, 288) + fp8_e8m0 scale
 #   qo_indptr:  (batch_size + 1,)                    int32
 #   kv_indptr:  (batch_size + 1,)                    int32
-#   config:     dict with MLA parameters
+#   config:     dict with MLA parameters (includes num_heads computed from tp)
 #
 # where qk_head_dim = kv_lora_rank + qk_rope_head_dim = 512 + 64 = 576
 #
@@ -33,4 +34,5 @@ class TestSpec(TypedDict):
     batchsize: int
     qseqlen: int
     kvseqlen: int
+    tp: int
     seed: int
diff --git a/problems/amd_202602/mixed-mla/task.yml b/problems/amd_202602/mixed-mla/task.yml
@@ -20,39 +20,40 @@ description: |
   persistent mode), which is ~2-3x faster than bf16 on MI355X.
 
   DeepSeek R1 forward_absorb MLA config:
-    - num_heads = 16 (query heads, after TP split)
+    - total_num_heads = 128 (query heads before TP split)
+    - num_heads = 128 // tp (query heads per device, tp=4 → 32, tp=8 → 16)
     - num_kv_heads = 1 (shared latent KV head)
     - kv_lora_rank = 512
     - qk_rope_head_dim = 64
     - qk_head_dim = 576 (kv_lora_rank + qk_rope_head_dim, absorbed q/k dim)
     - v_head_dim = 512 (= kv_lora_rank, output dim)
     - sm_scale = 1/sqrt(576)
     - dtype: q=bfloat16
-    - decode only (q_seq_len=1, kv_seq_len up to 8k)
+    - q_seq_len = 1 or 4, kv_seq_len up to 8k
 
   KV buffer format (forward_absorb):
     - Full 576 dims are used as keys (for Q@K^T score computation)
     - First 512 dims (kv_lora_rank) are used as values (for output computation)
 
   Input tuple: (q, kv_data, qo_indptr, kv_indptr, config)
-    - q: (total_q, 16, 576) bfloat16 — absorbed query
+    - q: (total_q, num_heads, 576) bfloat16 — absorbed query
     - kv_data: dict with three KV cache formats:
         kv_data["bf16"]  — Tensor (total_kv, 1, 576) bfloat16
         kv_data["fp8"]   — (Tensor, Tensor): kv_buffer fp8 + scalar scale
         kv_data["mxfp4"] — (Tensor, Tensor): kv_buffer fp4x2 + fp8_e8m0 scale
     - qo_indptr: (batch_size+1,) int32 — query segment pointers
     - kv_indptr: (batch_size+1,) int32 — KV segment pointers
-    - config: dict with MLA parameters
+    - config: dict with MLA parameters (includes num_heads computed from tp)
 
   Return:
-    - attention output: (total_q, 16, 512) bfloat16
+    - attention output: (total_q, num_heads, 512) bfloat16
 
   Key optimization opportunities:
     1. Use mxfp4 KV cache for even lower memory bandwidth (4x savings over bf16)
        - Fuse dequantization with attention to skip bf16 materialization
     2. Custom kernel with tighter memory access patterns
-    3. MQA: 1 KV head shared across 16 query heads — minimize redundant memory loads
-    4. Decode: q_seq_len=1, kv_seq_len up to 8k — memory-bound workload
+    3. MQA: 1 KV head shared across num_heads query heads — minimize redundant memory loads
+    4. q_seq_len=1 or 4, kv_seq_len up to 8k — memory-bound workload
     5. Variable-length batching: indptr-based segmented attention
     6. Split K/V from buffer: full 576 dims for keys, first 512 dims for values
 
@@ -69,27 +70,29 @@ benchmark_timeout: 900
 ranked_timeout: 1200
 
 tests:
-  # bs=4
-  - {"batchsize": 4, "qseqlen": 1, "kvseqlen": 1024, "seed": 4220}
-  # bs=32
-  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 1024, "seed": 5412}
-  # bs=64
-  - {"batchsize": 64, "qseqlen": 1, "kvseqlen": 8192, "seed": 1360}
-  # bs=256
-  - {"batchsize": 256, "qseqlen": 1, "kvseqlen": 8192, "seed": 9826}
+  # bs=4, tp=8
+  - {"batchsize": 4, "qseqlen": 1, "kvseqlen": 1024, "tp": 8, "seed": 4220}
+  - {"batchsize": 4, "qseqlen": 4, "kvseqlen": 1024, "tp": 8, "seed": 4231}
+  # bs=32, tp=4
+  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 1024, "tp": 4, "seed": 5412}
+  - {"batchsize": 32, "qseqlen": 4, "kvseqlen": 8192, "tp": 4, "seed": 5423}
+  # bs=128, tp=8
+  - {"batchsize": 128, "qseqlen": 1, "kvseqlen": 8192, "tp": 8, "seed": 7816}
+  - {"batchsize": 128, "qseqlen": 4, "kvseqlen": 8192, "tp": 4, "seed": 7827}
 
 benchmarks:
-  # bs=4
-  - {"batchsize": 4, "qseqlen": 1, "kvseqlen": 1024, "seed": 4217}
-  - {"batchsize": 4, "qseqlen": 1, "kvseqlen": 8192, "seed": 4220}
-  # bs=32
-  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 1024, "seed": 5412}
-  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 8192, "seed": 5415}
-  # bs=64
-  - {"batchsize": 64, "qseqlen": 1, "kvseqlen": 1024, "seed": 1357}
-  - {"batchsize": 64, "qseqlen": 1, "kvseqlen": 8192, "seed": 1360}
-  # bs=256
-  - {"batchsize": 256, "qseqlen": 1, "kvseqlen": 1024, "seed": 9823}
-  - {"batchsize": 256, "qseqlen": 1, "kvseqlen": 8192, "seed": 9826}
+  # bs=4, tp=4
+  - {"batchsize": 4, "qseqlen": 1, "kvseqlen": 1024, "tp": 4, "seed": 4237}
+  - {"batchsize": 4, "qseqlen": 4, "kvseqlen": 8192, "tp": 4, "seed": 4251}
+  # bs=32, tp=8
+  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 8192, "tp": 8, "seed": 5415}
+  - {"batchsize": 32, "qseqlen": 4, "kvseqlen": 1024, "tp": 8, "seed": 5420}
+  # bs=32, tp=4
+  - {"batchsize": 32, "qseqlen": 1, "kvseqlen": 1024, "tp": 4, "seed": 5432}
+  - {"batchsize": 32, "qseqlen": 4, "kvseqlen": 8192, "tp": 4, "seed": 5443}
+  # bs=128, tp=8
+  - {"batchsize": 128, "qseqlen": 1, "kvseqlen": 8192, "tp": 8, "seed": 7816}
+  - {"batchsize": 128, "qseqlen": 4, "kvseqlen": 8192, "tp": 8, "seed": 7824}
+
 
 ranking_by: "geom"
diff --git a/problems/amd_202602/moe-mxfp4/task.yml b/problems/amd_202602/moe-mxfp4/task.yml
@@ -65,22 +65,18 @@ description: |
 
   d_hidden_pad and d_expert_pad are the dimensions padded to 256-alignment for the CK kernel.
 
-  **Known issue:** The reference submission (which calls aiter's fused_moe) is non-deterministic
-  on MI355X — it does not pass correctness checks against itself. This appears to be an aiter
-  fused_moe kernel bug on gfx950. Submissions will be evaluated on benchmark performance only
-  until this is resolved.
-
   The ranking criteria is the geometric mean of the benchmark results.
 
   ```
   The AITER reference performance is (E includes shared expert, top_k = routed + shared):
     bs     E  d_hidden  d_expert  top_k  time[us]
-     4   257      7168       256      9     46.9
-    64   257      7168       256      9    187.7
-   256   257      7168       256      9    245.7
-    64    33      7168      2048      9    220.6
-   256    33      7168      2048      9    276.4
-  1024    33      7168      2048      9    572.2
+    16   257      7168       256      9    152.7
+   128   257      7168       256      9    239.0
+   512   257      7168       256      9    336.5
+    16    33      7168       512      9    106.2
+   128    33      7168       512      9    141.1
+   512    33      7168       512      9    225.0
+   512    33      7168      2048      9    380.4
   ```
 
   Input:
@@ -112,16 +108,18 @@ ranked_timeout: 840
 ranking_by: "geom"
 
 tests:
-  - {"dhidden": 4096, "dexpert": 1024, "nroutedexperts": 16, "nexpertspertoken": 4, "nsharedexperts": 1, "bs": 8, "seed": 9371}
+  - {"dhidden": 4096, "dexpert": 1024, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 8, "seed": 9371}
   - {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 32, "seed": 2291}
   - {"dhidden": 4096, "dexpert": 1536, "nroutedexperts": 64, "nexpertspertoken": 6, "nsharedexperts": 1, "bs": 128, "seed": 81934}
 
 benchmarks:
-  # EP off (all 257 experts on 1 GPU): E=257, top_k=9 (8 routed + 1 shared)
-  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 4, "seed": 9371}
-  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 64, "seed": 2291}
-  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 256, "seed": 81934}
-  # EP on (EP=8, 33 experts per GPU): E=33, top_k=9 (8 routed + 1 shared)
-  - {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 64, "seed": 2291}
-  - {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 256, "seed": 81934}
-  - {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 1024, "seed": 81934}
+  # TP=8
+  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 16, "seed": 9371}
+  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 128, "seed": 2291}
+  - {"dhidden": 7168, "dexpert": 256, "nroutedexperts": 256, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 512, "seed": 81934}
+  # TP=4
+  - {"dhidden": 7168, "dexpert": 512, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 16, "seed": 2291}
+  - {"dhidden": 7168, "dexpert": 512, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 128, "seed": 81934}
+  - {"dhidden": 7168, "dexpert": 512, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 512, "seed": 81934}
+  # EP on
+  - {"dhidden": 7168, "dexpert": 2048, "nroutedexperts": 32, "nexpertspertoken": 8, "nsharedexperts": 1, "bs": 512, "seed": 81934}
diff --git a/problems/amd_202602/mxfp4-mm/task.yml b/problems/amd_202602/mxfp4-mm/task.yml
@@ -36,11 +36,9 @@ description: |
     4 2880   512  8.198
    16 2112  7168 20.873
    32 4096   512  9.462
+   32 2880   512  9.173
    64 7168  2048 12.738
-   64 2880   512  9.873
-  128 2112  7168 27.284
   256 3072  1536 12.219
-  256 7168  2048 13.506
   ```
 config:
   main: "eval.py"
@@ -55,8 +53,6 @@ benchmarks:
   - {"m": 4, "n": 2880, "k": 512, "seed": 4565}
   - {"m": 16, "n": 2112, "k": 7168, "seed": 15}
   - {"m": 32, "n": 4096, "k": 512, "seed": 457}
+  - {"m": 32, "n": 2880, "k": 512, "seed": 54}
   - {"m": 64, "n": 7168, "k": 2048, "seed": 687}
-  - {"m": 64, "n": 2880, "k": 512, "seed": 54}
-  - {"m": 128, "n": 2112, "k": 7168, "seed": 24}
   - {"m": 256, "n": 3072, "k": 1536, "seed": 7856}
-  - {"m": 256, "n": 7168, "k": 2048, "seed": 223}
