minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Spec] Enable FlashInfer autotune for spec draft (#29595)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-07-01 13:36:07 -07:00
committed by GitHub
parent 4a8e76805c
commit 8f0d320d31
16 changed files with 313 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/runner/base_runner.py
+11 -153
View File
@@ -15,12 +15,9 @@
from __future__ import annotations
import datetime
import hashlib
import inspect
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any, Optional, Tuple
@@ -43,6 +40,10 @@ from sglang.srt.model_executor.forward_batch_info import (
PPProxyTensors,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner.flashinfer_autotune import (
run_flashinfer_autotune_forward,
should_run_flashinfer_autotune,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.spec_info import create_dummy_verify_input
from sglang.srt.utils import (
@@ -213,7 +214,7 @@ class BaseRunner(ABC):
self._pre_initialize_flashinfer_allreduce_workspace()
if self._should_run_flashinfer_autotune():
if should_run_flashinfer_autotune(self.model_runner):
buffers, batch_size = self._autotune_buffers()
assert (
buffers is not None
@@ -250,83 +251,6 @@ class BaseRunner(ABC):
dtype=mr.dtype,
)
def _should_run_flashinfer_autotune(self) -> bool:
"""Check if flashinfer autotune should be run."""
mr = self.model_runner
if mr.server_args.disable_flashinfer_autotune:
return False
# CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
# be autotuned -- its _dummy_run would dispatch more tokens per rank than
# SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
# Read server_args directly to avoid depending on initialize_moe_config()
# having already populated the MoE backend globals.
if (
mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
and mr.server_args.moe_a2a_backend == "deepep"
):
return False
backend_str = mr.server_args.moe_runner_backend
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend
moe_needs_autotune = backend_str in [
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
"flashinfer_mxfp4",
"flashinfer_cutedsl",
"flashinfer_cutlass",
]
from sglang.srt.layers.quantization.fp4_utils import (
get_fp4_gemm_runner_backend,
)
model_uses_fp4 = mr.model_config.quantization in (
"modelopt_fp4",
"modelopt_mixed",
)
fp4_gemm_needs_autotune = model_uses_fp4 and (
get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
)
from sglang.srt.layers.quantization.fp8_utils import (
get_fp8_gemm_runner_backend,
)
from sglang.srt.utils import is_sm100_supported
model_uses_modelopt_fp8 = mr.model_config.quantization in (
"modelopt",
"modelopt_fp8",
"modelopt_mixed",
)
# Online MXFP8 (microscaling) linears dispatch to flashinfer's
# ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
# exercise correctly -- it triggers an illegal memory access inside the
# mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
# tuning, so skip autotune for these models.
model_uses_mxfp8 = "mxfp8" in (mr.model_config.quantization or "")
fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
or (model_uses_modelopt_fp8 and is_sm100_supported())
)
if not (
moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune
):
return False
major, _ = torch.cuda.get_device_capability()
if major < 9:
return False
if mr.spec_algorithm.is_speculative():
return not mr.is_draft_worker
return True
def _flashinfer_autotune(self, *, buffers, batch_size):
"""Run flashinfer autotune.
@@ -335,76 +259,11 @@ class BaseRunner(ABC):
Supplied by warmup() (the decode runner's captured buffers when a graph
runner exists; a freshly-allocated dummy set in the eager path).
"""
from flashinfer.autotuner import autotune
from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
def forward_fn():
self._dummy_run(batch_size=batch_size, buffers=buffers)
mr = self.model_runner
cache_path = self._flashinfer_autotune_cache_path()
if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
autotune_cache = cache_path
logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
else:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
runs_dir = cache_path.parent / "runs"
runs_dir.mkdir(parents=True, exist_ok=True)
autotune_cache = (
runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
)
logger.info(
"Running FlashInfer autotune (cache reuse DISABLED via "
"SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
autotune_cache,
)
# Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
# calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
mr.forward_stream.wait_stream(torch.cuda.current_stream())
with torch.get_device_module(mr.device).stream(mr.forward_stream):
with (
torch.inference_mode(),
autotune(True, cache=str(autotune_cache)),
autotune_dummy_run_mode(),
):
self._dummy_run(batch_size=batch_size, buffers=buffers)
torch.cuda.current_stream().wait_stream(mr.forward_stream)
logger.info("FlashInfer autotune completed.")
def _flashinfer_autotune_cache_path(self) -> Path:
import flashinfer
mr = self.model_runner
major, minor = torch.cuda.get_device_capability(mr.device)
arch = f"sm{major}{minor}"
flashinfer_version = getattr(flashinfer, "__version__", "unknown")
server_args = mr.server_args
model_key = "|".join(
[
str(server_args.model_path),
str(mr.dtype),
str(server_args.quantization),
str(server_args.moe_runner_backend),
str(mr.tp_size),
str(mr.pp_size),
str(mr.dp_size),
str(mr.moe_ep_size),
str(mr.model_config.hf_config.__class__.__name__),
]
)
cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
cache_dir = (
Path(envs.SGLANG_CACHE_DIR.get())
/ "flashinfer"
/ "autotune"
/ flashinfer_version
/ arch
/ cache_key
)
cache_dir.mkdir(parents=True, exist_ok=True)
return (
cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
)
run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)
def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_bs: int = 1):
"""Allocate one static decode-buffer set for a dummy forward, sized to
@@ -413,10 +272,9 @@ class BaseRunner(ABC):
The PP-parallel DeepGEMM warmup sweeps batch sizes far larger than any
runner's max_bs (up to ~n_sms*block_m), so no pre-allocated runner buffer
set fits; it builds one here and hands it to _dummy_run (reused across the
sweep; _dummy_run slices it per shape). The flashinfer autotune does NOT
use this -- it reuses an existing runner's buffers via _autotune_buffers
(the eager input registry, or the decode cuda-graph runner's captured
buffers).
sweep; _dummy_run slices it per shape). Eager FlashInfer autotune also
allocates decode-shaped scratch buffers here. Decode cuda-graph autotune
reuses the captured runner buffers instead.
"""
mr = self.model_runner
return _allocate_decode_buffers(
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+15 -5
View File
@@ -71,6 +71,9 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
freeze_gc,
get_batch_sizes_to_capture,
)
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
@@ -853,15 +856,22 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
with canary_ctx:
shape_key = self._make_graph_key(bs, stream_idx, variant_label)
post_warmup_hook = getattr(
self.model_runner.attn_backend,
"on_after_cuda_graph_warmup",
None,
)
maybe_flashinfer_autotune_speculative_draft(
self,
run_once,
post_warmup_hook=post_warmup_hook,
skip_logits=False,
)
self.backend.capture_one(
shape_key,
run_once,
dummies=None,
post_warmup_hook=getattr(
self.model_runner.attn_backend,
"on_after_cuda_graph_warmup",
None,
),
post_warmup_hook=post_warmup_hook,
)
def recapture_if_needed(self, forward_batch: ForwardBatch):
python/sglang/srt/model_executor/runner/flashinfer_autotune.py
+224
View File
@@ -0,0 +1,224 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from __future__ import annotations
import contextlib
import datetime
import hashlib
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Optional
import torch
from sglang.srt.environ import envs
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.model_executor.runner.base_runner import BaseRunner
logger = logging.getLogger(__name__)
def should_run_flashinfer_autotune(
model_runner: ModelRunner, *, for_speculative_draft: bool = False
) -> bool:
"""Check if flashinfer autotune should be run."""
mr = model_runner
if mr.device != "cuda":
return False
if mr.server_args.disable_flashinfer_autotune:
return False
# CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
# be autotuned -- its _dummy_run would dispatch more tokens per rank than
# SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
# Read server_args directly to avoid depending on initialize_moe_config()
# having already populated the MoE backend globals.
if (
mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
and mr.server_args.moe_a2a_backend == "deepep"
):
return False
backend_str = mr.server_args.moe_runner_backend
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend
moe_needs_autotune = backend_str in [
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
"flashinfer_mxfp4",
"flashinfer_cutedsl",
"flashinfer_cutlass",
]
from sglang.srt.layers.quantization.fp4_utils import (
get_fp4_gemm_runner_backend,
)
model_quantization = mr.model_config.quantization
model_uses_fp4 = model_quantization in (
"modelopt_fp4",
"modelopt_mixed",
)
fp4_gemm_needs_autotune = model_uses_fp4 and (
get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
)
from sglang.srt.layers.quantization.fp8_utils import (
get_fp8_gemm_runner_backend,
)
from sglang.srt.utils import is_sm100_supported
model_uses_modelopt_fp8 = model_quantization in (
"modelopt",
"modelopt_fp8",
"modelopt_mixed",
)
# Online MXFP8 (microscaling) linears dispatch to flashinfer's
# ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
# exercise correctly -- it triggers an illegal memory access inside the
# mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
# tuning, so skip autotune for these models.
model_uses_mxfp8 = "mxfp8" in (model_quantization or "")
fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
or (model_uses_modelopt_fp8 and is_sm100_supported())
)
if not (moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune):
return False
if torch.cuda.get_device_capability()[0] < 9:
return False
if mr.spec_algorithm.is_speculative():
return mr.is_draft_worker if for_speculative_draft else not mr.is_draft_worker
return True
def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
import flashinfer
mr = model_runner
major, minor = torch.cuda.get_device_capability(mr.device)
arch = f"sm{major}{minor}"
flashinfer_version = getattr(flashinfer, "__version__", "unknown")
server_args = mr.server_args
model_key_parts = [
str(server_args.model_path),
str(mr.dtype),
str(server_args.quantization),
str(server_args.moe_runner_backend),
str(mr.tp_size),
str(mr.pp_size),
str(mr.dp_size),
str(mr.moe_ep_size),
str(mr.model_config.hf_config.__class__.__name__),
]
if mr.is_draft_worker:
model_key_parts.append(f"draft_quant={mr.model_config.quantization}")
model_key = "|".join(model_key_parts)
cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
cache_dir = (
Path(envs.SGLANG_CACHE_DIR.get())
/ "flashinfer"
/ "autotune"
/ flashinfer_version
/ arch
/ cache_key
)
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
@contextlib.contextmanager
def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool):
from flashinfer.autotuner import autotune
mr = model_runner
cache_path = flashinfer_autotune_cache_path(mr)
if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
autotune_cache = cache_path
logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
else:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
runs_dir = cache_path.parent / "runs"
runs_dir.mkdir(parents=True, exist_ok=True)
autotune_cache = runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
logger.info(
"Running FlashInfer autotune (cache reuse DISABLED via "
"SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
autotune_cache,
)
# Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
# calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
mr.forward_stream.wait_stream(torch.cuda.current_stream())
with torch.get_device_module(mr.device).stream(mr.forward_stream):
maybe_skip_logits = contextlib.nullcontext()
if skip_logits:
from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
maybe_skip_logits = autotune_dummy_run_mode()
with torch.inference_mode(), autotune(
True, cache=str(autotune_cache)
), maybe_skip_logits:
yield
torch.cuda.current_stream().wait_stream(mr.forward_stream)
logger.info("FlashInfer autotune completed.")
def run_flashinfer_autotune_forward(
model_runner: ModelRunner, forward_fn: Callable[[], None], *, skip_logits: bool
) -> None:
"""Run flashinfer autotune forward."""
with flashinfer_autotune_context(model_runner, skip_logits=skip_logits):
forward_fn()
def maybe_flashinfer_autotune_speculative_draft(
runner: BaseRunner,
forward_fn: Callable[[], None],
*,
post_warmup_hook: Optional[Callable[[], None]] = None,
skip_logits: bool = False,
) -> None:
"""Run speculative draft flashinfer autotune."""
mr = runner.model_runner
phase_key = f"{runner.__class__.__module__}.{runner.__class__.__qualname__}"
tuned_phases = getattr(mr, "_flashinfer_spec_draft_autotuned_phases", None)
if tuned_phases is None:
tuned_phases = set()
mr._flashinfer_spec_draft_autotuned_phases = tuned_phases
if phase_key in tuned_phases:
return
if (
not mr.spec_algorithm.is_speculative()
or not mr.is_draft_worker
or not should_run_flashinfer_autotune(mr, for_speculative_draft=True)
):
return
def run_and_reset():
forward_fn()
if post_warmup_hook is not None:
post_warmup_hook()
run_flashinfer_autotune_forward(mr, run_and_reset, skip_logits=skip_logits)
tuned_phases.add(phase_key)
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+13 -3
View File
@@ -28,6 +28,9 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -440,13 +443,20 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
forward_batch.mark_forward_metadata_ready()
self.deepep_adapter.capture(is_extend_in_batch=False)
shape_key = self._make_graph_key(num_seqs)
post_warmup_hook = getattr(
self.draft_attn_backend, "on_after_cuda_graph_warmup", None
)
maybe_flashinfer_autotune_speculative_draft(
self,
run_once,
post_warmup_hook=post_warmup_hook,
skip_logits=False,
)
self.backend.capture_one(
shape_key,
run_once,
dummies=None,
post_warmup_hook=getattr(
self.draft_attn_backend, "on_after_cuda_graph_warmup", None
),
post_warmup_hook=post_warmup_hook,
)
def _postprocess_output_to_raw_bs(self, out, raw_bs):
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+15 -5
View File
@@ -28,6 +28,9 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -432,15 +435,22 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
)
with canary_ctx:
shape_key = self._make_graph_key(bs)
post_warmup_hook = getattr(
self.draft_extend_attn_backend,
"on_after_cuda_graph_warmup",
None,
)
maybe_flashinfer_autotune_speculative_draft(
self,
run_once,
post_warmup_hook=post_warmup_hook,
skip_logits=False,
)
self.backend.capture_one(
shape_key,
run_once,
dummies=None,
post_warmup_hook=getattr(
self.draft_extend_attn_backend,
"on_after_cuda_graph_warmup",
None,
),
post_warmup_hook=post_warmup_hook,
)
def execute(self, forward_batch: ForwardBatch):
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+13 -3
View File
@@ -25,6 +25,9 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -321,13 +324,20 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
)
self.deepep_adapter.capture(is_extend_in_batch=False)
shape_key = self._make_graph_key(request_bs)
post_warmup_hook = getattr(
self.draft_attn_backend, "on_after_cuda_graph_warmup", None
)
maybe_flashinfer_autotune_speculative_draft(
self,
run_once,
post_warmup_hook=post_warmup_hook,
skip_logits=False,
)
self.backend.capture_one(
shape_key,
run_once,
dummies=None,
post_warmup_hook=getattr(
self.draft_attn_backend, "on_after_cuda_graph_warmup", None
),
post_warmup_hook=post_warmup_hook,
)
finally:
self.draft_attn_backend.token_to_kv_pool = saved_backend_pool
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+13 -3
View File
@@ -52,6 +52,9 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -369,13 +372,20 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
self.deepep_adapter.capture(is_extend_in_batch=True)
shape_key = self._make_graph_key(bs)
post_warmup_hook = getattr(
self.attn_backend, "on_after_cuda_graph_warmup", None
)
maybe_flashinfer_autotune_speculative_draft(
self,
run_once,
post_warmup_hook=post_warmup_hook,
skip_logits=False,
)
self.backend.capture_one(
shape_key,
run_once,
dummies=None,
post_warmup_hook=getattr(
self.attn_backend, "on_after_cuda_graph_warmup", None
),
post_warmup_hook=post_warmup_hook,
)
def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput):
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+1
View File
@@ -276,6 +276,7 @@ class TinyModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = sliding_window_size is not None
self.is_local_attention_model = sliding_window_size is not None
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+1
View File
@@ -242,6 +242,7 @@ class TinyDSAModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.attention_chunk_size = None
self.sliding_window_size = None
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+1
View File
@@ -265,6 +265,7 @@ class TinyDSV4ModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py
+1
View File
@@ -274,6 +274,7 @@ class TinyDualChunkModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.attention_chunk_size = None
self.sliding_window_size = None
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+1
View File
@@ -178,6 +178,7 @@ class TinyGDNModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py
+1
View File
@@ -184,6 +184,7 @@ class TinyKDAModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+1
View File
@@ -187,6 +187,7 @@ class TinyLightningModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+1
View File
@@ -275,6 +275,7 @@ class TinyMamba2ModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+1
View File
@@ -178,6 +178,7 @@ class TinyMLAModelConfig:
self.is_encoder_decoder = False
self.is_multimodal = False
self.is_generation = True
self.quantization = None
self.is_hybrid_swa = False
self.is_local_attention_model = False
self.attention_chunk_size = None