Allow FP32 split partials in prepared dense MLA

.lil/changes/b12x-dense-mla-fp32-partials.json

@@ -0,0 +1,18 @@
+{
+  "schema": "local-inference-release-change/v1",
+  "id": "b12x-dense-mla-fp32-partials",
+  "category": "feature",
+  "summary": "Allow dense MLA split partials to remain FP32 until the output merge",
+  "models": [
+    "Kimi-K3"
+  ],
+  "compatibility": "Set Caps.partial_dtype=torch.float32; torch.bfloat16 remains the default.",
+  "details": [
+    "FP32 doubles partial-output scratch without changing KV dtype, output dtype or split-selection policy."
+  ],
+  "pull_requests": [],
+  "requires": [],
+  "authors": [
+    "voipmonitor"
+  ]
+}

b12x/attention/dense_mla/_kernel.py

@@ -104,6 +104,7 @@ def _signature(binding: Binding) -> tuple[object, ...]:
         scratch.head_dim,
         scratch.v_head_dim,
         scratch.window_size,
+        str(scratch.partial_dtype),
         tuple(int(value) for value in binding.q.stride()),
         tuple(int(value) for value in binding.kv_cache.stride()),
         tuple(int(value) for value in binding.output.stride()),
@@ -115,6 +116,24 @@ def _signature(binding: Binding) -> tuple[object, ...]:
     )
 
 
+def _partial_fp32(scratch) -> bool:
+    return scratch.partial_output is not None and scratch.partial_dtype == torch.float32
+
+
+def _partial_cute_dtype(scratch):
+    """Element type of the split-partial operand handed to the kernels.
+
+    A one-split plan has no partial buffer; the forward kernel then receives
+    the BF16 output as a never-written placeholder.
+    """
+    return cutlass.Float32 if _partial_fp32(scratch) else cutlass.BFloat16
+
+
+def _partial_key_fields(scratch) -> tuple[object, ...]:
+    # BF16 partials keep the established compile keys.
+    return (key_field("partial_dtype", "fp32"),) if _partial_fp32(scratch) else ()
+
+
 @dataclass(frozen=True)
 class _ForwardLaunch:
     entry: DenseMlaForwardKernel
@@ -203,7 +222,7 @@ def _forward_launch(binding: Binding) -> _ForwardLaunch:
             dynamic_layout=True,
         ),
         _to_cute(final_lse, cutlass.Float32, align=4),
-        _to_cute(partial_output, cutlass.BFloat16, align=16),
+        _to_cute(partial_output, _partial_cute_dtype(scratch), align=16),
         _to_cute(partial_lse, cutlass.Float32, align=4),
         _to_cute(kv_scale, cutlass.Float32, align=4),
         _to_cute(q_scale, cutlass.Float32, align=4),
@@ -234,6 +253,7 @@ def _forward_launch(binding: Binding) -> _ForwardLaunch:
         key_field("v_head_dim", scratch.v_head_dim),
         key_field("window_size", scratch.window_size),
         key_field("record_stride_bytes", layout.record_stride_bytes),
