[BUG][enflame/gcu300] index_put grid.x 无上限分块崩溃（M > 4×65535 = 262140，batch≥2 采样触发）

[JosephNew]

现象

燧原 enflame（S60 / GCU300）上起 Qwen3.6-27B 服务，采样阶段崩溃、服务整体 SIGQUIT：

triton.runtime.errors.OutOfResources: out of resource: grid.x, Required: 124155, Hardware limit: 65535
SIGQUIT received. signum=None, frame=None. It usually means one child failed.

非 OOM（显存充足），是 triton kernel 的 grid.x 超 GCU300 的硬件上限。

环境

• FlagGems 5.4.0-rc2 HEAD ad89ec47 / FlagTree 0.7.0-rc2-triton3.6 803ea98b（enflame backend）/ sglang-plugin-FL 1c84db70
• Qwen3.6-27B（vocab_size = 248320，generation_config.json 里 top_k=20 / top_p=0.95），tp=4，--attention-backend fa3，黑名单 isin,_unique2
• 机器 S60 8 卡（GCU300），torch 2.11.0+cpu（设备名 gcu）

具体问题行定位

问题行：src/flag_gems/runtime/backend/_enflame/gcu300/ops/index_put.py:248

# generate_index_put_wrapper() 内部，line 237 起
def grid(meta):
    block_size0, block_size1 = mmu_safe_index_put_block_sizes(
        meta['BLOCK_SIZE0'], meta['BLOCK_SIZE1'],
        input_stride, input_shape, element_size,
    )
    return (
        triton.cdiv(M, block_size0),            # ← line 248：grid.x，没有任何上限
        min(triton.cdiv(N, block_size1), 255),  # ← line 250：grid.y 有 255 兜底
    )

生成的实例（indices_len=2, inp_rank=2 这一档）：

/root/.flaggems/code_cache/index_put_inp_rank_2_indices_len_2_index_rank_1.py:114
        return (
            triton.cdiv(M, block_size0),        # ← 114：grid.x 裸奔
            triton.cdiv(N, block_size1),        # ← 115
        )

