[diffusion] feat: layerwise NVTX markers for Nsight Systems profiling (#25683)

Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
mispa-ms
2026-05-26 11:24:27 +08:00
committed by GitHub
co-authored by Mick
parent 1953565ba1
commit 3142278c5f
14 changed files with 942 additions and 128 deletions
+21 -15
View File
@@ -178,12 +178,15 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
# we inject a stand-in module rather than letting patch() trigger
# the real import.
fake_module = MagicMock()
with patch(
"sglang.srt.mem_cache.registry.envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get",
return_value=True,
), patch.dict(
"sys.modules",
{"sglang.srt.mem_cache.radix_cache_cpp": fake_module},
with (
patch(
"sglang.srt.mem_cache.registry.envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get",
return_value=True,
),
patch.dict(
"sys.modules",
{"sglang.srt.mem_cache.radix_cache_cpp": fake_module},
),
):
result = default_radix_cache_factory(ctx)
fake_module.RadixCacheCpp.assert_called_once_with(
@@ -196,15 +199,18 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
# Shim both factory imports — each transitively loads sgl_kernel.
fake_components = MagicMock()
fake_radix = MagicMock()
with patch(
"sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get",
return_value=True,
), patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
with (
patch(
"sglang.srt.mem_cache.registry.envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get",
return_value=True,
),
patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
),
):
result = default_radix_cache_factory(ctx)
fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params)