fix: reorder DSA indexer dual-stream ops to avoid CUDA graph stream explosion (#30025)
Co-authored-by: Khoa Pham <264503018+kpham-sgl@users.noreply.github.com>
This commit is contained in:
@@ -668,7 +668,7 @@ class Envs:
|
|||||||
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
|
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
|
||||||
SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True)
|
SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True)
|
||||||
SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
|
SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
|
||||||
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(True)
|
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
|
||||||
|
|
||||||
# sgl-kernel
|
# sgl-kernel
|
||||||
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
|
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
|
||||||
|
|||||||
@@ -739,6 +739,13 @@ class Indexer(MultiPlatformOp):
|
|||||||
|
|
||||||
current_stream.wait_stream(self.alt_stream)
|
current_stream.wait_stream(self.alt_stream)
|
||||||
self.alt_stream.wait_stream(current_stream)
|
self.alt_stream.wait_stream(current_stream)
|
||||||
|
q_fp8, weights = fused_q_indexer_rope_first_quant(
|
||||||
|
q.contiguous(),
|
||||||
|
weights_raw,
|
||||||
|
q_scale_gate,
|
||||||
|
self._indexer_cos_sin_cache,
|
||||||
|
positions,
|
||||||
|
)
|
||||||
with torch.cuda.stream(self.alt_stream):
|
with torch.cuda.stream(self.alt_stream):
|
||||||
self._fused_k_prepare_and_store(
|
self._fused_k_prepare_and_store(
|
||||||
key,
|
key,
|
||||||
@@ -749,14 +756,6 @@ class Indexer(MultiPlatformOp):
|
|||||||
out_cache_loc=out_cache_loc,
|
out_cache_loc=out_cache_loc,
|
||||||
)
|
)
|
||||||
|
|
||||||
q_fp8, weights = fused_q_indexer_rope_first_quant(
|
|
||||||
q.contiguous(),
|
|
||||||
weights_raw,
|
|
||||||
q_scale_gate,
|
|
||||||
self._indexer_cos_sin_cache,
|
|
||||||
positions,
|
|
||||||
)
|
|
||||||
|
|
||||||
current_stream.wait_stream(self.alt_stream)
|
current_stream.wait_stream(self.alt_stream)
|
||||||
return q_fp8, weights
|
return q_fp8, weights
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user