@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c37f86e572

[Bugfix] GLM-5.3-Flash: fp8 plan dtype on SM90 sparse MLA, and right-size the indexer prefill workspace (#55222)

Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Mikhail Kostryukov <drakosha81@proton.me> Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>

drakosha committed Sep 30, 2026 at 01:36 UTC c37f86e5721f089a27c8e4358c30e435805fcaab

5 files changed +84 -19

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

+23 -8

@@ -208,12 +208,14 @@ def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_dec

208 assert builder.state.plan_calls[0][1].tolist() == expected_lens

209

210

211 -def test_plan_uses_state_params(monkeypatch):

211 +@pytest.mark.parametrize("kv_dtype", [torch.bfloat16, torch.float8_e4m3fn])

212 +def test_plan_uses_state_params(monkeypatch, kv_dtype):

213 """The NoPE/rope dims and scale live on the builder state, not the layer.

214

215 Without a device batch layout plan() takes exact per-row KV lengths and

216 rebuilds the schedule on every call; the indptrs are always full-size

216 - with zero-query padding rows past num_tokens.

217 + with zero-query padding rows past num_tokens. An fp8 cache plans as

218 + float8_e4m3fn.

219 """

220 wrapper = FakeWrapper()

221 plan_info = [0] * sm90_mod._PLAN_INFO_LEN

@@ -226,7 +228,7 @@ def test_plan_uses_state_params(monkeypatch):

228 state.device = torch.device("cpu")

229 state.wrapper = wrapper

230 state.num_heads = 4

229 - state.kv_dtype = torch.bfloat16

231 + state.kv_dtype = kv_dtype

232 state.kv_lora_rank = HEAD

233 state.qk_rope_head_dim = 64

234 state.sm_scale = 576**-0.5

@@ -249,7 +251,7 @@ def test_plan_uses_state_params(monkeypatch):

251 assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)

252 assert scale == 576**-0.5

253 assert kwargs["q_data_type"] == torch.bfloat16

252 - assert kwargs["kv_data_type"] == torch.bfloat16

254 + assert kwargs["kv_data_type"] == kv_dtype

255

256 # Replanning a smaller batch must clear the previous rows' lengths.

257 state.plan(1, torch.tensor([TOPK], dtype=torch.int32), None, None)

@@ -339,6 +341,14 @@ def test_kv_lens_host_empty():

341 _NO_KPOOL = object()

342

343

344 +@pytest.mark.parametrize(

345 + "spec_dtype,expected_kv_dtype",

346 + [

347 + (torch.bfloat16, torch.bfloat16),

348 + (torch.uint8, torch.float8_e4m3fn),

349 + ],

350 + ids=["bf16", "fp8"],

351 +)

352 @pytest.mark.parametrize(

353 "index_kpool,prefill_lens",

354 [

@@ -348,9 +358,12 @@ _NO_KPOOL = object()

358 ],

359 ids=["kpool4", "kpool_none", "no_kpool_attr"],

360 )

351 -def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens):

361 +def test_builder_kpool_from_model_config(

362 + monkeypatch, index_kpool, prefill_lens, spec_dtype, expected_kv_dtype

363 +):

364 """The builder took kpool from the KV cache spec, whose tokens_per_state is

353 - 1, so the tail pool was never read."""

365 + 1, so the tail pool was never read. The state plans the cache's logical

