@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1a085dadf2

[ROCm][K3] Extend FP8 asm MLA prefill to non-divisor small head counts (#51040)

Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

xiaohuguo2023 committed Aug 26, 2026 at 18:47 UTC 1a085dadf254b7cb2b5f904747a5750256d413fa

2 files changed +249 -45

tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py

+104 -21

@@ -43,8 +43,20 @@ pytestmark = pytest.mark.skipif(

43 )

45 # DeepSeek-style MLA head dims after kv_b_proj decompression (q/k carry the

46 -# nope+rope dims, v carries v_head_dim). FP8 prefill requires 16-aligned heads.

47 -NUM_HEADS = 16

46 +# nope+rope dims, v carries v_head_dim).

47 +#

48 +# The kernel requires 16-aligned heads, so _mla_fp8_prefill_attn replicate-pads

49 +# q/k/v up to get_fp8_prefill_num_heads(num_heads) and slices the output back.

50 +# Cover all three regimes:

51 +# 12 -> padded to 16 (non-divisor, tile+slice). Live case: a K3 rank at TP8.

52 +# 16 -> no padding, output aliases the caller's buffer.

53 +# 24 -> padded to 32. No current model and TP reaches this band (96 heads

54 +# would need TP4, which does not fit in 288 GiB; 128-head models give

55 +# 128/64/32/16/8), so this case exists to keep the general path honest

56 +# for future architectures: it is the only one of the three where

57 +# padding above 16 actually happens, so it is what proves the

58 +# replicate-pad/slice-back argument holds for a target other than 16.

59 +HEAD_COUNTS = [12, 16, 24]

60 QK_NOPE_HEAD_DIM = 128

61 QK_ROPE_HEAD_DIM = 64

62 QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM # 192

@@ -56,6 +68,10 @@ SCALE = 1.0 / math.sqrt(QK_HEAD_DIM)

68 # elements.

69 ATOL, RTOL = 1e-1, 5e-2

71 +# The assembly kernel writes bf16 through a raw output pointer. Keep its tensors

72 +# and workspace reservation on one dtype.

73 +ATTN_OUT_DTYPE = torch.bfloat16

74 +

76 @pytest.fixture(autouse=True)

77 def _workspace_manager():

@@ -73,14 +89,14 @@ def _workspace_manager():

89 reset_workspace_manager()

76 -def _make_impl():

92 +def _make_impl(num_heads: int):

93 """Minimal AiterMLAImpl exposing only what _mla_fp8_prefill_attn reads."""

94 from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1

96 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl

98 impl = object.__new__(AiterMLAImpl)

83 - impl.num_heads = NUM_HEADS

99 + impl.num_heads = num_heads

100 impl.v_head_dim = V_HEAD_DIM

101 impl.scale = SCALE

102 impl._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd

@@ -88,8 +104,13 @@ def _make_impl():

104 return impl

105

106

91 -def _build_prefill_metadata(seq_lens: list[int], device: torch.device):

92 - """Build fp8_prefill_* metadata via the real builder (no hand-rolling)."""

107 +def _build_prefill_metadata(

108 + seq_lens: list[int],

109 + device: torch.device,

110 + num_heads: int,

111 + attn_out_dtype: torch.dtype = ATTN_OUT_DTYPE,

112 +):

113 + """Build metadata with the production builder and output dtype."""

114 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder

115

116 qo_indptr_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)

@@ -99,12 +120,13 @@ def _build_prefill_metadata(seq_lens: list[int], device: torch.device):

120 max_q = max(seq_lens)

121

122 builder = object.__new__(AiterMLAMetadataBuilder)

102 - builder.num_heads = NUM_HEADS

123 + builder.num_heads = num_heads

124 builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM)

125 builder._init_fp8_prefill_ps_buffers(

126 max_num_reqs=len(seq_lens),

127 max_prefill_qlen=max_q,

128 max_num_batched_tokens=total_q,

129 + attn_out_dtype=attn_out_dtype,

130 device=device,

131 )

132

@@ -115,7 +137,7 @@ def _build_prefill_metadata(seq_lens: list[int], device: torch.device):

137 metadata = SimpleNamespace(prefill=prefill, num_decodes=0)

138 common = SimpleNamespace(query_start_loc_cpu=qo_indptr_cpu)

139 builder._build_fp8_prefill_ps_metadata(metadata, common)

118 - return metadata, total_q

140 + return metadata, total_q, builder

141

142

143 def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:

