biondizzle
e1a642452a
1. DeepseekV4MLAAttention.__init__ had a hard assertion that the attention backend MUST be FlashMLA. On Blackwell, FlashMLA doesn't work but we bypass it via _attention_impl_blackwell(). Added _is_blackwell flag to skip FlashMLA-specific init (fp8_ds_mla cache format conversion). 2. Added VLLM_NVFP4_GEMM_BACKEND=cutedsl env var to docker-compose.yml to force CuTeDSL kernel selection for NVFP4 linear layers. 3. Updated register_cutedsl_kernel.py to also register CuTeDSL in _NVFP4_BACKEND_TO_KERNEL dict (for the env var override path).
47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
#!/usr/bin
|
|
# Patch vLLM's linear kernel __init__.py to register the CuTeDSL NVFP4 kernel.
|
|
# This inserts our kernel at the TOP of the _POSSIBLE_NVFP4_KERNELS list,
|
|
# so it gets selected first on Blackwell GPUs.
|
|
|
|
import sys
|
|
|
|
def patch_init(path):
|
|
with open(path, 'r') as f:
|
|
content = f.read()
|
|
|
|
# Add import after the existing flashinfer import block
|
|
import_line = (
|
|
"from vllm.model_executor.kernels.linear.nvfp4.cutedsl import (\n"
|
|
" CuTeDSLNvFp4LinearKernel,\n"
|
|
")\n"
|
|
)
|
|
# Insert after the marlin import block
|
|
marker = "from vllm.model_executor.kernels.linear.nvfp4.marlin import ("
|
|
if "CuTeDSLNvFp4LinearKernel" in content:
|
|
print("CuTeDSL kernel already registered, skipping")
|
|
return
|
|
idx = content.find(marker)
|
|
if idx == -1:
|
|
print("ERROR: Could not find marlin import marker")
|
|
sys.exit(1)
|
|
# Find end of marlin import block
|
|
end = content.find("\n\n", idx)
|
|
content = content[:end] + "\n" + import_line + content[end:]
|
|
|
|
# Insert CuTeDSLNvFp4LinearKernel at TOP of _POSSIBLE_NVFP4_KERNELS CUDA list
|
|
old = " PlatformEnum.CUDA: [\n FlashInferCutlassNvFp4LinearKernel,"
|
|
new = " PlatformEnum.CUDA: [\n CuTeDSLNvFp4LinearKernel,\n FlashInferCutlassNvFp4LinearKernel,"
|
|
content = content.replace(old, new)
|
|
|
|
# Also add to _NVFP4_BACKEND_TO_KERNEL so VLLM_NVFP4_GEMM_BACKEND=cutedsl works
|
|
old_backend = ' "emulation": EmulationNvFp4LinearKernel,\n}'
|
|
new_backend = ' "emulation": EmulationNvFp4LinearKernel,\n "cutedsl": CuTeDSLNvFp4LinearKernel,\n}'
|
|
content = content.replace(old_backend, new_backend)
|
|
|
|
with open(path, 'w') as f:
|
|
f.write(content)
|
|
print("Patched CuTeDSL NVFP4 kernel into", path)
|
|
|
|
if __name__ == "__main__":
|
|
patch_init(sys.argv[1])
|