366 + dtype: uint8 fp8 storage becomes float8_e4m3fn."""

367 monkeypatch.setattr(

368 sm90_mod.FlashInferMLASparseMetadataBuilder,

369 "__init__",

@@ -358,14 +371,15 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens)

371 )

372

373 class _RecordingState:

361 - def __init__(self, *_args, **kwargs):

374 + def __init__(self, *args, **kwargs):

375 + self.kv_dtype = args[2]

376 self.index_topk = kwargs["index_topk"]

377 self.index_kpool = kwargs["index_kpool"]

378

379 monkeypatch.setattr(sm90_mod, "_SM90State", _RecordingState)

380 impl, _ = make_impl(64)

381 spec = MLAAttentionSpec(

368 - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=576, dtype=torch.bfloat16

382 + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=576, dtype=spec_dtype

383 )

384 assert spec.tokens_per_state == 1

385 hf_config = SimpleNamespace(index_topk=2048)

@@ -384,6 +398,7 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens)

398 builder = FlashInferMLASparseSM90Builder(

399 spec, ["attn"], vllm_config, torch.device("cpu")

400 )

401 + assert builder.state.kv_dtype == expected_kv_dtype

402 # req0: prefill chunk ending at context 42295; req1: context <= topk.

403 cam = SimpleNamespace(

404 num_reqs=2,

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

+41

@@ -24,6 +24,7 @@ from vllm.v1.attention.backends.mla.indexer import (

24 DeepseekV4IndexerBackend,

25 DeepseekV32IndexerMetadataBuilder,

26 DeepseekV41IndexerBackend,

27 + get_max_prefill_buffer_size,

28 )

29 from vllm.v1.attention.backends.mla.sparse_utils import (

30 ConvertReqIndexToGlobalIndexKernel,

@@ -377,3 +378,43 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_num_states():

378 device=device,

379 )

380 torch.testing.assert_close(valid_slots, expected)

381 +

382 +

383 +@pytest.mark.parametrize("compress_ratio", [1, 4])

384 +def test_indexer_prefill_budget_matches_compressed_workspace(compress_ratio):

385 + """The chunker budget is in compressed rows, like the K-gather workspace."""

386 + max_model_len = 1024

387 + kv_cache_spec = MLAAttentionSpec(

388 + block_size=256,

389 + num_kv_heads=1,

390 + head_size=128,

391 + dtype=torch.bfloat16,

392 + tokens_per_state=compress_ratio,

393 + )

394 + vllm_config = create_vllm_config(max_model_len=max_model_len)

395 + max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, max_model_len)

396 + block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size)

397 + builder = DeepseekV32IndexerMetadataBuilder(

398 + kv_cache_spec=kv_cache_spec,

399 + layer_names=["dummy"],

400 + vllm_config=vllm_config,

401 + device=torch.device("cpu"),

402 + block_table_width=block_table_width,

403 + )

404 +

405 + workspace_rows = get_max_prefill_buffer_size(vllm_config) // compress_ratio

406 + assert builder.max_prefill_buffer_size == workspace_rows

407 +

408 + # Overflows the workspace unless the chunker splits.

409 + num_reqs = 3 * (workspace_rows // (max_model_len // compress_ratio)) + 1

410 + compressed_seq_lens = torch.full((num_reqs,), max_model_len // compress_ratio)

411 + query_lens = torch.ones(num_reqs, dtype=torch.int64)

412 + chunks = builder._split_indexer_prefill_chunks(

413 + compressed_seq_lens,

414 + query_lens,

415 + builder.max_prefill_buffer_size,

416 + max_logits_bytes=10**15,

417 + )

418 + assert len(chunks) > 1

419 + for req_slice, _ in chunks:

420 + assert int(compressed_seq_lens[req_slice].sum()) <= workspace_rows

vllm/models/glm5next/common/attention.py

+3 -1

@@ -312,7 +312,9 @@ class Indexer(nn.Module):

312 self.prefix = prefix

313 from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size

314

315 - self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)

315 + self.max_total_seq_len = (

316 + get_max_prefill_buffer_size(vllm_config) // self.index_kpool

317 + )

318 self.indexer_op = SparseAttnIndexerKpool(

319 self.k_cache,

320 self.quant_block_size,

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

+7 -1

@@ -443,10 +443,16 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

443 assert topk_indices_buffer is not None

444 hf_config = vllm_config.model_config.hf_text_config

445 assert hf_config.index_topk is not None

446 + # fp8 KV cache is stored as uint8; plan() needs the logical dtype.

447 + kv_plan_dtype = (

448 + torch.float8_e4m3fn

449 + if kv_cache_spec.dtype == torch.uint8

450 + else kv_cache_spec.dtype

451 + )

452 self.state = _SM90State(

453 device,

454 impl.num_heads,

449 - kv_cache_spec.dtype,

455 + kv_plan_dtype,

456 vllm_config.scheduler_config.max_num_batched_tokens,

457 topk_indices_buffer.shape[1],

458 kv_lora_rank=impl.kv_lora_rank,

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

+10 -9

@@ -901,8 +901,17 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

901 self.use_pcp = self.pcp_world_size > 1

902 self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0

903 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

904 + # KV compression (DeepseekV4). Default to 1 for no compression.

905 + self.compress_ratio = 1

906 + if isinstance(self.kv_cache_spec, MLAAttentionSpec):

907 + assert isinstance(self.kv_cache_spec.tokens_per_state, int)

908 + self.compress_ratio = self.kv_cache_spec.tokens_per_state

909 # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.

905 - self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)

910 + # Counted in compressed rows, like the chunker's seq_lens and the

911 + # workspace.

912 + self.max_prefill_buffer_size = (

913 + get_max_prefill_buffer_size(self.vllm_config) // self.compress_ratio

914 + )

915 self.num_speculative_tokens = (

916 self.vllm_config.speculative_config.num_speculative_tokens

917 if self.vllm_config.speculative_config

@@ -984,14 +993,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

993 (self.num_sms + 1, 2), dtype=torch.int32, device=self.device

994 )

995

987 - # KV compression. Default to 1 for no compression.

988 - self.compress_ratio = 1

989 - # Get compress_ratio for DeepseekV4 support

990 - if isinstance(self.kv_cache_spec, MLAAttentionSpec):

991 - # MLA compression is a whole number of tokens per state (fractions

992 - # are whisper block pooling and never reach MLA).

993 - assert isinstance(self.kv_cache_spec.tokens_per_state, int)

994 - self.compress_ratio = self.kv_cache_spec.tokens_per_state

996 if self.dcp_world_size > 1 and self.compress_ratio > 1:

997 raise NotImplementedError(

998 "DCP is not supported with sparse indexer KV compression "