@@ -128,34 +150,95 @@ def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tenso

150 return out.transpose(0, 1) # [S, H, Dv]

151

152

153 +@pytest.mark.parametrize("num_heads", HEAD_COUNTS)

154 @pytest.mark.parametrize("seq_len", [128, 512])

155 @torch.inference_mode()

133 -def test_fp8_prefill_matches_reference(seq_len: int) -> None:

156 +def test_fp8_prefill_matches_reference(seq_len: int, num_heads: int) -> None:

157 device = torch.device("cuda")

158 torch.manual_seed(0)

159

137 - metadata, total_q = _build_prefill_metadata([seq_len], device)

138 -

139 - q = torch.randn(

140 - total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device

141 - )

142 - k = torch.randn(

143 - total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device

144 - )

145 - v = torch.randn(total_q, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=device)

160 + metadata, total_q, _ = _build_prefill_metadata([seq_len], device, num_heads)

161 +

162 + # Lock the workspace, as GPUModelRunner does after warmup, so the forward

163 + # cannot quietly grow it. The builder reserves scratch at num_head_k and

164 + # the forward requests it at its own padded width; if a future change makes

165 + # the forward's width exceed the builder's, that is an under-reservation

166 + # which is invisible on an unlocked workspace (it just grows) and fatal in

167 + # a real serve. Locking turns it into a failure here.

168 + #

169 + # Note this does not catch the reverse -- a builder wider than the forward

170 + # merely over-reserves -- so it would not have flagged the pre-PR

171 + # max(16, num_heads) sizing, which over-reserves 2.8x at 24 heads.

172 + from vllm.v1.worker.workspace import lock_workspace

173 +

174 + lock_workspace()

175 +

176 + qkv_kwargs = dict(dtype=ATTN_OUT_DTYPE, device=device)

177 + q = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs)

178 + k = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs)

179 + v = torch.randn(total_q, num_heads, V_HEAD_DIM, **qkv_kwargs)

180 out = torch.zeros(

147 - total_q, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device=device

181 + total_q, num_heads * V_HEAD_DIM, dtype=ATTN_OUT_DTYPE, device=device

182 )

183

150 - impl = _make_impl()

184 + impl = _make_impl(num_heads)

185 impl._mla_fp8_prefill_attn(q, k, v, metadata, out)

186

187 out_ref = _reference(q, k, v)

188

189 assert torch.isfinite(out).all()

190 torch.testing.assert_close(

157 - out.view(total_q, NUM_HEADS, V_HEAD_DIM).float(),

191 + out.view(total_q, num_heads, V_HEAD_DIM).float(),

192 out_ref.float(),

193 atol=ATOL,

194 rtol=RTOL,

195 )

196 +

197 +

198 +@pytest.mark.parametrize("num_heads", HEAD_COUNTS)

199 +def test_fp8_prefill_metadata_width_matches_forward(num_heads: int) -> None:

200 + """The PS metadata must be built for the head count the kernel is handed.

201 +

202 + ``_init_fp8_prefill_ps_buffers``/``build()`` size the work and reduce maps

203 + from ``num_head_k``, while ``_mla_fp8_prefill_attn`` pads q/k/v to its own

204 + target. Divergence is not a numerical bug in either direction -- maps built

205 + narrower than the kernel width still cover the real heads, since padding

206 + only ever appends duplicate heads above them. It is a sizing bug: narrower

207 + maps mean a lower head alignment, hence more partial tiles and a larger

208 + reservation, while wider ones under-reserve the forward's scratch. Pin the

209 + two to one helper rather than to independent expressions.

210 + """

211 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper

212 +

213 + width = AiterMLAHelper.get_fp8_prefill_num_heads(num_heads)

214 + assert width % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0

215 + assert width >= num_heads

216 + assert width - num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS

217 + # Every count the enable gate accepts must have a usable padded width.

218 + assert AiterMLAHelper.is_valid_num_heads(num_heads)

219 +

220 + # The builder must actually size its maps at that width. Assert against the

221 + # buffer it allocated rather than recomputing the width here, or the test

222 + # passes no matter what _init_fp8_prefill_ps_buffers does.

223 + #

224 + # fp8_ps_reduce_partial_map is sized by the partial-tile count, which is

225 + # gcd-driven and so is a sharp probe of the width used: a lower head

226 + # alignment yields *more* tiles. At 24 heads, sizing at the unpadded count

227 + # gives 4x the tiles (64 vs 16) and ~2.8x the scratch reservation (193.6

