Skip to content

Clarification on the Figure 3 baseline and plans for an optimized vLLM MISA kernel #43

Description

@dahaobenhao

Thank you for releasing the MISA implementation. We are evaluating the one-stage patch and would like to clarify the performance baseline used by Figure 3, as well as the roadmap for the vLLM integration.

From the current repository (eaf4dde), our reading is:

  • In MISA_preprint/misa_vllm_patch_1stage/misa_vllm_patch/test_speed.py, deepseek_sparse_attention_indexer() computes the dense all-head baseline with misa_mqa_sparse_head_return_logits_interface(). The script therefore compares two standalone TileLang paths using synthetic BF16 Q/K/weights.
  • The Figure 3 script does not call the vLLM prefill baseline in indexers.py.
  • In the actual vLLM patch, BaselineSparseAttnIndexer (use_misa=false) calls vLLM's fp8_mqa_logits() when DeepGEMM is available.
  • The current MISA prefill path (use_misa=true) calls misa_mqa_logits(), which converts the gathered FP8 Q/K data to BF16, performs pooling/routing/head selection, and then uses the TileLang sparse-head scoring kernel. Token selection is subsequently performed with torch.topk().

As an additional integration observation—not a reproduction of the paper's H200 Figure 3 result—we measured the following configuration:

  • vLLM 0.17.0 with the released one-stage plugin
  • GLM-5.2-FP8
  • 2 x 8 NVIDIA H100, TP=16
  • Batch size 1, 128K prefill, chunked prefill size 8192
  • Final chunk shape: M=8192, N=131072
  • Eager execution; use_misa=false versus use_misa=true; 8 active MISA heads

At the complete per-layer indexer boundary, including cache handling, scoring, MISA head routing, and token top-k, we observed:

Path Indexer latency Full prefill latency
Plugin baseline (use_misa=false, DeepGEMM available) 11.686 ms 13.978 s
MISA (use_misa=true) 24.310 ms 15.892 s

Under this different model/hardware/integration setting, MISA was 2.08x slower at the complete indexer boundary and 13.7% slower end to end.

We understand that this is not directly comparable to Figure 3: the GPU is H100 rather than H200, the model is GLM-5.2-FP8, and the measured boundary includes vLLM cache handling and token top-k.

Could you please clarify:

  1. Does the "DSA original indexer kernel" baseline reported for Figure 3 refer specifically to the dense all-head TileLang path in test_speed.py, rather than vLLM's FP8 fp8_mqa_logits() / DeepGEMM production path?
  2. Do you plan to provide an optimized FP8/fused MISA kernel for vLLM that avoids the full FP8-to-BF16 conversion and reduces the pooling/router/head-selection/top-k overhead?
  3. Do you have a recommended script and configuration for comparing MISA against vLLM's DeepGEMM-backed indexer, ideally including end-to-end prefill measurements?

We would be happy to test an optimized path or help evaluate a patch. Thanks!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions