diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 8dc082b08bbc..40bbd0267205 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -191,12 +191,14 @@ 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() @@ -204,7 +206,7 @@ def test_plan_uses_state_params(monkeypatch): 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(): diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index 6bd9ec44a152..8737d04e94fe 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 830bc52ab939..61a5e7d943d0 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/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,