/root/.flaggems/code_cache/index_put_inp_rank_2_indices_len_2_index_rank_1.py:118
    _index_put_jit_function[grid](              # ← 118：这里抛 OutOfResources

完整调用链（真机 traceback，行号已核对）：

sglang/srt/layers/sampler.py:470        probs_sort[...] = 0.0          ← aten::index_put_
  flag_gems/.../_enflame/gcu300/ops/index_put.py:480   _index_put_func(inp, tensor_indices, values, accumulate)
  flag_gems/.../_enflame/gcu300/ops/index_put.py:332   return overload(*args, **kwargs)
  /root/.flaggems/code_cache/index_put_inp_rank_2_indices_len_2_index_rank_1.py:118  _index_put_jit_function[grid](
  triton/runtime/autotuner.py:509         return self.fn.run(*args, **kwargs)
  triton/backends/enflame/backend.py:722  raise OutOfResources(grid_0, 65535, "grid.x")

根因

生成的 wrapper 里 grid.x = cdiv(M, BLOCK_SIZE0) 没有做上限分块，而 BLOCK_SIZE0 只有 tune config 里的唯一候选 4。

• src/flag_gems/runtime/backend/_enflame/gcu300/tune_configs.yaml → index_put.block_size0: [4]
• 于是 M = indices[0].numel() > 4 × 65535 = 262140 时必然崩

mmu_safe_index_put_block_sizes（utils/shape_utils.py:86）的两条 clamp 在这里都进不去：它要求索引跨度 index_stride × (shape[0]-1) × element_size ≥ 512MiB，而采样张量 (B, 248320) float32 的跨度只有几 MB，差三个数量级 —— 所以 block_size0 保持 4，不会被压成 1。

M = 496619 是怎么来的（服务里那一发）

sglang SamplingBatchInfo：

is_all_greedy = all(r.sampling_params.top_k <= 1 for r in reqs)

只要 batch 里混进一个非 greedy 请求，整批都走 pytorch 采样器、整批都做这行掩码赋值（sampler.py:470）：

probs_sort[torch.arange(0, V).view(1, -1) >= top_ks.view(-1, 1)] = 0.0

• 一行 temperature=0（sglang 归一化成 top_k=1）→ 掩掉 248320 - 1 = 248319 个
• 一行吃模型默认 top_k=20 → 掩掉 248320 - 20 = 248300 个
• 合计 M = 496619 → grid.x = cdiv(496619, 4) = 124155

⚠️ 触发条件是 batch ≥ 2。单请求 M 最大就是 vocab_size = 248320 → grid.x = 62080，不会崩。

分析方式

1. 拿到崩点：起服务 + 两条并发请求（见下），从 traceback 直接读到 index_put.py:480 → 生成的 _index_put_wrapper:118。
2. 读生成器定位公式：index_put.py 的 kernel 是 codegen 出来的，看 generate_index_put_wrapper()（line 216 起）就能看到 grid 的计算式，对比 grid.y 有 min(..., 255)、grid.x 什么都没有。
3. 扫阈值验证：用 torch.index_put_ 单独扫 M，实测边界精确到 1 —— M = 262140 → PASS、M = 262141 → Required: 65536 CRASH，与 4 × 65535 完全吻合。
4. 对照组：同样 shape 不开 flag_gems（另起进程）全部 PASS → 崩的是 flag_gems 的 kernel，不是 torch。
5. 崩溃物证（日志丢了也能用）：@triton.autotune 调优时任何 config 启动失败，_bench 返回 inf，而 Triton 照样把结果落盘，文件 mtime 就是崩溃时刻。容器里 grep -rl '"Infinity"' /root/.triton/cache/*/*.autotune.json 只有一条：

/root/.triton/cache/ITVV3YVFYCUPXKV6TKUAPLCCNQOE452XSZLQ2MA6PYWT6NGVIL3Q/_index_put_jit_function.autotune.json
mtime = 2026-09-18 06:04:47   （= 请求死掉的同一秒）

{"key": [496619, 1, "torch.float32", "torch.int32", "torch.int32", "torch.float32"],
 "configs_timings": [
   [{"kwargs": {"BLOCK_SIZE0": 4, "BLOCK_SIZE1": 4096}, ...}, [Infinity, Infinity, Infinity]],
   [{"kwargs": {"BLOCK_SIZE0": 4, "BLOCK_SIZE1": 2048}, ...}, [Infinity, Infinity, Infinity]]]}

key[0] = 496619（= M）、BLOCK_SIZE0 = 4，与上面的推导逐字吻合。

最小化测试脚本

自包含，不需要模型、不需要起服务：

import torch
import flag_gems

flag_gems.enable()  # 必须显式 enable，import 不注册算子

VOCAB = 248320  # Qwen3.6-27B vocab_size
DEVICE = flag_gems.device  # 燧原上是 "gcu"


def probe(top_ks):
    """top_ks: 逐行的 top_k，复刻 sglang 的掩码赋值 probs_sort[mask] = 0.0"""
    batch_size = len(top_ks)
    probs = torch.rand(batch_size, VOCAB, device=DEVICE, dtype=torch.float32)
    column = torch.arange(0, VOCAB, device=DEVICE).view(1, -1)
    keep = torch.tensor(top_ks, device=DEVICE, dtype=torch.int64).view(-1, 1)
    mask = column >= keep
    number_of_indices = int(mask.sum().item())
    grid_x = (number_of_indices + 3) // 4  # BLOCK_SIZE0 = 4
    print("top_ks=%-10s M=%-8d 预测 grid.x=%-8d" % (top_ks, number_of_indices, grid_x),
          end="")
    probs[mask] = 0.0  # <<< 崩在这里
    print("-> PASS")


for top_ks in ([1], [1, 1], [1, 20]):
    try:
        probe(top_ks)
    except Exception as error:
        print("-> %s: %s" % (type(error).__name__, str(error).splitlines()[0]))

GCU300 实测输出：

top_ks=[1]        M=248319   预测 grid.x=62080   -> PASS
top_ks=[1, 1]     M=496638   预测 grid.x=124160  -> OutOfResources: out of resource: grid.x, Required: 124160, Hardware limit: 65535.
top_ks=[1, 20]    M=496619   预测 grid.x=124155  -> OutOfResources: out of resource: grid.x, Required: 124155, Hardware limit: 65535.

端到端（起服务）

关键在于两条请求必须并发，才能进同一个 decode batch：

• 起服务：TOPS_VISIBLE_DEVICES=4,5,6,7 TORCH_GCU_ENABLE_INT64_AND_UINT64=1 ENABLE_I64_CHECK=0 TORCHDYNAMO_DISABLE=1 FLAGCX_PATH=/sgl-workspace/FlagCX SGLANG_FL_DIST_BACKEND=flagcx SGLANG_FL_FLAGOS_BLACKLIST=isin,_unique2 python3 -m sglang.launch_server --model-path /models/Qwen3.6-27B --tp 4 --mem-fraction-static 0.6 --attention-backend fa3 --reasoning-parser qwen3 --max-running-requests 64
• 请求 A：{"temperature": 0, "max_tokens": 512} → top_k=1
• 请求 B：{"temperature": 0.7, "max_tokens": 512} → 吃模型默认 top_k=20
• 两条同时发 → M = 496619 → 约 30 秒内 SIGQUIT

影响

M = indices[0].numel() > 262140 的 index_put_ 调用一律崩，触发面不限于采样：

• 采样路径：batch ≥ 2 且掩码覆盖面大（top_k / top_p / min_p 任一非默认，长 vocab 模型很容易踩）
• 任何 2-D 及以上、索引元素数 > 262140 的 tensor[idx] = value

建议修复

同目录 _enflame/gcu300/ops/index.py 里已经有正确做法可以直接照抄 —— index.py:278 起有
_GCU_MAX_GRID = 65535 / _GCU_MAX_GRID_Y = 255 的 chunk 拆分 + kernel 内循环 + 断言，
index_put.py 缺的就是这一套。或退一步：grid.x 也加 cap，
kernel 内用 for pid0 in tl.range(tl.program_id(0), num_blocks_m, tl.num_programs(0)) 补循环。