fix(mla): support BF16 and FP8 DeepSeek V4 dual-cache prefill

b12x/attention/_shared/mla/prefill.py

@@ -11,8 +11,9 @@
   * DSV4 single-cache: topk in {512, 1024, 2048} (FP8-QK) or 128 (BF16-QK)
     with 8-aligned heads split into a paired-head MG prefix plus optional
     single-group tails.
-  * DSV4 dual-cache (extra/indexed tokens): topk==128, heads % 8 == 0,
-    pbs_extra in {2, 64} (BF16-QK), using the same head partitioning.
+  * DSV4 dual-cache (extra/indexed tokens): topk in {128, 512} (BF16-QK) or
+    {1024, 2048} (FP8-QK), heads % 8 == 0, using the same head
+    partitioning, and pbs_extra in {2, 64}.
   * GLM_NSA: topk in {512, 1024, 2048}
   * GLM_NEXT: topk in {512, 1024, 2048, 2051, 2112}; 2112 is an
     alignment-only container whose per-row ``topk_length`` remains 2051.
@@ -431,12 +432,11 @@ def _run_partitioned_mg(
             scale_format=ScaleFormat.UE8M0_BYTE,
         )
 
-    # ── DSV4 dual-cache (has_extra) -> MG (BF16-QK), with strip-and-raise. ──────
-    # FI ships DSV4 dual-cache as topk==128, BF16-QK. 8-aligned head counts split
-    # into a paired prefix plus optional 16/8-head single-group tails. Everything
-    # else RAISEs (the decode-reuse has_extra body has been removed -- no fallback).
+    # DSV4 dual-cache uses BF16-QK for the 128-token text window and the
+    # 512-token image window, and FP8-QK for 1024/2048-wide main sections.
+    # Both paths retain the main/extra union in one online softmax.
     if has_extra:
