Skip to content

Commit 4dfbd4c

Browse files
AndyKong2020AndyKong2020
authored
Fix conv1d update state shift (sgl-project#428)
* fix(causal_conv1d_update): pre-shift conv_state to match rolling-buffer convention The kernel's in-place sliding window only writes back to positions [VAL .. state_len-1] (where VAL = state_len - seq_len), leaving the older prefix state[0..VAL-1] stale. This diverges from the vLLM/SGLang rolling-buffer convention where the final state should equal [old_state[seq_len:state_len], x[0:seq_len]]. Fix at the host integration layer by pre-shifting state[0..VAL-1] := state[seq_len..seq_len+VAL-1] before launching the kernel. The kernel only reads/writes positions [VAL..state_len-1], so the pre-shifted prefix is preserved through kernel execution and the final conv_state matches the reference. Pad-slot entries are filtered out, and the no-indices path is handled symmetrically. Verified by tests/python/sgl_kernel_npu/test_conv1d_update.py: output and conv_state both reach 100% match against the vLLM reference (previously state was 66.70%). * test(causal_conv1d_update): scale up NPU test config Bump BSZ 1->32, SEQ_LEN 1->4, HIDDEN 4096->2048, CACHE_LEN 10->32 to exercise the multi-batch / multi-token spec-decoding path that previously masked the conv_state rolling-buffer mismatch. * fix(causal_conv1d_update): tile along dim to avoid kernel UB overflow The kernel allocates ~26*dim bytes of UB per AIV core (width*dim weight plus several dim*sizeof operand/accumulator buffers). On Atlas A3 the 192KB UB budget is exceeded once dim is roughly above 7500, which manifests as a vector core "ub address out of bounds" exception. Real workloads such as Qwen3-Next call the op with conv_dim well past that, crashing in hybrid_linear_attn_backend.forward_extend. Fix at the host integration layer without touching the kernel: * Extract the kernel-launch / tiling-data / cache-management code into a per-tile lambda parameterized by tile_dim. * When dim exceeds kMaxDimTile (4096), slice x / weight / conv_state / bias along the dim axis, materialize contiguous tile copies (the kernel still requires contiguous inputs), launch the kernel per tile, and stitch y / conv_state back into the full tensors. * The pre-shift that maintains the rolling-buffer convention runs once on the full conv_state and is independent of tiling. Verified locally with a sweep over HIDDEN in {4096, 6144, 8192, 10240, 12288}: bf16 max output diff is 6.25e-2 (one ULP at magnitude) and conv_state matches exactly against the vLLM reference at every dim. * fix(causal_conv1d_update): make op graph-capture compatible 1. Add mutation annotation Tensor(a!) for conv_state in op schema so graph mode correctly tracks the in-place update. 2. Replace masked_select (dynamic-shape output) and data-dependent branch in pre-shift logic with clamp_min(0), ensuring all tensor shapes are statically known during graph capture. * test(causal_conv1d_update): adjust test config parameters * fix(causal_conv1d_update): replace aclrtMemcpy with CopyTensorHostToDevice for stream ordering Raw aclrtMemcpy bypasses OpCommand task queue, breaking stream ordering between pre-shift PyTorch ops and kernel launch in graph/non-blocking mode (precision errors without ASCEND_LAUNCH_BLOCKING). Use CopyTensorHostToDevice (pin_memory + .to(device)) for tiling cache miss paths so all device ops go through the same dispatch queue. * fix(causal_conv1d_update): fix stream ordering and reduce peak memory 1. Replace EXEC_KERNEL_CMD (RunOpApi) with OpCommand.SetCustomHandler pattern (same as lora ops) to ensure kernel launch goes through the same OpCommand task queue as pre-shift PyTorch ops. RunOpApi may bypass the task queue, causing stream ordering issues in graph / non-blocking mode (precision errors without ASCEND_LAUNCH_BLOCKING). 2. Optimize pre-shift memory: slice first then index_select to allocate only [B, val_shift, dim] instead of full [B, state_len, dim]. Reduces peak memory for large batch sizes. * Revert "fix(causal_conv1d_update): fix stream ordering and reduce peak memory" This reverts commit 2f0fe83. * fix(causal_conv1d_update): read padSlotId in kernel tiling parse Kernel Init() and ParseTilingData() did not read padSlotId from the tiling buffer, leaving it uninitialised. When the garbage value happened to match a valid stateIdx the corresponding batch was silently skipped, producing wrong output. The probability grows with batch size, explaining the "large BSZ → bad precision" symptom. * style: apply clang-format to causal_conv1d_update
1 parent d64de59 commit 4dfbd4c

5 files changed

Lines changed: 146 additions & 67 deletions

File tree

‎csrc/causal_conv1d_update/op_host/causal_conv1d_update.cpp‎

Lines changed: 138 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,31 @@ HOST_API at::Tensor causal_conv1d_update_impl(const at::Tensor &x, const at::Ten
103103
// Create output tensor
104104
at::Tensor y = at::empty_like(x);
105105

106+
// Pre-shift conv_state to maintain the rolling-buffer convention used by the
107+
// vLLM/SGLang reference. The kernel only writes positions [VAL .. state_len-1]
108+
// (where VAL = state_len - seq_len), leaving the older "scratch" prefix stale.
109+
// We shift state[0..VAL-1] := state[seq_len..seq_len+VAL-1] beforehand so that
110+
// after the kernel completes, the full conv_state matches the rolling window
111+
// [old[seq_len:state_len], x[0:seq_len]].
112+
const int64_t val_shift = state_len - seq_len;
113+
if (val_shift > 0) {
114+
if (has_indices) {
115+
// Map pad_slot_id (e.g. -1) to index 0 so that index_select stays
116+
// in-bounds. Shifting a pad slot is harmless because its data is
117+
// meaningless, and duplicate writes to index 0 are all identical
118+
// (same source row, same shift), so the result is correct.
119+
auto cs_indices_long = conv_state_indices.to(at::kLong);
120+
auto safe_indices = cs_indices_long.clamp_min(0);
121+
auto cs_view = conv_state.index_select(0, safe_indices); // [B, state_len, dim]
122+
auto src = cs_view.slice(1, seq_len, seq_len + val_shift).clone();
123+
cs_view.slice(1, 0, val_shift).copy_(src);
124+
conv_state.index_copy_(0, safe_indices, cs_view);
125+
} else {
126+
auto src = conv_state.slice(1, seq_len, seq_len + val_shift).clone();
127+
conv_state.slice(1, 0, val_shift).copy_(src);
128+
}
129+
}
130+
106131
auto ascendc_platform = platform_ascendc::PlatformAscendCManager::GetInstance();
107132
int32_t max_aiv_core = static_cast<int32_t>(ascendc_platform->GetCoreNumAiv());
108133
int32_t block_dim = std::min(max_aiv_core, static_cast<int32_t>(batch));
@@ -111,69 +136,121 @@ HOST_API at::Tensor causal_conv1d_update_impl(const at::Tensor &x, const at::Ten
111136
}
112137
int32_t workspace_size = static_cast<int32_t>(ascendc_platform->GetLibApiWorkSpaceSize());
113138

114-
// 1. Prepare Tiling Data Struct
115-
CausalConv1dUpdateTilingData tiling_data;
116-
SGLang::CausalConv1dUpdate::ComputeTilingData(batch, seq_len, dim, width, state_len, has_indices, has_bias,
117-
has_num_accept, has_query_loc, activation_mode, pad_slot_id,
118-
block_dim, tiling_data);
119-
120-
int32_t tilingSize = (sizeof(CausalConv1dUpdateTilingData) + PADDING_BYTE - 1) / PADDING_BYTE * PADDING_BYTE;
121-
at::Tensor tilingTensor;
122-
123-
// 2. Hash computation
124-
CausalConv1dUpdateTilingKey key{.batch = batch,
125-
.seqLen = seq_len,
126-
.dim = dim,
127-
.width = width,
128-
.stateLen = state_len,
129-
.hasIndices = has_indices ? 1 : 0,
130-
.hasBias = has_bias ? 1 : 0,
131-
.hasNumAccept = has_num_accept ? 1 : 0,
132-
.hasQueryLoc = has_query_loc ? 1 : 0,
133-
.activationMode = activation_mode ? 1 : 0,
134-
.padSlotId = pad_slot_id};
135-
uint64_t hashValue = CausalConv1dUpdateTilingKeyHash{}(key);
136-
137-
// 3. cache management
138-
static auto globalTilingBuffer =
139-
at::empty({tilingSize * MAX_CAPTURE_NUM}, at::TensorOptions().dtype(at::kByte).device(x.options().device()));
140-
141-
if (captureMap.find(hashValue) != captureMap.end()) {
142-
tilingTensor = at::from_blob(globalTilingBuffer.data_ptr<uint8_t>() + (tilingSize * captureMap[hashValue]),
143-
tilingSize, at::kByte);
144-
} else if (actualCaptureNum >= MAX_CAPTURE_NUM) {
145-
static auto tilingBuffer =
146-
at::empty({tilingSize}, at::TensorOptions().dtype(at::kByte).device(x.options().device()));
147-
aclrtMemcpy(tilingBuffer.data_ptr<uint8_t>(), sizeof(CausalConv1dUpdateTilingData), &tiling_data,
148-
sizeof(CausalConv1dUpdateTilingData), ACL_MEMCPY_HOST_TO_DEVICE);
149-
tilingTensor = at::from_blob(tilingBuffer.data_ptr<uint8_t>(), tilingSize, at::kByte);
150-
} else {
151-
captureMap[hashValue] = actualCaptureNum;
152-
aclrtMemcpy(globalTilingBuffer.data_ptr<uint8_t>() + actualCaptureNum * tilingSize,
153-
sizeof(CausalConv1dUpdateTilingData), &tiling_data, sizeof(CausalConv1dUpdateTilingData),
154-
ACL_MEMCPY_HOST_TO_DEVICE);
155-
actualCaptureNum++;
156-
tilingTensor = at::from_blob(globalTilingBuffer.data_ptr<uint8_t>() + (tilingSize * captureMap[hashValue]),
157-
tilingSize, at::kByte);
139+
// Per-tile launch helper. Captures everything needed except the tensors and the
140+
// current tile_dim, so we can call it once for the whole tensor or once per tile
141+
// when dim exceeds the kernel's UB budget.
142+
auto launch_for_tile = [&](const at::Tensor &x_t, const at::Tensor &weight_t, const at::Tensor &conv_state_t,
143+
const at::Tensor &bias_t, at::Tensor &y_t, int64_t tile_dim) {
144+
// 1. Prepare Tiling Data Struct
145+
CausalConv1dUpdateTilingData tiling_data;
146+
SGLang::CausalConv1dUpdate::ComputeTilingData(batch, seq_len, tile_dim, width, state_len, has_indices, has_bias,
147+
has_num_accept, has_query_loc, activation_mode, pad_slot_id,
148+
block_dim, tiling_data);
149+
150+
int32_t tilingSize = (sizeof(CausalConv1dUpdateTilingData) + PADDING_BYTE - 1) / PADDING_BYTE * PADDING_BYTE;
151+
at::Tensor tilingTensor;
152+
153+
// 2. Hash computation
154+
CausalConv1dUpdateTilingKey key{.batch = batch,
155+
.seqLen = seq_len,
156+
.dim = tile_dim,
157+
.width = width,
158+
.stateLen = state_len,
159+
.hasIndices = has_indices ? 1 : 0,
160+
.hasBias = has_bias ? 1 : 0,
161+
.hasNumAccept = has_num_accept ? 1 : 0,
162+
.hasQueryLoc = has_query_loc ? 1 : 0,
163+
.activationMode = activation_mode ? 1 : 0,
164+
.padSlotId = pad_slot_id};
165+
uint64_t hashValue = CausalConv1dUpdateTilingKeyHash{}(key);
166+
167+
// Helper: wrap tiling_data in a CPU tensor and copy to device via
168+
// torch_npu dispatch so the H2D transfer is properly ordered on the
169+
// same task queue as the preceding pre-shift ops and the subsequent
170+
// kernel launch. Using raw aclrtMemcpy here would bypass OpCommand
171+
// and break stream ordering in graph / non-blocking mode.
172+
auto copyTilingToDevice = [&]() {
173+
auto cpuTiling = at::empty({tilingSize}, at::kByte);
174+
std::memcpy(cpuTiling.data_ptr(), &tiling_data, sizeof(CausalConv1dUpdateTilingData));
175+
return TorchNpuHelper::CopyTensorHostToDevice(cpuTiling);
176+
};
177+
178+
// 3. cache management
179+
static auto globalTilingBuffer = at::empty({tilingSize * MAX_CAPTURE_NUM},
180+
at::TensorOptions().dtype(at::kByte).device(x_t.options().device()));
181+
182+
if (captureMap.find(hashValue) != captureMap.end()) {
183+
tilingTensor = at::from_blob(globalTilingBuffer.data_ptr<uint8_t>() + (tilingSize * captureMap[hashValue]),
184+
tilingSize, at::kByte);
185+
} else if (actualCaptureNum >= MAX_CAPTURE_NUM) {
186+
// Overflow: no room in cache, create a one-shot device copy.
187+
tilingTensor = copyTilingToDevice();
188+
} else {
189+
// New capture: copy to device and cache in globalTilingBuffer.
190+
captureMap[hashValue] = actualCaptureNum;
191+
auto deviceTiling = copyTilingToDevice();
192+
// D2D copy into the pre-allocated cache buffer (ordered via dispatch).
193+
globalTilingBuffer.slice(0, actualCaptureNum * tilingSize, actualCaptureNum * tilingSize + tilingSize)
194+
.copy_(deviceTiling);
195+
actualCaptureNum++;
196+
tilingTensor = at::from_blob(globalTilingBuffer.data_ptr<uint8_t>() + (tilingSize * captureMap[hashValue]),
197+
tilingSize, at::kByte);
198+
}
199+
200+
// 4. Create workspace
201+
auto workspace_tensor =
202+
at::empty({workspace_size}, at::TensorOptions().dtype(at::kByte).device(x_t.options().device()));
203+
204+
// 5. Launch kernel
205+
if (dtype == at::kBFloat16) {
206+
EXEC_KERNEL_CMD(causal_conv1d_update_bfloat16_t, block_dim, x_t, weight_t, conv_state_t,
207+
has_indices ? conv_state_indices : at::empty(0, x_t.options()),
208+
has_bias ? bias_t : at::empty(0, x_t.options()),
209+
has_num_accept ? num_accepted_tokens : at::empty(0, x_t.options()),
210+
has_query_loc ? query_start_loc : at::empty(0, x_t.options()), y_t, workspace_tensor,
211+
tilingTensor);
212+
} else {
213+
EXEC_KERNEL_CMD(causal_conv1d_update_half, block_dim, x_t, weight_t, conv_state_t,
214+
has_indices ? conv_state_indices : at::empty(0, x_t.options()),
215+
has_bias ? bias_t : at::empty(0, x_t.options()),
216+
has_num_accept ? num_accepted_tokens : at::empty(0, x_t.options()),
217+
has_query_loc ? query_start_loc : at::empty(0, x_t.options()), y_t, workspace_tensor,
218+
tilingTensor);
219+
}
220+
};
221+
222+
// The kernel allocates ~26*dim bytes of UB (width*dim weight + several dim*sizeof
223+
// operand/accum buffers). Atlas A3 has 192KB UB per AIV core, so dim much beyond
224+
// ~7500 overflows ("ub address out of bounds" / vector core exception). Models
225+
// like Qwen3-Next have a much larger conv_dim, so we tile along the dim axis
226+
// here at the integration layer and call the kernel per tile.
227+
constexpr int64_t kMaxDimTile = 4096;
228+
if (dim <= kMaxDimTile) {
229+
launch_for_tile(x, weight, conv_state, bias, y, dim);
230+
return y;
158231
}
159232

160-
// 4. Create workspace
161-
auto workspace_tensor =
162-
at::empty({workspace_size}, at::TensorOptions().dtype(at::kByte).device(x.options().device()));
163-
164-
// 5. Launch kernel
165-
if (dtype == at::kBFloat16) {
166-
EXEC_KERNEL_CMD(causal_conv1d_update_bfloat16_t, block_dim, x, weight, conv_state,
167-
has_indices ? conv_state_indices : at::empty(0, x.options()),
168-
has_bias ? bias : at::empty(0, x.options()),
169-
has_num_accept ? num_accepted_tokens : at::empty(0, x.options()),
170-
has_query_loc ? query_start_loc : at::empty(0, x.options()), y, workspace_tensor, tilingTensor);
171-
} else {
172-
EXEC_KERNEL_CMD(causal_conv1d_update_half, block_dim, x, weight, conv_state,
173-
has_indices ? conv_state_indices : at::empty(0, x.options()),
174-
has_bias ? bias : at::empty(0, x.options()),
175-
has_num_accept ? num_accepted_tokens : at::empty(0, x.options()),
176-
has_query_loc ? query_start_loc : at::empty(0, x.options()), y, workspace_tensor, tilingTensor);
233+
const int64_t num_tiles = (dim + kMaxDimTile - 1) / kMaxDimTile;
234+
for (int64_t t = 0; t < num_tiles; ++t) {
235+
const int64_t tile_start = t * kMaxDimTile;
236+
const int64_t tile_end = std::min(tile_start + kMaxDimTile, dim);
237+
const int64_t tile_dim = tile_end - tile_start;
238+
239+
// Slice along the dim axis. Slicing yields non-contiguous views; the kernel
240+
// requires contiguous inputs, so we materialize tile copies and write the
241+
// results back after the kernel returns.
242+
auto x_tile = x.slice(2, tile_start, tile_end).contiguous();
243+
auto weight_tile = weight.slice(1, tile_start, tile_end).contiguous();
244+
auto conv_state_tile = conv_state.slice(2, tile_start, tile_end).contiguous();
245+
auto bias_tile = has_bias ? bias.slice(0, tile_start, tile_end).contiguous() : bias;
246+
auto y_tile = at::empty_like(x_tile);
247+
248+
launch_for_tile(x_tile, weight_tile, conv_state_tile, bias_tile, y_tile, tile_dim);
249+
250+
// Stitch the tile outputs back into the full tensors. conv_state is updated
251+
// in-place by the kernel on the tile copy, so we must copy it back.
252+
y.slice(2, tile_start, tile_end).copy_(y_tile);
253+
conv_state.slice(2, tile_start, tile_end).copy_(conv_state_tile);
177254
}
178255

179256
return y;

‎csrc/causal_conv1d_update/op_kernel/causal_conv1d_update.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ __aicore__ inline void CausalConv1dUpdate<T>::Init(GM_ADDR x, GM_ADDR weight, GM
122122
tilingData_.hasNumAccept = tiling_data->hasNumAccept;
123123
tilingData_.hasQueryLoc = tiling_data->hasQueryLoc;
124124
tilingData_.activationMode = tiling_data->activationMode;
125+
tilingData_.padSlotId = tiling_data->padSlotId;
125126
this->ParseCoreBlocks(tilingData_, blockIdx_, batchNum_);
126127

127128
// alloc TQue

‎csrc/causal_conv1d_update/op_kernel/causal_conv1d_update_base.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ __aicore__ inline void CausalConv1dUpdateBase<T>::ParseTilingData(
5555
runTilingData.hasNumAccept = tilingData->hasNumAccept;
5656
runTilingData.hasQueryLoc = tilingData->hasQueryLoc;
5757
runTilingData.activationMode = tilingData->activationMode;
58+
runTilingData.padSlotId = tilingData->padSlotId;
5859
}
5960

6061
template <typename T>

‎csrc/pytorch_extensions.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ TORCH_LIBRARY_FRAGMENT(npu, m)
109109
m.def("triangular_inverse(Tensor x) -> Tensor");
110110

111111
m.def(
112-
"causal_conv1d_update(Tensor x, Tensor weight, Tensor conv_state, "
112+
"causal_conv1d_update(Tensor x, Tensor weight, Tensor(a!) conv_state, "
113113
"Tensor conv_state_indices, Tensor? bias=None, Tensor? num_accepted_tokens=None, "
114114
"Tensor? query_start_loc=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor");
115115

‎tests/python/sgl_kernel_npu/test_conv1d_update.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ def test_correctness_fixed():
113113
# --- Config ---
114114
BSZ = 8
115115
HIDDEN_SIZE = 4096
116-
SEQ_LEN = 2
116+
SEQ_LEN = 3
117117
KERNEL_SIZE = 3
118118
CACHE_LEN = 65
119119
DTYPE = torch.bfloat16
@@ -243,11 +243,11 @@ def test_npu_causal_conv1d_update():
243243
return
244244

245245
# --- Config ---
246-
BSZ = 1
247-
HIDDEN_SIZE = 4096 # Keep the hidden size moderate so the test stays fast.
248-
SEQ_LEN = 1
246+
BSZ = 32
247+
HIDDEN_SIZE = 12288 # Keep the hidden size moderate so the test stays fast.
248+
SEQ_LEN = 4
249249
KERNEL_SIZE = 4
250-
CACHE_LEN = 10
250+
CACHE_LEN = 32
251251
# conv_state must be large enough to hold (width - 1) + (seq_len - 1) values.
252252
CONV_STATE_LEN = (
253253
KERNEL_SIZE - 1 + SEQ_LEN - 1

0 commit comments

Comments
 (0)