Disable breakable CUDA graph for NemotronH (#33130)
This commit is contained in:
@@ -129,6 +129,14 @@ def is_deepseek_v4(config) -> bool:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def is_nemotron_h(config) -> bool:
|
||||||
|
return _hf_arch(config) in (
|
||||||
|
"NemotronHForCausalLM",
|
||||||
|
"NemotronHPuzzleForCausalLM",
|
||||||
|
"NemotronHForCausalLMMTP",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
|
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
|
||||||
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
|
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
|
||||||
return config.index_head_dim
|
return config.index_head_dim
|
||||||
|
|||||||
@@ -4426,7 +4426,11 @@ class ServerArgs:
|
|||||||
memory-saver rejection in its own __init__; config-time rules can be
|
memory-saver rejection in its own __init__; config-time rules can be
|
||||||
added here as they're discovered.
|
added here as they're discovered.
|
||||||
"""
|
"""
|
||||||
from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
|
from sglang.srt.configs.model_config import (
|
||||||
|
is_deepseek_dsa,
|
||||||
|
is_deepseek_v4,
|
||||||
|
is_nemotron_h,
|
||||||
|
)
|
||||||
|
|
||||||
rules = [
|
rules = [
|
||||||
# MLA prefill under BCG takes forward_mha, which has no eager
|
# MLA prefill under BCG takes forward_mha, which has no eager
|
||||||
@@ -4437,6 +4441,13 @@ class ServerArgs:
|
|||||||
lambda: self.use_mla_backend()
|
lambda: self.use_mla_backend()
|
||||||
and not is_deepseek_dsa(self.get_model_config().hf_config),
|
and not is_deepseek_dsa(self.get_model_config().hf_config),
|
||||||
),
|
),
|
||||||
|
# NemotronH's hybrid Mamba2 prefill is not BCG-safe: the mamba
|
||||||
|
# state-track write is not wired into the captured buffers, so a
|
||||||
|
# replay can commit a cache slot it never wrote.
|
||||||
|
(
|
||||||
|
"NemotronH (hybrid Mamba2 prefill)",
|
||||||
|
lambda: is_nemotron_h(self.get_model_config().hf_config),
|
||||||
|
),
|
||||||
# DSV4 is BCG-compatible but introduces heavy memory pressure: the
|
# DSV4 is BCG-compatible but introduces heavy memory pressure: the
|
||||||
# c4 indexer scratch is pinned in the capture pool and OOMs. Disable.
|
# c4 indexer scratch is pinned in the capture pool and OOMs. Disable.
|
||||||
(
|
(
|
||||||
|
|||||||
Reference in New Issue
Block a user