+        *_partial_key_fields(scratch),
         tensor_key(
             "q",
             binding.q,
@@ -271,12 +291,16 @@ def _merge_launch(binding: Binding) -> _MergeLaunch | None:
         return None
     assert scratch.partial_output is not None
     assert scratch.partial_lse is not None
-    entry = DenseMlaMergeKernel(scratch.num_splits, scratch.v_head_dim)
+    entry = DenseMlaMergeKernel(
+        scratch.num_splits,
+        scratch.v_head_dim,
+        partial_fp32=_partial_fp32(scratch),
+    )
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
     args = (
         _to_cute(
             scratch.partial_output,
-            cutlass.BFloat16,
+            _partial_cute_dtype(scratch),
             align=16,
         ),
         _to_cute(scratch.partial_lse, cutlass.Float32, align=4),
@@ -296,6 +320,7 @@ def _merge_launch(binding: Binding) -> _MergeLaunch | None:
         key_field("num_splits", scratch.num_splits),
         key_field("heads", scratch.num_q_heads),
         key_field("v_head_dim", scratch.v_head_dim),
+        *_partial_key_fields(scratch),
         tensor_key(
             "partial_output",
             scratch.partial_output,

b12x/attention/dense_mla/_merge.py

@@ -1,4 +1,9 @@
-"""Base-2 split reduction for standalone dense MLA."""
+"""Base-2 split reduction for standalone dense MLA.
+
+Split partials are normalized per split and stored as BF16 (two roundings
+between the FP32 accumulators and the merged BF16 output) or as FP32 (one
+rounding, at the merged output).
+"""
 
 from __future__ import annotations
 
@@ -11,6 +16,7 @@
     bfloat2_to_float2_scaled,
     get_ptr_as_int64,
     ld_global_v2_u32,
+    ld_global_v4_f32,
     pack_f32x2_to_bfloat2,
     st_global_v2_u32,
 )
@@ -19,11 +25,12 @@
 
 
 class DenseMlaMergeKernel:
-    """Merge normalized BF16 split partials and emit natural-log LSE."""
+    """Merge normalized split partials and emit natural-log LSE."""
 
-    def __init__(self, num_splits: int, value_dim: int):
+    def __init__(self, num_splits: int, value_dim: int, *, partial_fp32: bool = False):
         self.num_splits = int(num_splits)
         self.value_dim = int(value_dim)
+        self.partial_fp32 = bool(partial_fp32)
 
     def _get_shared_storage_cls(self):
         """Return the per-CTA split-weight storage.
@@ -166,21 +173,30 @@ def kernel(
                         * Int64(self.num_splits)
                         + Int64(split)
                     ) * Int64(self.value_dim) + Int64(dimension)
-                    packed0, packed1 = ld_global_v2_u32(
-                        get_ptr_as_int64(partial_output, partial_offset)
-                    )
-                    value0, value1 = bfloat2_to_float2_scaled(
-                        packed0,
-                        weight,
-                    )
-                    value2, value3 = bfloat2_to_float2_scaled(
-                        packed1,
-                        weight,
-                    )
-                    accumulator[0] += value0
-                    accumulator[1] += value1
-                    accumulator[2] += value2
-                    accumulator[3] += value3
+                    if cutlass.const_expr(self.partial_fp32):
+                        value0, value1, value2, value3 = ld_global_v4_f32(
+                            get_ptr_as_int64(partial_output, partial_offset)
+                        )
+                        accumulator[0] += value0 * weight
+                        accumulator[1] += value1 * weight
+                        accumulator[2] += value2 * weight
+                        accumulator[3] += value3 * weight
+                    else:
+                        packed0, packed1 = ld_global_v2_u32(
+                            get_ptr_as_int64(partial_output, partial_offset)
+                        )
+                        value0, value1 = bfloat2_to_float2_scaled(
+                            packed0,
+                            weight,
+                        )
+                        value2, value3 = bfloat2_to_float2_scaled(
+                            packed1,
+                            weight,
+                        )
+                        accumulator[0] += value0
+                        accumulator[1] += value1
+                        accumulator[2] += value2
+                        accumulator[3] += value3
                 split += Int32(1)
 
             inverse = Float32(inverse_storage[0])

b12x/attention/dense_mla/_preparation.py

@@ -289,6 +289,7 @@ def _query(caps: Caps, abi: FrozenMapping) -> DenseMlaQuery:
         window_size=caps.window_size, use_cuda_graph=caps.use_cuda_graph,
         max_page_table_width=caps.max_page_table_width,
         num_cache_pages=caps.num_cache_pages, abi=abi,
+        partial_dtype=str(caps.partial_dtype).removeprefix("torch."),
     )
 
 
@@ -303,6 +304,7 @@ def _caps(query: DenseMlaQuery, config: DenseMlaConfig, device: torch.device) ->
         head_dim=query.qk_head_dim, v_head_dim=query.v_head_dim,
         physical_record_width=query.physical_record_width,
         window_size=query.window_size, use_cuda_graph=query.use_cuda_graph,
+        partial_dtype=_dtype(query.partial_dtype),
         budget=__import__("b12x.attention.dense_mla.planner", fromlist=["Budget"]).Budget(max_splits=config.max_splits),
     )
 

b12x/attention/dense_mla/_scratch.py

@@ -76,9 +76,16 @@ class Caps:
     window_size: int | None = None
     use_cuda_graph: bool = False
     budget: Budget | None = None
+    # FP32 storage avoids BF16 rounding of normalized split partials before
+    # the merge. The merged output retains its BF16 dtype.
+    partial_dtype: torch.dtype = torch.bfloat16
 
     def __post_init__(self) -> None:
         object.__setattr__(self, "device", _canonical_device(self.device))
+        if self.partial_dtype not in (torch.bfloat16, torch.float32):
+            raise TypeError(
+                "dense MLA partial_dtype must be torch.bfloat16 or torch.float32"
+            )
         if self.mode not in ("decode", "extend", "verify"):
             raise ValueError(f"unsupported dense MLA mode {self.mode!r}")
         if self.dtype != torch.bfloat16:
@@ -214,7 +221,7 @@ def _dense_mla_scratch_layout(
             * caps.num_q_heads
             * num_splits
             * caps.v_head_dim
-            * dtype_nbytes(torch.bfloat16)
+            * dtype_nbytes(caps.partial_dtype)
         )
         cursor = align_up(cursor, SCRATCH_ALIGN_BYTES)
         partial_lse_offset_bytes = cursor
@@ -269,6 +276,7 @@ class Scratch:
     chunks_per_split: int
     query_tile: int
     use_cuda_graph: bool
+    partial_dtype: torch.dtype
     partial_output: torch.Tensor | None
     partial_lse: torch.Tensor | None
     final_lse: torch.Tensor
@@ -550,7 +558,7 @@ def _materialize(
                 layout.num_splits,
                 caps.v_head_dim,
             ),
-            dtype=torch.bfloat16,
+            dtype=caps.partial_dtype,
         )
         partial_lse, _ = materialize_scratch_view(
             scratch_storage,
@@ -594,6 +602,7 @@ def _materialize(
         chunks_per_split=layout.chunks_per_split,
         query_tile=query_tile,
         use_cuda_graph=caps.use_cuda_graph,
+        partial_dtype=caps.partial_dtype,
         partial_output=partial_output,
         partial_lse=partial_lse,
         final_lse=final_lse,

b12x/attention/dense_mla/_tuning.py

@@ -37,6 +37,7 @@ class DenseMlaQuery:
     max_page_table_width: int
     num_cache_pages: int
     abi: FrozenMapping
+    partial_dtype: str = "bfloat16"
     exhaustive: bool = field(default_factory=capture_exhaustive_search)
 
     def to_dict(self) -> dict[str, object]:
@@ -57,6 +58,7 @@ def to_dict(self) -> dict[str, object]:
             "max_page_table_width": self.max_page_table_width,
             "num_cache_pages": self.num_cache_pages,
             "abi": self.abi.to_dict(),
+            "partial_dtype": self.partial_dtype,
             "exhaustive": self.exhaustive,
         }
 
@@ -151,7 +153,7 @@ def _tuning_parameters(query: DenseMlaQuery, device: DeviceIdentity | None):
 
 TUNING = TuningContract(
     component_id="attention.mla",
-    query_schema_version=5,
+    query_schema_version=6,
     config_schema_version=1,
     query_fields=frozenset(
         {
@@ -171,6 +173,7 @@ def _tuning_parameters(query: DenseMlaQuery, device: DeviceIdentity | None):
             "max_page_table_width",
             "num_cache_pages",
             "abi",
+            "partial_dtype",
             "exhaustive",
         }
     ),

tests/attention/test_dense_mla.py

@@ -949,3 +949,92 @@ def test_fp8_page_ids_past_int32_scaled_offset_match_reference() -> None:
         expected_output,
         expected_lse,
     )
+
+
+def test_partial_dtype_accepts_only_bf16_and_fp32() -> None:
+    device = require_b12x()
+    common = dict(
+        device=device,
+        mode="decode",
+        kv_dtype=FP8,
+        num_q_heads=HEADS,
+        page_size=64,
+        max_total_q=1,
+        max_batch=1,
+        max_cache_tokens=4096,
+        max_page_table_width=64,
+        num_cache_pages=64,
+    )
+    assert dense_mla.Caps(**common).partial_dtype == torch.bfloat16
+    assert (
+        dense_mla.Caps(**common, partial_dtype=torch.float32).partial_dtype
+        == torch.float32
+    )
+    with pytest.raises(TypeError, match="partial_dtype"):
+        dense_mla.Caps(**common, partial_dtype=torch.float16)
+
+
+@pytest.mark.parametrize("live_tokens", [2_817, 10_240, 40_960])
+def test_fp32_split_partials_reduce_error_against_reference(live_tokens: int) -> None:
+    """Float32 partials remove the BF16 rounding between split and merge."""
+    device = require_b12x()
+    torch.manual_seed(20260924)
+    heads, page_size, max_tokens = 16, 64, 65_536
+    width = max_tokens // page_size
+    q_float = torch.randn(1, heads, QK_DIM, device=device) * 0.3
+    cache_float = torch.randn(width, page_size, QK_DIM, device=device) * 0.3
+    q_scale = (q_float.abs().max() / 400).reshape(1).float()
+    kv_scale = (cache_float.abs().max() / 400).reshape(1).float()
+    q = (q_float / q_scale).to(FP8)
+    cache = (cache_float / kv_scale).to(FP8)
+    page_table = torch.randperm(width, device=device).to(torch.int32).reshape(1, width)
+    cache_seqlens = torch.tensor([live_tokens], dtype=torch.int32, device=device)
+    cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    expected, expected_lse = dense_mla.reference(
+        q,
+        cache,
+        page_table,
+        cache_seqlens,
+        cu_seqlens_q,
+        q_scale=q_scale,
+        kv_scale=kv_scale,
+    )
+    errors = {}
+    for partial_dtype in (torch.bfloat16, torch.float32):
+        plan = dense_mla.plan(
+            dense_mla.Caps(
+                device=device,
+                mode="decode",
+                kv_dtype=FP8,
+                num_q_heads=heads,
+                page_size=page_size,
+                max_total_q=1,
+                max_batch=1,
+                max_cache_tokens=max_tokens,
+                max_page_table_width=width,
+                num_cache_pages=width,
+                use_cuda_graph=True,
+                partial_dtype=partial_dtype,
+            )
+        )
+        output = torch.empty(1, heads, VALUE_DIM, dtype=torch.bfloat16, device=device)
+        binding = _bind(
+            plan,
+            q=q,
+            kv_cache=cache,
+            output=output,
+            page_table=page_table,
+            cache_seqlens=cache_seqlens,
+            cu_seqlens_q=cu_seqlens_q,
+            q_scale=q_scale,
+            kv_scale=kv_scale,
+        )
+        assert binding.scratch.num_splits > 1
+        assert binding.scratch.partial_output.dtype == partial_dtype
+        actual, lse = dense_mla.run(binding=binding)
+        _assert_matches(actual, lse, expected, expected_lse)
+        errors[partial_dtype] = float(
+            ((actual.float() - expected).norm() / expected.norm()).item()
+        )
+    print(f"live {live_tokens}: relative L2 error {errors}")
+    assert errors[torch.float32] < errors[torch.bfloat16]