228 + # vs 68.6 MiB at batch=1/qlen=512) -- not wrong numerically (the 24-wide

229 + # maps still cover the real heads) but the reason to pin the widths

230 + # together.

231 + from aiter import get_ps_metadata_info_v1

232 +

233 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import _FP8_PREFILL_TILE_Q

234 +

235 + def _tiles(nhk: int) -> int:

236 + return get_ps_metadata_info_v1(

237 + batch_size=1,

238 + num_head_k=nhk,

239 + max_qlen=512,

240 + qlen_granularity=_FP8_PREFILL_TILE_Q,

241 + )[5][0]

242 +

243 + _, _, builder = _build_prefill_metadata([512], torch.device("cuda"), num_heads)

244 + assert builder.fp8_ps_reduce_partial_map.numel() == _tiles(width)

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+145 -24

@@ -414,9 +414,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

414 device=device,

415 )

416

417 - # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.

418 - self._fp8_prefill_enabled = (

419 - _fp8_mla_prefill_supported() and self.num_heads % 16 == 0

417 + # The assembly prefill requires FP8 KV, bf16 output, and 16-aligned

418 + # heads. It writes bf16 through a raw output pointer, so fp16 must use

419 + # the standard prefill path. Head counts that are not a multiple of 16

420 + # are replicate-padded up to one in _mla_fp8_prefill_attn, so the gate

421 + # is the same head-count predicate the decode path uses.

422 + self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (

423 + kv_cache_dtype_str == "fp8"

424 + and vllm_config.model_config.dtype == torch.bfloat16

425 + and AiterMLAHelper.is_valid_num_heads(self.num_heads)

426 )

427 if self._fp8_prefill_enabled:

428 max_prefill_qlen = min(

@@ -427,6 +433,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

433 max_num_reqs,

434 max_prefill_qlen,

435 vllm_config.scheduler_config.max_num_batched_tokens,

436 + vllm_config.model_config.dtype,

437 device,

438 )

439

@@ -444,6 +451,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

451 max_num_reqs: int,

452 max_prefill_qlen: int,

453 max_num_batched_tokens: int,

454 + attn_out_dtype: torch.dtype,

455 device: torch.device,

456 ) -> None:

457 """Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.

@@ -466,13 +474,30 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

474 is bounded by this budget rather than by a single request's

475 ``max_prefill_qlen`` — concurrent requests can sum to more than

476 ``max_model_len`` when ``max_model_len < max_num_batched_tokens``.

477 + attn_out_dtype: Dtype of the attention output buffer, used to size

478 + the padded-head output scratch (small head counts only).

479 device: Target device for the buffers.

480 """

481 from aiter import get_ps_metadata_info_v1

482

483 # After kv_b_proj decompression, K has num_heads heads (same as Q).

484 # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.

475 - num_head_k = self.num_heads

485 + # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are

486 + # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS

487 + # metadata for that same padded count so the work/reduce maps and the

488 + # scratch reservations describe the width the kernel is handed.

489 + #

490 + # This was previously max(16, num_heads), which agrees with the helper

491 + # at every head count reachable today (12 and 16 both give 16) and

492 + # differs only above 16: at 24 it leaves num_head_k=24 while the

493 + # forward pads to 32. That is not a correctness bug -- the 24-wide work

494 + # maps still cover head-tiles 0..23, which are the real heads -- but it

495 + # is expensive, because a lower head alignment yields more partial

496 + # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.

497 + # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512

498 + # (measured on gfx950). Sizing both from one helper keeps the widths

499 + # equal and takes the cheaper tiling.

500 + num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)

501 v_head_dim = self.mla_dims.v_head_dim

502 # gqa_ratio = 1

503 # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)

@@ -516,7 +541,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

541 from vllm.v1.worker.workspace import current_workspace_manager

542

543 max_num_partial_tiles = reduce_partial_map_size

519 - current_workspace_manager().get_simultaneous(

544 + reservations: list[tuple[tuple[int, ...], torch.dtype]] = [

545 (

546 (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),

547 torch.float32,

@@ -526,7 +551,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

551 torch.float32,

552 ),

553 ((max_num_batched_tokens, num_head_k), torch.float32),

529 - )

554 + ]

555 + if self.num_heads < num_head_k:

556 + # Padded head counts also take their kernel output buffer from the

557 + # workspace: the caller's [total_q, num_heads * v_head_dim] output

558 + # cannot back a num_head_k-head view (see _mla_fp8_prefill_attn).

