Revert mixed-mla to original config (no tp parameter)

The AMD-AIM tp parameter changes (num_heads=128//tp, qseqlen=4)
made the test cases ~16x slower due to more heads and larger queries.
Revert to the original 16-head decode-only config that runs within
the test timeout.

Author: Mark Saroufim

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, num_heads, 576) bfloat16 — absorbed query (num_heads = 128 // tp)
+  q:       (total_q, 16, 576) bfloat16 — absorbed query
   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
 # ---------------------------------------------------------------------------
-TOTAL_NUM_HEADS = 128
+NUM_HEADS = 16
 NUM_KV_HEADS = 1
 KV_LORA_RANK = 512
 QK_ROPE_HEAD_DIM = 64
@@ -285,23 +285,17 @@ def _aiter_mla_decode(
 # generate_input / ref_kernel / check_implementation
 # ---------------------------------------------------------------------------
 
-def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, tp: int, seed: int) -> input_t:
+def generate_input(batchsize: int, qseqlen: int, kvseqlen: 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)
 
@@ -310,7 +304,7 @@ def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, tp: int, seed: i
 
     # 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
 
@@ -338,7 +332,7 @@ def generate_input(batchsize: int, qseqlen: int, kvseqlen: int, tp: int, seed: i
 
     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,

problems/amd_202602/mixed-mla/submission.py

@@ -4,8 +4,7 @@
 Implement custom_kernel() to beat the aiter a8w8 reference (fp8 Q + fp8 KV).
 
 DeepSeek R1 forward_absorb MLA config:
-  total_num_heads  = 128    (query heads before TP split)
-  num_heads        = 128 // tp  (query heads per device, tp=4 → 32, tp=8 → 16)
+  num_heads        = 16     (query heads, after TP split)
   num_kv_heads     = 1      (shared latent KV head)
   kv_lora_rank     = 512    (latent dim)
   qk_rope_head_dim = 64     (RoPE dim)
@@ -18,24 +17,24 @@
   - First 512 dims (kv_lora_rank) used as values (for output computation)
 
 Input tuple:
-  q:          (total_q, num_heads, 576) bfloat16 — absorbed query
+  q:          (total_q, 16, 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 (includes num_heads computed from tp)
+  config:     dict with MLA parameters
 
 Output:
-  attention output: (total_q, num_heads, 512) bfloat16
+  attention output: (total_q, 16, 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 num_heads query heads — minimize redundant memory loads
+  3. MQA: 1 KV head shared across 16 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
 """

problems/amd_202602/mixed-mla/task.py

@@ -5,14 +5,13 @@
 #
 # 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 (includes num_heads computed from tp)
+#   config:     dict with MLA parameters
 #
 # where qk_head_dim = kv_lora_rank + qk_rope_head_dim = 512 + 64 = 576
 #
@@ -34,5 +33,4 @@ class TestSpec(TypedDict):
     batchsize: int
     qseqlen: int
     kvseqlen: int
-    tp: int
     seed: int

problems/amd_202602/mixed-mla/task.yml

@@ -20,40 +20,39 @@ description: |
   persistent mode), which is ~2-3x faster than bf16 on MI355X.
 
   DeepSeek R1 forward_absorb MLA config:
-    - total_num_heads = 128 (query heads before TP split)
-    - num_heads = 128 // tp (query heads per device, tp=4 → 32, tp=8 → 16)
+    - num_heads = 16 (query heads, after TP split)
     - 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
-    - q_seq_len = 1 or 4, kv_seq_len up to 8k
+    - decode only (q_seq_len=1, 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, num_heads, 576) bfloat16 — absorbed query
+    - q: (total_q, 16, 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 (includes num_heads computed from tp)
+    - config: dict with MLA parameters
 
   Return:
-    - attention output: (total_q, num_heads, 512) bfloat16
+    - attention output: (total_q, 16, 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 num_heads query heads — minimize redundant memory loads
-    4. q_seq_len=1 or 4, kv_seq_len up to 8k — memory-bound workload
+    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
     5. Variable-length batching: indptr-based segmented attention
     6. Split K/V from buffer: full 576 dims for keys, first 512 dims for values
 
@@ -70,29 +69,27 @@ benchmark_timeout: 900
 ranked_timeout: 1200
 
 tests:
-  # 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}
+  # 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}
 
 benchmarks:
-  # 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}
-
+  # 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}
 
 ranking_by: "geom"