kvcache-ai/sglang

Fork 0

mirror of https://github.com/kvcache-ai/sglang.git synced 2026-11-05 01:07:39 +00:00

[PR #77] fix(kt): restore integer Marlin weights without gradients #78

New issue

Open

opened 2026-10-29 21:07:01 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:01 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/77
Author: @zupengwang
Created: 8/9/2026
Status: 🔄 Open

Base: main ← Head: fix/kt-restore-integer-weights

📝 Commits (1)

f95c539 fix(kt): restore integer weights without gradients

📊 Changes

2 files changed (+36 additions, -1 deletions)

View changed files

📝 python/sglang/srt/layers/moe/kt_ep_wrapper.py (+2 -1)
➕ test/srt/test_kt_ep_wrapper.py (+34 -0)

📄 Description

Motivation

SharedFullContext._restore_raw_attrs() reconstructs raw Marlin weight tensors as torch.nn.Parameter objects. For integer weights such as the int32 Marlin layout, the default requires_grad=True is invalid and raises:

RuntimeError: Only Tensors of floating point and complex dtype can require gradients

This is the failure reported after hybrid GPU prefill in kvcache-ai/ktransformers#2107.

Modifications

Restore raw weight parameters with requires_grad=False, matching inference-only model weights and supporting integer dtypes.
Add a focused CUDA regression test that exercises the actual SharedFullContext._restore_raw_attrs() path.

Accuracy Tests

Tested on an RTX 3090 with PyTorch 2.9.1+cu128:

python -m pytest -q test/srt/test_kt_ep_wrapper.py

# Before
1 failed: RuntimeError: Only Tensors of floating point and complex dtype can require gradients

# After
1 passed in 19.61s

The test verifies the restored parameter's shape, int32 dtype, CUDA device, and requires_grad=False. This change does not alter weight values or model-forward math.

Additional checks:

python -m py_compile python/sglang/srt/layers/moe/kt_ep_wrapper.py test/srt/test_kt_ep_wrapper.py
python -m black --check test/srt/test_kt_ep_wrapper.py
python -m isort --check-only test/srt/test_kt_ep_wrapper.py
python -m ruff check --select=F401,F821 python/sglang/srt/layers/moe/kt_ep_wrapper.py test/srt/test_kt_ep_wrapper.py
git diff --check

All passed.

Benchmarking and Profiling

Not applicable. The change only controls gradient metadata while recreating a parameter after prefill; it does not change the inference compute path.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations. (Not applicable: no user-facing behavior or configuration changes.)
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (Not applicable: no model math or performance path changes.)
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/77 **Author:** [@zupengwang](https://github.com/zupengwang) **Created:** 8/9/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/kt-restore-integer-weights` --- ### 📝 Commits (1) - [`f95c539`](https://github.com/kvcache-ai/sglang/commit/f95c539207c0dd00324f0abf6903745f5715ccdd) fix(kt): restore integer weights without gradients ### 📊 Changes **2 files changed** (+36 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (+2 -1) ➕ `test/srt/test_kt_ep_wrapper.py` (+34 -0) </details> ### 📄 Description  ## Motivation `SharedFullContext._restore_raw_attrs()` reconstructs raw Marlin weight tensors as `torch.nn.Parameter` objects. For integer weights such as the `int32` Marlin layout, the default `requires_grad=True` is invalid and raises: ```text RuntimeError: Only Tensors of floating point and complex dtype can require gradients ``` This is the failure reported after hybrid GPU prefill in kvcache-ai/ktransformers#2107. ## Modifications - Restore raw weight parameters with `requires_grad=False`, matching inference-only model weights and supporting integer dtypes. - Add a focused CUDA regression test that exercises the actual `SharedFullContext._restore_raw_attrs()` path. ## Accuracy Tests Tested on an RTX 3090 with PyTorch 2.9.1+cu128: ```text python -m pytest -q test/srt/test_kt_ep_wrapper.py # Before 1 failed: RuntimeError: Only Tensors of floating point and complex dtype can require gradients # After 1 passed in 19.61s ``` The test verifies the restored parameter's shape, `int32` dtype, CUDA device, and `requires_grad=False`. This change does not alter weight values or model-forward math. Additional checks: ```text python -m py_compile python/sglang/srt/layers/moe/kt_ep_wrapper.py test/srt/test_kt_ep_wrapper.py python -m black --check test/srt/test_kt_ep_wrapper.py python -m isort --check-only test/srt/test_kt_ep_wrapper.py python -m ruff check --select=F401,F821 python/sglang/srt/layers/moe/kt_ep_wrapper.py test/srt/test_kt_ep_wrapper.py git diff --check ``` All passed. ## Benchmarking and Profiling Not applicable. The change only controls gradient metadata while recreating a parameter after prefill; it does not change the inference compute path. ## Checklist - [x] Format your code according to the [Format code with pre-commit](https://docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-commit). - [x] Add unit tests according to the [Run and add unit tests](https://docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-tests). - [ ] Update documentation according to [Write documentations](https://docs.sglang.io/developer_guide/contribution_guide.html#write-documentations). (Not applicable: no user-facing behavior or configuration changes.) - [ ] Provide accuracy and speed benchmark results according to [Test the accuracy](https://docs.sglang.io/developer_guide/contribution_guide.html#test-the-accuracy) and [Benchmark the speed](https://docs.sglang.io/developer_guide/contribution_guide.html#benchmark-the-speed). (Not applicable: no model math or performance path changes.) - [x] Follow the SGLang code style [guidance](https://docs.sglang.io/developer_guide/contribution_guide.html#code-style-guidance). ## Review Process 1. Ping Merge Oncalls to start the PR flow. See the [PR Merge Process](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md#pull-request-merge-process). 2. Get approvals from [CODEOWNERS](https://github.com/sgl-project/sglang/blob/main/.github/CODEOWNERS) and other reviewers. 3. Trigger CI tests with [comments](https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests) or contact authorized users to do so. - `/tag-run-ci-label`, `/rerun-failed-ci`, `/tag-and-rerun-ci` 4. After green CI and required approvals, ask Merge Oncalls to merge. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:01 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/sglang#78

No description provided.

Rows
Columns