Upload 9 files

added_tokens.json

@@ -0,0 +1,10 @@
+{
+  "<|execute_end|>": 73444,
+  "<|execute_start|>": 73443,
+  "<|fim_middle|>": 73446,
+  "<|fim_prefix|>": 73445,
+  "<|fim_suffix|>": 73447,
+  "<|im_end|>": 73440,
+  "<|im_start|>": 73441,
+  "<|tool_call|>": 73442
+}

cis_pooling.py

@@ -0,0 +1,118 @@
+import triton
+import triton.language as tl
+import torch
+
+MAX_LEN = 32768
+
+@triton.jit
+def nosa_mean_pool_kernel(
+    cis_ptr,            # [N, H]
+    cu_seqlens_ptr,     # int32 [B]
+    result_ptr,         # [H, N, M]
+    cis_stride_n,        # int32
+    cis_stride_h,        # int32
+    result_stride_h,     # int32
+    result_stride_n,     # int32
+    result_stride_m,     # int32
+    N, 
+    H,
+    M, 
+    kernel_size: tl.constexpr,
+    stride,
+    MAX_LEN: tl.constexpr,
+):
+    # grid: (H, B, M)
+    tidx_h = tl.program_id(0)  # head
+    tidx_b = tl.program_id(1)  # batch idx
+    tidx_m = tl.program_id(2)  # window idx
+
+
+    batch_start = tl.load(cu_seqlens_ptr + tidx_b)
+    batch_end = tl.load(cu_seqlens_ptr + tidx_b + 1)
+
+    block_idx = tl.arange(0, kernel_size)
+
+    beg_pos = cis_ptr + tidx_h * cis_stride_h + (batch_start + tidx_m * stride) * cis_stride_n
+
+    block_cis_ptrs = beg_pos + block_idx * cis_stride_n
+    mask = (block_idx + tidx_m * stride) < (batch_end - batch_start)
+    block_scores = tl.load(
+        block_cis_ptrs, 
+        mask=mask,
+        other=0.0,
+    )
+
+    # 对block_scores做平均值，注意mask要对, 分母上是mask的有效元素数
+    val_cnt = tl.sum(mask.to(tl.int32), axis=0)
+    acc = tl.sum(block_scores, axis=0) / val_cnt
+    
+    if tidx_m * stride + kernel_size <= batch_end - batch_start:
+        write_pos = result_ptr + tidx_h * result_stride_h + batch_start * result_stride_n + tidx_m * result_stride_m
+        write_idx = tl.arange(0, MAX_LEN)
+        write_ptrs = write_pos + write_idx * result_stride_n
+        tl.store(write_ptrs, acc, mask=write_idx < batch_end - batch_start)
+
+def nosa_mean_pooling(cis_score, cu_seqlens, max_seqlen, kernel_size=32, stride=16):
+    """
+    cis_score: [N, H] (torch.Tensor, float32/bfloat16/float16都行，但triton里先用float32)
+    cu_seqlens: [B+1] (torch.int32)
+    """
+    assert kernel_size == 32 and stride == 16
+
+    N, H = cis_score.shape
+    B = cu_seqlens.numel() - 1
+    M = max_seqlen // stride - 1  # 每个batch最大窗口数
+    M = max(M, 0) # bug fix
+    assert max_seqlen < MAX_LEN, f"Please increate MAX_LEN, MAX_LEN: {MAX_LEN}, max_seqlen: {max_seqlen}"
+
+    result = torch.zeros((H, N, M), dtype=cis_score.dtype, device=cis_score.device)
+
+    grid = (H, B, M)
+    nosa_mean_pool_kernel[grid](
+        cis_score,
+        cu_seqlens,
+        result,
+        cis_score.stride(0),
+        cis_score.stride(1),
+        result.stride(0),
+        result.stride(1),
+        result.stride(2),
+        N, H, M, kernel_size, stride, MAX_LEN
+    )
+
+    return result
+
+
+def main():
+    torch.manual_seed(0)
+    device = "cuda"
+
+    # 模拟数据
+    B = 2
+    H = 4
+    lens = [67, 1432]  # 每个 batch 的长度
+    cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(lens), dim=0)), dtype=torch.int32, device=device)
+    N = cu_seqlens[-1].item()
+    max_seqlen = max(lens)
+
+    cis_score = torch.randn(N, H, device=device, dtype=torch.bfloat16)
+
+    # Triton 版本
+    result = nosa_mean_pooling(cis_score, cu_seqlens, max_seqlen, kernel_size=32, stride=16)
+
+    # PyTorch baseline: 对每个 batch 做 pooling 然后广播
+    M = max_seqlen // 16 - 1
+    baseline = torch.zeros((H, N, M), device=device, dtype=torch.bfloat16)
+    for b in range(B):
+        start, end = cu_seqlens[b].item(), cu_seqlens[b+1].item()
+        seq = cis_score[start:end].T.unsqueeze(0)  # [1, H, L]
+        pooled = torch.nn.functional.avg_pool1d(seq, kernel_size=32, stride=16)  # [1, H, m]
+        pooled = pooled.squeeze(0)  # [H, m]
+        baseline[:, start:end, :pooled.size(-1)] = pooled.unsqueeze(1).expand(H, end-start, pooled.size(-1))
+
+    # 检查差异
+    max_diff = (result - baseline).abs().max()
+    print("Triton vs PyTorch max diff:", max_diff.item())
+
+if __name__ == "__main__":
+    main()

