From d887c66fb2fd7d0f5ea7e8ee91eb3382c604e22f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 09:16:12 +0000 Subject: [PATCH] Allow FP32 dense MLA split partials in prepared plans Carry partial_dtype through dense MLA declarations, tuning queries, scratch and kernels. FP32 avoids intermediate BF16 rounding and doubles partial-output scratch. Defaults and split selection remain unchanged. Validation: Four host dtype contract cases pass. Five dense MLA implementation ASTs match the qualified e3e4fff4cf source. No GPU test was run in this audit environment. Co-authored-by: Martin Vit Co-authored-by: Claude Opus 5.5 Co-authored-by: Codex --- .../changes/b12x-dense-mla-fp32-partials.json | 18 ++++ b12x/attention/dense_mla/_kernel.py | 31 ++++++- b12x/attention/dense_mla/_merge.py | 52 +++++++---- b12x/attention/dense_mla/_preparation.py | 2 + b12x/attention/dense_mla/_scratch.py | 13 ++- b12x/attention/dense_mla/_tuning.py | 5 +- tests/attention/test_dense_mla.py | 89 +++++++++++++++++++ 7 files changed, 186 insertions(+), 24 deletions(-) create mode 100644 .lil/changes/b12x-dense-mla-fp32-partials.json diff --git a/.lil/changes/b12x-dense-mla-fp32-partials.json b/.lil/changes/b12x-dense-mla-fp32-partials.json new file mode 100644 index 000000000..89fd264fe --- /dev/null +++ b/.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" + ] +} diff --git a/b12x/attention/dense_mla/_kernel.py b/b12x/attention/dense_mla/_kernel.py index 7d737891c..de78e00c4 100644 --- a/b12x/attention/dense_mla/_kernel.py +++ b/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, diff --git a/b12x/attention/dense_mla/_merge.py b/b12x/attention/dense_mla/_merge.py index 5f26c5aa2..5989027b4 100644 --- a/b12x/attention/dense_mla/_merge.py +++ b/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]) diff --git a/b12x/attention/dense_mla/_preparation.py b/b12x/attention/dense_mla/_preparation.py index b42373ff1..188052e50 100644 --- a/b12x/attention/dense_mla/_preparation.py +++ b/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), ) diff --git a/b12x/attention/dense_mla/_scratch.py b/b12x/attention/dense_mla/_scratch.py index 37235c1d4..abed7850c 100644 --- a/b12x/attention/dense_mla/_scratch.py +++ b/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, diff --git a/b12x/attention/dense_mla/_tuning.py b/b12x/attention/dense_mla/_tuning.py index 40bb341c4..86d6a5463 100644 --- a/b12x/attention/dense_mla/_tuning.py +++ b/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", } ), diff --git a/tests/attention/test_dense_mla.py b/tests/attention/test_dense_mla.py index 82d591ba7..c356ddc9c 100644 --- a/tests/attention/test_dense_mla.py +++ b/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]