559 + reservations.append(

560 + ((max_num_batched_tokens, num_head_k, v_head_dim), attn_out_dtype)

561 + )

562 + current_workspace_manager().get_simultaneous(*reservations)

563

564 logger.info(

565 "FP8 MLA prefill PS buffers allocated "

@@ -566,7 +599,22 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

599 kv_indptr_cpu = qo_indptr_cpu.clone()

600 seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32)

601

569 - num_head_k = self.num_heads

602 + # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are

603 + # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS

604 + # metadata for that same padded count so the work/reduce maps and the

605 + # scratch reservations describe the width the kernel is handed.

606 + #

607 + # This was previously max(16, num_heads), which agrees with the helper

608 + # at every head count reachable today (12 and 16 both give 16) and

609 + # differs only above 16: at 24 it leaves num_head_k=24 while the

610 + # forward pads to 32. That is not a correctness bug -- the 24-wide work

611 + # maps still cover head-tiles 0..23, which are the real heads -- but it

612 + # is expensive, because a lower head alignment yields more partial

613 + # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.

614 + # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512

615 + # (measured on gfx950). Sizing both from one helper keeps the widths

616 + # equal and takes the cheaper tiling.

617 + num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)

618 # gqa_ratio = 1

619 # qhead_granularity = max(gqa_ratio, 1)

620 # qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity

@@ -927,8 +975,39 @@ class AiterMLAHelper:

975 return -(-num_heads // m) * m

976

977 @staticmethod

930 - def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:

931 - m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)

978 + def get_fp8_prefill_num_heads(num_heads: int) -> int:

979 + """Head count the FP8 PS asm prefill runs at: the next multiple of 16.

980 +

981 + Deliberately *not* ``get_actual_mla_num_heads``. That one carves out

982 + native H24 when ``_aiter_mla_native_h24_supported()``, which probes the

983 + asm *decode* reducer and metadata. The prefill is a different kernel

984 + pair (``mla_prefill_ps_asm_fwd`` + ``mla_reduce_v1``) with no such

985 + probe, so 24 heads pad to 32 here even on a native-H24 build.

986 +

987 + The PS metadata in ``_init_fp8_prefill_ps_buffers``/``build()`` must be

988 + sized with this same function, or the work/reduce maps describe a

989 + different head count than the kernel is handed.

990 +

991 + This function itself has no upper bound; the ceiling comes from the

992 + ``is_valid_num_heads`` gate, which rejects counts above

993 + ``_AITER_MAX_PADDED_MLA_HEADS`` (128) unless they are already

994 + 16-aligned. That bound was established for the asm *decode* padding, so

995 + an architecture with, say, 136 heads per rank would be refused the

996 + prefill here for a reason that was never measured against this kernel

997 + pair. Revisit the constant rather than special-casing prefill.

998 + """

999 + m = AiterMLAHelper._AITER_MIN_MLA_HEADS

