Cheng Wan
d13d5c03ab
`--attention-backend` is one field of three: a launch that sets only
`--prefill-attention-backend` or `--decode-attention-backend` leaves the base
field at `None`. Seven decisions read that base field alone and therefore
answered from a field the operator never set. `attention_backends()` is the
pair with the base-field fallback already applied, so each site now asks it for
the half it actually needs:
- `inkling_common/attn` assembles backend-specific kwargs (rel_bias / score
mods) and gates its fused prologue; the backend those describe is the one
`self.attn` dispatches to, so `serving_attention_backend()` selects the pair
member by `forward_batch.forward_mode`, mirroring
`HybridAttnBackend._select_backend` exactly -- draft-extend routes through
the prefill branch like the dispatcher does -- and preferring the
runner-stamped pair, so a draft runner answers with its own backend. That
preference only works if every backend that can enter a ForwardContext
carries the stamp, so `DraftBackendFactory._create_backend` now stamps its
products with the backend it resolved (draft override first), and the
draft-extend conv-sidecar wrapper copies the wrapped backend's stamp -- the
replacement backends the spec workers install had no stamp at all and fell
back to the target's configured pair.
- The chunked-prefix-cache gate is a *prefill* feature -> prefill half. Reading
the base field switched the feature off for every prefill-only configuration.
- `init_deterministic_inference_config` maps *prefill* knobs
(SPLIT_TILE / PREFILL_TRUNCATION_ALIGN) -> prefill half; the map missed and
left truncation unset.
- `two_batch_overlap` computes extend positions -> prefill half.
- mrope's interleaved-rope kernel runs in both phases -> both halves must
support triton. This one is not conservative when it misreads:
`support_triton(None)` answers **True**, so a `--prefill-attention-backend
torch_native` launch took the triton path.
- The req-to-token writer has one caller, `alloc_for_extend` -> prefill half;
its fallback pays several `.item()` syncs per request, so gating it on the
decode half too would send every extend of a mixed launch through the slow
path. `get_last_loc` (the spec-decode allocator's helper) keeps the
both-halves reading: verify tokens are served by either half depending on
`speculative_attention_mode`.
- The flashinfer version floor is a guard; it never fired for a launch that
pinned flashinfer through a split field.
One more site the census found is not converted here: `gpt_oss` derives its
`sinks` parameter dtype from the backend, and a single parameter dtype cannot
serve a split pair (FA4 asserts bfloat16, trtllm_mha consumes float32), so
that one is a behaviour question rather than a config-source one and is fixed
in its own PR.
`test_split_attention_backend_decisions.py` pins the callable decisions by
calling them under a split-only publish, and pins the remaining ones
statically -- the file/why map fails if any of them goes back to the base field
(reverse-verified). It also asserts the `support_triton(None) is True` trap the
sweep exists for.
The stamp comes from the constructor, not the request: every factory leaf
answers ("effective_name", backend), because several map entries do not build
what their key says -- cutedsl_mla draft-extend builds the trtllm-mla backend,
"nsa" is a deprecated alias building dsa, and the hybrid-linear entries pick
fa3/intel_amx/triton by host, which no static rename table can express (a
review catch: on Blackwell the alias stamp reached Inkling's per-forward
kwargs assembly, which asserts a concrete kernel name, and crashed the first
draft-extend forward). The stamping is pinned by unit tests, not only by a
spec e2e: removing the child-stamping loop, stamping an alias from a leaf, or
dropping the wrapper copy goes red (reverse-verified), and a static guard
walks the factory source asserting no leaf answers an alias name. The child loop states its contract explicitly --
`create_decode_backend` passes `stamps_children=True` because its products
are per-step containers by construction, so a container without
`attn_backends` raises instead of being silently skipped by a defensive
probe. The `_version` invalidation names its contract (autograd's in-place
counter: private, chosen because it is the only per-tensor signal that ticks
on copy_-style updates; removal fails loudly). The version-floor guard's file
joins the pair-reader ratchet, and the one runner-seed chain read sharing the
backend's __init__ (`speculative_eagle_topk`) reads the spec bag.
532 lines
19 KiB
Python
532 lines
19 KiB
Python
from sglang.srt.runtime_context import attention_backends, get_spec
|
|
from sglang.srt.utils.common import (
|
|
cpu_has_amx_support,
|
|
is_blackwell,
|
|
is_cpu,
|
|
is_hip,
|
|
is_musa,
|
|
is_npu,
|
|
)
|
|
|
|
|
|
def _assert_draft_needs_no_conv_sidecar(draft_model_runner) -> None:
|
|
"""Refuse a multi-step draft decode backend for a draft with conv layers."""
|
|
from sglang.srt.configs.inkling import InklingMMConfig, InklingModelConfig
|
|
|
|
if isinstance(
|
|
draft_model_runner.model_config.hf_config,
|
|
(InklingModelConfig, InklingMMConfig),
|
|
):
|
|
raise NotImplementedError(
|
|
"Inkling's draft model runs its own short convs, which need the "
|
|
"conv-state sidecar the multi-step draft decode backend cannot carry. "
|
|
"Use --enable-multi-layer-eagle."
|
|
)
|
|
|
|
|
|
class DraftBackendFactory:
|
|
def __init__(
|
|
self,
|
|
draft_model_runner,
|
|
topk: int,
|
|
speculative_num_steps: int,
|
|
seed_dsa_topk_from_draft_extend: bool = False,
|
|
):
|
|
self.draft_model_runner = draft_model_runner
|
|
self.topk = topk
|
|
self.speculative_num_steps = speculative_num_steps
|
|
self.seed_dsa_topk_from_draft_extend = seed_dsa_topk_from_draft_extend
|
|
# The draft runner's own backend, not the process-wide config.
|
|
self.draft_attn_backend = draft_model_runner.draft_attention_backend
|
|
|
|
def _create_backend(
|
|
self,
|
|
backend_name: str,
|
|
backend_map: dict,
|
|
error_template: str,
|
|
stamps_children: bool = False,
|
|
):
|
|
# The split pair with the base-backend fallback already applied.
|
|
prefill_backend, decode_backend = attention_backends()
|
|
configured = (
|
|
decode_backend
|
|
if backend_name == "decode_attention_backend"
|
|
else prefill_backend
|
|
)
|
|
backend_type = self.draft_attn_backend or configured
|
|
|
|
if backend_type not in backend_map:
|
|
raise ValueError(error_template.format(backend_type=backend_type))
|
|
|
|
stamp, backend = backend_map[backend_type]()
|
|
if backend is not None:
|
|
backend.prefill_attention_backend_str = stamp
|
|
backend.decode_attention_backend_str = stamp
|
|
if stamps_children:
|
|
for child in backend.attn_backends:
|
|
child.prefill_attention_backend_str = stamp
|
|
child.decode_attention_backend_str = stamp
|
|
return backend
|
|
|
|
def create_decode_backend(self):
|
|
# No multi-step draft backend for steps=0 (nospec) or steps=1.
|
|
if self.speculative_num_steps <= 1:
|
|
return None
|
|
|
|
# Returns a per-step CONTAINER, not an AttentionBackend, so
|
|
# attn_backend_wrapper_for_draft_extend cannot give it a conv sidecar.
|
|
_assert_draft_needs_no_conv_sidecar(self.draft_model_runner)
|
|
|
|
backend_map = {
|
|
"flashinfer": self._create_flashinfer_decode_backend,
|
|
"triton": self._create_triton_decode_backend,
|
|
"intel_amx": self._create_intel_amx_decode_backend,
|
|
"aiter": self._create_aiter_decode_backend,
|
|
"fa3": self._create_fa3_decode_backend,
|
|
"hybrid_linear_attn": self._create_hybrid_linear_attn_decode_backend,
|
|
"flashmla": self._create_flashmla_decode_backend,
|
|
"trtllm_mha": self._create_trtllm_mha_decode_backend,
|
|
"trtllm_mla": self._create_trtllm_mla_decode_backend,
|
|
"cutedsl_mla": self._create_cutedsl_mla_decode_backend,
|
|
"tokenspeed_mla": self._create_tokenspeed_mla_decode_backend,
|
|
"dsa": self._create_dsa_decode_backend,
|
|
"nsa": self._create_dsa_decode_backend, # Deprecated alias for "dsa"
|
|
"ascend": self._create_ascend_decode_backend,
|
|
"fa4": self._create_fa4_decode_backend,
|
|
"dsv4": self._create_dsv4_decode_backend,
|
|
}
|
|
|
|
return self._create_backend(
|
|
"decode_attention_backend",
|
|
backend_map,
|
|
"EAGLE is not supported in decode attention backend {backend_type}",
|
|
stamps_children=True,
|
|
)
|
|
|
|
def create_draft_extend_backend(self):
|
|
backend_map = {
|
|
"flashinfer": self._create_flashinfer_prefill_backend,
|
|
"triton": self._create_triton_prefill_backend,
|
|
"intel_amx": self._create_intel_amx_prefill_backend,
|
|
"aiter": self._create_aiter_prefill_backend,
|
|
"fa3": self._create_fa3_prefill_backend,
|
|
"hybrid_linear_attn": self._create_hybrid_linear_attn_prefill_backend,
|
|
"flashmla": self._create_flashmla_prefill_backend,
|
|
"trtllm_mha": self._create_trtllm_mha_prefill_backend,
|
|
"trtllm_mla": self._create_trtllm_mla_prefill_backend,
|
|
# cute-dsl MLA only supports decode; draft-extend falls back to trtllm-gen.
|
|
"cutedsl_mla": self._create_trtllm_mla_prefill_backend,
|
|
"tokenspeed_mla": self._create_tokenspeed_mla_prefill_backend,
|
|
"dsa": self._create_dsa_prefill_backend,
|
|
"nsa": self._create_dsa_prefill_backend, # Deprecated alias for "dsa"
|
|
"ascend": self._create_ascend_prefill_backend,
|
|
"fa4": self._create_fa4_prefill_backend,
|
|
"dsv4": self._create_dsv4_prefill_backend,
|
|
}
|
|
backend_name = (
|
|
"decode_attention_backend"
|
|
if get_spec().speculative_attention_mode == "decode"
|
|
else "prefill_attention_backend"
|
|
)
|
|
backend = self._create_backend(
|
|
backend_name,
|
|
backend_map,
|
|
"EAGLE is not supported in attention backend {backend_type}",
|
|
)
|
|
# A draft with conv layers of its own (Inkling) needs its sidecar here too.
|
|
from sglang.srt.layers.attention.attention_registry import (
|
|
attn_backend_wrapper_for_draft_extend,
|
|
)
|
|
|
|
wrapped = attn_backend_wrapper_for_draft_extend(
|
|
self.draft_model_runner, backend
|
|
)
|
|
if wrapped is not backend and wrapped is not None and backend is not None:
|
|
wrapped.prefill_attention_backend_str = (
|
|
backend.prefill_attention_backend_str
|
|
)
|
|
wrapped.decode_attention_backend_str = backend.decode_attention_backend_str
|
|
return wrapped
|
|
|
|
def _create_dsa_decode_backend(self):
|
|
from sglang.srt.layers.attention.dsa_backend import (
|
|
DeepseekSparseAttnMultiStepBackend,
|
|
)
|
|
|
|
return (
|
|
"dsa",
|
|
DeepseekSparseAttnMultiStepBackend(
|
|
self.draft_model_runner,
|
|
self.topk,
|
|
self.speculative_num_steps,
|
|
seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
|
|
),
|
|
)
|
|
|
|
def _create_dsa_prefill_backend(self):
|
|
from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
|
|
|
|
return (
|
|
"dsa",
|
|
DeepseekSparseAttnBackend(
|
|
self.draft_model_runner,
|
|
skip_prefill=False,
|
|
seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
|
|
),
|
|
)
|
|
|
|
def _create_flashinfer_decode_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
from sglang.srt.layers.attention.flashinfer_backend import (
|
|
FlashInferMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"flashinfer",
|
|
FlashInferMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
else:
|
|
from sglang.srt.layers.attention.flashinfer_mla_backend import (
|
|
FlashInferMLAMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"flashinfer",
|
|
FlashInferMLAMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_triton_decode_backend(self):
|
|
from sglang.srt.layers.attention.triton_backend import (
|
|
TritonMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"triton",
|
|
TritonMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_intel_amx_decode_backend(self):
|
|
from sglang.srt.layers.attention.intel_amx_backend import (
|
|
IntelAMXMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"intel_amx",
|
|
IntelAMXMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_hybrid_linear_attn_decode_backend(self):
|
|
if is_cpu() and cpu_has_amx_support():
|
|
return self._create_intel_amx_decode_backend()
|
|
if is_blackwell():
|
|
return self._create_triton_decode_backend()
|
|
return self._create_fa3_decode_backend()
|
|
|
|
def _create_hybrid_linear_attn_prefill_backend(self):
|
|
if is_cpu() and cpu_has_amx_support():
|
|
return self._create_intel_amx_prefill_backend()
|
|
if is_blackwell():
|
|
return self._create_triton_prefill_backend()
|
|
return self._create_fa3_prefill_backend()
|
|
|
|
def _create_aiter_decode_backend(self):
|
|
from sglang.srt.layers.attention.aiter_backend import AiterMultiStepDraftBackend
|
|
|
|
return (
|
|
"aiter",
|
|
AiterMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_fa_decode_backend(self, fa_impl_ver: int = 3):
|
|
if not is_musa():
|
|
from sglang.srt.layers.attention.flashattention_backend import (
|
|
FlashAttentionMultiStepBackend,
|
|
)
|
|
else:
|
|
from sglang.srt.hardware_backend.musa.attention.flashattention_backend import (
|
|
MusaFlashAttentionMultiStepBackend as FlashAttentionMultiStepBackend,
|
|
)
|
|
|
|
return (
|
|
f"fa{fa_impl_ver}",
|
|
FlashAttentionMultiStepBackend(
|
|
self.draft_model_runner,
|
|
self.topk,
|
|
self.speculative_num_steps,
|
|
fa_impl_ver=fa_impl_ver,
|
|
),
|
|
)
|
|
|
|
def _create_fa3_decode_backend(self):
|
|
return self._create_fa_decode_backend(fa_impl_ver=3)
|
|
|
|
def _create_fa4_decode_backend(self):
|
|
return self._create_fa_decode_backend(fa_impl_ver=4)
|
|
|
|
def _create_flashmla_decode_backend(self):
|
|
from sglang.srt.layers.attention.flashmla_backend import (
|
|
FlashMLAMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"flashmla",
|
|
FlashMLAMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_trtllm_mha_decode_backend(self):
|
|
from sglang.srt.layers.attention.trtllm_mha_backend import (
|
|
TRTLLMHAAttnMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"trtllm_mha",
|
|
TRTLLMHAAttnMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_trtllm_mla_decode_backend(self, backend: str = "trtllm-gen"):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
raise ValueError(
|
|
"trtllm_mla backend requires MLA model (use_mla_backend=True)."
|
|
)
|
|
|
|
from sglang.srt.layers.attention.trtllm_mla_backend import (
|
|
TRTLLMMLAMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"trtllm_mla",
|
|
TRTLLMMLAMultiStepDraftBackend(
|
|
self.draft_model_runner,
|
|
self.topk,
|
|
self.speculative_num_steps,
|
|
backend=backend,
|
|
),
|
|
)
|
|
|
|
def _create_cutedsl_mla_decode_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
raise ValueError(
|
|
"cutedsl_mla backend requires MLA model (use_mla_backend=True)."
|
|
)
|
|
|
|
from sglang.srt.layers.attention.cutedsl_mla_backend import (
|
|
CuteDslMLAMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"cutedsl_mla",
|
|
CuteDslMLAMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_tokenspeed_mla_decode_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
raise ValueError(
|
|
"tokenspeed_mla backend requires MLA model (use_mla_backend=True)."
|
|
)
|
|
|
|
from sglang.srt.layers.attention.tokenspeed_mla_backend import (
|
|
TokenspeedMLAMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"tokenspeed_mla",
|
|
TokenspeedMLAMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_ascend_decode_backend(self):
|
|
from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
|
|
AscendAttnMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"ascend",
|
|
AscendAttnMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_dsv4_decode_backend(self):
|
|
# Decode here is the EAGLE multi-step draft decode path.
|
|
if is_npu():
|
|
from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import (
|
|
DeepseekV4AscendMultiStepDraftBackend,
|
|
)
|
|
|
|
return (
|
|
"dsv4",
|
|
DeepseekV4AscendMultiStepDraftBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
elif is_hip():
|
|
from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
|
|
DeepseekV4MultiStepBackend,
|
|
)
|
|
else:
|
|
from sglang.srt.layers.attention.deepseek_v4_backend import (
|
|
DeepseekV4MultiStepBackend,
|
|
)
|
|
|
|
return (
|
|
"dsv4",
|
|
DeepseekV4MultiStepBackend(
|
|
self.draft_model_runner, self.topk, self.speculative_num_steps
|
|
),
|
|
)
|
|
|
|
def _create_flashinfer_prefill_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
from sglang.srt.layers.attention.flashinfer_backend import (
|
|
FlashInferAttnBackend,
|
|
)
|
|
|
|
return (
|
|
"flashinfer",
|
|
FlashInferAttnBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
else:
|
|
from sglang.srt.layers.attention.flashinfer_mla_backend import (
|
|
FlashInferMLAAttnBackend,
|
|
)
|
|
|
|
return (
|
|
"flashinfer",
|
|
FlashInferMLAAttnBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_triton_prefill_backend(self):
|
|
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
|
|
|
|
return (
|
|
"triton",
|
|
TritonAttnBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_intel_amx_prefill_backend(self):
|
|
from sglang.srt.layers.attention.intel_amx_backend import IntelAMXAttnBackend
|
|
|
|
return ("intel_amx", IntelAMXAttnBackend(self.draft_model_runner))
|
|
|
|
def _create_aiter_prefill_backend(self):
|
|
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
|
|
|
|
return ("aiter", AiterAttnBackend(self.draft_model_runner, skip_prefill=False))
|
|
|
|
def _create_fa_prefill_backend(self, fa_impl_ver: int = 3):
|
|
if not is_musa():
|
|
from sglang.srt.layers.attention.flashattention_backend import (
|
|
FlashAttentionBackend,
|
|
)
|
|
else:
|
|
from sglang.srt.hardware_backend.musa.attention.flashattention_backend import (
|
|
MusaFlashAttentionBackend as FlashAttentionBackend,
|
|
)
|
|
return (
|
|
f"fa{fa_impl_ver}",
|
|
FlashAttentionBackend(
|
|
self.draft_model_runner, skip_prefill=False, fa_impl_ver=fa_impl_ver
|
|
),
|
|
)
|
|
|
|
def _create_fa3_prefill_backend(self):
|
|
return self._create_fa_prefill_backend(fa_impl_ver=3)
|
|
|
|
def _create_fa4_prefill_backend(self):
|
|
return self._create_fa_prefill_backend(fa_impl_ver=4)
|
|
|
|
def _create_trtllm_mha_prefill_backend(self):
|
|
from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
|
|
|
|
return (
|
|
"trtllm_mha",
|
|
TRTLLMHAAttnBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_trtllm_mla_prefill_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
raise ValueError(
|
|
"trtllm_mla backend requires MLA model (use_mla_backend=True)."
|
|
)
|
|
|
|
from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
|
|
|
|
return (
|
|
"trtllm_mla",
|
|
TRTLLMMLABackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_tokenspeed_mla_prefill_backend(self):
|
|
if not self.draft_model_runner.use_mla_backend:
|
|
raise ValueError(
|
|
"tokenspeed_mla backend requires MLA model (use_mla_backend=True)."
|
|
)
|
|
|
|
from sglang.srt.layers.attention.tokenspeed_mla_backend import (
|
|
TokenspeedMLABackend,
|
|
)
|
|
|
|
return (
|
|
"tokenspeed_mla",
|
|
TokenspeedMLABackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_ascend_prefill_backend(self):
|
|
from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
|
|
AscendAttnBackend,
|
|
)
|
|
|
|
return ("ascend", AscendAttnBackend(self.draft_model_runner))
|
|
|
|
def _create_flashmla_prefill_backend(self):
|
|
from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend
|
|
|
|
return (
|
|
"flashmla",
|
|
FlashMLABackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
|
|
def _create_dsv4_prefill_backend(self):
|
|
# On NPU the "dsv4" backend resolves to the Ascend V4 subclass; its
|
|
# draft-extend path uses the registered DSV4 prefill backend.
|
|
if is_npu():
|
|
from sglang.srt.layers.attention.attention_registry import (
|
|
ATTENTION_BACKENDS,
|
|
)
|
|
|
|
return ("dsv4", ATTENTION_BACKENDS["dsv4"](self.draft_model_runner))
|
|
elif is_hip():
|
|
from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
|
|
DeepseekV4HipRadixBackend,
|
|
)
|
|
|
|
return (
|
|
"dsv4",
|
|
DeepseekV4HipRadixBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|
|
from sglang.srt.layers.attention.deepseek_v4_backend import (
|
|
DeepseekV4AttnBackend,
|
|
)
|
|
|
|
return (
|
|
"dsv4",
|
|
DeepseekV4AttnBackend(self.draft_model_runner, skip_prefill=False),
|
|
)
|