config.json

@@ -0,0 +1,35 @@
+{
+  "_name_or_path": "openbmb/CPM-2B",
+  "architectures": [
+    "SparseLlamaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoModelForCausalLM": "modeling_llama_long_infllmv2.SparseLlamaForCausalLM"
+  },
+  "bos_token_id": 1,
+  "eos_token_id": [2,73440],
+  "pad_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.1,
+  "intermediate_size": 16384,
+  "head_dim": 128,
+  "max_position_embeddings": 32768,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "model_type": "llama",
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "rope_type": "longrope",
+    "attention_factor": 1.0,
+    "long_factor": [0.9977997200264581, 1.014658295992452, 1.0349680404997148, 1.059429246056193, 1.0888815016813513, 1.1243301355211495, 1.166977103606075, 1.2182568066927284, 1.2798772354275727, 1.3538666751582975, 1.4426259039919596, 1.5489853358570191, 1.6762658237220625, 1.8283407612492941, 2.0096956085876183, 2.225478927469756, 2.481536379650452, 2.784415934557119, 3.1413289096347365, 3.560047844772632, 4.048719380066383, 4.615569542115128, 5.2684819496549835, 6.014438591970396, 6.858830049237097, 7.804668263503327, 8.851768731513417, 9.99600492938444, 11.228766118181639, 12.536757560834843, 13.902257701387796, 15.303885189125953, 16.717837610115794, 18.119465097853947, 19.484965238406907, 20.792956681060105, 22.02571786985731, 23.16995406772833, 24.217054535738416, 25.16289275000465, 26.007284207271347, 26.753240849586767, 27.40615325712662, 27.973003419175363, 28.461674954469114, 28.880393889607006, 29.237306864684626, 29.540186419591297, 29.79624387177199, 30.01202719065413, 30.193382037992453, 30.34545697551969, 30.47273746338473, 30.579096895249787, 30.66785612408345, 30.741845563814174, 30.80346599254902, 30.85474569563567, 30.897392663720595, 30.932841297560394, 30.962293553185553, 30.986754758742034, 31.007064503249293, 31.02392307921529],
+    "short_factor": [0.9977997200264581, 1.014658295992452, 1.0349680404997148, 1.059429246056193, 1.0888815016813513, 1.1243301355211495, 1.166977103606075, 1.2182568066927284, 1.2798772354275727, 1.3538666751582975, 1.4426259039919596, 1.5489853358570191, 1.6762658237220625, 1.8283407612492941, 2.0096956085876183, 2.225478927469756, 2.481536379650452, 2.784415934557119, 3.1413289096347365, 3.560047844772632, 4.048719380066383, 4.615569542115128, 5.2684819496549835, 6.014438591970396, 6.858830049237097, 7.804668263503327, 8.851768731513417, 9.99600492938444, 11.228766118181639, 12.536757560834843, 13.902257701387796, 15.303885189125953, 16.717837610115794, 18.119465097853947, 19.484965238406907, 20.792956681060105, 22.02571786985731, 23.16995406772833, 24.217054535738416, 25.16289275000465, 26.007284207271347, 26.753240849586767, 27.40615325712662, 27.973003419175363, 28.461674954469114, 28.880393889607006, 29.237306864684626, 29.540186419591297, 29.79624387177199, 30.01202719065413, 30.193382037992453, 30.34545697551969, 30.47273746338473, 30.579096895249787, 30.66785612408345, 30.741845563814174, 30.80346599254902, 30.85474569563567, 30.897392663720595, 30.932841297560394, 30.962293553185553, 30.986754758742034, 31.007064503249293, 31.02392307921529],
+    "original_max_position_embeddings": 32768
+  },
+  "rope_theta": 10000.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.36.0",
+  "use_cache": true,
+  "vocab_size": 73448
+}

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": [
+    2,
+    73440
+  ],
+  "pad_token_id": 2,
+  "temperature": 0.8,
+  "top_p": 0.8,
+  "transformers_version": "4.46.1"
+}

special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "additional_special_tokens": [
+    "<|im_end|>",
+    "<|im_start|>",
+    "<|tool_call|>",
+    "<|execute_start|>",
+    "<|execute_end|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>"
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb74d51116831c3bf65db812c553f94ab0c88dcf97a5bbb37e3504f6d359c530
+size 1181204

tokenizer_config.json

@@ -0,0 +1,117 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73440": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73441": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73442": {
+      "content": "<|tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73443": {
+      "content": "<|execute_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73444": {
+      "content": "<|execute_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73445": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73446": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73447": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_end|>",
+    "<|im_start|>",
+    "<|tool_call|>",
+    "<|execute_start|>",
+    "<|execute_end|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>"
+  ],
+  "bos_token": "<s>",
+  "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% if enable_thinking is defined and enable_thinking is false %}{{ '<think>\n\n</think>\n' }}{% endif %}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}