1000 + return -(-num_heads // m) * m

1001 +

1002 + @staticmethod

1003 + def get_mla_padded_q(

1004 + num_heads: int, q: torch.Tensor, target_heads: int | None = None

1005 + ) -> torch.Tensor:

1006 + m = (

1007 + target_heads

1008 + if target_heads is not None

1009 + else AiterMLAHelper.get_actual_mla_num_heads(num_heads)

1010 + )

1011 if num_heads == m:

1012 return q

1013 if m % num_heads == 0:

@@ -1036,10 +1115,14 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1115 self.flash_attn_varlen_func = flash_attn_varlen_func

1116

1117 # FP8 MLA prefill kernel imports (lazy, only when enabled).

1039 - # Auto-enabled on gfx950 when AITER ships the kernels.

1040 - # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.

1041 - self._fp8_prefill_enabled = (

1042 - _fp8_mla_prefill_supported() and self.num_heads % 16 == 0

1118 + # Auto-enabled on gfx950 when AITER ships the kernels. Only runs when the

1119 + # KV cache is FP8. Head counts that are not a multiple of 16 are

1120 + # replicate-padded up to one (see _mla_fp8_prefill_attn).

1121 + from vllm.utils.torch_utils import is_quantized_kv_cache

1122 +

1123 + self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (

1124 + is_quantized_kv_cache(kv_cache_dtype)

1125 + and AiterMLAHelper.is_valid_num_heads(self.num_heads)

1126 )

1127 if self._fp8_prefill_enabled:

1128 from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1

@@ -1082,7 +1165,36 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1165

1166 fp8_dtype = current_platform.fp8_dtype()

1167 total_q = q.shape[0]

1085 - nhead = self.num_heads

1168 + # PS asm prefill + mla_reduce_v1 require 16-aligned heads, and the PS

1169 + # metadata is built for get_fp8_prefill_num_heads(num_heads). For head

1170 + # counts that are not a multiple of 16 (K3 = 12/rank at TP8)

1171 + # replicate-pad q/k/v up to that count, then slice the output back to

1172 + # the real head count.

1173 + #

1174 + # Counts above 16 that are not multiples of 16 (24, 40, ...) are

1175 + # handled by the same code but are not reached by any current model

1176 + # and TP that fits: 96 heads would need TP4, whose weights exceed a

1177 + # 288 GiB GPU, and 128-head models land on 128/64/32/16/8. The path is

1178 + # kept general for future architectures and its numerics are covered by

1179 + # test_fp8_prefill_matches_reference[num_heads=24]. Note the cost is

1180 + # (padded - real)/real extra FLOPs and q/k/v bytes, which is worst just

1181 + # above a multiple of 16 (17 heads pad to 32); a future arch landing

1182 + # there should measure against the flash_attn_varlen_func fallback

1183 + # rather than assume the asm path wins.

1184 + #

1185 + # Exact, not approximate: after kv_b_proj gqa_ratio is 1, so q, k and v

1186 + # all carry num_heads heads and attention is independent per head.

1187 + # Padding all three identically makes padded head j a duplicate of real

1188 + # head j % num_heads, so the real heads [0:num_heads] are bit-identical

1189 + # to the unpadded result. Same argument as the decode path; only the

1190 + # target width differs.

1191 + _real_nhead = self.num_heads

1192 + nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)

1193 + _pad = nhead != _real_nhead

1194 + if _pad:

1195 + q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q, nhead)

1196 + k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k, nhead)

1197 + v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v, nhead)

1198 v_head_dim = self.v_head_dim

1199 tile_q = _FP8_PREFILL_TILE_Q

1200

@@ -1106,19 +1218,25 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1218 num_partial_tiles = attn_metadata.fp8_prefill_num_partial_tiles

1219 assert num_partial_tiles is not None

1220

1109 - # Reuse the caller's output buffer to skip the per-call alloc + copy.

1110 - # The ASM and reduce kernels both write to a [total_q, nhead, v_head_dim]

1111 - # view, which aliases the [total_q, nhead * v_head_dim] storage of out.

1112 - out_3d = out.view(total_q, nhead, v_head_dim)

1113 -

1114 - # Per-call scratch (logits, attn_lse, final_lse) is served from the

1115 - # workspace manager so allocator churn in the prefill hot path is

1116 - # bounded after warmup, matching the pattern in PR #41002.

1117 - logits, attn_lse, final_lse = current_workspace_manager().get_simultaneous(

1221 + # Per-call scratch is served from the workspace manager so allocator

1222 + # churn in the prefill hot path is bounded after warmup, matching the

1223 + # pattern in PR #41002. The builder reserves the maximum shape of every

1224 + # tensor requested here before the workspace is locked.

1225 + scratch: list[tuple[tuple[int, ...], torch.dtype]] = [

1226 ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),

1227 ((num_partial_tiles * tile_q, nhead), torch.float32),

1228 ((total_q, nhead), torch.float32),

1121 - )

1229 + ]

1230 + if _pad:

1231 + # The ASM and reduce kernels write a [total_q, nhead, v_head_dim]

1232 + # buffer. With unpadded heads that aliases the caller's

1233 + # [total_q, nhead * v_head_dim] output, so write straight into it;

1234 + # padded heads do not fit that storage and need their own buffer.

1235 + scratch.append(((total_q, nhead, v_head_dim), out.dtype))

1236 +

1237 + workspace = current_workspace_manager()

1238 + logits, attn_lse, final_lse, *pad_out = workspace.get_simultaneous(*scratch)

1239 + out_3d = pad_out[0] if _pad else out.view(total_q, nhead, v_head_dim)

1240

1241 # Phase 1: persistent-scheduling assembly prefill kernel.

1242 self._mla_prefill_ps_asm_fwd(

@@ -1156,6 +1274,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1274 final_lse,

1275 )

1276

1277 + if _pad:

1278 + out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])

1279 +

1280 def forward_mha(

1281 self,

1282 q: torch.Tensor,