Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion python/sglang/srt/models/inkling_common/attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -364,7 +364,8 @@ def _project_qkvr(

def _fused_attn_prologue_verify(self, q, k, v, forward_batch, log_scaling_tau=None):
"""Fused target-verify {k/v sconv + save_windows + qk-norm (+ KV store)}
(jit_kernel/inkling_attn_prologue.py); returns ``(q, k, v, did_store)``.
(kernels/ops/model/inkling/inkling_attn_prologue.py); returns
``(q, k, v, did_store)``.

The fused kernel writes raw bf16 KV, so it only does the store when the
KV pool is bf16: full layers at ``out_cache_loc`` in the full pool,
Expand Down
12 changes: 7 additions & 5 deletions python/sglang/srt/models/inkling_common/kernels/comm.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def _ar_jit():
lazy so importing comm.py doesn't pull in the JIT machinery)."""
if not is_cuda():
return None
from sglang.jit_kernel import inkling_all_reduce
from sglang.kernels.ops.model.inkling import inkling_all_reduce

return inkling_all_reduce

Expand All @@ -103,7 +103,7 @@ def _ar_jit():
def _ar_fused_jit():
if not is_cuda():
return None
from sglang.jit_kernel import inkling_ar_fused
from sglang.kernels.ops.model.inkling import inkling_ar_fused

return inkling_ar_fused

Expand Down Expand Up @@ -240,7 +240,8 @@ def ar_sconv_norm_fusable(
"""True when a decode {all-reduce -> sconv -> add+RMSNorm} chain
(attn-side: wo_ud AR -> attn_sconv -> mlp_norm; MoE-side: MoE AR ->
mlp_sconv -> next attn_norm)
can run as the single fused kernel (jit_kernel/inkling_ar_fused.py). Must be
can run as the single fused kernel
(kernels/ops/model/inkling/inkling_ar_fused.py). Must be
evaluated identically by the producing layer (MoE ``reduce=False``) and the
consuming layer/tail -- it is a pure function of per-forward state."""
if not is_cuda():
Expand Down Expand Up @@ -674,7 +675,7 @@ def all_gather_hidden(input: torch.Tensor, group: GroupCoordinator) -> torch.Ten
def _ar_ssconv_jit():
if not is_cuda():
return None
from sglang.jit_kernel import inkling_ar_scattered_sconv
from sglang.kernels.ops.model.inkling import inkling_ar_scattered_sconv

return inkling_ar_scattered_sconv

Expand All @@ -688,7 +689,8 @@ def scattered_ar_sconv_fusable(
) -> bool:
"""True when an extend {reduce_scatter_hidden -> sconv(shard) ->
all_gather_hidden} chain can run as the single fused v3/v3b-style kernel
(jit_kernel/inkling_ar_scattered_sconv.py). Pure function of per-forward
(kernels/ops/model/inkling/inkling_ar_scattered_sconv.py). Pure function of
per-forward
state -- the producing layer (reduce=False) and the consuming site must
evaluate it identically."""
if not is_cuda():
Expand Down
Loading