-        if model_type == ModelType.DSV4 and int(topk) == 128:
+        if model_type == ModelType.DSV4 and int(topk) in (128, 512):
             return _run_partitioned_mg(
                 compute_mode=ComputeMode.BF16,
                 model_type=ModelType.DSV4,
@@ -447,10 +447,28 @@ def _run_partitioned_mg(
                 extra_page_block_size=extra_page_block_size,
                 stride_extra_kv_block=stride_extra_kv_block,
             )
+        if (
+            _mg_enabled
+            and model_type == ModelType.DSV4
+            and compute_mode == ComputeMode.FP8
+            and scale_format == ScaleFormat.UE8M0_BYTE
+            and topk in (1024, 2048)
+        ):
+            return _run_partitioned_mg(
+                compute_mode=ComputeMode.FP8,
+                model_type=ModelType.DSV4,
+                scale_format=ScaleFormat.UE8M0_BYTE,
+                extra_kv_cache=extra_kv_cache,
+                extra_indices=extra_indices,
+                extra_topk_length=extra_topk_length,
+                extra_page_block_size=extra_page_block_size,
+                stride_extra_kv_block=stride_extra_kv_block,
+            )
         raise ValueError(
             f"DSV4 dual-cache prefill (heads={heads}, topk={topk}, "
             f"pbs_extra={int(extra_page_block_size)}) requires MG dispatch; only "
-            "DSV4 topk==128 with heads divisible by 8 is supported. "
+            "DSV4 topk in {128,512} (BF16-QK) or topk in {1024,2048} (FP8-QK) "
+            "with heads divisible by 8 is supported. "
             "No decode-reuse fallback."
         )
 
@@ -464,7 +482,8 @@ def _run_partitioned_mg(
         "Supported (MG) shapes: single-cache heads%8==0; "
         "DSV4 single-cache topk in {512, 1024, 2048} (FP8) or 128 "
         "(BF16-QK, heads%8==0); "
-        "DSV4 dual-cache topk==128 with heads%8==0 and pbs_extra in {2, 64}; "
+        "DSV4 dual-cache topk in {128, 512} (BF16-QK) or topk in {1024, 2048} "
+        "(FP8-QK), heads%8==0, pbs_extra in {2, 64}; "
         "GLM_NSA topk in {512, 1024, 2048}; GLM_NEXT topk in "
         "{512, 1024, 2048, 2051, 2112}; "
         "NVFP4 (GLM-family, scale_format=2) topk in {128, 512, 1024, 2048}; "

b12x/attention/_shared/mla/prefill_mg.py

@@ -332,6 +332,7 @@ def s2_qk_rope_global_mg_dsv4(
     lane: Int32,
     page_block_size: Int32,
     stride_kv_block: Int64,
+    base_delta: Int64,
     *,
     d_rope: cutlass.Constexpr,
     q_rope_stride: cutlass.Constexpr,
@@ -355,7 +356,7 @@ def s2_qk_rope_global_mg_dsv4(
     a_col = (lane >> Int32(4)) * Int32(8)
     entry = warp_first_cand + gid
     idx = _ld_global_index_i32(index_base_ptr, entry)
-    rope_base = _dsv4_rope_base_off(idx, page_block_size, stride_kv_block)
+    rope_base = _dsv4_rope_base_off(idx, page_block_size, stride_kv_block) + base_delta
 
     for ks in cutlass.range_constexpr(d_rope // 16):
         ko = Int32(ks) * Int32(16)
@@ -2919,10 +2920,21 @@ def _body(
                     if ci >= num_main_tiles:
                         index_base_ptr = get_ptr_as_int64(extra_row, split_cand_start)
 
-                # MAIN rope geometry used by the non-dual FP8 / GLM arms.
+                # Keep the rope operand anchored to the main-cache tensor. For
+                # dual-cache EXTRA tiles, a scalar byte delta re-points that
+                # operand at the extra cache without dynamically rebinding a
+                # cute.Tensor inside this runtime section switch.
                 rope_cache = kv_cache_u8
                 rope_pbs = Int32(self.page_block_size)
                 rope_stride = stride_kv_block
+                rope_delta = Int64(0)
+                if cutlass.const_expr(has_extra):
+                    if ci >= num_main_tiles:
+                        rope_pbs = Int32(self.pbs_extra)
+                        rope_stride = stride_extra_kv_block
+                        rope_delta = get_ptr_as_int64(
+                            extra_kv_cache_u8, Int64(0)
+                        ) - get_ptr_as_int64(kv_cache_u8, Int64(0))
 
                 acc0 = [
                     [
@@ -3173,6 +3185,7 @@ def _body(
                             lane,
                             rope_pbs,
                             rope_stride,
+                            rope_delta,
                             d_rope=t.d_rope,
                             q_rope_stride=L.q_rope_stride,
                             n_hg=n_hg,

tests/attention/test_attention_mla_sm120.py

@@ -1115,14 +1115,15 @@ def test_unified_decode_dual_cache_matches_extra_ref(
 
 
 @torch.inference_mode()
-def test_unified_prefill_dual_cache_80_heads_split_tail_matches_extra_ref() -> None:
+@pytest.mark.parametrize("topk", [128, 512])
+def test_unified_prefill_dual_cache_80_heads_split_tail_matches_extra_ref(topk: int) -> None:
     """DSV4 dual-cache prefill heads=80 uses the split MG path (64-head paired
     prefix + 16-head tail) and matches the PyTorch extra-cache oracle."""
     device = require_b12x_sparse_mla()
     from b12x.attention._shared.mla.kernel import run_unified_prefill
 
     num_heads = 80
-    topk, extra_topk, pbs_extra = 128, 128, 2
+    extra_topk, pbs_extra = 128, 2
     main_blocks = 16
     case = dsv4_extra_ref.make_dsv4_extra_decode_case(
         num_heads=num_heads,
@@ -1183,6 +1184,86 @@ def test_unified_prefill_dual_cache_80_heads_split_tail_matches_extra_ref() -> N
     assert (got - exp).abs().max().item() < 2e-2
 
 
+@torch.inference_mode()
+@pytest.mark.parametrize("num_heads", [16, 32])
+@pytest.mark.parametrize("topk", [512, 1024, 2048])
+@pytest.mark.parametrize("pbs_extra", [2, 64])
+def test_unified_prefill_dual_cache_matches_extra_ref(
+    num_heads: int, topk: int, pbs_extra: int
+) -> None:
+    """DSV4 BF16/FP8 dual-cache prefill uses the extra cache for K-RoPE."""
+    device = require_b12x_sparse_mla()
+    from b12x.attention._shared.mla.kernel import run_unified_prefill
+
+    extra_topk = 128
+    main_blocks = (topk + _DSV4_PAGE - 1) // _DSV4_PAGE
+    case = dsv4_extra_ref.make_dsv4_extra_decode_case(
+        num_heads=num_heads,
+        topk=topk,
+        extra_topk=extra_topk,
+        num_tokens=1,
+        num_blocks=main_blocks,
+        page_block_size=_DSV4_PAGE,
+        pbs_extra=pbs_extra,
+        invalidate_half=False,
+        with_sink=False,
+        device=device,
+        seed=9_200 + num_heads,
+    )
+    q = case["q"].contiguous()
+    main_cache = _repack_dsv4_to_compressed(case["kv_cache"], _DSV4_PAGE, main_blocks)
+    extra_blocks = case["extra_kv_cache"].shape[0]
+    extra_cache = _repack_dsv4_to_compressed(
+        case["extra_kv_cache"], pbs_extra, extra_blocks
+    )
+    main_indices = case["topk_indices"].contiguous()
+    extra_indices = case["extra_indices"].contiguous()
+    main_lengths = torch.full((1,), topk, dtype=torch.int32, device=device)
+    extra_lengths = torch.full((1,), extra_topk, dtype=torch.int32, device=device)
+
+    expected, expected_lse = dsv4_extra_ref.dsv4_extra_decode_reference(
+        q,
+        case["kv_cache"],
+        main_indices,
+        case["sm_scale"],
+        case["extra_kv_cache"],
+        extra_indices,
+        page_block_size=_DSV4_PAGE,
+        pbs_extra=pbs_extra,
+        topk_length=main_lengths,
+        extra_topk_length=extra_lengths,
+        main_kv_dequant=case["kv_dequant"],
+        extra_kv_dequant=case["extra_kv_dequant"],
+    )
+    output, lse = run_unified_prefill(
+        q=q,
+        kv_cache=main_cache,
+        topk_indices=main_indices,
+        topk_length=main_lengths,
+        sm_scale=case["sm_scale"],
+        page_block_size=_DSV4_PAGE,
+        extra_kv_cache=extra_cache,
+        extra_indices=extra_indices,
+        extra_topk_length=extra_lengths,
+        extra_page_block_size=pbs_extra,
+    )
+    torch.cuda.synchronize()
+    got = output[0].float()
+    expected = expected[0].float()
+    assert torch.isfinite(got).all()
+    assert torch.count_nonzero(got).item() > 0
+    assert torch.isfinite(lse).all()
+    torch.testing.assert_close(
+        lse.float(),
+        expected_lse.float(),
+        atol=6.0e-2,
+        rtol=2.0e-2,
+    )
+    cos = _cosine(got, expected)
+    assert cos > 0.999, f"DSV4 dual-cache prefill topk={topk} heads={num_heads} O cos={cos}"
+    assert (got - expected).abs().max().item() < 2e-2
+
+
 @torch.inference_mode()
 def test_unified_prefill_dsv4_valid_hpb_8_matches_prefill_ref() -> None:
     """DSV4 prefill heads=8 uses a single MG group with VALID_HPB=8 and must not