[ROCm] Enable DeepSeek-V4 on gfx11 (#47017)
This PR enables DeepSeek-V4 checkpoints on ROCm gfx11/RDNA devices. Signed-off-by: Letian Kang <Letian.Kang@amd.com>
PikaPikachu committed
Aug 12, 2026 at 08:55 UTC
466855a2bfc57e58f7e2d03ee2deba79ceab2617
3 files changed
+24
-2
vllm/_aiter_ops.py
+16
-1
@@ -2025,12 +2025,27 @@ class rocm_aiter_ops:
2025
2026
@staticmethod
2027
def register_ops_once() -> None:
2028
+ global _OPS_REGISTERED
2029
+
2030
if not (
2031
is_aiter_found_and_supported() or is_aiter_found_and_supported_on_rdna4()
2032
):
2033
+ if not current_platform.is_rocm():
2034
+ return
2035
+
2036
+ from vllm.platforms.rocm import on_gfx11
2037
+
2038
+ if on_gfx11() and not _OPS_REGISTERED:
2039
+ direct_register_custom_op(
2040
+ op_name="rocm_aiter_sparse_attn_indexer",
2041
+ op_func=rocm_aiter_sparse_attn_indexer,
2042
+ mutates_args=["topk_indices_buffer"],
2043
+ fake_impl=rocm_aiter_sparse_attn_indexer_fake,
2044
+ dispatch_key=current_platform.dispatch_key,
2045
+ )
2046
+ _OPS_REGISTERED = True
2047
return
2048
2033
- global _OPS_REGISTERED
2049
if not _OPS_REGISTERED:
2050
# register all the custom ops here
2051
direct_register_custom_op(
vllm/model_executor/layers/sparse_attn_indexer.py
+7
-1
@@ -850,7 +850,13 @@ class SparseAttnIndexer(CustomOp):
850
assert isinstance(q_quant, torch.Tensor), (
851
"AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
852
)
853
- if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():
853
+ from vllm.platforms.rocm import on_gfx11
854
+
855
+ if (
856
+ rocm_aiter_ops.is_enabled()
857
+ or rocm_aiter_ops.is_rdna_aiter_enabled()
858
+ or on_gfx11()
859
+ ):
860
return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
861
hidden_states,
862
_encode_layer_name(self.k_cache.prefix),
vllm/platforms/rocm.py
+1
@@ -520,6 +520,7 @@ class RocmPlatform(Platform):
520
"deepseek_v4_fp8",
521
"compressed-tensors",
522
"fbgemm_fp8",
523
+ "inc",
524
"quark",
525
"mxfp4",
526
"mxfp8",