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

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

@@ -191,20 +191,22 @@ def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_dec
     assert builder.state.plan_calls[0][1].tolist() == expected_lens
 
 
-def test_plan_uses_state_params(monkeypatch):
+@pytest.mark.parametrize("kv_dtype", [torch.bfloat16, torch.float8_e4m3fn])
+def test_plan_uses_state_params(monkeypatch, kv_dtype):
     """The NoPE/rope dims and scale live on the builder state, not the layer.
 
     plan() takes exact per-row KV lengths; the schedule is rebuilt on every
     call (contexts grow between steps) and the indptrs are always full-size
-    with zero-query padding rows past num_tokens.
+    with zero-query padding rows past num_tokens. An fp8 cache plans as
+    float8_e4m3fn while the query side stays bfloat16.
     """
     impl, rows = make_impl(64, "auto")
     wrapper = FakeWrapper()
     state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
     state.device = torch.device("cpu")
     state.wrapper = wrapper
     state.num_heads = 4
-    state.kv_dtype = torch.bfloat16
+    state.kv_dtype = kv_dtype
     state.kv_lora_rank = HEAD
     state.qk_rope_head_dim = 64
     state.sm_scale = 576**-0.5
@@ -226,7 +228,24 @@ def test_plan_uses_state_params(monkeypatch):
     assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)
     assert scale == 576**-0.5
     assert kwargs["q_data_type"] == torch.bfloat16
-    assert kwargs["kv_data_type"] == torch.bfloat16
+    assert kwargs["kv_data_type"] == kv_dtype
+
+
+@pytest.mark.parametrize(
+    "spec_dtype,expected",
+    [
+        (torch.uint8, torch.float8_e4m3fn),
+        (torch.float8_e4m3fn, torch.float8_e4m3fn),
+        (torch.bfloat16, torch.bfloat16),
+    ],
+)
+def test_plan_dtype_translates_fp8_storage(spec_dtype, expected):
+    """An fp8 cache is allocated as uint8 but planned as float8_e4m3fn.
+
+    The wrapper rejects uint8, so a cache spec carrying the storage dtype has
+    to be translated before it reaches plan(); other dtypes pass through.
+    """
+    assert FlashInferMLASparseSM90Builder._plan_dtype(spec_dtype) == expected
 
 
 def test_kv_lens_host_formula():

vllm/models/glm5next/common/attention.py

@@ -300,7 +300,12 @@ def __init__(
         self.prefix = prefix
         from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
 
-        self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
+        # Right-size the indexer prefill workspace: it is sized in tokens but
+        # this indexer KV is pool-granular (compress_ratio == index_kpool),
+        # same as deepseek_v4/attention.py already does. Thread #53906.
+        self.max_total_seq_len = (
+            get_max_prefill_buffer_size(vllm_config) // self.index_kpool
+        )
         self.indexer_op = SparseAttnIndexerKpool(
             self.k_cache,
             self.quant_block_size,

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

@@ -266,6 +266,16 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
 
     metadata_cls = FlashInferMLASparseSM90Metadata
 
+    @staticmethod
+    def _plan_dtype(spec_dtype: torch.dtype) -> torch.dtype:
+        """Dtype plan() must be given for a cache of ``spec_dtype``.
+
+        An fp8 KV cache is allocated as uint8 storage and run() views it as
+        float8_e4m3fn; plan() has to be told the same dtype, the wrapper
+        rejects uint8.
+        """
+        return torch.float8_e4m3fn if spec_dtype == torch.uint8 else spec_dtype
+
     def __init__(
         self,
         kv_cache_spec: "AttentionSpec",
@@ -289,7 +299,7 @@ def __init__(
         self.state = _SM90State(
             device,
             impl.num_heads,
-            kv_cache_spec.dtype,
+            self._plan_dtype(kv_cache_spec.dtype),
             vllm_config.scheduler_config.max_num_batched_tokens,
             topk_indices_buffer.shape[1],
             kv_lora_rank=impl.kv_lora_rank,