Skip to content

[Build][SM70] Make TokenSpeed MLA optional for V100 builds #409

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ga-it wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Build][SM70] Make TokenSpeed MLA optional for V100 builds #409

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion docs/design/attention_backends.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,9 @@ hardware and configuration.

> **‡** Automatic selection tries FlashAttention first. On Blackwell
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
> TokenSpeed MLA. On other GPUs, only FlashAttention is considered. TokenSpeed
> MLA is optional; its package must use an `apache-tvm-ffi` version compatible
> with the installed TileLang release.

### Decode Backends

Expand Down
4 changes: 2 additions & 2 deletions requirements/cuda.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@ fastsafetensors >= 0.3.3
nvidia-cutlass-dsl[cu13]==4.7.0
quack-kernels>=0.3.3

# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.2.5
# TokenSpeed MLA is optional. Version 0.2.5 pins apache-tvm-ffi==0.1.13,
# which cannot coexist with the TileLang version pinned above.

# Humming kernels for quantization gemm
humming-kernels[cu13]==0.1.2
5 changes: 3 additions & 2 deletions vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,9 @@ def supports_compute_capability(cls, device_capability: "DeviceCapability") -> b
return device_capability.major == 10

_INSTALL_HINT = (
"tokenspeed_mla package is not installed. "
"Install it with: `uv pip install tokenspeed-mla`"
"tokenspeed_mla is an optional package and is not installed. "
"Its apache-tvm-ffi requirement must be compatible with the installed "
"TileLang version."
)

@classmethod
Expand Down
5 changes: 3 additions & 2 deletions vllm/v1/attention/backends/mla/tokenspeed_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,8 +101,9 @@ def supports_combination(
import tokenspeed_mla # noqa: F401
except ImportError:
return (
"tokenspeed_mla package is not installed. "
"Install it with: `uv pip install tokenspeed-mla`"
"tokenspeed_mla is an optional package and is not installed. "
"Its apache-tvm-ffi requirement must be compatible with the "
"installed TileLang version."
)

# tokenspeed_mla CuTe DSL kernel is shape-specialized for DeepSeek R1
Expand Down
Loading